SERVER-41725 Make BtreeAccessMethod generate keys using KeyString

1 files changed, 30 insertions, 21 deletions

diff --git a/src/mongo/db/storage/key_string.h b/src/mongo/db/storage/key_string.h
index aefa78c75a2..48529162659 100644
--- a/src/mongo/db/storage/key_string.h
+++ b/src/mongo/db/storage/key_string.h
@@ -354,6 +354,8 @@ enum DecimalContinuationMarker {
     kDCMHasContinuationLargerThanDoubleRoundedUpTo15Digits = 0x3
 };
 
+using StringTransformFn = std::function<std::string(StringData)>;
+
 template <class BufferT>
 class BuilderBase {
 public:
@@ -424,7 +426,7 @@ public:
      */
     template <typename T = BufferT>
     typename std::enable_if<std::is_same<T, BufBuilder>::value, Value>::type release() {
-        _prepareForRelease();
+        _doneAppending();
         _transition(BuildState::kReleased);
         return {version, _typeBits, static_cast<size_t>(_buffer.len()), _buffer.release()};
     }
@@ -434,7 +436,7 @@ public:
      * buffer.
      */
     Value getValueCopy() {
-        _prepareForRelease();
+        _doneAppending();
         invariant(_state == BuildState::kEndAdded || _state == BuildState::kAppendedRecordID ||
                   _state == BuildState::kAppendedTypeBits);
         BufBuilder newBuf(_buffer.len());
@@ -444,7 +446,11 @@ public:
 
     void appendRecordId(RecordId loc);
     void appendTypeBits(const TypeBits& bits);
-    void appendBSONElement(const BSONElement& elem);
+
+    /*
+     * Function 'f' will be applied to all string elements contained in 'elem'.
+     */
+    void appendBSONElement(const BSONElement& elem, const StringTransformFn& f = nullptr);
 
     /**
      * Resets to an empty state.
@@ -519,15 +525,15 @@ private:
     void _appendDate(Date_t val, bool invert);
     void _appendTimestamp(Timestamp val, bool invert);
     void _appendOID(OID val, bool invert);
-    void _appendString(StringData val, bool invert);
+    void _appendString(StringData val, bool invert, const StringTransformFn& f);
     void _appendSymbol(StringData val, bool invert);
     void _appendCode(StringData val, bool invert);
     void _appendCodeWString(const BSONCodeWScope& val, bool invert);
     void _appendBinData(const BSONBinData& val, bool invert);
     void _appendRegex(const BSONRegEx& val, bool invert);
     void _appendDBRef(const BSONDBRef& val, bool invert);
-    void _appendArray(const BSONArray& val, bool invert);
-    void _appendObject(const BSONObj& val, bool invert);
+    void _appendArray(const BSONArray& val, bool invert, const StringTransformFn& f);
+    void _appendObject(const BSONObj& val, bool invert, const StringTransformFn& f);
     void _appendNumberDouble(const double num, bool invert);
     void _appendNumberLong(const long long num, bool invert);
     void _appendNumberInt(const int num, bool invert);
@@ -538,10 +544,13 @@ private:
      *              if NULL, not included in encoding
      *              if not NULL, put in after type, before value
      */
-    void _appendBsonValue(const BSONElement& elem, bool invert, const StringData* name);
+    void _appendBsonValue(const BSONElement& elem,
+                          bool invert,
+                          const StringData* name,
+                          const StringTransformFn& f);
 
     void _appendStringLike(StringData str, bool invert);
-    void _appendBson(const BSONObj& obj, bool invert);
+    void _appendBson(const BSONObj& obj, bool invert, const StringTransformFn& f);
     void _appendSmallDouble(double value, DecimalContinuationMarker dcm, bool invert);
     void _appendLargeDouble(double value, DecimalContinuationMarker dcm, bool invert);
     void _appendInteger(const long long num, bool invert);
@@ -560,7 +569,7 @@ private:
 
     void _appendBytes(const void* source, size_t bytes, bool invert);
 
-    void _prepareForRelease() {
+    void _doneAppending() {
         if (_state == BuildState::kAppendingBSONElements) {
             _appendDiscriminator(_discriminator);
         }
@@ -631,39 +640,39 @@ struct isKeyString<BuilderBase<BufferT>> : public std::true_type {};
 template <>
 struct isKeyString<Value> : public std::true_type {};
 
-template <class T>
+template <class T, class U>
 inline typename std::enable_if<isKeyString<T>::value, bool>::type operator<(const T& lhs,
-                                                                            const T& rhs) {
+                                                                            const U& rhs) {
     return lhs.compare(rhs) < 0;
 }
 
-template <class T>
+template <class T, class U>
 inline typename std::enable_if<isKeyString<T>::value, bool>::type operator<=(const T& lhs,
-                                                                             const T& rhs) {
+                                                                             const U& rhs) {
     return lhs.compare(rhs) <= 0;
 }
 
-template <class T>
+template <class T, class U>
 inline typename std::enable_if<isKeyString<T>::value, bool>::type operator==(const T& lhs,
-                                                                             const T& rhs) {
+                                                                             const U& rhs) {
     return lhs.compare(rhs) == 0;
 }
 
-template <class T>
+template <class T, class U>
 inline typename std::enable_if<isKeyString<T>::value, bool>::type operator>(const T& lhs,
-                                                                            const T& rhs) {
+                                                                            const U& rhs) {
     return lhs.compare(rhs) > 0;
 }
 
-template <class T>
+template <class T, class U>
 inline typename std::enable_if<isKeyString<T>::value, bool>::type operator>=(const T& lhs,
-                                                                             const T& rhs) {
+                                                                             const U& rhs) {
     return lhs.compare(rhs) >= 0;
 }
 
-template <class T>
+template <class T, class U>
 inline typename std::enable_if<isKeyString<T>::value, bool>::type operator!=(const T& lhs,
-                                                                             const T& rhs) {
+                                                                             const U& rhs) {
     return !(lhs == rhs);
 }