From 218036b2a46b8bed6f46145c3b10c176f47a2b8f Mon Sep 17 00:00:00 2001
From: Aaron <aaron@10gen.com>
Date: Mon, 1 Mar 2010 10:07:01 -0800
Subject: [PATCH] SERVER-580 add putBack() to db client cursor

---
 client/dbclient.cpp     |  8 ++++++++
 client/dbclient.h       |  8 +++++++-
 dbtests/clienttests.cpp | 34 +++++++++++++++++++++++++++++++++-
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/client/dbclient.cpp b/client/dbclient.cpp
index 1079d6e8ca9..e3eb7a054e9 100644
--- a/client/dbclient.cpp
+++ b/client/dbclient.cpp
@@ -821,6 +821,9 @@ namespace mongo {
 
     /** If true, safe to call next().  Requests more from server if necessary. */
     bool DBClientCursor::more() {
+        if ( !_putBack.empty() )
+            return true;
+        
         if (haveLimit && pos >= nToReturn)
             return false;
 
@@ -836,6 +839,11 @@ namespace mongo {
 
     BSONObj DBClientCursor::next() {
         assert( more() );
+        if ( !_putBack.empty() ) {
+            BSONObj ret = _putBack.top();
+            _putBack.pop();
+            return ret;
+        }
         pos++;
         BSONObj o(data);
         data += o.objsize();
diff --git a/client/dbclient.h b/client/dbclient.h
index 787ca0a6232..5b958796846 100644
--- a/client/dbclient.h
+++ b/client/dbclient.h
@@ -210,7 +210,7 @@ namespace mongo {
             if you want to exhaust whatever data has been fetched to the client already but 
             then perhaps stop.
         */
-        bool moreInCurrentBatch() { return pos < nReturned; }
+        bool moreInCurrentBatch() { return !_putBack.empty() || pos < nReturned; }
 
         /** next
 		   @return next object in the result cursor.
@@ -219,6 +219,11 @@ namespace mongo {
            if you do not want to handle that yourself, call nextSafe().
         */
         BSONObj next();
+        
+        /** 
+            restore an object previously returned by next() to the cursor
+         */
+        void putBack( const BSONObj &o ) { _putBack.push( o.getOwned() ); }
 
 		/** throws AssertionException if get back { $err : ... } */
         BSONObj nextSafe() {
@@ -317,6 +322,7 @@ namespace mongo {
         int opts;
         int batchSize;
         auto_ptr<Message> m;
+        stack< BSONObj > _putBack;
 
         int resultFlags;
         long long cursorId;
diff --git a/dbtests/clienttests.cpp b/dbtests/clienttests.cpp
index 5ff4dfd9cfc..fb97397aa3e 100644
--- a/dbtests/clienttests.cpp
+++ b/dbtests/clienttests.cpp
@@ -111,7 +111,38 @@ namespace ClientTests {
             ASSERT_EQUALS( 1111, c->itcount() );
         }
     };
-    
+
+    class PushBack : public Base {
+    public:
+        PushBack() : Base( "PushBack" ) {}
+        void run() {
+            for( int i = 0; i < 10; ++i )
+                db.insert( ns(), BSON( "i" << i ) );
+            auto_ptr< DBClientCursor > c = db.query( ns(), Query().sort( BSON( "i" << 1 ) ) );
+            BSONObj o = c->next();
+            ASSERT( c->more() );
+            ASSERT( c->moreInCurrentBatch() );
+            c->putBack( o );
+            ASSERT( c->more() );
+            ASSERT( c->moreInCurrentBatch() );
+            o = c->next();
+            BSONObj o2 = c->next();
+            BSONObj o3 = c->next();
+            c->putBack( o3 );
+            c->putBack( o2 );
+            c->putBack( o );
+            for( int i = 0; i < 10; ++i ) {
+                o = c->next();
+                ASSERT_EQUALS( i, o[ "i" ].number() );
+            }
+            ASSERT( !c->more() );
+            ASSERT( !c->moreInCurrentBatch() );
+            c->putBack( o );
+            ASSERT( c->more() );
+            ASSERT( c->moreInCurrentBatch() );
+            ASSERT_EQUALS( 1, c->itcount() );
+        }
+    };
 
     class All : public Suite {
     public:
@@ -123,6 +154,7 @@ namespace ClientTests {
             add<ReIndex>();
             add<ReIndex2>();
             add<CS_10>();
+            add<PushBack>();
         }
         
     } all;