SERVER-41725 Make BtreeAccessMethod generate keys using KeyString

4 files changed, 116 insertions, 0 deletions

diff --git a/src/mongo/db/query/collation/SConscript b/src/mongo/db/query/collation/SConscript
index 1a2c36d3ace..4304071e676 100644
--- a/src/mongo/db/query/collation/SConscript
+++ b/src/mongo/db/query/collation/SConscript
@@ -109,6 +109,7 @@ icuEnv.CppUnitTest(
         "collator_interface_mock_test.cpp",
     ],
     LIBDEPS=[
+        '$BUILD_DIR/mongo/db/storage/key_string',
         "collator_factory_mock",
         "collator_icu",
         "collator_interface",
diff --git a/src/mongo/db/query/collation/collation_index_key_test.cpp b/src/mongo/db/query/collation/collation_index_key_test.cpp
index 20a788d7df4..1f184778fe2 100644
--- a/src/mongo/db/query/collation/collation_index_key_test.cpp
+++ b/src/mongo/db/query/collation/collation_index_key_test.cpp
@@ -34,12 +34,36 @@
 #include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/bson/json.h"
 #include "mongo/db/query/collation/collator_interface_mock.h"
+#include "mongo/db/storage/key_string.h"
 #include "mongo/unittest/unittest.h"
 
 namespace {
 
 using namespace mongo;
 
+void assertKeyStringCollatorOutput(const CollatorInterfaceMock& collator,
+                                   const BSONObj& dataObj,
+                                   const BSONObj& expected) {
+    KeyString::Builder ks(KeyString::Version::kLatestVersion, KeyString::ALL_ASCENDING);
+    ks.appendBSONElement(dataObj.firstElement(), [&](StringData stringData) {
+        return collator.getComparisonString(stringData);
+    });
+
+    ASSERT_EQ(
+        ks.getValueCopy(),
+        KeyString::Builder(KeyString::Version::kLatestVersion, expected, KeyString::ALL_ASCENDING));
+}
+
+void assertKeyStringCollatorThrows(const CollatorInterfaceMock& collator, const BSONObj& dataObj) {
+    KeyString::Builder ks(KeyString::Version::kLatestVersion, KeyString::ALL_ASCENDING);
+    ASSERT_THROWS_CODE(ks.appendBSONElement(dataObj.firstElement(),
+                                            [&](StringData stringData) {
+                                                return collator.getComparisonString(stringData);
+                                            }),
+                       AssertionException,
+                       ErrorCodes::CannotBuildIndexKeys);
+}
+
 TEST(CollationIndexKeyTest, IsCollatableTypeShouldBeTrueForString) {
     BSONObj obj = BSON("foo"
                        << "string");
@@ -79,6 +103,16 @@ TEST(CollationIndexKeyTest, CollationAwareAppendReversesStringWithReverseMockCol
                            << "gnirts"));
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendReversesStringWithReverseMockCollator) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("foo"
+                           << "string");
+    assertKeyStringCollatorOutput(collator,
+                                  dataObj,
+                                  BSON(""
+                                       << "gnirts"));
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlySerializesEmptyComparisonKey) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObjBuilder builder;
@@ -94,6 +128,16 @@ TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlySerializesEmptyComparis
     ASSERT_BSONOBJ_EQ(out.obj(), expectedObj);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendCorrectlySerializesEmptyComparisonKey) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObjBuilder builder;
+    builder.append("foo", StringData());
+
+    BSONObjBuilder expectedBuilder;
+    expectedBuilder.append("", StringData());
+    assertKeyStringCollatorOutput(collator, builder.obj(), expectedBuilder.obj());
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlySerializesWithEmbeddedNullByte) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObjBuilder builder;
@@ -109,6 +153,16 @@ TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlySerializesWithEmbeddedN
     ASSERT_BSONOBJ_EQ(out.obj(), expectedObj);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendCorrectlySerializesWithEmbeddedNullByte) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObjBuilder builder;
+    builder.append("foo", "a\0b"_sd);
+
+    BSONObjBuilder expectedBuilder;
+    expectedBuilder.append("", "b\0a"_sd);
+    assertKeyStringCollatorOutput(collator, builder.obj(), expectedBuilder.obj());
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesSimpleEmbeddedObject) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj dataObj = BSON("" << BSON("a"
@@ -121,6 +175,15 @@ TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesSimpleEmbeddedO
     ASSERT_BSONOBJ_EQ(out.obj(), expected);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendCorrectlyReversesSimpleEmbeddedObject) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("" << BSON("a"
+                                      << "!foo"));
+    BSONObj expected = BSON("" << BSON("a"
+                                       << "oof!"));
+    assertKeyStringCollatorOutput(collator, dataObj, expected);
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesSimpleEmbeddedArray) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj dataObj = BSON("" << BSON_ARRAY("foo"
@@ -133,6 +196,15 @@ TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesSimpleEmbeddedA
     ASSERT_BSONOBJ_EQ(out.obj(), expected);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendCorrectlyReversesSimpleEmbeddedArray) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("" << BSON_ARRAY("foo"
+                                            << "bar"));
+    BSONObj expected = BSON("" << BSON_ARRAY("oof"
+                                             << "rab"));
+    assertKeyStringCollatorOutput(collator, dataObj, expected);
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesComplexNesting) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj dataObj = fromjson(
@@ -149,6 +221,19 @@ TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesComplexNesting)
     ASSERT_BSONOBJ_EQ(out.obj(), expected);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendCorrectlyReversesComplexNesting) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = fromjson(
+        "{ '' : [{'a': 'ha', 'b': 2},"
+        "'bar',"
+        "{'c': 2, 'd': 'ah', 'e': 'abc', 'f': ['cba', 'xyz']}]})");
+    BSONObj expected = fromjson(
+        "{ '' : [{'a': 'ah', 'b': 2},"
+        "'rab',"
+        "{'c': 2, 'd': 'ha', 'e': 'cba', 'f': ['abc', 'zyx']}]})");
+    assertKeyStringCollatorOutput(collator, dataObj, expected);
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendThrowsIfSymbol) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj dataObj = BSON("" << BSONSymbol("mySymbol"));
@@ -159,6 +244,12 @@ TEST(CollationIndexKeyTest, CollationAwareAppendThrowsIfSymbol) {
         ErrorCodes::CannotBuildIndexKeys);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendThrowsIfSymbol) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("" << BSONSymbol("mySymbol"));
+    assertKeyStringCollatorThrows(collator, dataObj);
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendDoesNotThrowOnSymbolIfNoCollation) {
     BSONObj dataObj = BSON("" << BSONSymbol("mySymbol"));
     BSONObj expected = BSON("" << BSONSymbol("mySymbol"));
@@ -179,6 +270,14 @@ TEST(CollationIndexKeyTest, CollationAwareAppendThrowsIfSymbolInsideObject) {
         ErrorCodes::CannotBuildIndexKeys);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendThrowsIfSymbolInsideObject) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("" << BSON("a"
+                                      << "foo"
+                                      << "b" << BSONSymbol("mySymbol")));
+    assertKeyStringCollatorThrows(collator, dataObj);
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendThrowsIfSymbolInsideArray) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj dataObj = BSON("" << BSON_ARRAY("foo" << BSONSymbol("mySymbol")));
@@ -189,4 +288,10 @@ TEST(CollationIndexKeyTest, CollationAwareAppendThrowsIfSymbolInsideArray) {
         ErrorCodes::CannotBuildIndexKeys);
 }
 
+TEST(CollationIndexKeyTest, KeyStringAppendThrowsIfSymbolInsideArray) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("" << BSON_ARRAY("foo" << BSONSymbol("mySymbol")));
+    assertKeyStringCollatorThrows(collator, dataObj);
+}
+
 }  // namespace
diff --git a/src/mongo/db/query/collation/collator_interface.cpp b/src/mongo/db/query/collation/collator_interface.cpp
index d000ae85527..ae2dedbcde9 100644
--- a/src/mongo/db/query/collation/collator_interface.cpp
+++ b/src/mongo/db/query/collation/collator_interface.cpp
@@ -40,4 +40,8 @@ void CollatorInterface::hash_combine(size_t& seed, StringData stringToHash) cons
     SimpleStringDataComparator::kInstance.hash_combine(seed, comparisonKey.getKeyData());
 }
 
+std::string CollatorInterface::getComparisonString(StringData stringData) const {
+    return getComparisonKey(stringData).getKeyData().toString();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/query/collation/collator_interface.h b/src/mongo/db/query/collation/collator_interface.h
index 7cba6aa2bb1..0e19137f8df 100644
--- a/src/mongo/db/query/collation/collator_interface.h
+++ b/src/mongo/db/query/collation/collator_interface.h
@@ -115,6 +115,12 @@ public:
     virtual ComparisonKey getComparisonKey(StringData stringData) const = 0;
 
     /**
+     * Returns the comparison key string for 'stringData', according to this collation. See
+     * ComparisonKey's comments for details.
+     */
+    std::string getComparisonString(StringData stringData) const;
+
+    /**
      * Returns whether this collation has the same matching and sorting semantics as 'other'.
      */
     bool operator==(const CollatorInterface& other) const {