SERVER-69969 Avoid pushing constant argument values on the VM stack

6 files changed, 104 insertions, 53 deletions

diff --git a/buildscripts/gdb/mongo_printers.py b/buildscripts/gdb/mongo_printers.py
index 76ac681c1aa..2fdd409ebaf 100644
--- a/buildscripts/gdb/mongo_printers.py
+++ b/buildscripts/gdb/mongo_printers.py
@@ -802,7 +802,7 @@ class SbeCodeFragmentPrinter(object):
                 offset = read_as_integer(cur_op, int_size)
                 cur_op += int_size
                 args = 'offset: ' + str(offset) + ', target: ' + hex(cur_op + offset)
-            elif op_name in ['pushConstVal', 'getFieldConst']:
+            elif op_name in ['pushConstVal', 'getFieldImm']:
                 tag = read_as_integer(cur_op, tag_size)
                 args = 'tag: ' + self.valuetags_lookup.get(tag, "unknown") + \
                     ', value: ' + hex(read_as_integer(cur_op + tag_size, value_size))
@@ -814,7 +814,7 @@ class SbeCodeFragmentPrinter(object):
                 args = 'convert to: ' + \
                     self.valuetags_lookup.get(read_as_integer(cur_op, tag_size), "unknown")
                 cur_op += tag_size
-            elif op_name in ['typeMatch']:
+            elif op_name in ['typeMatchImm']:
                 args = 'mask: ' + hex(read_as_integer(cur_op, uint32_size))
                 cur_op += uint32_size
             elif op_name in ['function', 'functionSmall']:
@@ -826,10 +826,10 @@ class SbeCodeFragmentPrinter(object):
                 args = 'builtin: ' + self.builtins_lookup.get(builtin_id, "unknown")
                 args += ' arity: ' + str(read_as_integer(cur_op + builtin_size, arity_size))
                 cur_op += (builtin_size + arity_size)
-            elif op_name in ['fillEmptyConst']:
+            elif op_name in ['fillEmptyImm']:
                 args = 'Instruction::Constants: ' + str(read_as_integer(cur_op, uint8_size))
                 cur_op += uint8_size
-            elif op_name in ['traverseFConst', 'traversePConst']:
+            elif op_name in ['traverseFImm', 'traversePImm']:
                 const_enum = read_as_integer(cur_op, uint8_size)
                 cur_op += uint8_size
                 args = \
diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index 19fd4cf83c4..d3457dd2ba8 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -605,6 +605,22 @@ vm::CodeFragment EFunction::compileDirect(CompileCtx& ctx) const {
         }
         vm::CodeFragment code;
 
+        // Optimize well known set of functions with constant arguments and generate their
+        // specialized variants.
+        if (_name == "typeMatch" && _nodes[1]->as<EConstant>()) {
+            auto [tag, val] = _nodes[1]->as<EConstant>()->getConstant();
+            if (tag == value::TypeTags::NumberInt64) {
+                auto mask = value::bitcastTo<int64_t>(val);
+                uassert(6996901,
+                        "Second argument to typeMatch() must be a 32-bit integer constant",
+                        mask >> 32 == 0 || mask >> 32 == -1);
+                code.append(_nodes[0]->compileDirect(ctx));
+                code.appendTypeMatch(mask);
+
+                return code;
+            }
+        }
+
         for (size_t idx = arity; idx-- > 0;) {
             code.append(_nodes[idx]->compileDirect(ctx));
         }
diff --git a/src/mongo/db/exec/sbe/sbe_test.cpp b/src/mongo/db/exec/sbe/sbe_test.cpp
index c3ffba02394..8a93c14bf28 100644
--- a/src/mongo/db/exec/sbe/sbe_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_test.cpp
@@ -469,29 +469,29 @@ TEST(SBEVM, CodeFragmentToStringArgs) {
         std::string toStringPattern{kAddrPattern};
 
         code.appendFillEmpty(vm::Instruction::Null);
-        toStringPattern += instrPattern("fillEmptyConst", "k: Null");
+        toStringPattern += instrPattern("fillEmptyImm", "k: Null");
         code.appendFillEmpty(vm::Instruction::False);
-        toStringPattern += instrPattern("fillEmptyConst", "k: False");
+        toStringPattern += instrPattern("fillEmptyImm", "k: False");
         code.appendFillEmpty(vm::Instruction::True);
-        toStringPattern += instrPattern("fillEmptyConst", "k: True");
+        toStringPattern += instrPattern("fillEmptyImm", "k: True");
 
         code.appendTraverseP(0xAA, vm::Instruction::Nothing);
         auto offsetP1 = 0xAA - code.instrs().size();
         toStringPattern +=
-            instrPattern("traversePConst", "k: Nothing, offset: " + std::to_string(offsetP1));
+            instrPattern("traversePImm", "k: Nothing, offset: " + std::to_string(offsetP1));
         code.appendTraverseP(0xAA, vm::Instruction::Int32One);
         auto offsetP2 = 0xAA - code.instrs().size();
         toStringPattern +=
-            instrPattern("traversePConst", "k: 1, offset: " + std::to_string(offsetP2));
+            instrPattern("traversePImm", "k: 1, offset: " + std::to_string(offsetP2));
         code.appendTraverseF(0xBB, vm::Instruction::True);
         auto offsetF = 0xBB - code.instrs().size();
         toStringPattern +=
-            instrPattern("traverseFConst", "k: True, offset: " + std::to_string(offsetF));
+            instrPattern("traverseFImm", "k: True, offset: " + std::to_string(offsetF));
 
         auto [tag, val] = value::makeNewString("Hello world!");
         value::ValueGuard guard{tag, val};
         code.appendGetField(tag, val);
-        toStringPattern += instrPattern("getFieldConst", "value: \"Hello world!\"");
+        toStringPattern += instrPattern("getFieldImm", "value: \"Hello world!\"");
 
         code.appendAdd();
         toStringPattern += instrPattern("add", "");
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index c1427ba6eee..7875ef7f499 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -108,16 +108,16 @@ int Instruction::stackOffset[Instruction::Tags::lastInstruction] = {
     -2,  // collCmp3w
 
     -1,  // fillEmpty
-    0,   // fillEmptyConst
+    0,   // fillEmptyImm
     -1,  // getField
-    0,   // getFieldConst
+    0,   // getFieldImm
     -1,  // getElement
     -1,  // collComparisonKey
     -1,  // getFieldOrElement
     -2,  // traverseP
-    0,   // traversePConst
+    0,   // traversePImm
     -2,  // traverseF
-    0,   // traverseFConst
+    0,   // traverseFImm
     -2,  // setField
     0,   // getArraySize
 
@@ -144,6 +144,7 @@ int Instruction::stackOffset[Instruction::Tags::lastInstruction] = {
     0,  // isMinKey
     0,  // isMaxKey
     0,  // isTimestamp
+    0,  // typeMatchImm
 
     0,  // function is special, the stack offset is encoded in the instruction itself
     0,  // functionSmall is special, the stack offset is encoded in the instruction itself
@@ -260,8 +261,8 @@ std::string CodeFragment::toString() const {
                 break;
             }
             // Instructions with other kinds of arguments.
-            case Instruction::traversePConst:
-            case Instruction::traverseFConst: {
+            case Instruction::traversePImm:
+            case Instruction::traverseFImm: {
                 auto k = readFromMemory<Instruction::Constants>(pcPointer);
                 pcPointer += sizeof(k);
                 auto offset = readFromMemory<int>(pcPointer);
@@ -269,13 +270,13 @@ std::string CodeFragment::toString() const {
                 ss << "k: " << Instruction::toStringConstants(k) << ", offset: " << offset;
                 break;
             }
-            case Instruction::fillEmptyConst: {
+            case Instruction::fillEmptyImm: {
                 auto k = readFromMemory<Instruction::Constants>(pcPointer);
                 pcPointer += sizeof(k);
                 ss << "k: " << Instruction::toStringConstants(k);
                 break;
             }
-            case Instruction::getFieldConst:
+            case Instruction::getFieldImm:
             case Instruction::pushConstVal: {
                 auto tag = readFromMemory<value::TypeTags>(pcPointer);
                 pcPointer += sizeof(tag);
@@ -303,6 +304,12 @@ std::string CodeFragment::toString() const {
                 ss << "tag: " << tag;
                 break;
             }
+            case Instruction::typeMatchImm: {
+                auto mask = readFromMemory<uint32_t>(pcPointer);
+                pcPointer += sizeof(mask);
+                ss << "mask: " << mask;
+                break;
+            }
             case Instruction::function:
             case Instruction::functionSmall: {
                 auto f = readFromMemory<Builtin>(pcPointer);
@@ -513,7 +520,7 @@ void CodeFragment::appendSimpleInstruction(Instruction::Tags tag) {
 
 void CodeFragment::appendFillEmpty(Instruction::Constants k) {
     Instruction i;
-    i.tag = Instruction::fillEmptyConst;
+    i.tag = Instruction::fillEmptyImm;
     adjustStackSimple(i);
 
     auto offset = allocateSpace(sizeof(Instruction) + sizeof(k));
@@ -530,7 +537,7 @@ void CodeFragment::appendGetField(value::TypeTags tag, value::Value val) {
     invariant(value::isString(tag));
 
     Instruction i;
-    i.tag = Instruction::getFieldConst;
+    i.tag = Instruction::getFieldImm;
     adjustStackSimple(i);
 
     auto offset = allocateSpace(sizeof(Instruction) + sizeof(tag) + sizeof(val));
@@ -630,7 +637,7 @@ void CodeFragment::appendIsRecordId() {
 
 void CodeFragment::appendTraverseP(int codePosition, Instruction::Constants k) {
     Instruction i;
-    i.tag = Instruction::traversePConst;
+    i.tag = Instruction::traversePImm;
     adjustStackSimple(i);
 
     auto size = sizeof(Instruction) + sizeof(codePosition) + sizeof(k);
@@ -645,7 +652,7 @@ void CodeFragment::appendTraverseP(int codePosition, Instruction::Constants k) {
 
 void CodeFragment::appendTraverseF(int codePosition, Instruction::Constants k) {
     Instruction i;
-    i.tag = Instruction::traverseFConst;
+    i.tag = Instruction::traverseFImm;
     adjustStackSimple(i);
 
     auto size = sizeof(Instruction) + sizeof(codePosition) + sizeof(k);
@@ -658,6 +665,18 @@ void CodeFragment::appendTraverseF(int codePosition, Instruction::Constants k) {
     offset += writeToMemory(offset, codeOffset);
 }
 
+void CodeFragment::appendTypeMatch(uint32_t mask) {
+    Instruction i;
+    i.tag = Instruction::typeMatchImm;
+    adjustStackSimple(i);
+
+    auto size = sizeof(Instruction) + sizeof(mask);
+    auto offset = allocateSpace(size);
+
+    offset += writeToMemory(offset, i);
+    offset += writeToMemory(offset, mask);
+}
+
 void CodeFragment::appendFunction(Builtin f, ArityType arity) {
     Instruction i;
     const bool isSmallArity = (arity <= std::numeric_limits<SmallArityType>::max());
@@ -5321,7 +5340,7 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
                     }
                     break;
                 }
-                case Instruction::fillEmptyConst: {
+                case Instruction::fillEmptyImm: {
                     auto k = readFromMemory<Instruction::Constants>(pcPointer);
                     pcPointer += sizeof(k);
 
@@ -5371,7 +5390,7 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
                     }
                     break;
                 }
-                case Instruction::getFieldConst: {
+                case Instruction::getFieldImm: {
                     auto tagField = readFromMemory<value::TypeTags>(pcPointer);
                     pcPointer += sizeof(tagField);
                     auto valField = readFromMemory<value::Value>(pcPointer);
@@ -5467,7 +5486,7 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
                     traverseP(code);
                     break;
                 }
-                case Instruction::traversePConst: {
+                case Instruction::traversePImm: {
                     auto k = readFromMemory<Instruction::Constants>(pcPointer);
                     pcPointer += sizeof(k);
 
@@ -5485,7 +5504,7 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
                     traverseF(code);
                     break;
                 }
-                case Instruction::traverseFConst: {
+                case Instruction::traverseFImm: {
                     auto k = readFromMemory<Instruction::Constants>(pcPointer);
                     pcPointer += sizeof(k);
 
@@ -5733,6 +5752,21 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
                     runTagCheck(value::TypeTags::Timestamp);
                     break;
                 }
+                case Instruction::typeMatchImm: {
+                    auto mask = readFromMemory<uint32_t>(pcPointer);
+                    pcPointer += sizeof(mask);
+
+                    auto [owned, tag, val] = getFromStack(0);
+                    if (tag != value::TypeTags::Nothing) {
+                        topStack(false,
+                                 value::TypeTags::Boolean,
+                                 value::bitcastFrom<bool>(getBSONTypeMask(tag) & mask));
+                    }
+                    if (owned) {
+                        value::releaseValue(tag, val);
+                    }
+                    break;
+                }
                 case Instruction::function:
                 case Instruction::functionSmall: {
                     auto f = readFromMemory<Builtin>(pcPointer);
diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h
index 3f6a959aaf9..db435950d33 100644
--- a/src/mongo/db/exec/sbe/vm/vm.h
+++ b/src/mongo/db/exec/sbe/vm/vm.h
@@ -274,16 +274,16 @@ struct Instruction {
         collCmp3w,
 
         fillEmpty,
-        fillEmptyConst,
+        fillEmptyImm,
         getField,
-        getFieldConst,
+        getFieldImm,
         getElement,
         collComparisonKey,
         getFieldOrElement,
         traverseP,  // traverse projection paths
-        traversePConst,
+        traversePImm,
         traverseF,  // traverse filter paths
-        traverseFConst,
+        traverseFImm,
         setField,
         getArraySize,  // number of elements
 
@@ -310,6 +310,7 @@ struct Instruction {
         isMinKey,
         isMaxKey,
         isTimestamp,
+        typeMatchImm,
 
         function,
         functionSmall,
@@ -422,12 +423,12 @@ struct Instruction {
                 return "collCmp3w";
             case fillEmpty:
                 return "fillEmpty";
-            case fillEmptyConst:
-                return "fillEmptyConst";
+            case fillEmptyImm:
+                return "fillEmptyImm";
             case getField:
                 return "getField";
-            case getFieldConst:
-                return "getFieldConst";
+            case getFieldImm:
+                return "getFieldImm";
             case getElement:
                 return "getElement";
             case collComparisonKey:
@@ -436,12 +437,12 @@ struct Instruction {
                 return "getFieldOrElement";
             case traverseP:
                 return "traverseP";
-            case traversePConst:
-                return "traversePConst";
+            case traversePImm:
+                return "traversePImm";
             case traverseF:
                 return "traverseF";
-            case traverseFConst:
-                return "traverseFConst";
+            case traverseFImm:
+                return "traverseFImm";
             case setField:
                 return "setField";
             case getArraySize:
@@ -488,6 +489,8 @@ struct Instruction {
                 return "isMaxKey";
             case isTimestamp:
                 return "isTimestamp";
+            case typeMatchImm:
+                return "typeMatchImm";
             case function:
                 return "function";
             case functionSmall:
@@ -783,6 +786,7 @@ public:
     void appendIsTimestamp() {
         appendSimpleInstruction(Instruction::isTimestamp);
     }
+    void appendTypeMatch(uint32_t mask);
     void appendFunction(Builtin f, ArityType arity);
     void appendJump(int jumpOffset);
     void appendJumpTrue(int jumpOffset);
diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.cpp b/src/mongo/db/query/sbe_stage_builder_helpers.cpp
index e9672bce394..4823b68b46b 100644
--- a/src/mongo/db/query/sbe_stage_builder_helpers.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_helpers.cpp
@@ -113,16 +113,20 @@ std::unique_ptr<sbe::EExpression> makeIsMember(std::unique_ptr<sbe::EExpression>
 
     return makeIsMember(std::move(input), std::move(arr), std::move(collatorVar));
 }
-
-std::unique_ptr<sbe::EExpression> generateNullOrMissing(const sbe::EVariable& var) {
-    return makeBinaryOp(sbe::EPrimBinary::logicOr,
-                        makeNot(makeFunction("exists", var.clone())),
+std::unique_ptr<sbe::EExpression> generateNullOrMissingExpr(const sbe::EExpression& expr) {
+    return makeFunction("fillEmpty",
                         makeFunction("typeMatch",
-                                     var.clone(),
+                                     expr.clone(),
                                      makeConstant(sbe::value::TypeTags::NumberInt64,
                                                   sbe::value::bitcastFrom<int64_t>(
                                                       getBSONTypeMask(BSONType::jstNULL) |
-                                                      getBSONTypeMask(BSONType::Undefined)))));
+                                                      getBSONTypeMask(BSONType::Undefined)))),
+                        sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Boolean,
+                                                   sbe::value::bitcastFrom<bool>(true)));
+}
+
+std::unique_ptr<sbe::EExpression> generateNullOrMissing(const sbe::EVariable& var) {
+    return generateNullOrMissingExpr(var);
 }
 
 std::unique_ptr<sbe::EExpression> generateNullOrMissing(const sbe::FrameId frameId,
@@ -132,14 +136,7 @@ std::unique_ptr<sbe::EExpression> generateNullOrMissing(const sbe::FrameId frame
 }
 
 std::unique_ptr<sbe::EExpression> generateNullOrMissing(std::unique_ptr<sbe::EExpression> arg) {
-    return makeBinaryOp(sbe::EPrimBinary::logicOr,
-                        makeNot(makeFunction("exists", arg->clone())),
-                        makeFunction("typeMatch",
-                                     arg->clone(),
-                                     makeConstant(sbe::value::TypeTags::NumberInt64,
-                                                  sbe::value::bitcastFrom<int64_t>(
-                                                      getBSONTypeMask(BSONType::jstNULL) |
-                                                      getBSONTypeMask(BSONType::Undefined)))));
+    return generateNullOrMissingExpr(*arg);
 }
 
 std::unique_ptr<sbe::EExpression> generateNonNumericCheck(const sbe::EVariable& var) {