SERVER-23689 make $pull respect the collation

8 files changed, 162 insertions, 28 deletions

diff --git a/src/mongo/bson/mutable/SConscript b/src/mongo/bson/mutable/SConscript
index ebac4bdf87e..831ea59560a 100644
--- a/src/mongo/bson/mutable/SConscript
+++ b/src/mongo/bson/mutable/SConscript
@@ -20,6 +20,7 @@ env.Library(
         'mutable_bson_test_utils.cpp',
     ],
     LIBDEPS=[
+        '$BUILD_DIR/mongo/db/query/collation/collator_interface_mock',
         '$BUILD_DIR/mongo/unittest/unittest',
         'mutable_bson',
     ],
diff --git a/src/mongo/bson/mutable/const_element-inl.h b/src/mongo/bson/mutable/const_element-inl.h
index 325df232b09..a7f3864ae17 100644
--- a/src/mongo/bson/mutable/const_element-inl.h
+++ b/src/mongo/bson/mutable/const_element-inl.h
@@ -165,17 +165,21 @@ inline SafeNum ConstElement::getValueSafeNum() const {
 }
 
 inline int ConstElement::compareWithElement(const ConstElement& other,
-                                            bool considerFieldName) const {
-    return _basis.compareWithElement(other, considerFieldName);
+                                            bool considerFieldName,
+                                            StringData::ComparatorInterface* comparator) const {
+    return _basis.compareWithElement(other, considerFieldName, comparator);
 }
 
 inline int ConstElement::compareWithBSONElement(const BSONElement& other,
-                                                bool considerFieldName) const {
-    return _basis.compareWithBSONElement(other, considerFieldName);
+                                                bool considerFieldName,
+                                                StringData::ComparatorInterface* comparator) const {
+    return _basis.compareWithBSONElement(other, considerFieldName, comparator);
 }
 
-inline int ConstElement::compareWithBSONObj(const BSONObj& other, bool considerFieldName) const {
-    return _basis.compareWithBSONObj(other, considerFieldName);
+inline int ConstElement::compareWithBSONObj(const BSONObj& other,
+                                            bool considerFieldName,
+                                            StringData::ComparatorInterface* comparator) const {
+    return _basis.compareWithBSONObj(other, considerFieldName, comparator);
 }
 
 inline void ConstElement::writeTo(BSONObjBuilder* builder) const {
diff --git a/src/mongo/bson/mutable/const_element.h b/src/mongo/bson/mutable/const_element.h
index 80ad1aac724..1997de64a32 100644
--- a/src/mongo/bson/mutable/const_element.h
+++ b/src/mongo/bson/mutable/const_element.h
@@ -88,12 +88,17 @@ public:
     inline bool isValueMaxKey() const;
     inline SafeNum getValueSafeNum() const;
 
-    inline int compareWithElement(const ConstElement& other, bool considerFieldName = true) const;
+    inline int compareWithElement(const ConstElement& other,
+                                  bool considerFieldName = true,
+                                  StringData::ComparatorInterface* comparator = nullptr) const;
 
     inline int compareWithBSONElement(const BSONElement& other,
-                                      bool considerFieldName = true) const;
+                                      bool considerFieldName = true,
+                                      StringData::ComparatorInterface* comparator = nullptr) const;
 
-    inline int compareWithBSONObj(const BSONObj& other, bool considerFieldName = true) const;
+    inline int compareWithBSONObj(const BSONObj& other,
+                                  bool considerFieldName = true,
+                                  StringData::ComparatorInterface* comparator = nullptr) const;
 
     inline void writeTo(BSONObjBuilder* builder) const;
     inline void writeArrayTo(BSONArrayBuilder* builder) const;
diff --git a/src/mongo/bson/mutable/document-inl.h b/src/mongo/bson/mutable/document-inl.h
index 3307567536e..ec06777e8a5 100644
--- a/src/mongo/bson/mutable/document-inl.h
+++ b/src/mongo/bson/mutable/document-inl.h
@@ -30,15 +30,19 @@
 namespace mongo {
 namespace mutablebson {
 
-inline int Document::compareWith(const Document& other, bool considerFieldName) const {
+inline int Document::compareWith(const Document& other,
+                                 bool considerFieldName,
+                                 StringData::ComparatorInterface* comparator) const {
     // We cheat and use Element::compareWithElement since we know that 'other' is a
     // Document and has a 'hidden' fieldname that is always indentical across all Document
     // instances.
-    return root().compareWithElement(other.root(), considerFieldName);
+    return root().compareWithElement(other.root(), considerFieldName, comparator);
 }
 
-inline int Document::compareWithBSONObj(const BSONObj& other, bool considerFieldName) const {
-    return root().compareWithBSONObj(other, considerFieldName);
+inline int Document::compareWithBSONObj(const BSONObj& other,
+                                        bool considerFieldName,
+                                        StringData::ComparatorInterface* comparator) const {
+    return root().compareWithBSONObj(other, considerFieldName, comparator);
 }
 
 inline void Document::writeTo(BSONObjBuilder* builder) const {
diff --git a/src/mongo/bson/mutable/document.cpp b/src/mongo/bson/mutable/document.cpp
index 612baf1a285..6bb159b0e7c 100644
--- a/src/mongo/bson/mutable/document.cpp
+++ b/src/mongo/bson/mutable/document.cpp
@@ -1485,7 +1485,9 @@ SafeNum Element::getValueSafeNum() const {
     }
 }
 
-int Element::compareWithElement(const ConstElement& other, bool considerFieldName) const {
+int Element::compareWithElement(const ConstElement& other,
+                                bool considerFieldName,
+                                StringData::ComparatorInterface* comparator) const {
     verify(ok());
     verify(other.ok());
 
@@ -1507,13 +1509,15 @@ int Element::compareWithElement(const ConstElement& other, bool considerFieldNam
     // TODO: Andy has suggested that this may not be legal since woCompare is not reflexive
     // in all cases.
     if (impl.hasValue(thisRep))
-        return -other.compareWithBSONElement(impl.getSerializedElement(thisRep), considerFieldName);
+        return -other.compareWithBSONElement(
+            impl.getSerializedElement(thisRep), considerFieldName, comparator);
 
     const Document::Impl& oimpl = other.getDocument().getImpl();
     const ElementRep& otherRep = oimpl.getElementRep(other.getIdx());
 
     if (oimpl.hasValue(otherRep))
-        return compareWithBSONElement(oimpl.getSerializedElement(otherRep), considerFieldName);
+        return compareWithBSONElement(
+            oimpl.getSerializedElement(otherRep), considerFieldName, comparator);
 
     // Leaf elements should always have a value, so we should only be dealing with Objects
     // or Arrays here.
@@ -1554,7 +1558,8 @@ int Element::compareWithElement(const ConstElement& other, bool considerFieldNam
         if (!otherIter.ok())
             return 1;
 
-        const int result = thisIter.compareWithElement(otherIter, considerChildFieldNames);
+        const int result =
+            thisIter.compareWithElement(otherIter, considerChildFieldNames, comparator);
         if (result != 0)
             return result;
 
@@ -1563,7 +1568,9 @@ int Element::compareWithElement(const ConstElement& other, bool considerFieldNam
     }
 }
 
-int Element::compareWithBSONElement(const BSONElement& other, bool considerFieldName) const {
+int Element::compareWithBSONElement(const BSONElement& other,
+                                    bool considerFieldName,
+                                    StringData::ComparatorInterface* comparator) const {
     verify(ok());
 
     const Document::Impl& impl = getDocument().getImpl();
@@ -1572,7 +1579,7 @@ int Element::compareWithBSONElement(const BSONElement& other, bool considerField
     // If we have a representation as a BSONElement, we can just use BSONElement::woCompare
     // to do the entire comparison.
     if (impl.hasValue(thisRep))
-        return impl.getSerializedElement(thisRep).woCompare(other, considerFieldName);
+        return impl.getSerializedElement(thisRep).woCompare(other, considerFieldName, comparator);
 
     // Leaf elements should always have a value, so we should only be dealing with Objects
     // or Arrays here.
@@ -1599,10 +1606,12 @@ int Element::compareWithBSONElement(const BSONElement& other, bool considerField
     const bool considerChildFieldNames =
         (impl.getType(thisRep) != mongo::Array) && (other.type() != mongo::Array);
 
-    return compareWithBSONObj(other.Obj(), considerChildFieldNames);
+    return compareWithBSONObj(other.Obj(), considerChildFieldNames, comparator);
 }
 
-int Element::compareWithBSONObj(const BSONObj& other, bool considerFieldName) const {
+int Element::compareWithBSONObj(const BSONObj& other,
+                                bool considerFieldName,
+                                StringData::ComparatorInterface* comparator) const {
     verify(ok());
 
     const Document::Impl& impl = getDocument().getImpl();
@@ -1624,7 +1633,7 @@ int Element::compareWithBSONObj(const BSONObj& other, bool considerFieldName) co
         if (otherVal.eoo())
             return 1;
 
-        const int result = thisIter.compareWithBSONElement(otherVal, considerFieldName);
+        const int result = thisIter.compareWithBSONElement(otherVal, considerFieldName, comparator);
         if (result != 0)
             return result;
 
diff --git a/src/mongo/bson/mutable/document.h b/src/mongo/bson/mutable/document.h
index 1efd9e95729..cca50443311 100644
--- a/src/mongo/bson/mutable/document.h
+++ b/src/mongo/bson/mutable/document.h
@@ -277,10 +277,14 @@ public:
     //
 
     /** Compare this Document to 'other' with the semantics of BSONObj::woCompare. */
-    inline int compareWith(const Document& other, bool considerFieldName = true) const;
+    inline int compareWith(const Document& other,
+                           bool considerFieldName = true,
+                           StringData::ComparatorInterface* comparator = nullptr) const;
 
     /** Compare this Document to 'other' with the semantics of BSONObj::woCompare. */
-    inline int compareWithBSONObj(const BSONObj& other, bool considerFieldName = true) const;
+    inline int compareWithBSONObj(const BSONObj& other,
+                                  bool considerFieldName = true,
+                                  StringData::ComparatorInterface* comparator = nullptr) const;
 
 
     //
diff --git a/src/mongo/bson/mutable/element.h b/src/mongo/bson/mutable/element.h
index 2c3a54c348b..ee0d6de3d45 100644
--- a/src/mongo/bson/mutable/element.h
+++ b/src/mongo/bson/mutable/element.h
@@ -350,7 +350,9 @@ public:
      *   Returns 0 if this == other either tautologically, or according to woCompare.
      *   Returns 1 if this > other according to BSONElement::woCompare
      */
-    int compareWithElement(const ConstElement& other, bool considerFieldName = true) const;
+    int compareWithElement(const ConstElement& other,
+                           bool considerFieldName = true,
+                           StringData::ComparatorInterface* comparator = nullptr) const;
 
     /** Compare this Element with BSONElement 'other'. You should not call this on the root
      *  Element of the Document because the root Element does not have a field name. Use
@@ -360,7 +362,9 @@ public:
      *   Returns 0 if this == other either tautologically, or according to woCompare.
      *   Returns 1 if this > other according to BSONElement::woCompare
      */
-    int compareWithBSONElement(const BSONElement& other, bool considerFieldName = true) const;
+    int compareWithBSONElement(const BSONElement& other,
+                               bool considerFieldName = true,
+                               StringData::ComparatorInterface* comparator = nullptr) const;
 
     /** Compare this Element, which must be an Object or an Array, with 'other'.
      *
@@ -368,7 +372,9 @@ public:
      *   Returns 0 if this object == other either tautologically, or according to woCompare.
      *   Returns 1 if this object > other according to BSONElement::woCompare
      */
-    int compareWithBSONObj(const BSONObj& other, bool considerFieldName = true) const;
+    int compareWithBSONObj(const BSONObj& other,
+                           bool considerFieldName = true,
+                           StringData::ComparatorInterface* comparator = nullptr) const;
 
 
     //
diff --git a/src/mongo/bson/mutable/mutable_bson_test.cpp b/src/mongo/bson/mutable/mutable_bson_test.cpp
index b1ea43ac213..0ec8dfdbcd6 100644
--- a/src/mongo/bson/mutable/mutable_bson_test.cpp
+++ b/src/mongo/bson/mutable/mutable_bson_test.cpp
@@ -32,9 +32,10 @@
 #include "mongo/base/status.h"
 #include "mongo/bson/bsonobj.h"
 #include "mongo/bson/mutable/algorithm.h"
-#include "mongo/bson/mutable/mutable_bson_test_utils.h"
 #include "mongo/bson/mutable/damage_vector.h"
+#include "mongo/bson/mutable/mutable_bson_test_utils.h"
 #include "mongo/db/json.h"
+#include "mongo/db/query/collation/collator_interface_mock.h"
 #include "mongo/platform/decimal128.h"
 #include "mongo/unittest/unittest.h"
 
@@ -3254,6 +3255,106 @@ TEST(DocumentComparison, SimpleComparisonWithDeserializedElements) {
     ASSERT_EQUALS(0, doc2.compareWith(doc1));
 }
 
+TEST(DocumentComparison, DocumentCompareWithRespectsCollation) {
+    mongo::CollatorInterfaceMock collator(mongo::CollatorInterfaceMock::MockType::kAlwaysEqual);
+    const mmb::Document doc1(mongo::fromjson("{a: 'foo'}"));
+    const mmb::Document doc2(mongo::fromjson("{a: 'bar'}"));
+    // Pass true to indicate that we should compare field names. The two documents should be unequal
+    // without the collator, but equal when using the "always equal" collator.
+    ASSERT_NE(0, doc1.compareWith(doc2, true));
+    ASSERT_EQ(0, doc1.compareWith(doc2, true, &collator));
+}
+
+TEST(DocumentComparison, DocumentCompareWithBSONObjRespectsCollation) {
+    mongo::CollatorInterfaceMock collator(mongo::CollatorInterfaceMock::MockType::kAlwaysEqual);
+    const mmb::Document doc1(mongo::fromjson("{a: 'foo'}"));
+    const mongo::BSONObj doc2 = mongo::fromjson("{a: 'bar'}");
+    // Pass true to indicate that we should compare field names. The two documents should be unequal
+    // without the collator, but equal when using the "always equal" collator.
+    ASSERT_NE(0, doc1.compareWithBSONObj(doc2, true));
+    ASSERT_EQ(0, doc1.compareWithBSONObj(doc2, true, &collator));
+}
+
+TEST(DocumentComparison, ElementCompareWithElementRespectsCollator) {
+    mongo::CollatorInterfaceMock collator(mongo::CollatorInterfaceMock::MockType::kAlwaysEqual);
+    const mmb::Document doc1(mongo::fromjson("{a: 'foo'}"));
+    const mmb::Document doc2(mongo::fromjson("{a: 'bar'}"));
+    const mmb::ConstElement element1 = doc1.root().leftChild();
+    const mmb::ConstElement element2 = doc2.root().leftChild();
+    // Pass true to indicate that we should compare field names. The two documents should be unequal
+    // without the collator, but equal when using the "always equal" collator.
+    ASSERT_NE(0, element1.compareWithElement(element2, true));
+    ASSERT_EQ(0, element1.compareWithElement(element2, true, &collator));
+}
+
+TEST(DocumentComparison, ElementCompareWithBSONElementRespectsCollator) {
+    mongo::CollatorInterfaceMock collator(mongo::CollatorInterfaceMock::MockType::kAlwaysEqual);
+    const mmb::Document doc1(mongo::fromjson("{a: 'foo'}"));
+    const mongo::BSONObj doc2 = mongo::fromjson("{a: 'bar'}");
+    const mmb::ConstElement element1 = doc1.root().leftChild();
+    const mongo::BSONElement element2 = doc2["a"];
+    // Pass true to indicate that we should compare field names. The two documents should be unequal
+    // without the collator, but equal when using the "always equal" collator.
+    ASSERT_NE(0, element1.compareWithBSONElement(element2, true));
+    ASSERT_EQ(0, element1.compareWithBSONElement(element2, true, &collator));
+}
+
+TEST(DocumentComparison, ElementCompareWithBSONObjRespectsCollator) {
+    mongo::CollatorInterfaceMock collator(mongo::CollatorInterfaceMock::MockType::kAlwaysEqual);
+    const mmb::Document doc1(mongo::fromjson("{b: {c: 'foo'}}"));
+    const mongo::BSONObj doc2 = mongo::fromjson("{c: 'bar'}");
+    const mmb::ConstElement element1 = doc1.root().leftChild();
+    // Pass true to indicate that we should compare field names. The two documents should be unequal
+    // without the collator, but equal when using the "always equal" collator.
+    ASSERT_NE(0, element1.compareWithBSONObj(doc2, true));
+    ASSERT_EQ(0, element1.compareWithBSONObj(doc2, true, &collator));
+}
+
+TEST(DocumentComparison, DocumentCompareWithRespectsCollationRecursively) {
+    mongo::CollatorInterfaceMock collator(mongo::CollatorInterfaceMock::MockType::kAlwaysEqual);
+    const mmb::Document doc1(mongo::fromjson("{a: [{b: 'foo'}, {b: 'bar'}]}"));
+    const mmb::Document doc2(mongo::fromjson("{a: [{b: 'notFoo'}, {b: 'notBar'}]}"));
+    // Pass true to indicate that we should compare field names. The two documents should be unequal
+    // without the collator, but equal when using the "always equal" collator.
+    ASSERT_NE(0, doc1.compareWith(doc2, true));
+    ASSERT_EQ(0, doc1.compareWith(doc2, true, &collator));
+}
+
+TEST(DocumentComparison, DocumentCompareWithRespectsCollationWithDeserializedElement) {
+    mongo::CollatorInterfaceMock collator(mongo::CollatorInterfaceMock::MockType::kAlwaysEqual);
+    mmb::Document doc1(mongo::fromjson("{a: ['foo', 'foo']}"));
+    mmb::Document doc2(mongo::fromjson("{a: ['bar', 'bar']}"));
+
+    // With the always equal collator, the documents should start out comparing equal.
+    ASSERT_EQ(0, doc1.compareWith(doc2, true, &collator));
+    ASSERT_EQ(0, doc2.compareWith(doc1, true, &collator));
+
+    // They should still be equal after causing deserialization of one of the leaf elements of
+    // 'doc1'.
+    {
+        mmb::Element elementA = doc1.root()["a"];
+        ASSERT_TRUE(elementA.ok());
+        mmb::Element elementA0 = elementA[0];
+        ASSERT_TRUE(elementA0.ok());
+        ASSERT_OK(elementA0.remove());
+        ASSERT_OK(elementA.pushBack(elementA0));
+        ASSERT_EQ(0, doc1.compareWith(doc2, true, &collator));
+        ASSERT_EQ(0, doc2.compareWith(doc1, true, &collator));
+    }
+
+    // And they should remain equal after doing the same to 'doc2'.
+    {
+        mmb::Element elementA = doc2.root()["a"];
+        ASSERT_TRUE(elementA.ok());
+        mmb::Element elementA0 = elementA[0];
+        ASSERT_TRUE(elementA0.ok());
+        ASSERT_OK(elementA0.remove());
+        ASSERT_OK(elementA.pushBack(elementA0));
+        ASSERT_EQ(0, doc1.compareWith(doc2, true, &collator));
+        ASSERT_EQ(0, doc2.compareWith(doc1, true, &collator));
+    }
+}
+
 TEST(UnorderedEqualityChecker, Identical) {
     const mongo::BSONObj b1 =
         mongo::fromjson("{ a : [1, 2, { 'a' : 'b', 'x' : 'y' } ], b : { x : 1, y : 2, z : 3 } }");