SERVER-40828 Add getter for expressions on projection nodes

(cherry picked from commit cf6cb438e0302be611bc9c2901919803e66eebba)

4 files changed, 60 insertions, 0 deletions

diff --git a/src/mongo/db/pipeline/parsed_add_fields.h b/src/mongo/db/pipeline/parsed_add_fields.h
index 96579ada1a8..469f6d79d9d 100644
--- a/src/mongo/db/pipeline/parsed_add_fields.h
+++ b/src/mongo/db/pipeline/parsed_add_fields.h
@@ -76,6 +76,10 @@ public:
         return TransformerType::kComputedProjection;
     }
 
+    const InclusionNode& getRoot() const {
+        return *_root;
+    }
+
     /**
      * Parses the addFields specification given by 'spec', populating internal data structures.
      */
diff --git a/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp b/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp
index 3be7c403e1c..18b6e243f09 100644
--- a/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp
+++ b/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp
@@ -65,6 +65,19 @@ void ProjectionNode::addExpressionForPath(const FieldPath& path,
     addOrGetChild(path.getFieldName(0).toString())->addExpressionForPath(path.tail(), expr);
 }
 
+boost::intrusive_ptr<Expression> ProjectionNode::getExpressionForPath(const FieldPath& path) const {
+    if (path.getPathLength() == 1) {
+        if (_expressions.find(path.getFieldName(0)) != _expressions.end()) {
+            return _expressions.at(path.getFieldName(0));
+        }
+        return nullptr;
+    }
+    if (auto child = getChild(path.getFieldName(0).toString())) {
+        return child->getExpressionForPath(path.tail());
+    }
+    return nullptr;
+}
+
 ProjectionNode* ProjectionNode::addOrGetChild(const std::string& field) {
     auto child = getChild(field);
     return child ? child : addChild(field);
diff --git a/src/mongo/db/pipeline/parsed_aggregation_projection_node.h b/src/mongo/db/pipeline/parsed_aggregation_projection_node.h
index 8671a289695..0ec444bc9f5 100644
--- a/src/mongo/db/pipeline/parsed_aggregation_projection_node.h
+++ b/src/mongo/db/pipeline/parsed_aggregation_projection_node.h
@@ -64,6 +64,12 @@ public:
     void addProjectionForPath(const FieldPath& path);
 
     /**
+     * Get the expression for the given path. Returns null if no expression for the given path is
+     * found.
+     */
+    boost::intrusive_ptr<Expression> getExpressionForPath(const FieldPath& path) const;
+
+    /**
      * Recursively adds 'path' into the tree as a computed field, creating any child nodes if
      * necessary.
      *
diff --git a/src/mongo/db/pipeline/parsed_aggregation_projection_test.cpp b/src/mongo/db/pipeline/parsed_aggregation_projection_test.cpp
index 6a76410e8e2..15efa442726 100644
--- a/src/mongo/db/pipeline/parsed_aggregation_projection_test.cpp
+++ b/src/mongo/db/pipeline/parsed_aggregation_projection_test.cpp
@@ -38,6 +38,7 @@
 #include "mongo/bson/json.h"
 #include "mongo/db/pipeline/document.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
+#include "mongo/db/pipeline/parsed_inclusion_projection.h"
 #include "mongo/db/pipeline/value.h"
 #include "mongo/unittest/unittest.h"
 
@@ -573,6 +574,42 @@ TEST(ParsedAggregationProjectionType, ShouldCoerceNumericsToBools) {
     }
 }
 
+TEST(ParsedAggregationProjectionType, GetExpressionForPathGetsTopLevelExpression) {
+    const boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+    auto projectObj = BSON("$add" << BSON_ARRAY(BSON("$const" << 1) << BSON("$const" << 3)));
+    auto expr = Expression::parseObject(expCtx, projectObj, expCtx->variablesParseState);
+    ParsedAggregationProjection::ProjectionPolicies defaultPolicies;
+    auto node = InclusionNode(defaultPolicies);
+    node.addExpressionForPath(FieldPath("key"), expr);
+    BSONObjBuilder bob;
+    ASSERT_EQ(expr, node.getExpressionForPath(FieldPath("key")));
+}
+
+TEST(ParsedAggregationProjectionType, GetExpressionForPathGetsCorrectTopLevelExpression) {
+    const boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+    auto correctObj = BSON("$add" << BSON_ARRAY(BSON("$const" << 1) << BSON("$const" << 3)));
+    auto incorrectObj = BSON("$add" << BSON_ARRAY(BSON("$const" << 2) << BSON("$const" << 4)));
+    auto correctExpr = Expression::parseObject(expCtx, correctObj, expCtx->variablesParseState);
+    auto incorrectExpr = Expression::parseObject(expCtx, incorrectObj, expCtx->variablesParseState);
+    ParsedAggregationProjection::ProjectionPolicies defaultPolicies;
+    auto node = InclusionNode(defaultPolicies);
+    node.addExpressionForPath(FieldPath("key"), correctExpr);
+    node.addExpressionForPath(FieldPath("other"), incorrectExpr);
+    BSONObjBuilder bob;
+    ASSERT_EQ(correctExpr, node.getExpressionForPath(FieldPath("key")));
+}
+
+TEST(ParsedAggregationProjectionType, GetExpressionForPathGetsNonTopLevelExpression) {
+    const boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+    auto projectObj = BSON("$add" << BSON_ARRAY(BSON("$const" << 1) << BSON("$const" << 3)));
+    auto expr = Expression::parseObject(expCtx, projectObj, expCtx->variablesParseState);
+    ParsedAggregationProjection::ProjectionPolicies defaultPolicies;
+    auto node = InclusionNode(defaultPolicies);
+    node.addExpressionForPath(FieldPath("key.second"), expr);
+    BSONObjBuilder bob;
+    ASSERT_EQ(expr, node.getExpressionForPath(FieldPath("key.second")));
+}
+
 }  // namespace
 }  // namespace parsed_aggregation_projection
 }  // namespace mongo