SERVER-59299: Flatten top-level nested $match stages in doOptimizeAt

(cherry picked from commit 4db5eceda2cff697f35c84cd08232bac8c33beec)

4 files changed, 61 insertions, 5 deletions

diff --git a/jstests/sharding/pipeline_length_limit.js b/jstests/sharding/pipeline_length_limit.js
index a9315549d92..8c4d49c9a95 100644
--- a/jstests/sharding/pipeline_length_limit.js
+++ b/jstests/sharding/pipeline_length_limit.js
@@ -108,6 +108,7 @@ function testLimits(testDB, lengthLimit) {
      {$group: {_id: "$_id"}},
      {$limit: 1},
      {$lookup: {from: collname, localField: "_id", foreignField: "_id", as: "foo"}},
+     {$match: {_id: {$exists: true}}},
      {$project: {_id: 1}},
      {$redact: "$$KEEP"},
      {$replaceWith: "$$ROOT"},
diff --git a/src/mongo/db/pipeline/document_source_match.cpp b/src/mongo/db/pipeline/document_source_match.cpp
index f5dc049df46..182444e0c4c 100644
--- a/src/mongo/db/pipeline/document_source_match.cpp
+++ b/src/mongo/db/pipeline/document_source_match.cpp
@@ -33,6 +33,8 @@
 
 #include <memory>
 
+#include "mongo/bson/bsonobj.h"
+#include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/db/exec/document_value/document.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/matcher/expression_algo.h"
@@ -379,7 +381,33 @@ bool DocumentSourceMatch::isTextQuery(const BSONObj& query) {
 }
 
 void DocumentSourceMatch::joinMatchWith(intrusive_ptr<DocumentSourceMatch> other) {
-    rebuild(BSON("$and" << BSON_ARRAY(_predicate << other->getQuery())));
+    BSONObjBuilder bob;
+    BSONArrayBuilder arrBob(bob.subarrayStart("$and"));
+
+    auto addPredicates = [&](const auto& predicates) {
+        if (predicates.isEmpty()) {
+            arrBob.append(predicates);
+        }
+
+        for (auto&& pred : predicates) {
+            // If 'pred' is an $and, add its children directly to the new top-level $and to avoid
+            // nesting $and's. Otherwise, add 'pred' itself as a child.
+            if (pred.fieldNameStringData() == "$and") {
+                for (auto& child : pred.Array()) {
+                    arrBob.append(child);
+                }
+            } else {
+                BSONObjBuilder childBob(arrBob.subobjStart());
+                childBob.append(pred);
+            }
+        }
+    };
+
+    addPredicates(_predicate);
+    addPredicates(other->_predicate);
+
+    arrBob.doneFast();
+    rebuild(bob.obj());
 }
 
 pair<intrusive_ptr<DocumentSourceMatch>, intrusive_ptr<DocumentSourceMatch>>
diff --git a/src/mongo/db/pipeline/document_source_match.h b/src/mongo/db/pipeline/document_source_match.h
index ecdda5dc274..1e9b9bb6635 100644
--- a/src/mongo/db/pipeline/document_source_match.h
+++ b/src/mongo/db/pipeline/document_source_match.h
@@ -107,7 +107,7 @@ public:
 
     /**
      * Attempts to combine with any subsequent $match stages, joining the query objects with a
-     * $and.
+     * $and and flattening top-level $and's in the process.
      */
     Pipeline::SourceContainer::iterator doOptimizeAt(Pipeline::SourceContainer::iterator itr,
                                                      Pipeline::SourceContainer* container) final;
diff --git a/src/mongo/db/pipeline/document_source_match_test.cpp b/src/mongo/db/pipeline/document_source_match_test.cpp
index 0ad0a0f1d32..45008a4f668 100644
--- a/src/mongo/db/pipeline/document_source_match_test.cpp
+++ b/src/mongo/db/pipeline/document_source_match_test.cpp
@@ -473,9 +473,7 @@ TEST_F(DocumentSourceMatchTest, MultipleMatchStagesShouldCombineIntoOne) {
     container.push_back(match3);
     match1->optimizeAt(container.begin(), &container);
     ASSERT_EQUALS(container.size(), 1U);
-    ASSERT_BSONOBJ_EQ(match1->getQuery(),
-                      fromjson("{'$and': [{'$and': [{a:1}, {b:1}]},"
-                               "{c:1}]}"));
+    ASSERT_BSONOBJ_EQ(match1->getQuery(), fromjson("{'$and': [{a:1}, {b:1}, {c:1}]}"));
 }
 
 TEST_F(DocumentSourceMatchTest, ShouldPropagatePauses) {
@@ -526,6 +524,35 @@ TEST_F(DocumentSourceMatchTest, ShouldCorrectlyJoinWithSubsequentMatch) {
     ASSERT_TRUE(match->getNext().isEOF());
 }
 
+TEST_F(DocumentSourceMatchTest, RepeatedJoinWithShouldNotNestAnds) {
+    auto match1 = DocumentSourceMatch::create(fromjson("{}"), getExpCtx());
+    Pipeline::SourceContainer container{
+        match1,
+        DocumentSourceMatch::create(fromjson("{}"), getExpCtx()),
+        DocumentSourceMatch::create(fromjson("{a: 1}"), getExpCtx()),
+        DocumentSourceMatch::create(fromjson("{b: 1}"), getExpCtx()),
+        DocumentSourceMatch::create(fromjson("{$and: [{c: 1}, {d: 1}]}"), getExpCtx()),
+        DocumentSourceMatch::create(fromjson("{$or: [{e: 1}, {f: 1}]}"), getExpCtx()),
+        DocumentSourceMatch::create(fromjson("{}"), getExpCtx()),
+        DocumentSourceMatch::create(
+            fromjson("{$and: [{g: 1}, {h: 1}], $or: [{i: 1}, {j: 1}], $and: [{k: 1}, {l: 1}]}"),
+            getExpCtx()),
+        DocumentSourceMatch::create(fromjson("{$and: [{m: 1}, {$and: [{n: 1}, {o: 1}]}]}"),
+                                    getExpCtx())};
+
+    // Call optimizeAt() repeatedly to trigger joinWith() behavior
+    for (size_t i = 0; i < 8; i++) {
+        match1->optimizeAt(container.begin(), &container);
+    }
+
+    ASSERT_EQUALS(container.size(), 1U);
+    ASSERT_BSONOBJ_EQ(
+        match1->getQuery(),
+        fromjson("{'$and': [{}, {}, {a: 1}, {b: 1}, {c: 1}, {d: 1}, {$or: [{e: 1}, {f: 1}]}, {}, "
+                 "{g: 1}, {h: 1}, {$or: [{i: 1}, {j: 1}]}, {k: 1}, {l: 1}, {m: 1}, {$and: [{n: 1}, "
+                 "{o: 1}]}]}"));
+}
+
 DEATH_TEST_REGEX_F(DocumentSourceMatchTest,
                    ShouldFailToDescendExpressionOnPathThatIsNotACommonPrefix,
                    "Invariant failure.*expression::isPathPrefixOf") {