Make labeled subobj specs easier in c++. Now you can do BSON( "a" << GT << 4 << LTE << 22.3 << NE << "foo" )

3 files changed, 158 insertions, 11 deletions

diff --git a/db/jsobj.cpp b/db/jsobj.cpp
index c1d87f00cf5..855e4e34dbf 100644
--- a/db/jsobj.cpp
+++ b/db/jsobj.cpp
@@ -1090,4 +1090,11 @@ namespace mongo {
 */
     }
 
+    Labeler::Label GT( "$gt" );
+    Labeler::Label GTE( "$gte" );
+    Labeler::Label LT( "$lt" );
+    Labeler::Label LTE( "$lte" );
+    Labeler::Label NE( "$ne" );
+    Labeler::Label IN( "$in" );
+    
 } // namespace mongo
diff --git a/db/jsobj.h b/db/jsobj.h
index b2edb64a2f2..85f78f30623 100644
--- a/db/jsobj.h
+++ b/db/jsobj.h
@@ -36,6 +36,7 @@ namespace mongo {
     class BSONObj;
     class Record;
     class BSONObjBuilder;
+    class BSONObjBuilderValueStream;
 
 #pragma pack(1)
 
@@ -730,12 +731,36 @@ namespace mongo {
 */
 #define BSON(x) (( BSONObjBuilder() << x ).obj())
 
+    class Labeler {
+    public:
+        struct Label {
+            Label( const char *l ) : l_( l ) {}
+            const char *l_;
+        };
+        Labeler( const Label &l, BSONObjBuilderValueStream *s ) : l_( l ), s_( s ) {}
+        template<class T>
+        BSONObjBuilder& operator<<( T value );
+    private:
+        const Label &l_;
+        BSONObjBuilderValueStream *s_;
+    };
+    
+    extern Labeler::Label GT;
+    extern Labeler::Label GTE;
+    extern Labeler::Label LT;
+    extern Labeler::Label LTE;
+    extern Labeler::Label NE;
+    extern Labeler::Label IN;
+    
     class BSONObjBuilderValueStream {
     public:
-        BSONObjBuilderValueStream( const char * fieldName , BSONObjBuilder * builder );
+        friend class Labeler;
+        BSONObjBuilderValueStream( BSONObjBuilder * builder );
 
         template<class T> 
         BSONObjBuilder& operator<<( T value );
+        
+        Labeler operator<<( const Labeler::Label &l );
 /*
         BSONObjBuilder& operator<<( const char * value );
         BSONObjBuilder& operator<<( const string& v ) { return (*this << v.c_str()); }
@@ -744,11 +769,15 @@ namespace mongo {
         BSONObjBuilder& operator<<( const unsigned long value ){ return (*this << (double)value); }
 */
 
+        void endField( const char *nextFieldName = 0 );
+        bool subobjStarted() const { return _fieldName; }
+        
     private:
         const char * _fieldName;
         BSONObjBuilder * _builder;
+        auto_ptr< BSONObjBuilder > _subobj;
     };
-
+    
     /**
        utility for creating a BSONObj
      */
@@ -985,15 +1014,23 @@ namespace mongo {
         }
 
         /** Stream oriented way to add field names and values. */
-        BSONObjBuilderValueStream operator<<(const char * name ) {
-            return BSONObjBuilderValueStream( name , this );
+        BSONObjBuilderValueStream &operator<<(const char * name ) {
+            if ( !s_.get() )
+                s_.reset( new BSONObjBuilderValueStream( this ) );
+            s_->endField( name );
+            return *s_;
         }
 
         /** Stream oriented way to add field names and values. */
-        BSONObjBuilderValueStream operator<<( string name ) {
-            return BSONObjBuilderValueStream( name.c_str() , this );
+        BSONObjBuilderValueStream &operator<<( string name ) {
+            return operator<<( name.c_str() );
         }
 
+        Labeler operator<<( const Labeler::Label &l ) {
+            massert( "Value stream expected", s_.get() );
+            massert( "No subobject started", s_->subobjStarted() );
+            return *s_ << l;
+        }
 
     private:
         // Append the provided arr object as an array.
@@ -1004,6 +1041,8 @@ namespace mongo {
         }
 
         char* _done() {
+            if ( s_.get() )
+                s_->endField();
             b.append((char) EOO);
             char *data = b.buf();
             *((int*)data) = b.len();
@@ -1011,6 +1050,7 @@ namespace mongo {
         }
 
         BufBuilder b;
+        auto_ptr< BSONObjBuilderValueStream > s_;
     };
 
 
@@ -1187,15 +1227,35 @@ namespace mongo {
         return false;
     }
 
-    inline BSONObjBuilderValueStream::BSONObjBuilderValueStream( const char * fieldName , BSONObjBuilder * builder ) {
-        _fieldName = fieldName;
+    inline BSONObjBuilderValueStream::BSONObjBuilderValueStream( BSONObjBuilder * builder ) {
+        _fieldName = 0;
         _builder = builder;
+        _subobj.reset( new BSONObjBuilder() );
     }
-
+    
     template<class T> inline 
     BSONObjBuilder& BSONObjBuilderValueStream::operator<<( T value ) { 
         _builder->append(_fieldName, value);
+        _fieldName = 0;
         return *_builder;
     }
 
+    inline Labeler BSONObjBuilderValueStream::operator<<( const Labeler::Label &l ) { 
+        return Labeler( l, this );
+    }
+
+    inline void BSONObjBuilderValueStream::endField( const char *nextFieldName ) {
+        if ( _fieldName ) {
+            _builder->append( _fieldName, _subobj->done() );
+            _subobj.reset( new BSONObjBuilder() );
+        }
+        _fieldName = nextFieldName;
+    }    
+
+    template<class T> inline
+    BSONObjBuilder& Labeler::operator<<( T value ) {
+        s_->_subobj->append( l_.l_, value );
+        return *s_->_builder;
+    }    
+        
 } // namespace mongo
diff --git a/dbtests/jsobjtests.cpp b/dbtests/jsobjtests.cpp
index 1199858f97c..5e201c8de4b 100644
--- a/dbtests/jsobjtests.cpp
+++ b/dbtests/jsobjtests.cpp
@@ -433,15 +433,89 @@ namespace JsobjTests {
                 OID a;
                 OID b;
                 
-                cout << endl;
                 a.init();
                 b.init( a.str() );
                 
                 ASSERT( a == b );
             }
         };
-    }
+    } // namespace OIDTests
+    
+    namespace ValueStreamTests {
+        
+        class LabelBase {
+        public:
+            void run() {
+                cout << expected().toString() << endl;
+                cout << actual().toString() << endl;
+                ASSERT( !expected().woCompare( actual() ) );
+            }
+        protected:
+            virtual BSONObj expected() = 0;
+            virtual BSONObj actual() = 0;
+        };
+        
+        class LabelBasic : public LabelBase {
+            BSONObj expected() {
+                return BSON( "a" << ( BSON( "$gt" << 1 ) ) );
+            }
+            BSONObj actual() {
+                return BSON( "a" << GT << 1 );
+            }            
+        };
+
+        class LabelShares : public LabelBase {
+            BSONObj expected() {
+                return BSON( "z" << "q" << "a" << ( BSON( "$gt" << 1 ) ) << "x" << "p" );
+            }
+            BSONObj actual() {
+                return BSON( "z" << "q" << "a" << GT << 1 << "x" << "p" );
+            }            
+        };        
+        
+        class LabelDouble : public LabelBase {
+            BSONObj expected() {
+                return BSON( "a" << ( BSON( "$gt" << 1 << "$lte" << "x" ) ) );
+            }
+            BSONObj actual() {
+                return BSON( "a" << GT << 1 << LTE << "x" );
+            }            
+        };        
+
+        class LabelDoubleShares : public LabelBase {
+            BSONObj expected() {
+                return BSON( "z" << "q" << "a" << ( BSON( "$gt" << 1 << "$lte" << "x" ) ) << "x" << "p" );
+            }
+            BSONObj actual() {
+                return BSON( "z" << "q" << "a" << GT << 1 << LTE << "x" << "x" << "p" );
+            }            
+        };        
 
+        class LabelMulti : public LabelBase {
+            BSONObj expected() {
+                return BSON( "z" << "q"
+                            << "a" << BSON( "$gt" << 1 << "$lte" << "x" )
+                            << "b" << BSON( "$ne" << 1 << "$ne" << "f" << "$ne" << 22.3 )
+                            << "x" << "p" );
+            }
+            BSONObj actual() {
+                return BSON( "z" << "q"
+                            << "a" << GT << 1 << LTE << "x"
+                            << "b" << NE << 1 << NE << "f" << NE << 22.3
+                            << "x" << "p" );
+            }            
+        };        
+        
+        class Unallowed {
+        public:
+            void run() {
+                ASSERT_EXCEPTION( BSON( GT << 4 ), MsgAssertionException );
+                ASSERT_EXCEPTION( BSON( "a" << 1 << GT << 4 ), MsgAssertionException );
+            }
+        };
+        
+    } // namespace ValueStreamTests
+    
     class All : public UnitTest::Suite {
     public:
         All() {
@@ -487,6 +561,12 @@ namespace JsobjTests {
             add< BSONObjTests::Validation::Fuzz >( .001 );
             add< OIDTests::init1 >();
             add< OIDTests::initParse1 >();
+            add< ValueStreamTests::LabelBasic >();
+            add< ValueStreamTests::LabelShares >();
+            add< ValueStreamTests::LabelDouble >();
+            add< ValueStreamTests::LabelDoubleShares >();
+            add< ValueStreamTests::LabelMulti >();
+            add< ValueStreamTests::Unallowed >();
         }
     };