SERVER-54129 Simplify generated SBE expressions for $mod, $add, and $split when some arguments are constants

2 files changed, 77 insertions, 0 deletions

diff --git a/jstests/aggregation/expressions/add.js b/jstests/aggregation/expressions/add.js
new file mode 100644
index 00000000000..cb9b146c84f
--- /dev/null
+++ b/jstests/aggregation/expressions/add.js
@@ -0,0 +1,31 @@
+// Confirm correctness of $add evaluation in find projection.
+(function() {
+"use strict";
+
+load("jstests/aggregation/extras/utils.js");  // For assertArrayEq.
+load("jstests/libs/sbe_util.js");             // For checkSBEEnabled.
+
+const isSBEEnabled = checkSBEEnabled(db);
+if (isSBEEnabled) {
+    // Override error-code-checking APIs. We only load this when SBE is explicitly enabled, because
+    // it causes failures in the parallel suites.
+    load("jstests/libs/sbe_assert_error_override.js");
+}
+
+const coll = db.expression_add;
+coll.drop();
+assert.commandWorked(coll.insert({a: NumberInt(2), b: NumberLong(3), c: 3.5}));
+
+const testCases = [
+    [["$a", "$b", "$c"], 8.5],
+    [["$a", "$b", null], null],
+    [["$a", "$b", 5], 10],
+    [[5, "$a", "$b"], 10],
+    [["$a", 5, "$b"], 10],
+    [["$a", 20, "$c", 10], 35.5],
+];
+for (const testCase of testCases) {
+    const [addExpr, expected] = testCase;
+    assert.eq(coll.findOne({}, {sum: {$add: addExpr}, _id: 0}), {sum: expected}, testCase);
+}
+})();
diff --git a/jstests/aggregation/expressions/split.js b/jstests/aggregation/expressions/split.js
index 3425e81ecc2..1c03a5c1e2a 100644
--- a/jstests/aggregation/expressions/split.js
+++ b/jstests/aggregation/expressions/split.js
@@ -6,7 +6,13 @@
 
 load("jstests/aggregation/extras/utils.js");  // For assertErrorCode and testExpression.
 load("jstests/libs/sbe_assert_error_override.js");
+load("jstests/libs/sbe_util.js");  // For checkSBEEnabled.
 
+// TODO SERVER-58095: When the classic engine is used, it will eagerly return null values, even
+// if some of its arguments are invalid. This is not the case when SBE is enabled because of the
+// order in which arguments are evaluated. In certain cases, errors will be thrown or the empty
+// string will be returned instead of null.
+const sbeEnabled = checkSBEEnabled(db);
 const coll = db.split;
 
 coll.drop();
@@ -42,8 +48,18 @@ testExpression(coll, {$split: [null, "abc"]}, null);
 // Ensure that $split produces null when given missing fields as input.
 testExpression(coll, {$split: ["$a", "a"]}, null);
 testExpression(coll, {$split: ["a", "$a"]}, null);
+testExpression(coll, {$split: ["$a", null]}, null);
+testExpression(coll, {$split: [null, "$a"]}, null);
 testExpression(coll, {$split: ["$missing", {$toLower: "$missing"}]}, null);
 
+// SBE expression translation will detect our empty string, whereas the classic engine will
+// detect that "$a" is missing and return null.
+if (sbeEnabled) {
+    testExpression(coll, {$split: ["", "$a"]}, [""]);
+} else {
+    testExpression(coll, {$split: ["", "$a"]}, null);
+}
+
 //
 // Error Code tests with constant-folding optimization.
 //
@@ -78,4 +94,34 @@ pipeline = {
     $project: {split: {$split: ["abc", ""]}}
 };
 assertErrorCode(coll, pipeline, 40087);
+
+const stringNumericArg = {
+    $split: [1, "$a"]
+};
+if (sbeEnabled) {
+    pipeline = {$project: {split: stringNumericArg}};
+    assertErrorCode(coll, pipeline, 40085);
+} else {
+    testExpression(coll, stringNumericArg, null);
+}
+
+const splitNumArg = {
+    $split: ["$b", 1]
+};
+if (sbeEnabled) {
+    pipeline = {$project: {split: splitNumArg}};
+    assertErrorCode(coll, pipeline, 40086);
+} else {
+    testExpression(coll, splitNumArg, null);
+}
+
+const emptyStringDelim = {
+    $split: ["$abc", ""]
+};
+if (sbeEnabled) {
+    pipeline = {$project: {split: emptyStringDelim}};
+    assertErrorCode(coll, pipeline, 40087);
+} else {
+    testExpression(coll, emptyStringDelim, null);
+}
 })();