SERVER-51540 Support mod expression in SBE

4 files changed, 273 insertions, 4 deletions

diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index 590cd01d804..de541b40d4b 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -96,6 +96,7 @@ env.CppUnitTest(
         'expressions/sbe_index_of_test.cpp',
         'expressions/sbe_is_member_builtin_test.cpp',
         'expressions/sbe_iso_date_to_parts_test.cpp',
+        'expressions/sbe_mod_expression_test.cpp',
         'expressions/sbe_set_expressions_test.cpp',
         'expressions/sbe_to_upper_to_lower_test.cpp',
         'expressions/sbe_trigonometric_expressions_test.cpp',
diff --git a/src/mongo/db/exec/sbe/expressions/sbe_mod_expression_test.cpp b/src/mongo/db/exec/sbe/expressions/sbe_mod_expression_test.cpp
new file mode 100644
index 00000000000..6ca58baf241
--- /dev/null
+++ b/src/mongo/db/exec/sbe/expressions/sbe_mod_expression_test.cpp
@@ -0,0 +1,231 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/exec/sbe/expression_test_base.h"
+
+namespace mongo::sbe {
+
+class SBEModExprTest : public EExpressionTestFixture {
+protected:
+    void runAndAssertExpression(const vm::CodeFragment* compiledExpr, double expectedVal) {
+        auto [tag, val] = runCompiledExpression(compiledExpr);
+        value::ValueGuard guard(tag, val);
+        ASSERT_EQUALS(value::TypeTags::NumberDouble, tag);
+        ASSERT_APPROX_EQUAL(value::bitcastTo<double>(val), expectedVal, 0.000001);
+    }
+
+    void runAndAssertExpression(const vm::CodeFragment* compiledExpr, Decimal128 expectedVal) {
+        auto [tag, val] = runCompiledExpression(compiledExpr);
+        value::ValueGuard guard(tag, val);
+
+        ASSERT_EQUALS(value::TypeTags::NumberDecimal, tag);
+        ASSERT(value::bitcastTo<Decimal128>(val)
+                   .subtract(expectedVal)
+                   .toAbs()
+                   .isLessEqual(Decimal128(".000001")));
+    }
+
+    void runAndAssertExpression(const vm::CodeFragment* compiledExpr, int32_t expectedVal) {
+        auto [tag, val] = runCompiledExpression(compiledExpr);
+        value::ValueGuard guard(tag, val);
+        ASSERT_EQUALS(value::TypeTags::NumberInt32, tag);
+        ASSERT_EQUALS(value::bitcastTo<int32_t>(val), expectedVal);
+    }
+
+    void runAndAssertExpression(const vm::CodeFragment* compiledExpr, int64_t expectedVal) {
+        auto [tag, val] = runCompiledExpression(compiledExpr);
+        value::ValueGuard guard(tag, val);
+        ASSERT_EQUALS(value::TypeTags::NumberInt64, tag);
+        ASSERT_EQUALS(value::bitcastTo<int64_t>(val), expectedVal);
+    }
+
+    void runAndAssertNothing(const vm::CodeFragment* compiledExpr) {
+        auto [tag, val] = runCompiledExpression(compiledExpr);
+        value::ValueGuard guard(tag, val);
+        ASSERT_EQUALS(value::TypeTags::Nothing, tag);
+    }
+
+    void runAndAssertThrows(const vm::CodeFragment* compiledExpr) {
+        ASSERT_THROWS_CODE(runCompiledExpression(compiledExpr), AssertionException, 4848403);
+    }
+};
+
+TEST_F(SBEModExprTest, ComputesMod) {
+    value::OwnedValueAccessor slotAccessor1, slotAccessor2;
+    auto argSlot1 = bindAccessor(&slotAccessor1);
+    auto argSlot2 = bindAccessor(&slotAccessor2);
+    auto modExpr = sbe::makeE<sbe::EFunction>(
+        "mod", sbe::makeEs(makeE<EVariable>(argSlot1), makeE<EVariable>(argSlot2)));
+    auto compiledExpr = compileExpression(*modExpr);
+
+    const int32_t i32Val = 16;
+    const int32_t i32Mod = 4;
+    const int64_t i64Val = 2147483648;
+    const int64_t i64Mod = 2147483649;
+    const Decimal128 decVal(123.4);
+    const Decimal128 decMod(43.2);
+    const double dblVal(9.9);
+    const double dblMod(2.3);
+
+    slotAccessor1.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Val));
+    slotAccessor2.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Mod));
+    runAndAssertExpression(compiledExpr.get(), i32Val % i32Mod);
+
+    slotAccessor1.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(i64Val));
+    slotAccessor2.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(i64Mod));
+    runAndAssertExpression(compiledExpr.get(), i64Val % i64Mod);
+
+    auto [tagArgDecVal, valArgDecVal] = value::makeCopyDecimal(decVal);
+    slotAccessor1.reset(tagArgDecVal, valArgDecVal);
+    auto [tagArgDecMod, valArgDecMod] = value::makeCopyDecimal(decMod);
+    slotAccessor2.reset(tagArgDecMod, valArgDecMod);
+    runAndAssertExpression(compiledExpr.get(), decVal.modulo(decMod));
+
+    slotAccessor1.reset(value::TypeTags::NumberDouble, value::bitcastFrom<double>(dblVal));
+    slotAccessor2.reset(value::TypeTags::NumberDouble, value::bitcastFrom<double>(dblMod));
+    runAndAssertExpression(compiledExpr.get(), std::fmod(dblVal, dblMod));
+}
+
+TEST_F(SBEModExprTest, ComputesModDifferentWidths) {
+    value::OwnedValueAccessor slotAccessor1, slotAccessor2;
+    auto argSlot1 = bindAccessor(&slotAccessor1);
+    auto argSlot2 = bindAccessor(&slotAccessor2);
+    auto modExpr = sbe::makeE<sbe::EFunction>(
+        "mod", sbe::makeEs(makeE<EVariable>(argSlot1), makeE<EVariable>(argSlot2)));
+    auto compiledExpr = compileExpression(*modExpr);
+
+    const int32_t i32Val = 16;
+    const int32_t i32Mod = 4;
+    const int64_t i64Val = 2147483648;
+    const int64_t i64Mod = 2147483649;
+
+    slotAccessor1.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(i64Val));
+    slotAccessor2.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Mod));
+    runAndAssertExpression(compiledExpr.get(), i64Val % i32Mod);
+
+    slotAccessor1.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Val));
+    slotAccessor2.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(i64Mod));
+    runAndAssertExpression(compiledExpr.get(), i32Val % i64Mod);
+}
+
+TEST_F(SBEModExprTest, ComputesNothingIfNotNumeric) {
+    value::OwnedValueAccessor slotAccessor1, slotAccessor2;
+    auto argSlot1 = bindAccessor(&slotAccessor1);
+    auto argSlot2 = bindAccessor(&slotAccessor2);
+    auto modExpr = sbe::makeE<sbe::EFunction>(
+        "mod", sbe::makeEs(makeE<EVariable>(argSlot1), makeE<EVariable>(argSlot2)));
+    auto compiledExpr = compileExpression(*modExpr);
+
+    const int32_t i32Val = 16;
+    const int32_t i32Mod = 4;
+
+    auto [tagStrArg1, valStrArg1] = value::makeNewString("abc");
+    slotAccessor1.reset(tagStrArg1, valStrArg1);
+    auto [tagStrArg2, valStrArg2] = value::makeNewString("xyz");
+    slotAccessor2.reset(tagStrArg2, valStrArg2);
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Val));
+    slotAccessor2.reset(tagStrArg2, valStrArg2);
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(tagStrArg1, valStrArg1);
+    slotAccessor2.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Mod));
+    runAndAssertNothing(compiledExpr.get());
+}
+
+TEST_F(SBEModExprTest, ComputesNothingIfNullOrMissing) {
+    value::OwnedValueAccessor slotAccessor1, slotAccessor2;
+    auto argSlot1 = bindAccessor(&slotAccessor1);
+    auto argSlot2 = bindAccessor(&slotAccessor2);
+    auto modExpr = sbe::makeE<sbe::EFunction>(
+        "mod", sbe::makeEs(makeE<EVariable>(argSlot1), makeE<EVariable>(argSlot2)));
+    auto compiledExpr = compileExpression(*modExpr);
+
+    const int32_t i32Val = 16;
+    const int32_t i32Mod = 4;
+
+    slotAccessor1.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Val));
+    slotAccessor2.reset(value::TypeTags::Nothing, 0);
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::Nothing, 0);
+    slotAccessor2.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Mod));
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::Nothing, 0);
+    slotAccessor2.reset(value::TypeTags::Nothing, 0);
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Val));
+    slotAccessor2.reset(value::TypeTags::Null, 0);
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::Null, 0);
+    slotAccessor2.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Mod));
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::Null, 0);
+    slotAccessor2.reset(value::TypeTags::Null, 0);
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::Null, 0);
+    slotAccessor2.reset(value::TypeTags::Nothing, 0);
+    runAndAssertNothing(compiledExpr.get());
+
+    slotAccessor1.reset(value::TypeTags::Nothing, 0);
+    slotAccessor2.reset(value::TypeTags::Null, 0);
+    runAndAssertNothing(compiledExpr.get());
+}
+
+TEST_F(SBEModExprTest, ErrorIfModRHSIsZero) {
+    value::OwnedValueAccessor slotAccessor1, slotAccessor2;
+    auto argSlot1 = bindAccessor(&slotAccessor1);
+    auto argSlot2 = bindAccessor(&slotAccessor2);
+    auto modExpr = sbe::makeE<sbe::EFunction>(
+        "mod", sbe::makeEs(makeE<EVariable>(argSlot1), makeE<EVariable>(argSlot2)));
+    auto compiledExpr = compileExpression(*modExpr);
+
+    const int32_t i32Val = 16;
+    const int32_t i32Mod = 0;
+    const Decimal128 decVal(123.4);
+    const Decimal128 decMod(0);
+
+    slotAccessor1.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Val));
+    slotAccessor2.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(i32Mod));
+    runAndAssertThrows(compiledExpr.get());
+
+    auto [tagArgDecVal, valArgDecVal] = value::makeCopyDecimal(decVal);
+    slotAccessor1.reset(tagArgDecVal, valArgDecVal);
+    auto [tagArgDecMod, valArgDecMod] = value::makeCopyDecimal(decMod);
+    slotAccessor2.reset(tagArgDecMod, valArgDecMod);
+    runAndAssertThrows(compiledExpr.get());
+}
+
+}  // namespace mongo::sbe
\ No newline at end of file
diff --git a/src/mongo/db/exec/sbe/vm/arith.cpp b/src/mongo/db/exec/sbe/vm/arith.cpp
index f663f5880d6..36d91a9f2ce 100644
--- a/src/mongo/db/exec/sbe/vm/arith.cpp
+++ b/src/mongo/db/exec/sbe/vm/arith.cpp
@@ -497,19 +497,19 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericMod(value::Type
                 return {false, value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(result)};
             }
             case value::TypeTags::NumberInt64: {
-                assertNonZero(numericCast<int32_t>(rhsTag, rhsValue) != 0);
+                assertNonZero(numericCast<int64_t>(rhsTag, rhsValue) != 0);
                 auto result = overflow::safeMod(numericCast<int64_t>(lhsTag, lhsValue),
                                                 numericCast<int64_t>(rhsTag, rhsValue));
                 return {false, value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(result)};
             }
             case value::TypeTags::NumberDouble: {
-                assertNonZero(numericCast<int32_t>(rhsTag, rhsValue) != 0);
+                assertNonZero(numericCast<double>(rhsTag, rhsValue) != 0);
                 auto result = fmod(numericCast<double>(lhsTag, lhsValue),
                                    numericCast<double>(rhsTag, rhsValue));
                 return {false, value::TypeTags::NumberDouble, value::bitcastFrom<double>(result)};
             }
             case value::TypeTags::NumberDecimal: {
-                assertNonZero(numericCast<Decimal128>(rhsTag, rhsValue).isZero());
+                assertNonZero(!numericCast<Decimal128>(rhsTag, rhsValue).isZero());
                 auto result = numericCast<Decimal128>(lhsTag, lhsValue)
                                   .modulo(numericCast<Decimal128>(rhsTag, rhsValue));
                 auto [tag, val] = value::makeCopyDecimal(result);
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index 0ca2094949f..04c97fc023b 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -1836,7 +1836,44 @@ public:
         unsupportedExpression("$meta");
     }
     void visit(ExpressionMod* expr) final {
-        unsupportedExpression(expr->getOpName());
+        auto frameId = _context->frameIdGenerator->generate();
+        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};
+
+        // If the rhs is a small integral double, convert it to int32 to match $mod MQL semantics.
+        auto numericConvert32 =
+            sbe::makeE<sbe::ENumericConvert>(rhsVar.clone(), sbe::value::TypeTags::NumberInt32);
+        auto rhsExpr = buildMultiBranchConditional(
+            CaseValuePair{
+                sbe::makeE<sbe::EPrimBinary>(
+                    sbe::EPrimBinary::logicAnd,
+                    sbe::makeE<sbe::ETypeMatch>(
+                        rhsVar.clone(), getBSONTypeMask(sbe::value::TypeTags::NumberDouble)),
+                    sbe::makeE<sbe::EPrimUnary>(
+                        sbe::EPrimUnary::logicNot,
+                        sbe::makeE<sbe::ETypeMatch>(
+                            lhsVar.clone(), getBSONTypeMask(sbe::value::TypeTags::NumberDouble)))),
+                sbe::makeE<sbe::EFunction>(
+                    "fillEmpty", sbe::makeEs(std::move(numericConvert32), rhsVar.clone()))},
+            rhsVar.clone());
+
+        auto modExpr = buildMultiBranchConditional(
+            CaseValuePair{sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr,
+                                                       generateNullOrMissing(lhsVar),
+                                                       generateNullOrMissing(rhsVar)),
+                          sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Null, 0)},
+            CaseValuePair{sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr,
+                                                       generateNonNumericCheck(lhsVar),
+                                                       generateNonNumericCheck(rhsVar)),
+                          sbe::makeE<sbe::EFail>(ErrorCodes::Error{5154000},
+                                                 "$mod only supports numeric types")},
+            sbe::makeE<sbe::EFunction>("mod", sbe::makeEs(lhsVar.clone(), std::move(rhsExpr))));
+
+        _context->pushExpr(
+            sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(modExpr)));
     }
     void visit(ExpressionMultiply* expr) final {
         auto arity = expr->getChildren().size();