SERVER-50430 Implement SBE support for $bsonSize expression

6 files changed, 152 insertions, 1 deletions

diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index a9461dc72a6..38bf760a8dc 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -94,6 +94,7 @@ env.CppUnitTest(
         'sbe_plan_stage_test.cpp',
         'sbe_sort_test.cpp',
         'sbe_test.cpp',
+        'expressions/sbe_bson_size_test.cpp',
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/db/concurrency/lock_manager',
diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index c950ac2228c..5c598445272 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -364,6 +364,7 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = {
     {"bitTestMask", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::bitTestMask, false}},
     {"bitTestPosition",
      BuiltinFn{[](size_t n) { return n == 3; }, vm::Builtin::bitTestPosition, false}},
+    {"bsonSize", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::bsonSize, false}},
 };
 
 /**
diff --git a/src/mongo/db/exec/sbe/expressions/sbe_bson_size_test.cpp b/src/mongo/db/exec/sbe/expressions/sbe_bson_size_test.cpp
new file mode 100644
index 00000000000..511bbbed7ad
--- /dev/null
+++ b/src/mongo/db/exec/sbe/expressions/sbe_bson_size_test.cpp
@@ -0,0 +1,109 @@
+/**
+ *    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 {
+using SBEBsonSizeTest = EExpressionTestFixture;
+
+TEST_F(SBEBsonSizeTest, ComputesSizeForBsonDocument) {
+    value::ViewOfValueAccessor slotAccessor;
+    auto argSlot = bindAccessor(&slotAccessor);
+    auto bsonSizeExpr =
+        sbe::makeE<sbe::EFunction>("bsonSize", sbe::makeEs(makeE<EVariable>(argSlot)));
+    auto compiledExpr = compileExpression(*bsonSizeExpr);
+
+    BSONObjBuilder objBuilder;
+    objBuilder.append("name", "Test string element");
+    objBuilder.append("age", 32);
+    objBuilder.append("citizen", true);
+    auto bsonObj = objBuilder.done();
+
+    slotAccessor.reset(value::TypeTags::bsonObject, value::bitcastFrom(bsonObj.objdata()));
+    auto [tag, val] = runCompiledExpression(compiledExpr.get());
+    value::ValueGuard guard(tag, val);
+
+    ASSERT_EQUALS(value::TypeTags::NumberInt32, tag);
+    ASSERT_EQUALS(value::bitcastTo<uint32_t>(val), bsonObj.objsize());
+}
+
+TEST_F(SBEBsonSizeTest, ComputesSizeForSbeObject) {
+    value::ViewOfValueAccessor slotAccessor;
+    auto argSlot = bindAccessor(&slotAccessor);
+    auto bsonSizeExpr =
+        sbe::makeE<sbe::EFunction>("bsonSize", sbe::makeEs(makeE<EVariable>(argSlot)));
+    auto compiledExpr = compileExpression(*bsonSizeExpr);
+
+    auto [tagArg1, valArg1] = value::makeNewString("Test string element");
+    auto [tagArg2, valArg2] = value::makeNewObject();
+    auto obj = value::getObjectView(valArg2);
+    obj->push_back("name", tagArg1, valArg1);
+    obj->push_back("age", value::TypeTags::NumberInt32, value::bitcastFrom(32));
+    obj->push_back("citizen", value::TypeTags::Boolean, value::bitcastFrom(true));
+    value::ValueGuard argGuard(tagArg2, valArg2);
+
+    slotAccessor.reset(value::TypeTags::Object, valArg2);
+    auto [tag, val] = runCompiledExpression(compiledExpr.get());
+    value::ValueGuard guard(tag, val);
+
+    ASSERT_EQUALS(value::TypeTags::NumberInt32, tag);
+    ASSERT_EQUALS(value::bitcastTo<uint32_t>(val), 54);
+}
+
+TEST_F(SBEBsonSizeTest, ReturnsNothingForNonObject) {
+    value::ViewOfValueAccessor slotAccessor;
+    auto argSlot = bindAccessor(&slotAccessor);
+    auto bsonSizeExpr =
+        sbe::makeE<sbe::EFunction>("bsonSize", sbe::makeEs(makeE<EVariable>(argSlot)));
+    auto compiledExpr = compileExpression(*bsonSizeExpr);
+
+    auto [tagArg1, valArg1] = value::makeNewString("Test string element");
+    value::ValueGuard guard1(tagArg1, valArg1);
+
+    auto [tagArg2, valArg2] = value::makeNewArray();
+    value::ValueGuard guard2(tagArg2, valArg2);
+    auto arr = value::getArrayView(valArg2);
+    arr->push_back(value::TypeTags::NumberInt32, value::bitcastFrom(32));
+    auto [tagArg3, valArg3] = value::copyValue(tagArg1, valArg1);
+    arr->push_back(tagArg3, valArg3);
+
+    std::vector<std::pair<value::TypeTags, value::Value>> testData = {
+        std::make_pair(value::TypeTags::NumberInt32, value::bitcastFrom(12789)),
+        std::make_pair(tagArg1, valArg1),
+        std::make_pair(tagArg2, valArg2)};
+
+    for (size_t i = 0; i < testData.size(); i++) {
+        slotAccessor.reset(testData[i].first, testData[i].second);
+        auto [tag, val] = runCompiledExpression(compiledExpr.get());
+        value::ValueGuard guard(tag, val);
+        ASSERT_EQUALS(value::TypeTags::Nothing, tag);
+    }
+}
+
+}  // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index 1ff189b1c37..b5890497f45 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -1196,6 +1196,22 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinBitTestMask(uin
     return {false, value::TypeTags::Boolean, result};
 }
 
+std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinBsonSize(uint8_t arity) {
+    auto [_, tagOperand, valOperand] = getFromStack(0);
+
+    if (tagOperand == value::TypeTags::Object) {
+        BSONObjBuilder objBuilder;
+        bson::convertToBsonObj(objBuilder, value::getObjectView(valOperand));
+        auto sz = objBuilder.done().objsize();
+        return {false, value::TypeTags::NumberInt32, value::bitcastFrom(sz)};
+    } else if (tagOperand == value::TypeTags::bsonObject) {
+        auto beginObj = value::getRawPointerView(valOperand);
+        auto sz = ConstDataView(beginObj).read<LittleEndian<uint32_t>>();
+        return {false, value::TypeTags::NumberInt32, value::bitcastFrom(sz)};
+    }
+    return {false, value::TypeTags::Nothing, 0};
+}
+
 std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builtin f,
                                                                           uint8_t arity) {
     switch (f) {
@@ -1229,6 +1245,8 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builti
             return builtinBitTestMask(arity);
         case Builtin::bitTestPosition:
             return builtinBitTestPosition(arity);
+        case Builtin::bsonSize:
+            return builtinBsonSize(arity);
     }
 
     MONGO_UNREACHABLE;
diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h
index cf8c4411953..a5197d17437 100644
--- a/src/mongo/db/exec/sbe/vm/vm.h
+++ b/src/mongo/db/exec/sbe/vm/vm.h
@@ -185,6 +185,7 @@ enum class Builtin : uint8_t {
     bitTestZero,      // test bitwise mask & value is zero
     bitTestMask,      // test bitwise mask & value is mask
     bitTestPosition,  // test BinData with a bit position list
+    bsonSize,         // implements $bsonSize
 };
 
 class CodeFragment {
@@ -422,6 +423,7 @@ private:
     std::tuple<bool, value::TypeTags, value::Value> builtinBitTestZero(uint8_t arity);
     std::tuple<bool, value::TypeTags, value::Value> builtinBitTestMask(uint8_t arity);
     std::tuple<bool, value::TypeTags, value::Value> builtinBitTestPosition(uint8_t arity);
+    std::tuple<bool, value::TypeTags, value::Value> builtinBsonSize(uint8_t arity);
 
     std::tuple<bool, value::TypeTags, value::Value> dispatchBuiltin(Builtin f, uint8_t arity);
 
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index 46b512aca6a..b17164df951 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -918,7 +918,27 @@ public:
         unsupportedExpression(expr->getOpName());
     }
     void visit(ExpressionBsonSize* expr) final {
-        unsupportedExpression(expr->getOpName());
+        // Build an expression which evaluates the size of a BSON document and validates the input
+        // argument.
+        // 1. If the argument is null or empty, return null.
+        // 2. Else, if the argument is a BSON document, return its size.
+        // 3. Else, raise an error.
+
+        auto frameId = _context->frameIdGenerator->generate();
+        auto binds = sbe::makeEs(_context->popExpr());
+        sbe::EVariable inputRef(frameId, 0);
+
+        auto bsonSizeExpr = sbe::makeE<sbe::EIf>(
+            generateNullOrMissing(frameId, 0),
+            sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Null, 0),
+            sbe::makeE<sbe::EIf>(
+                sbe::makeE<sbe::EFunction>("isObject", sbe::makeEs(inputRef.clone())),
+                sbe::makeE<sbe::EFunction>("bsonSize", sbe::makeEs(inputRef.clone())),
+                sbe::makeE<sbe::EFail>(ErrorCodes::Error{5043001},
+                                       "$bsonSize requires a document input")));
+
+        _context->pushExpr(
+            sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(bsonSizeExpr)));
     }
     void visit(ExpressionCeil* expr) final {
         unsupportedExpression(expr->getOpName());