SERVER-8046 Add support for cloning elements to mutable

6 files changed, 199 insertions, 11 deletions

diff --git a/src/mongo/bson/mutable/const_element-inl.h b/src/mongo/bson/mutable/const_element-inl.h
index 2e928ce53ac..8da6fc1b8f3 100644
--- a/src/mongo/bson/mutable/const_element-inl.h
+++ b/src/mongo/bson/mutable/const_element-inl.h
@@ -176,6 +176,11 @@ namespace mutablebson {
         return _basis.toString();
     }
 
+    template<typename Builder>
+    inline void ConstElement::writeElement(Builder* builder, const StringData* fieldName) const {
+        return _basis.writeElement(builder, fieldName);
+    }
+
     inline bool operator==(const ConstElement& l, const ConstElement& r) {
         return l._basis == r._basis;
     }
@@ -200,5 +205,6 @@ namespace mutablebson {
         return !(l == r);
     }
 
+
 } // namespace mutablebson
 } // namespace mongo
diff --git a/src/mongo/bson/mutable/const_element.h b/src/mongo/bson/mutable/const_element.h
index 1253a0566eb..553e2c89079 100644
--- a/src/mongo/bson/mutable/const_element.h
+++ b/src/mongo/bson/mutable/const_element.h
@@ -93,6 +93,11 @@ namespace mutablebson {
         friend bool operator==(const ConstElement&, const ConstElement&);
 
     private:
+        friend class Document;
+
+        template<typename Builder>
+        inline void writeElement(Builder* builder, const StringData* fieldName = NULL) const;
+
         Element _basis;
     };
 
diff --git a/src/mongo/bson/mutable/document.cpp b/src/mongo/bson/mutable/document.cpp
index f93c0ae3258..79973cbf3e3 100644
--- a/src/mongo/bson/mutable/document.cpp
+++ b/src/mongo/bson/mutable/document.cpp
@@ -1104,11 +1104,8 @@ namespace mutablebson {
 
         if (thisRep->serialized) {
             // For leaf elements we just create a new Element with the current value and
-            // replace.
-            dassert(impl.hasValue(*thisRep));
-            Element replacement = getDocument().makeElementWithNewFieldName(
-                newName, impl.getSerializedElement(*thisRep));
-            // NOTE: This call will also invalidate thisRep.
+            // replace. Note that the 'setValue' call below will invalidate thisRep.
+            Element replacement = _doc->makeElementWithNewFieldName(newName, *this);
             setValue(&replacement);
         } else {
             // The easy case: just update what our field name offset refers to.
@@ -1756,7 +1753,7 @@ namespace mutablebson {
     } // namespace
 
     template<typename Builder>
-    void Element::writeElement(Builder* builder) const {
+    void Element::writeElement(Builder* builder, const StringData* fieldName) const {
         // No need to verify(ok()) since we are only called from methods that have done so.
         dassert(ok());
 
@@ -1764,10 +1761,15 @@ namespace mutablebson {
         const ElementRep& thisRep = impl.getElementRep(_repIdx);
 
         if (thisRep.serialized) {
-            builder->append(impl.getSerializedElement(thisRep));
+            BSONElement element = impl.getSerializedElement(thisRep);
+            if (fieldName)
+                builder->appendAs(element, *fieldName);
+            else
+                builder->append(element);
         } else {
             const BSONType type = impl.getType(thisRep);
-            SubBuilder<Builder> subBuilder(builder, type, impl.getFieldName(thisRep));
+            const StringData subName = fieldName ? *fieldName : impl.getFieldName(thisRep);
+            SubBuilder<Builder> subBuilder(builder, type, subName);
             if (type == mongo::Array) {
                 BSONArrayBuilder child_builder(subBuilder.buffer);
                 writeChildren(&child_builder);
@@ -2092,7 +2094,7 @@ namespace mutablebson {
         case EOO:
             verify(false);
         case NumberDouble:
-            return makeElementDouble(fieldName,value._numberDouble());
+            return makeElementDouble(fieldName, value._numberDouble());
         case String:
             return makeElementString(fieldName,
                                      StringData(value.valuestr(), value.valuestrsize() - 1));
@@ -2160,6 +2162,40 @@ namespace mutablebson {
         }
     }
 
+    Element Document::makeElement(ConstElement element) {
+        return makeElement(element, NULL);
+    }
+
+    Element Document::makeElementWithNewFieldName(const StringData& fieldName,
+                                                  ConstElement element) {
+        return makeElement(element, &fieldName);
+    }
+
+    Element Document::makeElement(ConstElement element, const StringData* fieldName) {
+        if (this == &element.getDocument()) {
+            // If the Element that we want to build from belongs to this Document, then we have
+            // to first copy it to the side, and then back in, since otherwise we might be
+            // attempting both read to and write from the underlying BufBuilder simultaneously,
+            // which will not work.
+            BSONObjBuilder builder;
+            element.writeElement(&builder, fieldName);
+            BSONObj built = builder.obj();
+            BSONElement newElement = built.firstElement();
+            return makeElement(newElement);
+        } else {
+            // If the Element belongs to another document, then we can just stream it into our
+            // builder. We still do need to dassert that the field name doesn't alias us
+            // somehow.
+            Impl& impl = getImpl();
+            if (fieldName)
+                dassert(impl.doesNotAlias(*fieldName));
+            BSONObjBuilder& builder = impl.leafBuilder();
+            const int leafRef = builder.len();
+            element.writeElement(&builder, fieldName);
+            return Element(this, impl.insertLeafElement(leafRef));
+        }
+    }
+
     Element Document::end() {
         return Element(this, kInvalidRepIdx);
     }
diff --git a/src/mongo/bson/mutable/document.h b/src/mongo/bson/mutable/document.h
index 375094c9111..8698cda6e61 100644
--- a/src/mongo/bson/mutable/document.h
+++ b/src/mongo/bson/mutable/document.h
@@ -367,6 +367,18 @@ namespace mutablebson {
          */
         Element makeElementSafeNum(const StringData& fieldName, SafeNum value);
 
+        /** Construct a new element with the same name, type, and value as the provided mutable
+         *  Element. The data is copied from the given Element. Unlike most methods in this
+         *  class the provided Element may be from a different Document.
+         */
+        Element makeElement(ConstElement elt);
+
+        /** Construct a new Element with the same type and value as the provided mutable
+         *  Element, but with a new field name. The data is copied from the given
+         *  Element. Unlike most methods in this class the provided Element may be from a
+         *  different Document.
+         */
+        Element makeElementWithNewFieldName(const StringData& fieldName, ConstElement elt);
 
         //
         // Accessors
@@ -393,6 +405,9 @@ namespace mutablebson {
         class Impl;
         inline Impl& getImpl();
         inline const Impl& getImpl() const;
+
+        Element makeElement(ConstElement element, const StringData* fieldName);
+
         const boost::scoped_ptr<Impl> _impl;
 
         // The root element of this document.
diff --git a/src/mongo/bson/mutable/element.h b/src/mongo/bson/mutable/element.h
index 4d33cd3be58..89ee4e0de2b 100644
--- a/src/mongo/bson/mutable/element.h
+++ b/src/mongo/bson/mutable/element.h
@@ -556,6 +556,8 @@ namespace mutablebson {
 
     private:
         friend class Document;
+        friend class ConstElement;
+
         friend bool operator==(const Element&, const Element&);
 
         inline Element(Document* doc, RepIdx repIdx);
@@ -567,7 +569,7 @@ namespace mutablebson {
         Status setValue(Element* newValue);
 
         template<typename Builder>
-        inline void writeElement(Builder* builder) const;
+        inline void writeElement(Builder* builder, const StringData* fieldName = NULL) const;
 
         template<typename Builder>
         inline void writeChildren(Builder* builder) const;
diff --git a/src/mongo/bson/mutable/mutable_bson_test.cpp b/src/mongo/bson/mutable/mutable_bson_test.cpp
index a9b859c93c0..4b412a9059e 100644
--- a/src/mongo/bson/mutable/mutable_bson_test.cpp
+++ b/src/mongo/bson/mutable/mutable_bson_test.cpp
@@ -1206,5 +1206,129 @@ namespace {
         ASSERT_EQUALS(obj.toString(), doc.toString());
     }
 
-} // namespace
+    TEST(Document, ElementCloningToDifferentDocument) {
+
+        const char initial[] = "{ a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6 ] }";
+
+        mmb::Document source(mongo::fromjson(initial));
+
+        // Dirty the 'd' node and parents.
+        source.root()["d"].pushBack(source.makeElementInt(mongo::StringData(), 7));
+
+        mmb::Document target;
+
+        mmb::Element newElement = target.makeElement(source.root()["d"]);
+        ASSERT_TRUE(newElement.ok());
+        mongo::Status status = target.root().pushBack(newElement);
+        ASSERT_OK(status);
+        const char* expected =
+            "{ d : [ 4, 5, 6, 7 ] }";
+        ASSERT_EQUALS(mongo::fromjson(expected), target);
+
+        newElement = target.makeElement(source.root()["b"]);
+        ASSERT_TRUE(newElement.ok());
+        status = target.root().pushBack(newElement);
+        ASSERT_OK(status);
+        expected =
+            "{ d : [ 4, 5, 6, 7 ], b : [ 1, 2, 3 ] }";
+        ASSERT_EQUALS(mongo::fromjson(expected), target);
+
+        newElement = target.makeElementWithNewFieldName("C", source.root()["c"]);
+        ASSERT_TRUE(newElement.ok());
+        status = target.root().pushBack(newElement);
+        ASSERT_OK(status);
+        expected =
+            "{ d : [ 4, 5, 6, 7 ], b : [ 1, 2, 3 ], C : { 'c' : 'c' } }";
+        ASSERT_EQUALS(mongo::fromjson(expected), target);
+    }
+
+    TEST(Document, ElementCloningToSameDocument) {
+
+        const char initial[] = "{ a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6 ] }";
+
+        mmb::Document doc(mongo::fromjson(initial));
+
+        // Dirty the 'd' node and parents.
+        doc.root()["d"].pushBack(doc.makeElementInt(mongo::StringData(), 7));
+
+        mmb::Element newElement = doc.makeElement(doc.root()["d"]);
+        ASSERT_TRUE(newElement.ok());
+        mongo::Status status = doc.root().pushBack(newElement);
+        ASSERT_OK(status);
+        const char* expected =
+            "{ "
+            " a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6, 7 ], "
+            " d : [ 4, 5, 6, 7 ] "
+            "}";
+        ASSERT_EQUALS(mongo::fromjson(expected), doc);
+
+        newElement = doc.makeElement(doc.root()["b"]);
+        ASSERT_TRUE(newElement.ok());
+        status = doc.root().pushBack(newElement);
+        ASSERT_OK(status);
+        expected =
+            "{ "
+            " a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6, 7 ], "
+            " d : [ 4, 5, 6, 7 ], "
+            " b : [ 1, 2, 3 ] "
+            "}";
+        ASSERT_EQUALS(mongo::fromjson(expected), doc);
+
+        newElement = doc.makeElementWithNewFieldName("C", doc.root()["c"]);
+        ASSERT_TRUE(newElement.ok());
+        status = doc.root().pushBack(newElement);
+        ASSERT_OK(status);
+        expected =
+            "{ "
+            " a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6, 7 ], "
+            " d : [ 4, 5, 6, 7 ], "
+            " b : [ 1, 2, 3 ], "
+            " C : { 'c' : 'c' } "
+            "}";
+        ASSERT_EQUALS(mongo::fromjson(expected), doc);
+    }
+
+    TEST(Document, RootCloningToDifferentDocument) {
+
+        const char initial[] = "{ a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6 ] }";
+
+        mmb::Document source(mongo::fromjson(initial));
+
+        // Dirty the 'd' node and parents.
+        source.root()["d"].pushBack(source.makeElementInt(mongo::StringData(), 7));
+
+        mmb::Document target;
+
+        mmb::Element newElement = target.makeElementWithNewFieldName("X", source.root());
+        mongo::Status status = target.root().pushBack(newElement);
+        ASSERT_OK(status);
+        const char expected[] =
+            "{ X : { a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6, 7 ] } }";
+
+        ASSERT_EQUALS(mongo::fromjson(expected), target);
+    }
+
+    TEST(Document, RootCloningToSameDocument) {
+
+        const char initial[] = "{ a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6 ] }";
+
+        mmb::Document doc(mongo::fromjson(initial));
+
+        // Dirty the 'd' node and parents.
+        doc.root()["d"].pushBack(doc.makeElementInt(mongo::StringData(), 7));
+
+        mmb::Element newElement = doc.makeElementWithNewFieldName("X", doc.root());
+        mongo::Status status = doc.root().pushBack(newElement);
+        ASSERT_OK(status);
+        const char expected[] =
+            "{ "
+            " a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6, 7 ], "
+            "X : { a : 1, b : [ 1, 2, 3 ], c : { 'c' : 'c' }, d : [ 4, 5, 6, 7 ] }"
+            "}";
+
+        ASSERT_EQUALS(mongo::fromjson(expected), doc);
+    }
+
+
+} // namespacem