SERVER-30991 Introduce MatchExpression::optimize().

This patch refactors CanonicalQuery::normalizeTree() so that the
normalization logic for each type of MatchExpression goes with the
class, rather than all the optimization rules getting bundled into one
huge else if chain.

We wanted something along the lines of an optimize() member function
that would optimize 'this' and return the optimized result (possibly
the same object). However, we also wanted unique_ptr semantics, so
that the optimize function creates a new tree that does not include
the original object, it autmotatically gets destroyed. There's no way
to specify a member function that accepts a unique_ptr 'this' value.

To get around that, we provide a getOptimizer() private function that
returns a function with the unique_ptr signature we want:
unique_ptr<MatchExpression> -> unique_ptr<MatchExpression>.

This way, we still get to replace our if else chain with virtual
dispatch, and we can maintain unique_ptr semantics for the
MatchExpression tree.

1 files changed, 41 insertions, 1 deletions

diff --git a/src/mongo/db/pipeline/document_source_lookup_test.cpp b/src/mongo/db/pipeline/document_source_lookup_test.cpp
index 03ea80c5d7d..66931114dbe 100644
--- a/src/mongo/db/pipeline/document_source_lookup_test.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup_test.cpp
@@ -678,7 +678,10 @@ TEST_F(DocumentSourceLookUpTest,
     auto subPipeline = lookupStage->getSubPipeline_forTest(DOC("_id" << 5));
     ASSERT(subPipeline);
 
-    // TODO SERVER-30991: $match within $facet should optimize to $const.
+    // TODO: The '$$var1' in this expression actually gets optimized to a constant expression with
+    // value 5, but the explain output does not reflect that change. SERVER-31292 will make that
+    // optimization visible here, so we will need to replace '$$var1' with a $const expression for
+    // this test to pass.
     auto expectedPipe =
         fromjson(str::stream() << "[{mock: {}}, {$match: {x:1}}, {$sort: {sortKey: {x: 1}}}, "
                                << sequentialCacheStageObj()
@@ -688,6 +691,43 @@ TEST_F(DocumentSourceLookUpTest,
     ASSERT_VALUE_EQ(Value(subPipeline->writeExplainOps(kExplain)), Value(BSONArray(expectedPipe)));
 }
 
+TEST_F(DocumentSourceLookUpTest, ExprEmbeddedInMatchExpressionShouldBeOptimized) {
+    auto expCtx = getExpCtx();
+    NamespaceString fromNs("test", "coll");
+    expCtx->setResolvedNamespace(fromNs, {fromNs, std::vector<BSONObj>{}});
+
+    // This pipeline includes a $match stage that itself includes a $expr expression.
+    auto docSource = DocumentSourceLookUp::createFromBson(
+        fromjson("{$lookup: {let: {var1: '$_id'}, pipeline: [{$match: {$expr: {$eq: "
+                 "['$_id','$$var1']}}}], from: 'coll', as: 'as'}}")
+            .firstElement(),
+        expCtx);
+
+    auto lookupStage = static_cast<DocumentSourceLookUp*>(docSource.get());
+    ASSERT(lookupStage);
+
+    lookupStage->injectMongodInterface(
+        std::shared_ptr<MockMongodInterface>(new MockMongodInterface({})));
+
+    auto subPipeline = lookupStage->getSubPipeline_forTest(DOC("_id" << 5));
+    ASSERT(subPipeline);
+
+    auto sources = subPipeline->getSources();
+    ASSERT_GTE(sources.size(), 2u);
+
+    // The first source is our mock data source, and the second should be the $match expression.
+    auto secondSource = *(++sources.begin());
+    auto& matchSource = dynamic_cast<const DocumentSourceMatch&>(*secondSource);
+
+    // Ensure that the '$$var' in the embedded expression got optimized to ExpressionConstant.
+    BSONObjBuilder builder;
+    matchSource.getMatchExpression()->serialize(&builder);
+    auto serializedMatch = builder.obj();
+    auto expectedMatch = fromjson("{$expr: {$eq: ['$_id', {$const: 5}]}}");
+
+    ASSERT_VALUE_EQ(Value(serializedMatch), Value(expectedMatch));
+}
+
 TEST_F(DocumentSourceLookUpTest,
        ShouldIgnoreLocalVariablesShadowingLetVariablesWhenFindingNonCorrelatedPrefix) {
     auto expCtx = getExpCtx();