Merge branch 'master' of github.com:mongodb/mongo

2 files changed, 111 insertions, 9 deletions

diff --git a/db/projection.cpp b/db/projection.cpp
index 519e3c5162f..85bb5463503 100644
--- a/db/projection.cpp
+++ b/db/projection.cpp
@@ -30,6 +30,9 @@ namespace mongo {
         while ( i.more() ){
             BSONElement e = i.next();
 
+            if ( ! e.isNumber() )
+                _hasNonSimple = true;
+
             if (e.type() == Object){
                 BSONObj obj = e.embeddedObject();
                 BSONElement e2 = obj.firstElement();
@@ -213,4 +216,73 @@ namespace mongo {
             }
         }
     }
+
+    Projection::KeyOnly* Projection::checkKey( const BSONObj& keyPattern ) const {
+        if ( _include ){
+            // if we default to including then we can't
+            // use an index because we don't know what we're missing
+            return 0;
+        }
+        
+        if ( _hasNonSimple )
+            return 0;
+        
+        if ( _includeID && keyPattern["_id"].eoo() )
+            return 0;
+
+        // at this point we know its all { x : 1 } style
+        
+        auto_ptr<KeyOnly> p( new KeyOnly() );
+
+        int got = 0;
+        BSONObjIterator i( keyPattern );
+        while ( i.more() ){
+            BSONElement k = i.next();
+            
+            if ( _source[k.fieldName()].type() ){
+                if ( ! _includeID && mongoutils::str::equals( k.fieldName() , "_id" ) ){
+                    p->addNo();
+                }
+                else {
+                    p->addYes( k.fieldName() );
+                    got++;
+                }
+            }
+            else if ( mongoutils::str::equals( "_id" , k.fieldName() ) && _includeID ){
+                p->addYes( "_id" );
+            }
+            else {
+                p->addNo();
+            }
+
+        }
+        
+        int need = _source.nFields();
+        if ( ! _includeID )
+            need--;
+
+        if ( got == need )
+            return p.release();
+
+        return 0;
+    }
+
+    BSONObj Projection::KeyOnly::hydrate( const BSONObj& key ) const {
+        assert( _include.size() == _names.size() );
+
+        BSONObjBuilder b( key.objsize() + _stringSize + 16 );
+
+        BSONObjIterator i(key);
+        unsigned n=0;
+        while ( i.more() ){
+            assert( n < _include.size() );
+            BSONElement e = i.next();
+            if ( _include[n] ){
+                b.appendAs( e , _names[n] );
+            }
+            n++;
+        }
+        
+        return b.obj();
+    }
 }
diff --git a/db/projection.h b/db/projection.h
index 65e27fc7a84..b64767e4094 100644
--- a/db/projection.h
+++ b/db/projection.h
@@ -27,13 +27,39 @@ namespace mongo {
      */
     class Projection {
     public:
-        Projection()
-            : _include(true)
-            , _special(false)
-            , _includeID(true)
-            , _skip(0)
-            , _limit(-1)
-        {}
+
+        class KeyOnly {
+        public:
+            
+            KeyOnly() : _stringSize(0){}
+
+            BSONObj hydrate( const BSONObj& key ) const;
+            
+            void addNo() { _add( false , "" ); }
+            void addYes( const string& name ) { _add( true , name ); }
+
+        private:
+            
+            void _add( bool b , const string& name ) { 
+                _include.push_back( b ); 
+                _names.push_back( name ); 
+                _stringSize += name.size(); 
+            }
+
+            vector<bool> _include; // one entry per field in key.  true iff should be in output
+            vector<string> _names; // name of field since key doesn't have names
+            
+            int _stringSize;
+        };
+
+        Projection() :
+            _include(true) ,
+            _special(false) ,
+            _includeID(true) ,
+            _skip(0) ,
+            _limit(-1) ,
+            _hasNonSimple(false){
+        }
         
         /**
          * called once per lifetime
@@ -59,10 +85,12 @@ namespace mongo {
 
         /**
          * @return if the key has all the information needed to return
+         *         return a new KeyOnly 
+         *         otherwise null
          *         NOTE: a key may have modified the actual data
-         *               which has to be handled above this
+         *               which has to be handled above this (arrays, geo)
          */
-        bool keyEnough( const BSONObj& keyPattern ) const;
+        KeyOnly* checkKey( const BSONObj& keyPattern ) const;
         
     private:
 
@@ -89,6 +117,8 @@ namespace mongo {
         // used for $slice operator
         int _skip;
         int _limit;
+
+        bool _hasNonSimple;
     };