SERVER-51540 Support mod expression in SBE

8 files changed, 352 insertions, 5 deletions

diff --git a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml
index ef61e08aee2..92535b6a866 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml
@@ -54,6 +54,7 @@ selector:
   - jstests/core/or2.js
   - jstests/core/or3.js
   - jstests/core/orj.js
+  - jstests/core/projection_expr_mod.js
   - jstests/core/ref.js
   - jstests/core/ref4.js
   - jstests/core/regex_limit.js
diff --git a/jstests/aggregation/sources/merge/mode_replace_insert.js b/jstests/aggregation/sources/merge/mode_replace_insert.js
index 07791658496..1975365ae06 100644
--- a/jstests/aggregation/sources/merge/mode_replace_insert.js
+++ b/jstests/aggregation/sources/merge/mode_replace_insert.js
@@ -1,7 +1,4 @@
 // Tests for the $merge stage with whenMatched: "replace" and whenNotMatched: "insert".
-// @tags: [
-//   sbe_incompatible,
-// ]
 (function() {
 "use strict";
 
diff --git a/jstests/core/projection_expr_mod.js b/jstests/core/projection_expr_mod.js
new file mode 100644
index 00000000000..e06c7b9224f
--- /dev/null
+++ b/jstests/core/projection_expr_mod.js
@@ -0,0 +1,79 @@
+// Confirm correctness of $mod evaluation in find projection.
+
+(function() {
+"use strict";
+
+load("jstests/aggregation/extras/utils.js");        // For assertArrayEq.
+load("jstests/libs/sbe_assert_error_override.js");  // Override error-code-checking APIs.
+
+const coll = db.projection_expr_mod;
+coll.drop();
+
+assert.commandWorked(coll.insertMany([
+    {n: "double % double", v: 138.5, m: 3.0},
+    {n: "double % long", v: 138.5, m: NumberLong(3)},
+    {n: "double % int", v: 138.5, m: NumberInt(3)},
+    {n: "int % double", v: NumberInt(8), m: 3.25},
+    {n: "int % double int", v: NumberInt(8), m: 3.0},
+    {n: "int % int", v: NumberInt(8), m: NumberInt(3)},
+    {n: "int % long", v: NumberInt(8), m: NumberLong(3)},
+    {n: "long % double", v: NumberLong(8), m: 3.25},
+    {n: "long % double int", v: NumberLong(8), m: 3.0},
+    {n: "long % double long", v: NumberLong(500000000000), m: 450000000000.0},
+    {n: "long % int", v: NumberLong(8), m: NumberInt(3)},
+    {n: "long % long", v: NumberLong(8), m: NumberLong(3)},
+    {n: "very long % very long", v: NumberLong(800000000000), m: NumberLong(300000000000)}
+]));
+
+const result = coll.find({}, {f: {$mod: ["$v", "$m"]}, _id: 0, n: 1}).toArray();
+const expectedResult = [
+    {n: "double % double", f: 0.5},
+    {n: "double % long", f: 0.5},
+    {n: "double % int", f: 0.5},
+    {n: "int % double", f: 1.5},
+    {n: "int % double int", f: 2},
+    {n: "int % int", f: 2},
+    {n: "int % long", f: NumberLong(2)},
+    {n: "long % double", f: 1.5},
+    {n: "long % double int", f: NumberLong(2)},
+    {n: "long % double long", f: 50000000000},
+    {n: "long % int", f: NumberLong(2)},
+    {n: "long % long", f: NumberLong(2)},
+    {n: "very long % very long", f: NumberLong(200000000000)}
+];
+assertArrayEq({actual: result, expected: expectedResult});
+
+//
+// Confirm error cases.
+//
+
+// Confirm mod by 0 fails in an expected manner.
+assert(coll.drop());
+assert.commandWorked(coll.insert({a: 10}));
+let error = assert.throws(() => coll.find({}, {f: {$mod: ["$a", 0]}, _id: 0, n: 1}).toArray());
+assert.commandFailedWithCode(error, 4848403);
+
+assert(coll.drop());
+assert.commandWorked(coll.insert({a: NumberInt(10)}));
+error =
+    assert.throws(() => coll.find({}, {f: {$mod: ["$a", NumberInt(0)]}, _id: 0, n: 1}).toArray());
+assert.commandFailedWithCode(error, 4848403);
+
+assert(coll.drop());
+assert.commandWorked(coll.insert({a: NumberLong(10)}));
+error =
+    assert.throws(() => coll.find({}, {f: {$mod: ["$a", NumberLong(0)]}, _id: 0, n: 1}).toArray());
+assert.commandFailedWithCode(error, 4848403);
+
+// Clear collection again and reset.
+assert(coll.drop());
+assert.commandWorked(coll.insert({a: 10}));
+
+// Confirm expected behavior for NaN and Infinity values.
+assert.eq(coll.findOne({}, {f: {$mod: ["$a", NaN]}, _id: 0}), {f: NaN});
+assert.eq(coll.findOne({}, {f: {$mod: ["$a", Infinity]}, _id: 0}), {f: 10});
+assert.eq(coll.findOne({}, {f: {$mod: ["$a", -Infinity]}, _id: 0}), {f: 10});
+assert.eq(coll.findOne({}, {f: {$mod: [Infinity, "$a"]}, _id: 0}), {f: NaN});
+assert.eq(coll.findOne({}, {f: {$mod: [-Infinity, "$a"]}, _id: 0}), {f: NaN});
+assert.eq(coll.findOne({}, {f: {$mod: [NaN, "$a"]}, _id: 0}), {f: NaN});
+})();
\ No newline at end of file
diff --git a/jstests/libs/sbe_assert_error_override.js b/jstests/libs/sbe_assert_error_override.js
index 8d6147bccd9..5ebfa28410f 100644
--- a/jstests/libs/sbe_assert_error_override.js
+++ b/jstests/libs/sbe_assert_error_override.js
@@ -30,6 +30,7 @@ const equivalentErrorCodesList = [
     [16007, 5066300],
     [16608, 4848401],
     [16609, 5073101],
+    [16610, 4848403],
     [16555, 5073102],
     [28680, 4903701],
     [28689, 5126701],
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index d6305e59c9c..b7a31d18696 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -97,6 +97,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..e859fd62860 100644
--- a/src/mongo/db/exec/sbe/vm/arith.cpp
+++ b/src/mongo/db/exec/sbe/vm/arith.cpp
@@ -509,7 +509,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericMod(value::Type
                 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 e91b33159b6..0559766a4bf 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -2160,7 +2160,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();