From 3cf4e0593c394dd7eb45d8000d76b5dc73a3f425 Mon Sep 17 00:00:00 2001
From: Justin Seyster <justin.seyster@mongodb.com>
Date: Thu, 28 Sep 2017 18:20:41 -0400
Subject: 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.
---
 .../db/pipeline/document_source_lookup_test.cpp    | 42 +++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

(limited to 'src/mongo/db/pipeline/document_source_lookup_test.cpp')

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();
-- 
cgit v1.2.1