SERVER-66545 Collator should be passed as an argument to "generateSortKey" builtin in SBE

5 files changed, 20 insertions, 21 deletions

diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index 653ad7b4490..7e0758b6e22 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -506,7 +506,7 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = {
     {"hash", BuiltinFn{kAnyNumberOfArgs, vm::Builtin::hash, false}},
     {"ftsMatch", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::ftsMatch, false}},
     {"generateSortKey",
-     BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::generateSortKey, false}},
+     BuiltinFn{[](size_t n) { return n == 2 || n == 3; }, vm::Builtin::generateSortKey, false}},
     {"tsSecond", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::tsSecond, false}},
     {"tsIncrement", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::tsIncrement, false}},
     {"typeMatch", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::typeMatch, false}},
diff --git a/src/mongo/db/exec/sbe/values/sort_spec.h b/src/mongo/db/exec/sbe/values/sort_spec.h
index 39d0e195fb7..af48415c539 100644
--- a/src/mongo/db/exec/sbe/values/sort_spec.h
+++ b/src/mongo/db/exec/sbe/values/sort_spec.h
@@ -39,31 +39,24 @@ namespace mongo::sbe::value {
  */
 class SortSpec {
 public:
-    SortSpec(const BSONObj& sortPattern, const CollatorInterface* collator)
-        : _sortPattern(sortPattern.getOwned()), _collator(collator), _keyGen(initKeyGen()) {}
-
-    SortSpec(const SortSpec& other)
-        : _sortPattern(other._sortPattern), _collator(other._collator), _keyGen(initKeyGen()) {}
+    SortSpec(const BSONObj& sortPattern)
+        : _sortPattern(sortPattern.getOwned()), _keyGen(initKeyGen()) {}
+    SortSpec(const SortSpec& other) : _sortPattern(other._sortPattern), _keyGen(initKeyGen()) {}
 
     SortSpec& operator=(const SortSpec&) = delete;
 
-    KeyString::Value generateSortKey(const BSONObj& obj) const;
+    KeyString::Value generateSortKey(const BSONObj& obj, const CollatorInterface* collator);
 
     const BSONObj& getPattern() const {
         return _sortPattern;
     }
 
-    const CollatorInterface* getCollator() const {
-        return _collator;
-    }
-
     size_t getApproximateSize() const;
 
 private:
     BtreeKeyGenerator initKeyGen() const;
 
     const BSONObj _sortPattern;
-    const CollatorInterface* const _collator;
     const BtreeKeyGenerator _keyGen;
 };
 }  // namespace mongo::sbe::value
diff --git a/src/mongo/db/exec/sbe/values/value.cpp b/src/mongo/db/exec/sbe/values/value.cpp
index 90b626e2a8d..cde56e31ba2 100644
--- a/src/mongo/db/exec/sbe/values/value.cpp
+++ b/src/mongo/db/exec/sbe/values/value.cpp
@@ -175,12 +175,12 @@ size_t PcreRegex::getApproximateSize() const {
     return sizeof(PcreRegex) + _pattern.size() + 1 + _options.size() + 1 + pcreSize;
 }
 
-KeyString::Value SortSpec::generateSortKey(const BSONObj& obj) const {
+KeyString::Value SortSpec::generateSortKey(const BSONObj& obj, const CollatorInterface* collator) {
     KeyStringSet keySet;
     SharedBufferFragmentBuilder allocator(KeyString::HeapBuilder::kHeapAllocatorDefaultBytes);
     const bool skipMultikey = false;
     MultikeyPaths* multikeyPaths = nullptr;
-    _keyGen.getKeys(allocator, obj, skipMultikey, &keySet, multikeyPaths);
+    _keyGen.getKeys(allocator, obj, skipMultikey, &keySet, multikeyPaths, collator);
 
     // When 'isSparse' is false, BtreeKeyGenerator::getKeys() is guaranteed to insert at least
     // one key into 'keySet', so this assertion should always be true.
@@ -212,12 +212,10 @@ BtreeKeyGenerator SortSpec::initKeyGen() const {
     auto version = KeyString::Version::kLatestVersion;
     auto ordering = Ordering::make(_sortPattern);
 
-    return {std::move(fields), std::move(fixed), isSparse, _collator, version, ordering};
+    return {std::move(fields), std::move(fixed), isSparse, version, ordering};
 }
 
 size_t SortSpec::getApproximateSize() const {
-    // _collator points to a block of memory that SortSpec doesn't own, so we don't acccount
-    // for the size of this block of memory here.
     auto size = sizeof(SortSpec);
     size += _sortPattern.isOwned() ? _sortPattern.objsize() : 0;
     size += _keyGen.getApproximateSize() - sizeof(_keyGen);
diff --git a/src/mongo/db/exec/sbe/values/value_printer.cpp b/src/mongo/db/exec/sbe/values/value_printer.cpp
index 12df709d5c4..78e655114a3 100644
--- a/src/mongo/db/exec/sbe/values/value_printer.cpp
+++ b/src/mongo/db/exec/sbe/values/value_printer.cpp
@@ -459,8 +459,6 @@ void ValuePrinter<T>::writeValueToStream(TypeTags tag, Value val, size_t depth)
         case TypeTags::sortSpec:
             stream << "SortSpec(";
             writeObjectToStream(getSortSpecView(val)->getPattern());
-            stream << ", ";
-            writeCollatorToStream(getSortSpecView(val)->getCollator());
             stream << ')';
             break;
         case TypeTags::indexBounds:
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index a8bd51ab7ba..ea111abda70 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -4045,7 +4045,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinGetRegexFlags(A
 }
 
 std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinGenerateSortKey(ArityType arity) {
-    invariant(arity == 2);
+    invariant(arity == 2 || arity == 3);
 
     auto [ssOwned, ssTag, ssVal] = getFromStack(0);
     auto [objOwned, objTag, objVal] = getFromStack(1);
@@ -4053,6 +4053,15 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinGenerateSortKey
         return {false, value::TypeTags::Nothing, 0};
     }
 
+    CollatorInterface* collator{nullptr};
+    if (arity == 3) {
+        auto [collatorOwned, collatorTag, collatorVal] = getFromStack(2);
+        if (collatorTag != value::TypeTags::collator) {
+            return {false, value::TypeTags::Nothing, 0};
+        }
+        collator = value::getCollatorView(collatorVal);
+    }
+
     auto ss = value::getSortSpecView(ssVal);
 
     auto obj = [objTag = objTag, objVal = objVal]() {
@@ -4069,7 +4078,8 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinGenerateSortKey
 
     return {true,
             value::TypeTags::ksValue,
-            value::bitcastFrom<KeyString::Value*>(new KeyString::Value(ss->generateSortKey(obj)))};
+            value::bitcastFrom<KeyString::Value*>(
+                new KeyString::Value(ss->generateSortKey(obj, collator)))};
 }
 
 std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinReverseArray(ArityType arity) {