SERVER-51468 [SBE] Fix bug with $in operator

2 files changed, 20 insertions, 9 deletions

diff --git a/jstests/core/in_with_mixed_values.js b/jstests/core/in_with_mixed_values.js
index c5877c3fdaf..a9ada4034fe 100644
--- a/jstests/core/in_with_mixed_values.js
+++ b/jstests/core/in_with_mixed_values.js
@@ -46,6 +46,8 @@ assert.commandWorked(coll.insert(docs));
 
 assert.eq(8, coll.find({x: {$in: [1, /^a/]}}).itcount());
 assert.eq(7, coll.find({x: {$in: [2, /^a/]}}).itcount());
+assert.eq(6, coll.find({x: {$in: [[1, 2], /^a/]}}).itcount());
+assert.eq(5, coll.find({x: {$in: [/^a/]}}).itcount());
 assert.eq(1, coll.find({'x.y': {$in: [1, 2]}}).itcount());
 assert.eq(2, coll.find({'x.y': {$in: [1, /^a/]}}).itcount());
 assert.eq(1, coll.find({'x': {$in: [{y: 1}, /^z/]}}).itcount());
diff --git a/src/mongo/db/query/sbe_stage_builder_filter.cpp b/src/mongo/db/query/sbe_stage_builder_filter.cpp
index aef0ccef3ae..dec1e9528a4 100644
--- a/src/mongo/db/query/sbe_stage_builder_filter.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_filter.cpp
@@ -1058,10 +1058,12 @@ public:
                 // FROM branch binds an array constant carrying the regex patterns to a slot. Then
                 // the inner branch executes 'regexMatch' once per regex.
                 auto [regexTag, regexVal] = sbe::value::copyValue(arrTag, arrVal);
-                auto regexFromStage = makeProject(EvalStage{},
-                                                  _context->planNodeId,
-                                                  regexArraySlot,
-                                                  sbe::makeE<sbe::EConstant>(regexTag, regexVal));
+
+                auto regexFromStage =
+                    makeProject(equalities.size() > 0 ? EvalStage{} : std::move(inputStage),
+                                _context->planNodeId,
+                                regexArraySlot,
+                                sbe::makeE<sbe::EConstant>(regexTag, regexVal));
 
                 auto regexInnerStage =
                     makeProject(EvalStage{},
@@ -1100,10 +1102,17 @@ public:
                         EvalStage{});
                     branches.emplace_back(regexOutputSlot, std::move(regexStage));
 
-                    return generateShortCircuitingLogicalOp(sbe::EPrimBinary::logicOr,
-                                                            std::move(branches),
-                                                            _context->planNodeId,
-                                                            _context->slotIdGenerator);
+                    auto [shortCircuitingExpr, shortCircuitingStage] =
+                        generateShortCircuitingLogicalOp(sbe::EPrimBinary::logicOr,
+                                                         std::move(branches),
+                                                         _context->planNodeId,
+                                                         _context->slotIdGenerator);
+
+                    inputStage = makeLoopJoin(std::move(inputStage),
+                                              std::move(shortCircuitingStage),
+                                              _context->planNodeId);
+
+                    return {std::move(shortCircuitingExpr), std::move(inputStage)};
                 }
 
                 return {regexOutputSlot, std::move(regexStage)};
@@ -1111,7 +1120,7 @@ public:
             generatePredicate(_context,
                               expr->path(),
                               std::move(makePredicate),
-                              LeafTraversalMode::kArrayElementsOnly);
+                              LeafTraversalMode::kArrayAndItsElements);
         }
     }