SERVER-61839 Fix SBE implementation of $exists and $type when a positional projection is used

2 files changed, 48 insertions, 14 deletions

diff --git a/jstests/core/positional_projection.js b/jstests/core/positional_projection.js
index 9b0b013d51a..5c16231a667 100644
--- a/jstests/core/positional_projection.js
+++ b/jstests/core/positional_projection.js
@@ -181,6 +181,18 @@ testSingleDocument(
 testSingleDocument({a: {$in: [2, 3]}}, {'b.$': 1}, {a: [1, 2, 3], b: [4, 5, 6]}, {b: [5]});
 testSingleDocument({$or: [{a: 2}, {a: 3}]}, {'b.$': 1}, {a: [1, 2, 3], b: [4, 5, 6]}, {b: [5]});
 
+// SERVER-61839: Test out some cases involving $exists and $type where we've had bugs in the past.
+testSingleDocument(
+    {a: {$elemMatch: {y: {$exists: true}}}}, {"a.$": 1}, {a: [{y: 1}, {y: 2}]}, {a: [{y: 1}]});
+
+testSingleDocument(
+    {a: {$elemMatch: {y: {$type: ["array"]}}}}, {"a.$": 1}, {a: [{y: 1}, {y: []}]}, {a: [{y: []}]});
+
+testSingleDocument({a: {$elemMatch: {y: {$type: ["array", "double"]}}}},
+                   {"a.$": 1},
+                   {a: [{y: 1}, {y: []}]},
+                   {a: [{y: 1}]});
+
 // Tests involving getMore. Test the $-positional operator across multiple batches.
 assert.commandWorked(coll.insert([
     {
diff --git a/src/mongo/db/query/sbe_stage_builder_filter.cpp b/src/mongo/db/query/sbe_stage_builder_filter.cpp
index 8e3c768e3d2..01bb3ccd8eb 100644
--- a/src/mongo/db/query/sbe_stage_builder_filter.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_filter.cpp
@@ -1370,17 +1370,26 @@ public:
     }
 
     void visit(const ExistsMatchExpression* expr) final {
-        auto makePredicate = [](sbe::value::SlotId inputSlot,
-                                EvalStage inputStage) -> EvalExprStagePair {
-            return {sbe::makeE<sbe::EFunction>("exists",
-                                               sbe::makeEs(sbe::makeE<sbe::EVariable>(inputSlot))),
-                    std::move(inputStage)};
+        const auto traversalMode = LeafTraversalMode::kDoNotTraverseLeaf;
+
+        auto makePredicate = [expr, context = _context](sbe::value::SlotId inputSlot,
+                                                        EvalStage inputStage) -> EvalExprStagePair {
+            auto resultExpr = sbe::makeE<sbe::EFunction>(
+                "exists", sbe::makeEs(sbe::makeE<sbe::EVariable>(inputSlot)));
+
+            // $exists is always applied to the leaf of the field path. For kDoNotTraverseLeaf mode,
+            // generatePredicate() does not convert the predicate value to state when generating
+            // traversal for leaf nodes of field path. For this reason, we need to perform this
+            // conversion manually.
+            if (expr->fieldRef() && !expr->fieldRef()->empty() &&
+                context->evalStack.topFrame().data().inputSlot) {
+                resultExpr = context->stateHelper.makeState(std::move(resultExpr));
+            }
+
+            return {std::move(resultExpr), std::move(inputStage)};
         };
 
-        generatePredicate(_context,
-                          expr->fieldRef(),
-                          std::move(makePredicate),
-                          LeafTraversalMode::kDoNotTraverseLeaf);
+        generatePredicate(_context, expr->fieldRef(), std::move(makePredicate), traversalMode);
     }
 
     void visit(const ExprMatchExpression* matchExpr) final {
@@ -1743,17 +1752,30 @@ public:
     void visit(const TwoDPtInAnnulusExpression* expr) final {}
 
     void visit(const TypeMatchExpression* expr) final {
-        auto makePredicate = [expr](sbe::value::SlotId inputSlot,
-                                    EvalStage inputStage) -> EvalExprStagePair {
+        const auto traversalMode = expr->typeSet().hasType(BSONType::Array)
+            ? LeafTraversalMode::kDoNotTraverseLeaf
+            : LeafTraversalMode::kArrayElementsOnly;
+
+        auto makePredicate =
+            [expr, traversalMode, context = _context](sbe::value::SlotId inputSlot,
+                                                      EvalStage inputStage) -> EvalExprStagePair {
             const MatcherTypeSet& ts = expr->typeSet();
             auto resultExpr = makeFillEmptyFalse(
                 sbe::makeE<sbe::ETypeMatch>(makeVariable(inputSlot), ts.getBSONTypeMask()));
+
+            // $type is always applied to the leaf of the field path. For kDoNotTraverseLeaf mode,
+            // generatePredicate() does not convert the predicate value to state when generating
+            // traversal for leaf nodes of field path. For this reason, we need to perform this
+            // conversion manually.
+            if (expr->fieldRef() && !expr->fieldRef()->empty() &&
+                context->evalStack.topFrame().data().inputSlot &&
+                traversalMode == LeafTraversalMode::kDoNotTraverseLeaf) {
+                resultExpr = context->stateHelper.makeState(std::move(resultExpr));
+            }
+
             return {std::move(resultExpr), std::move(inputStage)};
         };
 
-        const auto traversalMode = expr->typeSet().hasType(BSONType::Array)
-            ? LeafTraversalMode::kDoNotTraverseLeaf
-            : LeafTraversalMode::kArrayElementsOnly;
         generatePredicate(_context, expr->fieldRef(), std::move(makePredicate), traversalMode);
     }