SERVER-21647 fix $rename to not re-order fields when the destination field is already present

3 files changed, 192 insertions, 0 deletions

diff --git a/src/mongo/bson/mutable/document.cpp b/src/mongo/bson/mutable/document.cpp
index 307177db043..2ed68db1ca3 100644
--- a/src/mongo/bson/mutable/document.cpp
+++ b/src/mongo/bson/mutable/document.cpp
@@ -1907,6 +1907,30 @@ Status Element::setValueSafeNum(const SafeNum value) {
     }
 }
 
+Status Element::setValueElement(ConstElement setFrom) {
+    verify(ok());
+
+    // Can't set to your own root element, since this would create a circular document.
+    if (_doc->root() == setFrom) {
+        return Status(ErrorCodes::IllegalOperation,
+                      "Attempt to set an element to its own document's root");
+    }
+
+    // Setting to self is a no-op.
+    //
+    // Setting the root is always an error so we want to fall through to the error handling in this
+    // case.
+    if (*this == setFrom && _repIdx != kRootRepIdx) {
+        return Status::OK();
+    }
+
+    Document::Impl& impl = getDocument().getImpl();
+    ElementRep thisRep = impl.getElementRep(_repIdx);
+    const StringData fieldName = impl.getFieldNameForNewElement(thisRep);
+    Element newValue = getDocument().makeElementWithNewFieldName(fieldName, setFrom);
+    return setValue(newValue._repIdx);
+}
+
 BSONType Element::getType() const {
     verify(ok());
     const Document::Impl& impl = getDocument().getImpl();
diff --git a/src/mongo/bson/mutable/element.h b/src/mongo/bson/mutable/element.h
index 68bf99dfc08..2c3a54c348b 100644
--- a/src/mongo/bson/mutable/element.h
+++ b/src/mongo/bson/mutable/element.h
@@ -473,6 +473,12 @@ public:
      */
     Status setValueSafeNum(const SafeNum value);
 
+    /** Set the value of this Element to the value from another Element.
+     *
+     * The name of this Element is not modified.
+     */
+    Status setValueElement(ConstElement setFrom);
+
 
     //
     // Accessors
diff --git a/src/mongo/bson/mutable/mutable_bson_test.cpp b/src/mongo/bson/mutable/mutable_bson_test.cpp
index f44afe2d9df..348f297c7b7 100644
--- a/src/mongo/bson/mutable/mutable_bson_test.cpp
+++ b/src/mongo/bson/mutable/mutable_bson_test.cpp
@@ -1461,6 +1461,168 @@ TEST(Document, SetValueBSONElementFieldNameHandling) {
     ASSERT_EQUALS(mongo::fromjson(outJson), doc.getObject());
 }
 
+TEST(Document, SetValueElementFromSeparateDocument) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc1(inObj);
+
+    mongo::BSONObj inObj2 = mongo::fromjson("{ b : 5 }");
+    const mmb::Document doc2(inObj2);
+
+    auto setTo = doc1.root().leftChild();
+    auto setFrom = doc2.root().leftChild();
+    ASSERT_OK(setTo.setValueElement(setFrom));
+
+    ASSERT_EQ(mongo::fromjson("{ a : 5 }"), doc1.getObject());
+
+    // Doc containing the 'setFrom' element should be unchanged.
+    ASSERT_EQ(inObj2, doc2.getObject());
+}
+
+TEST(Document, SetValueElementIsNoopWhenSetToSelf) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc(inObj);
+
+    auto element = doc.root().leftChild();
+    ASSERT_OK(element.setValueElement(element));
+
+    ASSERT_EQ(inObj, doc.getObject());
+}
+
+TEST(Document, SetValueElementIsNoopWhenSetToSelfFromCopy) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc(inObj);
+
+    auto element = doc.root().leftChild();
+    auto elementCopy = element;
+    ASSERT_OK(element.setValueElement(elementCopy));
+
+    ASSERT_EQ(inObj, doc.getObject());
+}
+
+TEST(Document, SetValueElementIsNoopWhenSetToSelfNonRootElement) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : { b : { c: 4 } } }");
+    mmb::Document doc(inObj);
+
+    auto element = doc.root().leftChild().leftChild().leftChild();
+    ASSERT_EQ("c", element.getFieldName());
+    ASSERT_OK(element.setValueElement(element));
+
+    ASSERT_EQ(inObj, doc.getObject());
+}
+
+TEST(Document, SetValueElementSetToNestedObject) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc1(inObj);
+
+    mongo::BSONObj inObj2 = mongo::fromjson("{ b : { c : 5, d : 6 } }");
+    const mmb::Document doc2(inObj2);
+
+    auto setTo = doc1.root().leftChild();
+    auto setFrom = doc2.root().leftChild();
+    ASSERT_OK(setTo.setValueElement(setFrom));
+
+    ASSERT_EQ(mongo::fromjson("{ a : { c : 5, d : 6 } }"), doc1.getObject());
+
+    // Doc containing the 'setFrom' element should be unchanged.
+    ASSERT_EQ(inObj2, doc2.getObject());
+}
+
+TEST(Document, SetValueElementNonRootElements) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : { b : 5, c : 6 } }");
+    mmb::Document doc1(inObj);
+
+    mongo::BSONObj inObj2 = mongo::fromjson("{ d : { e : 8, f : 9 } }");
+    const mmb::Document doc2(inObj2);
+
+    auto setTo = doc1.root().leftChild().rightChild();
+    ASSERT_EQ("c", setTo.getFieldName());
+    auto setFrom = doc2.root().leftChild().leftChild();
+    ASSERT_EQ("e", setFrom.getFieldName());
+    ASSERT_OK(setTo.setValueElement(setFrom));
+
+    ASSERT_EQ(mongo::fromjson("{ a : { b : 5, c : 8 } }"), doc1.getObject());
+
+    // Doc containing the 'setFrom' element should be unchanged.
+    ASSERT_EQ(inObj2, doc2.getObject());
+}
+
+TEST(Document, SetValueElementSetRootToSelfErrors) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc(inObj);
+
+    auto element = doc.root();
+    ASSERT_NOT_OK(element.setValueElement(element));
+    ASSERT_EQ(inObj, doc.getObject());
+}
+
+TEST(Document, SetValueElementSetRootToAnotherDocRootErrors) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc1(inObj);
+
+    mongo::BSONObj inObj2 = mongo::fromjson("{ b : 5 }");
+    const mmb::Document doc2(inObj2);
+
+    auto setTo = doc1.root();
+    auto setFrom = doc2.root();
+    ASSERT_NOT_OK(setTo.setValueElement(setFrom));
+
+    ASSERT_EQ(inObj, doc1.getObject());
+    ASSERT_EQ(inObj2, doc2.getObject());
+}
+
+TEST(Document, SetValueElementSetRootToNotRootInSelfErrors) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc(inObj);
+
+    auto setTo = doc.root();
+    auto setFrom = doc.root().leftChild();
+    ASSERT_NOT_OK(setTo.setValueElement(setFrom));
+    ASSERT_EQ(inObj, doc.getObject());
+}
+
+TEST(Document, SetValueElementSetRootToNotRootInAnotherDocErrors) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : 4 }");
+    mmb::Document doc1(inObj);
+
+    mongo::BSONObj inObj2 = mongo::fromjson("{ b : 5 }");
+    const mmb::Document doc2(inObj2);
+
+    auto setTo = doc1.root();
+    auto setFrom = doc2.root().leftChild();
+    ASSERT_NOT_OK(setTo.setValueElement(setFrom));
+
+    ASSERT_EQ(inObj, doc1.getObject());
+    ASSERT_EQ(inObj2, doc2.getObject());
+}
+
+TEST(Document, SetValueElementSetToOwnRootErrors) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : { b : 4 } }");
+    mmb::Document doc(inObj);
+
+    auto setTo = doc.root().leftChild().leftChild();
+    ASSERT_EQ("b", setTo.getFieldName());
+    auto setFrom = doc.root();
+
+    ASSERT_NOT_OK(setTo.setValueElement(setFrom));
+    ASSERT_EQ(inObj, doc.getObject());
+}
+
+TEST(Document, SetValueElementSetToOtherDocRoot) {
+    mongo::BSONObj inObj = mongo::fromjson("{ a : { b : 4 } }");
+    mmb::Document doc1(inObj);
+
+    mongo::BSONObj inObj2 = mongo::fromjson("{ c : 5 } }");
+    mmb::Document doc2(inObj2);
+
+    auto setTo = doc1.root().leftChild().leftChild();
+    ASSERT_EQ("b", setTo.getFieldName());
+    auto setFrom = doc2.root();
+
+    ASSERT_OK(setTo.setValueElement(setFrom));
+    ASSERT_EQ(mongo::fromjson("{ a : { b : { c : 5 } } }"), doc1.getObject());
+    ASSERT_EQ(inObj2, doc2.getObject());
+}
+
 TEST(Document, CreateElementWithEmptyFieldName) {
     mmb::Document doc;
     mmb::Element noname = doc.makeElementObject(mongo::StringData());