SERVER-49742 ExpressionAdd in SBE should use doubleDoubleSum built-in when called with more than two operands

4 files changed, 116 insertions, 6 deletions

diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index 5c598445272..6b517f293b5 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -359,7 +359,7 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = {
     {"addToArray", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::addToArray, true}},
     {"addToSet", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::addToSet, true}},
     {"doubleDoubleSum",
-     BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::doubleDoubleSum, true}},
+     BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::doubleDoubleSum, false}},
     {"bitTestZero", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::bitTestZero, false}},
     {"bitTestMask", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::bitTestMask, false}},
     {"bitTestPosition",
@@ -402,6 +402,7 @@ static stdx::unordered_map<std::string, InstrFn> kInstrFunctions = {
      InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsNumber, false}},
     {"isBinData",
      InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsBinData, false}},
+    {"isDate", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsDate, false}},
     {"sum", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendSum, true}},
     {"min", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendMin, true}},
     {"max", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendMax, true}},
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index b5890497f45..ba7c849431b 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -96,6 +96,7 @@ int Instruction::stackOffset[Instruction::Tags::lastInstruction] = {
     0,  // isString
     0,  // isNumber
     0,  // isBinData
+    0,  // isDate
     0,  // typeMatch
 
     0,  // function is special, the stack offset is encoded in the instruction itself
@@ -314,6 +315,10 @@ void CodeFragment::appendIsBinData() {
     appendSimpleInstruction(Instruction::isBinData);
 }
 
+void CodeFragment::appendIsDate() {
+    appendSimpleInstruction(Instruction::isDate);
+}
+
 void CodeFragment::appendTypeMatch(uint32_t typeMask) {
     Instruction i;
     i.tag = Instruction::typeMatch;
@@ -1814,6 +1819,18 @@ std::tuple<uint8_t, value::TypeTags, value::Value> ByteCode::run(const CodeFragm
                     }
                     break;
                 }
+                case Instruction::isDate: {
+                    auto [owned, tag, val] = getFromStack(0);
+
+                    if (tag != value::TypeTags::Nothing) {
+                        topStack(false, value::TypeTags::Boolean, tag == value::TypeTags::Date);
+                    }
+
+                    if (owned) {
+                        value::releaseValue(tag, val);
+                    }
+                    break;
+                }
                 case Instruction::typeMatch: {
                     auto typeMask = value::readFromMemory<uint32_t>(pcPointer);
                     pcPointer += sizeof(typeMask);
diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h
index a5197d17437..e4590a79c71 100644
--- a/src/mongo/db/exec/sbe/vm/vm.h
+++ b/src/mongo/db/exec/sbe/vm/vm.h
@@ -149,6 +149,7 @@ struct Instruction {
         isString,
         isNumber,
         isBinData,
+        isDate,
         typeMatch,
 
         function,
@@ -259,6 +260,7 @@ public:
     void appendIsString();
     void appendIsNumber();
     void appendIsBinData();
+    void appendIsDate();
     void appendTypeMatch(uint32_t typeMask);
     void appendFunction(Builtin f, uint8_t arity);
     void appendJump(int jumpOffset);
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index b17164df951..22c03198b9b 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -883,13 +883,103 @@ public:
         _context->pushExpr(
             sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(absExpr)));
     }
+
     void visit(ExpressionAdd* expr) final {
-        _context->ensureArity(2);
-        auto rhs = _context->popExpr();
-        auto lhs = _context->popExpr();
-        _context->pushExpr(
-            sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::add, std::move(lhs), std::move(rhs)));
+        size_t arity = expr->getChildren().size();
+        _context->ensureArity(arity);
+        auto frameId = _context->frameIdGenerator->generate();
+
+
+        auto generateNotNumberOrDate = [frameId](const sbe::value::SlotId slotId) {
+            sbe::EVariable var{frameId, slotId};
+            return sbe::makeE<sbe::EPrimBinary>(
+                sbe::EPrimBinary::logicAnd,
+                sbe::makeE<sbe::EPrimUnary>(
+                    sbe::EPrimUnary::logicNot,
+                    sbe::makeE<sbe::EFunction>("isNumber", sbe::makeEs(var.clone()))),
+                sbe::makeE<sbe::EPrimUnary>(
+                    sbe::EPrimUnary::logicNot,
+                    sbe::makeE<sbe::EFunction>("isDate", sbe::makeEs(var.clone()))));
+        };
+
+        if (arity == 2) {
+            auto rhs = _context->popExpr();
+            auto lhs = _context->popExpr();
+            auto binds = sbe::makeEs(std::move(lhs), std::move(rhs));
+            sbe::EVariable lhsVar{frameId, 0};
+            sbe::EVariable rhsVar{frameId, 1};
+
+            auto addExpr = sbe::makeE<sbe::EIf>(
+                sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr,
+                                             generateNullOrMissing(frameId, 0),
+                                             generateNullOrMissing(frameId, 1)),
+                sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Null, 0),
+                sbe::makeE<sbe::EIf>(
+                    sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr,
+                                                 generateNotNumberOrDate(0),
+                                                 generateNotNumberOrDate(1)),
+                    sbe::makeE<sbe::EFail>(
+                        ErrorCodes::Error{4974201},
+                        "only numbers and dates are allowed in an $add expression"),
+                    sbe::makeE<sbe::EIf>(
+                        sbe::makeE<sbe::EPrimBinary>(
+                            sbe::EPrimBinary::logicAnd,
+                            sbe::makeE<sbe::EFunction>("isDate", sbe::makeEs(lhsVar.clone())),
+                            sbe::makeE<sbe::EFunction>("isDate", sbe::makeEs(rhsVar.clone()))),
+                        sbe::makeE<sbe::EFail>(ErrorCodes::Error{4974202},
+                                               "only one date allowed in an $add expression"),
+                        sbe::makeE<sbe::EPrimBinary>(
+                            sbe::EPrimBinary::add, lhsVar.clone(), rhsVar.clone()))));
+
+            _context->pushExpr(
+                sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(addExpr)));
+        } else {
+            std::vector<std::unique_ptr<sbe::EExpression>> binds;
+            for (size_t i = 0; i < arity; i++) {
+                binds.push_back(_context->popExpr());
+            }
+            std::reverse(std::begin(binds), std::end(binds));
+
+            std::vector<std::unique_ptr<sbe::EExpression>> checkExprsNull;
+            std::vector<std::unique_ptr<sbe::EExpression>> checkExprsNotNumberOrDate;
+            std::vector<std::unique_ptr<sbe::EExpression>> argVars;
+            for (size_t idx = 0; idx < arity; idx++) {
+                checkExprsNull.push_back(generateNullOrMissing(frameId, idx));
+                checkExprsNotNumberOrDate.push_back(generateNotNumberOrDate(idx));
+                argVars.push_back(sbe::makeE<sbe::EVariable>(frameId, idx));
+            }
+
+            using iter_t = std::vector<std::unique_ptr<sbe::EExpression>>::iterator;
+            auto checkNullAllArguments =
+                std::accumulate(std::move_iterator<iter_t>(checkExprsNull.begin() + 1),
+                                std::move_iterator<iter_t>(checkExprsNull.end()),
+                                std::move(checkExprsNull.front()),
+                                [](auto&& acc, auto&& ex) {
+                                    return sbe::makeE<sbe::EPrimBinary>(
+                                        sbe::EPrimBinary::logicOr, std::move(acc), std::move(ex));
+                                });
+            auto checkNotNumberOrDateAllArguments =
+                std::accumulate(std::move_iterator<iter_t>(checkExprsNotNumberOrDate.begin() + 1),
+                                std::move_iterator<iter_t>(checkExprsNotNumberOrDate.end()),
+                                std::move(checkExprsNotNumberOrDate.front()),
+                                [](auto&& acc, auto&& ex) {
+                                    return sbe::makeE<sbe::EPrimBinary>(
+                                        sbe::EPrimBinary::logicOr, std::move(acc), std::move(ex));
+                                });
+            auto addExpr = sbe::makeE<sbe::EIf>(
+                std::move(checkNullAllArguments),
+                sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Null, 0),
+                sbe::makeE<sbe::EIf>(
+                    std::move(checkNotNumberOrDateAllArguments),
+                    sbe::makeE<sbe::EFail>(
+                        ErrorCodes::Error{4974203},
+                        "only numbers and dates are allowed in an $add expression"),
+                    sbe::makeE<sbe::EFunction>("doubleDoubleSum", std::move(argVars))));
+            _context->pushExpr(
+                sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(addExpr)));
+        }
     }
+
     void visit(ExpressionAllElementsTrue* expr) final {
         unsupportedExpression(expr->getOpName());
     }