SERVER-41721 Make IndexAccessMethod::getKeys output a KeyString set

1 files changed, 282 insertions, 128 deletions

diff --git a/src/mongo/db/index/s2_key_generator_test.cpp b/src/mongo/db/index/s2_key_generator_test.cpp
index 93fc8ac545d..9e158f4cb87 100644
--- a/src/mongo/db/index/s2_key_generator_test.cpp
+++ b/src/mongo/db/index/s2_key_generator_test.cpp
@@ -49,11 +49,12 @@ using namespace mongo;
 
 namespace {
 
-std::string dumpKeyset(const BSONObjSet& objs) {
+std::string dumpKeyset(const KeyStringSet& keyStrings) {
     std::stringstream ss;
     ss << "[ ";
-    for (BSONObjSet::iterator i = objs.begin(); i != objs.end(); ++i) {
-        ss << i->toString() << " ";
+    for (auto& keyString : keyStrings) {
+        auto key = KeyString::toBson(keyString, Ordering::make(BSONObj()));
+        ss << key.toString() << " ";
     }
     ss << "]";
 
@@ -76,17 +77,14 @@ std::string dumpMultikeyPaths(const MultikeyPaths& multikeyPaths) {
     return ss.str();
 }
 
-bool assertKeysetsEqual(const BSONObjSet& expectedKeys, const BSONObjSet& actualKeys) {
+bool assertKeysetsEqual(const KeyStringSet& expectedKeys, const KeyStringSet& actualKeys) {
     if (expectedKeys.size() != actualKeys.size()) {
         log() << "Expected: " << dumpKeyset(expectedKeys) << ", "
               << "Actual: " << dumpKeyset(actualKeys);
         return false;
     }
 
-    if (!std::equal(expectedKeys.begin(),
-                    expectedKeys.end(),
-                    actualKeys.begin(),
-                    SimpleBSONObjComparator::kInstance.makeEqualTo())) {
+    if (!std::equal(expectedKeys.begin(), expectedKeys.end(), actualKeys.begin())) {
         log() << "Expected: " << dumpKeyset(expectedKeys) << ", "
               << "Actual: " << dumpKeyset(actualKeys);
         return false;
@@ -120,14 +118,21 @@ long long getCellID(int x, int y, bool multiPoint = false) {
     const CollatorInterface* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     // There's no need to compute the prefixes of the indexed fields that cause the index to be
     // multikey when computing the cell id of the geo field.
     MultikeyPaths* multikeyPaths = nullptr;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &keys, multikeyPaths);
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &keys,
+                                     multikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
 
     ASSERT_EQUALS(1U, keys.size());
-    return (*keys.begin()).firstElement().Long();
+    auto key = KeyString::toBson(*keys.begin(), Ordering::make(BSONObj()));
+    return key.firstElement().Long();
 }
 
 TEST(S2KeyGeneratorTest, GetS2KeysFromSubobjectWithArrayOfGeoAndNonGeoSubobjects) {
@@ -141,16 +146,30 @@ TEST(S2KeyGeneratorTest, GetS2KeysFromSubobjectWithArrayOfGeoAndNonGeoSubobjects
     CollatorInterfaceMock* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(
-        genKeysFrom, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << 1 << "" << getCellID(0, 0)));
-    expectedKeys.insert(BSON("" << 1 << "" << getCellID(3, 3)));
-    expectedKeys.insert(BSON("" << 2 << "" << getCellID(0, 0)));
-    expectedKeys.insert(BSON("" << 2 << "" << getCellID(3, 3)));
+    ExpressionKeysPrivate::getS2Keys(genKeysFrom,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString1(KeyString::Version::kLatestVersion,
+                                      BSON("" << 1 << "" << getCellID(0, 0)),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString2(KeyString::Version::kLatestVersion,
+                                      BSON("" << 1 << "" << getCellID(3, 3)),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString3(KeyString::Version::kLatestVersion,
+                                      BSON("" << 2 << "" << getCellID(0, 0)),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString4(KeyString::Version::kLatestVersion,
+                                      BSON("" << 2 << "" << getCellID(3, 3)),
+                                      Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{
+        keyString1.release(), keyString2.release(), keyString3.release(), keyString4.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{{1U}, {1U}}, actualMultikeyPaths);
@@ -167,15 +186,26 @@ TEST(S2KeyGeneratorTest, GetS2KeysFromArrayOfNonGeoSubobjectsWithArrayValues) {
     CollatorInterfaceMock* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(
-        genKeysFrom, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << 1 << "" << getCellID(0, 0)));
-    expectedKeys.insert(BSON("" << 2 << "" << getCellID(0, 0)));
-    expectedKeys.insert(BSON("" << 3 << "" << getCellID(0, 0)));
+    ExpressionKeysPrivate::getS2Keys(genKeysFrom,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString1(KeyString::Version::kLatestVersion,
+                                      BSON("" << 1 << "" << getCellID(0, 0)),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString2(KeyString::Version::kLatestVersion,
+                                      BSON("" << 2 << "" << getCellID(0, 0)),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString3(KeyString::Version::kLatestVersion,
+                                      BSON("" << 3 << "" << getCellID(0, 0)),
+                                      Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString1.release(), keyString2.release(), keyString3.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{{0U, 1U}, std::set<size_t>{}}, actualMultikeyPaths);
@@ -190,16 +220,27 @@ TEST(S2KeyGeneratorTest, GetS2KeysFromMultiPointInGeoField) {
     CollatorInterfaceMock* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(
-        genKeysFrom, keyPattern, params, &actualKeys, &actualMultikeyPaths);
+    ExpressionKeysPrivate::getS2Keys(genKeysFrom,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
 
     const bool multiPoint = true;
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << 1 << "" << getCellID(0, 0, multiPoint)));
-    expectedKeys.insert(BSON("" << 1 << "" << getCellID(1, 0, multiPoint)));
-    expectedKeys.insert(BSON("" << 1 << "" << getCellID(1, 1, multiPoint)));
+    KeyString::HeapBuilder keyString1(KeyString::Version::kLatestVersion,
+                                      BSON("" << 1 << "" << getCellID(0, 0, multiPoint)),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString2(KeyString::Version::kLatestVersion,
+                                      BSON("" << 1 << "" << getCellID(1, 0, multiPoint)),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString3(KeyString::Version::kLatestVersion,
+                                      BSON("" << 1 << "" << getCellID(1, 1, multiPoint)),
+                                      Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString1.release(), keyString2.release(), keyString3.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, {0U}}, actualMultikeyPaths);
@@ -213,13 +254,21 @@ TEST(S2KeyGeneratorTest, CollationAppliedToNonGeoStringFieldAfterGeoField) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "gnirts"));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << getCellID(0, 0) << ""
+                                             << "gnirts"),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, std::set<size_t>{}},
@@ -234,14 +283,22 @@ TEST(S2KeyGeneratorTest, CollationAppliedToNonGeoStringFieldBeforeGeoField) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON(""
-                             << "gnirts"
-                             << "" << getCellID(0, 0)));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON(""
+                                          << "gnirts"
+                                          << "" << getCellID(0, 0)),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, std::set<size_t>{}},
@@ -256,15 +313,23 @@ TEST(S2KeyGeneratorTest, CollationAppliedToAllNonGeoStringFields) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON(""
-                             << "gnirts"
-                             << "" << getCellID(0, 0) << ""
-                             << "2gnirts"));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON(""
+                                          << "gnirts"
+                                          << "" << getCellID(0, 0) << ""
+                                          << "2gnirts"),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(
@@ -280,13 +345,21 @@ TEST(S2KeyGeneratorTest, CollationAppliedToNonGeoStringFieldWithMultiplePathComp
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "gnirts"));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << getCellID(0, 0) << ""
+                                             << "gnirts"),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, std::set<size_t>{}},
@@ -301,15 +374,25 @@ TEST(S2KeyGeneratorTest, CollationAppliedToStringsInArray) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "gnirts"));
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "2gnirts"));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString1(KeyString::Version::kLatestVersion,
+                                      BSON("" << getCellID(0, 0) << ""
+                                              << "gnirts"),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString2(KeyString::Version::kLatestVersion,
+                                      BSON("" << getCellID(0, 0) << ""
+                                              << "2gnirts"),
+                                      Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString1.release(), keyString2.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, {0U}}, actualMultikeyPaths);
@@ -324,27 +407,42 @@ TEST(S2KeyGeneratorTest, CollationAppliedToStringsInAllArrays) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "gnirts"
-                                << ""
-                                << "cba"));
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "gnirts"
-                                << ""
-                                << "fed"));
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "2gnirts"
-                                << ""
-                                << "cba"));
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "2gnirts"
-                                << ""
-                                << "fed"));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString1(KeyString::Version::kLatestVersion,
+                                      BSON("" << getCellID(0, 0) << ""
+                                              << "gnirts"
+                                              << ""
+                                              << "cba"),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString2(KeyString::Version::kLatestVersion,
+                                      BSON("" << getCellID(0, 0) << ""
+                                              << "gnirts"
+                                              << ""
+                                              << "fed"),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString3(KeyString::Version::kLatestVersion,
+                                      BSON("" << getCellID(0, 0) << ""
+                                              << "2gnirts"
+                                              << ""
+                                              << "cba"),
+                                      Ordering::make(BSONObj()));
+    KeyString::HeapBuilder keyString4(KeyString::Version::kLatestVersion,
+                                      BSON("" << getCellID(0, 0) << ""
+                                              << "2gnirts"
+                                              << ""
+                                              << "fed"),
+                                      Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{
+        keyString1.release(), keyString2.release(), keyString3.release(), keyString4.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, {0U}, {0U}}, actualMultikeyPaths);
@@ -358,12 +456,20 @@ TEST(S2KeyGeneratorTest, CollationDoesNotAffectNonStringFields) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << "" << 5));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << getCellID(0, 0) << "" << 5),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, std::set<size_t>{}},
@@ -378,14 +484,22 @@ TEST(S2KeyGeneratorTest, CollationAppliedToStringsInNestedObjects) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     ExpressionParams::initialize2dsphereParams(infoObj, &collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << BSON("c"
-                                        << "gnirts")));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << getCellID(0, 0) << ""
+                                             << BSON("c"
+                                                     << "gnirts")),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, std::set<size_t>{}},
@@ -400,13 +514,21 @@ TEST(S2KeyGeneratorTest, NoCollation) {
     const CollatorInterface* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << ""
-                                << "string"));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << getCellID(0, 0) << ""
+                                             << "string"),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, std::set<size_t>{}},
@@ -421,12 +543,20 @@ TEST(S2KeyGeneratorTest, EmptyArrayForLeadingFieldIsConsideredMultikey) {
     const CollatorInterface* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << BSONUndefined << "" << getCellID(0, 0)));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << BSONUndefined << "" << getCellID(0, 0)),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{{0U}, std::set<size_t>{}}, actualMultikeyPaths);
@@ -440,12 +570,20 @@ TEST(S2KeyGeneratorTest, EmptyArrayForTrailingFieldIsConsideredMultikey) {
     const CollatorInterface* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << getCellID(0, 0) << "" << BSONUndefined));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << getCellID(0, 0) << "" << BSONUndefined),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{std::set<size_t>{}, {0U}}, actualMultikeyPaths);
@@ -459,12 +597,20 @@ TEST(S2KeyGeneratorTest, SingleElementTrailingArrayIsConsideredMultikey) {
     const CollatorInterface* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << 99 << "" << getCellID(0, 0)));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << 99 << "" << getCellID(0, 0)),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{{1U}, std::set<size_t>{}}, actualMultikeyPaths);
@@ -478,12 +624,20 @@ TEST(S2KeyGeneratorTest, MidPathSingleElementArrayIsConsideredMultikey) {
     const CollatorInterface* collator = nullptr;
     ExpressionParams::initialize2dsphereParams(infoObj, collator, &params);
 
-    BSONObjSet actualKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet actualKeys;
     MultikeyPaths actualMultikeyPaths;
-    ExpressionKeysPrivate::getS2Keys(obj, keyPattern, params, &actualKeys, &actualMultikeyPaths);
-
-    BSONObjSet expectedKeys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
-    expectedKeys.insert(BSON("" << 99 << "" << getCellID(0, 0)));
+    ExpressionKeysPrivate::getS2Keys(obj,
+                                     keyPattern,
+                                     params,
+                                     &actualKeys,
+                                     &actualMultikeyPaths,
+                                     KeyString::Version::kLatestVersion,
+                                     Ordering::make(BSONObj()));
+
+    KeyString::HeapBuilder keyString(KeyString::Version::kLatestVersion,
+                                     BSON("" << 99 << "" << getCellID(0, 0)),
+                                     Ordering::make(BSONObj()));
+    KeyStringSet expectedKeys{keyString.release()};
 
     assertKeysetsEqual(expectedKeys, actualKeys);
     assertMultikeyPathsEqual(MultikeyPaths{{0U}, std::set<size_t>{}}, actualMultikeyPaths);