SERVER-49943 [SBE] Fix edge case bug with 'find({a.b.c: ..})'

2 files changed, 32 insertions, 5 deletions

diff --git a/jstests/core/arrayfind10.js b/jstests/core/arrayfind10.js
index 6f416c80286..6a3243d96eb 100644
--- a/jstests/core/arrayfind10.js
+++ b/jstests/core/arrayfind10.js
@@ -62,6 +62,27 @@ runWithAndWithoutIndex({"a.b": 1}, () => {
 
 assert(t.drop());
 
+assert.commandWorked(
+    t.insert([{"a": [{"b": [{"c": [5, 7]}]}]}, {"a": [{"b": []}, {"b": [{"c": [5, 7]}]}]}]));
+
+runWithAndWithoutIndex({"a.b": 1}, () => {
+    assert(arrayEq(t.find({"a.b.c": 5}, {_id: 0}).toArray(),
+                   [{"a": [{"b": [{"c": [5, 7]}]}]}, {"a": [{"b": []}, {"b": [{"c": [5, 7]}]}]}]));
+    assert(arrayEq(t.find({"a.b.c": {$eq: 5}}, {_id: 0}).toArray(),
+                   [{"a": [{"b": [{"c": [5, 7]}]}]}, {"a": [{"b": []}, {"b": [{"c": [5, 7]}]}]}]));
+
+    assert(arrayEq(t.find({"a.b.c": {$in: [5]}}, {_id: 0}).toArray(),
+                   [{"a": [{"b": [{"c": [5, 7]}]}]}, {"a": [{"b": []}, {"b": [{"c": [5, 7]}]}]}]));
+
+    assert(arrayEq(t.find({"a.b.c": {$size: 2}}, {_id: 0}).toArray(),
+                   [{"a": [{"b": [{"c": [5, 7]}]}]}, {"a": [{"b": []}, {"b": [{"c": [5, 7]}]}]}]));
+
+    assert(arrayEq(t.find({"a.b.c": {$elemMatch: {$eq: 5}}}, {_id: 0}).toArray(),
+                   [{"a": [{"b": [{"c": [5, 7]}]}]}, {"a": [{"b": []}, {"b": [{"c": [5, 7]}]}]}]));
+});
+
+assert(t.drop());
+
 assert.commandWorked(t.insert([
     {a: 1},
     {a: 2},
diff --git a/src/mongo/db/query/sbe_stage_builder_filter.cpp b/src/mongo/db/query/sbe_stage_builder_filter.cpp
index 97f82bdc812..aef0ccef3ae 100644
--- a/src/mongo/db/query/sbe_stage_builder_filter.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_filter.cpp
@@ -249,20 +249,26 @@ EvalExprStagePair generatePathTraversal(EvalStage inputStage,
         projectEvalExpr(std::move(innerExpr), std::move(innerBranch), planNodeId, slotIdGenerator);
 
     // Generate the traverse stage for the current nested level.
-    auto outputSlot = slotIdGenerator->generate();
+    auto traverseOutputSlot = slotIdGenerator->generate();
     auto outputStage =
         makeTraverse(std::move(fromBranch),
                      std::move(innerBranch),  // NOLINT(bugprone-use-after-move)
                      fieldSlot,
-                     outputSlot,
+                     traverseOutputSlot,
                      innerSlot,
                      sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr,
-                                                  sbe::makeE<sbe::EVariable>(outputSlot),
+                                                  sbe::makeE<sbe::EVariable>(traverseOutputSlot),
                                                   sbe::makeE<sbe::EVariable>(innerSlot)),
-                     sbe::makeE<sbe::EVariable>(outputSlot),
+                     sbe::makeE<sbe::EVariable>(traverseOutputSlot),
                      planNodeId,
                      1);
 
+    auto outputSlot = slotIdGenerator->generate();
+    outputStage = makeProject(std::move(outputStage),
+                              planNodeId,
+                              outputSlot,
+                              makeFillEmptyFalse(sbe::makeE<sbe::EVariable>(traverseOutputSlot)));
+
     if (isLeafField && mode == LeafTraversalMode::kArrayAndItsElements) {
         // For the last level, if 'mode' == kArrayAndItsElements and getField() returns an array we
         // need to apply the predicate both to the elements of the array _and_ to the array itself.
@@ -275,7 +281,7 @@ EvalExprStagePair generatePathTraversal(EvalStage inputStage,
 
         outputExpr = sbe::makeE<sbe::EPrimBinary>(
             sbe::EPrimBinary::logicOr,
-            makeFillEmptyFalse(sbe::makeE<sbe::EVariable>(outputSlot)),
+            sbe::makeE<sbe::EVariable>(outputSlot),
             sbe::makeE<sbe::EPrimBinary>(
                 sbe::EPrimBinary::logicAnd,
                 makeFillEmptyFalse(sbe::makeE<sbe::EFunction>(