SERVER-51549 Support expression $reverseArray in SBE

8 files changed, 259 insertions, 3 deletions

diff --git a/jstests/aggregation/bugs/reverseArray.js b/jstests/aggregation/bugs/reverseArray.js
index cf80c040171..4853188d555 100644
--- a/jstests/aggregation/bugs/reverseArray.js
+++ b/jstests/aggregation/bugs/reverseArray.js
@@ -1,7 +1,8 @@
 // SERVER-23029 added a new expression, $reverseArray, which consumes an array or a nullish value
 // and produces either the reversed version of that array, or null. In this test file, we check the
 // behavior and error cases.
-load("jstests/aggregation/extras/utils.js");  // For assertErrorCode.
+load("jstests/libs/sbe_assert_error_override.js");  // Override error-code-checking APIs.
+load("jstests/aggregation/extras/utils.js");        // For assertErrorCode.
 
 (function() {
 "use strict";
@@ -11,9 +12,18 @@ coll.drop();
 
 // We need a document to flow through the pipeline, even though we don't care what fields it
 // has.
-coll.insert({});
+coll.insert({
+    nullField: null,
+    undefField: undefined,
+    embedded: [[1, 2], [3, 4]],
+    singleElem: [1],
+    normal: [1, 2, 3],
+    num: 1,
+    empty: []
+});
 
 assertErrorCode(coll, [{$project: {reversed: {$reverseArray: 1}}}], 34435);
+assertErrorCode(coll, [{$project: {reversed: {$reverseArray: "$num"}}}], 34435);
 
 var res = coll.aggregate([{$project: {reversed: {$reverseArray: {$literal: [1, 2]}}}}]);
 var output = res.toArray();
@@ -29,4 +39,59 @@ var res = coll.aggregate([{$project: {reversed: {$reverseArray: "$notAField"}}}]
 var output = res.toArray();
 assert.eq(1, output.length);
 assert.eq(output[0].reversed, null);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: {$literal: [[1, 2], [3, 4]]}}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, [[3, 4], [1, 2]]);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: "$embedded"}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, [[3, 4], [1, 2]]);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: {$literal: null}}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, null);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: "$nullField"}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, null);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: {$literal: undefined}}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, null);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: "$undefField"}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, null);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: {$literal: [1]}}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, [1]);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: "$singleElem"}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, [1]);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: "$normal"}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, [3, 2, 1]);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: {$literal: []}}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, []);
+
+var res = coll.aggregate([{$project: {reversed: {$reverseArray: "$empty"}}}]);
+var output = res.toArray();
+assert.eq(1, output.length);
+assert.eq(output[0].reversed, []);
 }());
diff --git a/jstests/libs/sbe_assert_error_override.js b/jstests/libs/sbe_assert_error_override.js
index 6d0bfd491c1..78bb6996b83 100644
--- a/jstests/libs/sbe_assert_error_override.js
+++ b/jstests/libs/sbe_assert_error_override.js
@@ -48,6 +48,7 @@ const equivalentErrorCodesList = [
     [28766, 4903706],
     [31034, 4848972],
     [31095, 4848972],
+    [34435, 5154901],
     [40066, 4934200],
     [40085, 5155402],
     [40086, 5155400],
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index 18d2c7819eb..31e94c6c5c1 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -134,6 +134,7 @@ env.CppUnitTest(
         'expressions/sbe_mod_expression_test.cpp',
         'expressions/sbe_regex_test.cpp',
         'expressions/sbe_replace_one_expression_test.cpp',
+        'expressions/sbe_reverse_array_builtin_test.cpp',
         'expressions/sbe_set_expressions_test.cpp',
         'expressions/sbe_shard_filter_builtin_test.cpp',
         'expressions/sbe_to_upper_to_lower_test.cpp',
diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index ba23b8e9bf1..8783b35d7fd 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -442,6 +442,7 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = {
     {"extractSubArray",
      BuiltinFn{[](size_t n) { return n == 2 || n == 3; }, vm::Builtin::extractSubArray, false}},
     {"isArrayEmpty", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::isArrayEmpty, false}},
+    {"reverseArray", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::reverseArray, false}},
     {"dateAdd", BuiltinFn{[](size_t n) { return n == 5; }, vm::Builtin::dateAdd, false}},
     {"hasNullBytes", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::hasNullBytes, false}},
 };
diff --git a/src/mongo/db/exec/sbe/expressions/sbe_reverse_array_builtin_test.cpp b/src/mongo/db/exec/sbe/expressions/sbe_reverse_array_builtin_test.cpp
new file mode 100644
index 00000000000..4d32ddf1597
--- /dev/null
+++ b/src/mongo/db/exec/sbe/expressions/sbe_reverse_array_builtin_test.cpp
@@ -0,0 +1,111 @@
+/**
+ *    Copyright (C) 2021-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/platform/basic.h"
+
+#include "mongo/db/exec/sbe/expression_test_base.h"
+#include "mongo/db/exec/sbe/values/bson.h"
+
+namespace mongo::sbe {
+
+class SBEBuiltinReverseArrayTest : public EExpressionTestFixture {
+protected:
+    using TypedValue = std::pair<value::TypeTags, value::Value>;
+
+    /**
+     * Assert that result of 'reverseArray(array)' is equal to 'expected'.
+     * NOTE: Values behind arguments and the return value of this function are owned by the caller.
+     */
+    void runAndAssertExpression(TypedValue array, TypedValue expected) {
+        // We do not copy array value on purpose. During the copy, order of elements in ArraySet
+        // may change. 'reverseArray' return value depends on the order of elements in the input
+        // array. After copy 'reverse' may return elements in a different order from the one
+        // expected by the caller.
+        value::ViewOfValueAccessor arraySlotAccessor;
+        auto arraySlot = bindAccessor(&arraySlotAccessor);
+        arraySlotAccessor.reset(array.first, array.second);
+        auto arrayExpr = makeE<EVariable>(arraySlot);
+
+        auto reverseArrayExpr = makeE<EFunction>("reverseArray", makeEs(std::move(arrayExpr)));
+        auto compiledExpr = compileExpression(*reverseArrayExpr);
+
+        auto actual = runCompiledExpression(compiledExpr.get());
+        value::ValueGuard actualGuard{actual};
+
+        auto [compareTag, compareValue] =
+            value::compareValue(actual.first, actual.second, expected.first, expected.second);
+        ASSERT_EQ(compareTag, value::TypeTags::NumberInt32);
+        ASSERT_EQ(compareValue, 0);
+    }
+};
+
+TEST_F(SBEBuiltinReverseArrayTest, Array) {
+    for (auto makeArrayFn : {makeBsonArray, makeArray}) {
+        auto testArray = makeArrayFn(BSON_ARRAY(1 << 2 << 3));
+        value::ValueGuard testArrayGuard{testArray};
+
+        auto expectedResult = makeArray(BSON_ARRAY(3 << 2 << 1));
+        value::ValueGuard expectedResultGuard{expectedResult};
+
+        runAndAssertExpression(testArray, expectedResult);
+    }
+}
+
+TEST_F(SBEBuiltinReverseArrayTest, ArraySet) {
+    // This test needs to do a bit more work to find the correct reversed order since ArraySet's
+    // internal order is determined by it's hash function and not the order that elements are added
+    // to it.
+    auto testArray = makeArraySet(BSON_ARRAY(1 << 2 << 3));
+    value::ValueGuard testArrayGuard{testArray};
+    value::ArrayEnumerator testEnumerator{testArray.first, testArray.second};
+
+    std::vector<std::pair<value::TypeTags, value::Value>> testArrayContents;
+    auto expectedResult = value::makeNewArray();
+    value::ValueGuard expectedResultGuard{expectedResult};
+    auto expectedResultView = value::getArrayView(expectedResult.second);
+
+    while (!testEnumerator.atEnd()) {
+        testArrayContents.push_back(testEnumerator.getViewOfValue());
+        testEnumerator.advance();
+    }
+
+    for (auto it = testArrayContents.rbegin(); it != testArrayContents.rend(); ++it) {
+        auto [copyTag, copyVal] = copyValue(it->first, it->second);
+        expectedResultView->push_back(copyTag, copyVal);
+    }
+
+    runAndAssertExpression(testArray, expectedResult);
+}
+
+TEST_F(SBEBuiltinReverseArrayTest, NotArray) {
+    runAndAssertExpression(makeNothing(), makeNothing());
+    runAndAssertExpression(makeInt32(123), makeNothing());
+}
+
+}  // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index 08b51da0f9f..3e061ef2f7c 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -2834,6 +2834,65 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinGetRegexFlags(A
     return {true, strType, strValue};
 }
 
+std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinReverseArray(ArityType arity) {
+    invariant(arity == 1);
+    auto [inputOwned, inputType, inputVal] = getFromStack(0);
+
+    if (!value::isArray(inputType)) {
+        return {false, value::TypeTags::Nothing, 0};
+    }
+
+    auto [resultTag, resultVal] = value::makeNewArray();
+    auto resultView = value::getArrayView(resultVal);
+    value::ValueGuard resultGuard{resultTag, resultVal};
+
+    if (inputType == value::TypeTags::Array) {
+        auto inputView = value::getArrayView(inputVal);
+        size_t inputSize = inputView->size();
+        resultView->reserve(inputSize);
+        for (size_t i = 0; i < inputSize; i++) {
+            auto [origTag, origVal] = inputView->getAt(inputSize - 1 - i);
+            auto [copyTag, copyVal] = copyValue(origTag, origVal);
+            resultView->push_back(copyTag, copyVal);
+        }
+
+        resultGuard.reset();
+        return {true, resultTag, resultVal};
+    } else if (inputType == value::TypeTags::bsonArray || inputType == value::TypeTags::ArraySet) {
+        value::ArrayEnumerator enumerator{inputType, inputVal};
+
+        // Using intermediate vector since bsonArray and ArraySet don't
+        // support reverse iteration.
+        std::vector<std::pair<value::TypeTags, value::Value>> inputContents;
+
+        if (inputType == value::TypeTags::ArraySet) {
+            // Reserve space to avoid resizing on push_back calls.
+            auto arraySetView = value::getArraySetView(inputVal);
+            inputContents.reserve(arraySetView->size());
+            resultView->reserve(arraySetView->size());
+        }
+
+        while (!enumerator.atEnd()) {
+            inputContents.push_back(enumerator.getViewOfValue());
+            enumerator.advance();
+        }
+
+        // Run through the array backwards and copy into the result array.
+        for (auto it = inputContents.rbegin(); it != inputContents.rend(); ++it) {
+            auto [copyTag, copyVal] = copyValue(it->first, it->second);
+            resultView->push_back(copyTag, copyVal);
+        }
+
+        resultGuard.reset();
+        return {true, resultTag, resultVal};
+    } else {
+        // Earlier in this function we bailed out if the `inputType` wasn't
+        // Array, ArraySet or bsonArray, so it should be impossible to reach
+        // this point.
+        MONGO_UNREACHABLE;
+    }
+}
+
 std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinDateAdd(ArityType arity) {
     invariant(arity == 5);
 
@@ -3021,6 +3080,8 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builti
             return builtinExtractSubArray(arity);
         case Builtin::isArrayEmpty:
             return builtinIsArrayEmpty(arity);
+        case Builtin::reverseArray:
+            return builtinReverseArray(arity);
         case Builtin::dateAdd:
             return builtinDateAdd(arity);
         case Builtin::hasNullBytes:
diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h
index baa6fc82732..5dfe25cc536 100644
--- a/src/mongo/db/exec/sbe/vm/vm.h
+++ b/src/mongo/db/exec/sbe/vm/vm.h
@@ -308,6 +308,7 @@ enum class Builtin : uint8_t {
     shardFilter,
     extractSubArray,
     isArrayEmpty,
+    reverseArray,
     dateAdd,
     hasNullBytes,
     getRegexPattern,
@@ -721,6 +722,7 @@ private:
     std::tuple<bool, value::TypeTags, value::Value> builtinShardFilter(ArityType arity);
     std::tuple<bool, value::TypeTags, value::Value> builtinExtractSubArray(ArityType arity);
     std::tuple<bool, value::TypeTags, value::Value> builtinIsArrayEmpty(ArityType arity);
+    std::tuple<bool, value::TypeTags, value::Value> builtinReverseArray(ArityType arity);
     std::tuple<bool, value::TypeTags, value::Value> builtinDateAdd(ArityType arity);
     std::tuple<bool, value::TypeTags, value::Value> builtinHasNullBytes(ArityType arity);
     std::tuple<bool, value::TypeTags, value::Value> builtinGetRegexPattern(ArityType arity);
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index 820035e8450..397fe5bdb0f 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -2307,7 +2307,21 @@ public:
         unsupportedExpression(expr->getOpName());
     }
     void visit(ExpressionReverseArray* expr) final {
-        unsupportedExpression(expr->getOpName());
+        auto frameId = _context->frameIdGenerator->generate();
+        auto binds = sbe::makeEs(_context->popExpr());
+        sbe::EVariable inputRef{frameId, 0};
+
+        auto argumentIsNotArray = makeNot(makeFunction("isArray", inputRef.clone()));
+        auto exprRevArr = buildMultiBranchConditional(
+            CaseValuePair{generateNullOrMissing(inputRef),
+                          makeConstant(sbe::value::TypeTags::Null, 0)},
+            CaseValuePair{std::move(argumentIsNotArray),
+                          sbe::makeE<sbe::EFail>(ErrorCodes::Error{5154901},
+                                                 "$reverseArray argument must be an array")},
+            makeFunction("reverseArray", inputRef.clone()));
+
+        _context->pushExpr(
+            sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(exprRevArr)));
     }
     void visit(ExpressionSlice* expr) final {
         unsupportedExpression(expr->getOpName());