3 files changed, 65 insertions, 10 deletions

diff --git a/src/mongo/bson/mutable/SConscript b/src/mongo/bson/mutable/SConscript
index 831ea59560a..cbad6d15552 100644
--- a/src/mongo/bson/mutable/SConscript
+++ b/src/mongo/bson/mutable/SConscript
@@ -43,6 +43,7 @@ env.CppUnitTest(
         'mutable_bson_algo_test.cpp',
     ],
     LIBDEPS=[
+        '$BUILD_DIR/mongo/db/query/collation/collator_interface_mock',
         'mutable_bson',
         'mutable_bson_test_utils',
     ],
diff --git a/src/mongo/bson/mutable/algorithm.h b/src/mongo/bson/mutable/algorithm.h
index 850f37d9f00..de1134e0162 100644
--- a/src/mongo/bson/mutable/algorithm.h
+++ b/src/mongo/bson/mutable/algorithm.h
@@ -31,6 +31,7 @@
 #include <cstddef>
 #include <vector>
 
+#include "mongo/base/string_data_comparator_interface.h"
 #include "mongo/bson/mutable/const_element.h"
 #include "mongo/bson/mutable/element.h"
 #include "mongo/util/mongoutils/str.h"
@@ -166,58 +167,70 @@ void deduplicateChildren(Element parent, EqualityComparator equal) {
 class woLess {
     // TODO: This should possibly derive from std::binary_function.
 public:
-    woLess(bool considerFieldName = true) : _considerFieldName(considerFieldName) {}
+    woLess(bool considerFieldName = true,
+           const StringData::ComparatorInterface* comparator = nullptr)
+        : _considerFieldName(considerFieldName), _comp(comparator) {}
 
     inline bool operator()(const ConstElement& left, const ConstElement& right) const {
-        return left.compareWithElement(right, _considerFieldName) < 0;
+        return left.compareWithElement(right, _considerFieldName, _comp) < 0;
     }
 
 private:
     const bool _considerFieldName;
+    const StringData::ComparatorInterface* _comp = nullptr;
 };
 
 /** A greater-than ordering for Elements that compares based on woCompare */
 class woGreater {
     // TODO: This should possibly derive from std::binary_function.
 public:
-    woGreater(bool considerFieldName = true) : _considerFieldName(considerFieldName) {}
+    woGreater(bool considerFieldName = true,
+              const StringData::ComparatorInterface* comparator = nullptr)
+        : _considerFieldName(considerFieldName), _comp(comparator) {}
 
     inline bool operator()(const ConstElement& left, const ConstElement& right) const {
-        return left.compareWithElement(right, _considerFieldName) > 0;
+        return left.compareWithElement(right, _considerFieldName, _comp) > 0;
     }
 
 private:
     const bool _considerFieldName;
+    const StringData::ComparatorInterface* _comp = nullptr;
 };
 
 /** An equality predicate for elements that compares based on woCompare */
 class woEqual {
     // TODO: This should possibly derive from std::binary_function.
 public:
-    woEqual(bool considerFieldName = true) : _considerFieldName(considerFieldName) {}
+    woEqual(bool considerFieldName = true,
+            const StringData::ComparatorInterface* comparator = nullptr)
+        : _considerFieldName(considerFieldName), _comp(comparator) {}
 
     inline bool operator()(const ConstElement& left, const ConstElement& right) const {
-        return left.compareWithElement(right, _considerFieldName) == 0;
+        return left.compareWithElement(right, _considerFieldName, _comp) == 0;
     }
 
 private:
     const bool _considerFieldName;
+    const StringData::ComparatorInterface* _comp = nullptr;
 };
 
 /** An equality predicate for elements that compares based on woCompare */
 class woEqualTo {
     // TODO: This should possibly derive from std::binary_function.
 public:
-    woEqualTo(const ConstElement& value, bool considerFieldName = true)
-        : _value(value), _considerFieldName(considerFieldName) {}
+    woEqualTo(const ConstElement& value,
+              bool considerFieldName = true,
+              const StringData::ComparatorInterface* comparator = nullptr)
+        : _value(value), _considerFieldName(considerFieldName), _comp(comparator) {}
 
     inline bool operator()(const ConstElement& elt) const {
-        return _value.compareWithElement(elt, _considerFieldName) == 0;
+        return _value.compareWithElement(elt, _considerFieldName, _comp) == 0;
     }
 
 private:
-    const ConstElement& _value;
+    const ConstElement _value;
     const bool _considerFieldName;
+    const StringData::ComparatorInterface* _comp = nullptr;
 };
 
 // NOTE: Originally, these truly were algorithms, in that they executed the loop over a
diff --git a/src/mongo/bson/mutable/mutable_bson_algo_test.cpp b/src/mongo/bson/mutable/mutable_bson_algo_test.cpp
index cdd85576ebf..ade055a543f 100644
--- a/src/mongo/bson/mutable/mutable_bson_algo_test.cpp
+++ b/src/mongo/bson/mutable/mutable_bson_algo_test.cpp
@@ -32,11 +32,13 @@
 #include "mongo/bson/mutable/document.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/basic.h"
 #include "mongo/unittest/unittest.h"
 
 namespace {
 
+using mongo::CollatorInterfaceMock;
 using mongo::Status;
 using namespace mongo::mutablebson;
 
@@ -329,4 +331,43 @@ TEST(FullNameTest, InsideArray) {
                   getFullName(doc.root().leftChild().leftChild().leftChild().rightSibling()));
 }
 
+TEST(WoLessTest, CollationAware) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    Document less(mongo::fromjson("{ x: 'cbc' }"));
+    Document greater(mongo::fromjson("{ x: 'abd' }"));
+
+    woLess comp(true, &collator);
+    ASSERT_TRUE(comp(less.root(), greater.root()));
+    ASSERT_FALSE(comp(greater.root(), less.root()));
+}
+
+TEST(WoGreaterTest, CollationAware) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    Document less(mongo::fromjson("{ x: 'cbc' }"));
+    Document greater(mongo::fromjson("{ x: 'abd' }"));
+
+    woGreater comp(true, &collator);
+    ASSERT_TRUE(comp(greater.root(), less.root()));
+    ASSERT_FALSE(comp(less.root(), greater.root()));
+}
+
+TEST(WoEqualTest, CollationAware) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kAlwaysEqual);
+    Document docA(mongo::fromjson("{ x: 'not' }"));
+    Document docB(mongo::fromjson("{ x: 'equal' }"));
+
+    woEqual comp(true, &collator);
+    ASSERT_TRUE(comp(docA.root(), docB.root()));
+    ASSERT_TRUE(comp(docB.root(), docA.root()));
+}
+
+TEST(WoEqualToTest, CollationAware) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kAlwaysEqual);
+    Document docA(mongo::fromjson("{ x: 'not' }"));
+    Document docB(mongo::fromjson("{ x: 'equal' }"));
+
+    woEqualTo comp(docA.root(), true, &collator);
+    ASSERT_TRUE(comp(docB.root()));
+}
+
 }  // namespace