SERVER-50766 Use kArrayElementsOnly path traversal mode where possible

2 files changed, 105 insertions, 16 deletions

diff --git a/jstests/core/expressions_matching_whole_array.js b/jstests/core/expressions_matching_whole_array.js
new file mode 100644
index 00000000000..23faf87fe31
--- /dev/null
+++ b/jstests/core/expressions_matching_whole_array.js
@@ -0,0 +1,56 @@
+// Test MQL expressions which can match against the whole array stored in the field.
+(function() {
+"use strict";
+
+const coll = db.expressions_matching_whole_array;
+coll.drop();
+
+function testSingleDocument(document, query, shouldMatch) {
+    assert.commandWorked(coll.insert(document));
+
+    const result = coll.find(query).toArray();
+    if (shouldMatch) {
+        assert.eq(result.length, 1);
+        delete result[0]._id;
+        assert.docEq(document, result[0]);
+    } else {
+        assert.eq(result.length, 0);
+    }
+
+    assert(coll.drop());
+}
+
+function assertMatches(document, query) {
+    testSingleDocument(document, query, true);
+}
+
+function assertDoesNotMatch(document, query) {
+    testSingleDocument(document, query, false);
+}
+
+// $type
+assertMatches({a: [1, 2, 3]}, {a: {$type: 'array'}});
+assertMatches({a: []}, {a: {$type: 'array'}});
+assertDoesNotMatch({a: 1}, {a: {$type: 'array'}});
+
+// $in
+assertMatches({a: [1, 2, 3]}, {a: {$in: [[1, 2, 3]]}});
+assertDoesNotMatch({a: [1, 2, 3]}, {a: {$in: [[1, 2]]}});
+
+assertMatches({a: []}, {a: {$in: [[]]}});
+assertDoesNotMatch({a: []}, {a: {$in: [[1]]}});
+
+// $where matches ONLY against whole arrays.
+assertMatches({a: [1]}, {$where: 'Array.isArray(this.a) && this.a.length == 1 && this.a[0] == 1'});
+assertDoesNotMatch({a: [1, 2, 3]}, {$where: 'this.a == 1'});
+
+// Comparison expressions match whole array only when RHS has array, MaxKey or MinKey types.
+assertMatches({a: [1, 2, 3]}, {a: [1, 2, 3]});
+assertDoesNotMatch({a: [1, 2, 3]}, {a: [1, 2]});
+
+assertMatches({a: [1, 2, 3]}, {a: {$gt: MinKey}});
+assertMatches({a: []}, {a: {$gt: MinKey}});
+
+assertMatches({a: [1, 2, 3]}, {a: {$lt: MaxKey}});
+assertMatches({a: []}, {a: {$lt: MaxKey}});
+}());
\ No newline at end of file
diff --git a/src/mongo/db/query/sbe_stage_builder_filter.cpp b/src/mongo/db/query/sbe_stage_builder_filter.cpp
index 9d9c8ebaefa..562b059dbb3 100644
--- a/src/mongo/db/query/sbe_stage_builder_filter.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_filter.cpp
@@ -482,7 +482,7 @@ EvalExprStagePair generatePathTraversal(EvalStage inputStage,
 void generatePredicate(MatchExpressionVisitorContext* context,
                        const FieldRef* path,
                        MakePredicateFn makePredicate,
-                       LeafTraversalMode mode = LeafTraversalMode::kArrayAndItsElements,
+                       LeafTraversalMode mode,
                        bool useCombinator = true) {
     auto& frame = context->evalStack.topFrame();
 
@@ -721,7 +721,23 @@ void generateComparison(MatchExpressionVisitorContext* context,
                 std::move(inputStage)};
     };
 
-    generatePredicate(context, expr->fieldRef(), std::move(makePredicate));
+    // A 'kArrayAndItsElements' traversal mode matches the following semantics: when the path we are
+    // comparing is a path to an array, the comparison is considered true if it evaluates to true
+    // for the array itself or for any of the array's elements.
+    // However, we use 'kArrayElementsOnly' for the general case, because the comparison with the
+    // array will almost always be false. There are two exceptions:
+    // 1) when the 'rhs' operand is an array and
+    // 2) when the 'rhs' operand is MinKey or MaxKey.
+    // In the former case, the comparison we would skip by using 'kArrayElementsOnly' mode is an
+    // array-to-array comparison that can return true. In the latter case, we are avoiding a
+    // potential bug where traversing the path to the empty array ([]) would prevent _any_
+    // comparison, meaning a comparison like {$gt: MinKey} would return false.
+    const auto& rhs = expr->getData();
+    const auto checkWholeArray = rhs.type() == BSONType::Array || rhs.type() == BSONType::MinKey ||
+        rhs.type() == BSONType::MaxKey;
+    const auto traversalMode = checkWholeArray ? LeafTraversalMode::kArrayAndItsElements
+                                               : LeafTraversalMode::kArrayElementsOnly;
+    generatePredicate(context, expr->fieldRef(), std::move(makePredicate), traversalMode);
 }
 
 /**
@@ -809,7 +825,8 @@ void generateBitTest(MatchExpressionVisitorContext* context,
                 std::move(inputStage)};
     };
 
-    generatePredicate(context, expr->fieldRef(), std::move(makePredicate));
+    generatePredicate(
+        context, expr->fieldRef(), std::move(makePredicate), LeafTraversalMode::kArrayElementsOnly);
 }
 
 // Each logical expression child is evaluated in a separate EvalFrame. Set up a new EvalFrame with a
@@ -1275,7 +1292,10 @@ public:
                     std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
+        generatePredicate(_context,
+                          expr->fieldRef(),
+                          std::move(makePredicate),
+                          LeafTraversalMode::kDoNotTraverseLeaf);
     }
 
     void visit(const ExprMatchExpression* matchExpr) final {
@@ -1334,6 +1354,7 @@ public:
 
         auto arrSet = sbe::value::getArraySetView(arrSetVal);
 
+        auto hasArray = false;
         auto hasNull = false;
         for (auto&& equality : equalities) {
             auto [tagView, valView] = sbe::bson::convertFrom(true,
@@ -1341,14 +1362,17 @@ public:
                                                              equality.rawdata() + equality.size(),
                                                              equality.fieldNameSize() - 1);
 
-            if (tagView == sbe::value::TypeTags::Null) {
-                hasNull = true;
-            }
+            hasNull |= tagView == sbe::value::TypeTags::Null;
+            hasArray |= sbe::value::isArray(tagView);
+
             // An ArraySet assumes ownership of it's values so we have to make a copy here.
             auto [tag, val] = sbe::value::copyValue(tagView, valView);
             arrSet->push_back(tag, val);
         }
 
+        const auto traversalMode = hasArray ? LeafTraversalMode::kArrayAndItsElements
+                                            : LeafTraversalMode::kArrayElementsOnly;
+
         // If the InMatchExpression doesn't carry any regex patterns, we can just check if the value
         // in bound to the inputSlot is a member of the equalities set.
         if (expr->getRegexes().size() == 0) {
@@ -1372,7 +1396,7 @@ public:
                         std::move(inputStage)};
             };
 
-            generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
+            generatePredicate(_context, expr->fieldRef(), std::move(makePredicate), traversalMode);
             return;
         } else {
             // If the InMatchExpression contains regex patterns, then we need to handle a regex-only
@@ -1476,10 +1500,7 @@ public:
 
                 return {regexOutputSlot, std::move(regexStage)};
             };
-            generatePredicate(_context,
-                              expr->fieldRef(),
-                              std::move(makePredicate),
-                              LeafTraversalMode::kArrayAndItsElements);
+            generatePredicate(_context, expr->fieldRef(), std::move(makePredicate), traversalMode);
         }
     }
     // The following are no-ops. The internal expr comparison match expression are produced
@@ -1556,7 +1577,10 @@ public:
                     std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
+        generatePredicate(_context,
+                          expr->fieldRef(),
+                          std::move(makePredicate),
+                          LeafTraversalMode::kArrayElementsOnly);
     }
 
     void visit(const NorMatchExpression* expr) final {
@@ -1611,7 +1635,10 @@ public:
             return {std::move(resultExpr), std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
+        generatePredicate(_context,
+                          expr->fieldRef(),
+                          std::move(makePredicate),
+                          LeafTraversalMode::kArrayElementsOnly);
     }
 
     void visit(const SizeMatchExpression* expr) final {
@@ -1631,7 +1658,10 @@ public:
                     std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
+        const auto traversalMode = expr->typeSet().hasType(BSONType::Array)
+            ? LeafTraversalMode::kDoNotTraverseLeaf
+            : LeafTraversalMode::kArrayElementsOnly;
+        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate), traversalMode);
     }
 
     void visit(const WhereMatchExpression* expr) final {
@@ -1647,7 +1677,10 @@ public:
             return {std::move(whereExpr), std::move(inputStage)};
         };
 
-        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate));
+        generatePredicate(_context,
+                          expr->fieldRef(),
+                          std::move(makePredicate),
+                          LeafTraversalMode::kDoNotTraverseLeaf);
     }
 
     void visit(const WhereNoOpMatchExpression* expr) final {}