SERVER-38695 Make QuerySolutionNode subclasses for projection fast-paths

1 files changed, 94 insertions, 41 deletions

diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index 491d6d99e5b..4e4d724b555 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -32,6 +32,7 @@
 
 #include <memory>
 
+#include "mongo/base/string_data.h"
 #include "mongo/bson/bsonobj_comparator_interface.h"
 #include "mongo/db/fts/fts_query.h"
 #include "mongo/db/jsobj.h"
@@ -50,6 +51,12 @@ class GeoNearExpression;
  */
 struct QuerySolutionNode {
     QuerySolutionNode() {}
+
+    /**
+     * Constructs a QuerySolutionNode with a single child.
+     */
+    QuerySolutionNode(std::unique_ptr<QuerySolutionNode> child) : children{child.release()} {}
+
     virtual ~QuerySolutionNode() {
         for (size_t i = 0; i < children.size(); ++i) {
             delete children[i];
@@ -520,40 +527,26 @@ struct IndexScanNode : public QuerySolutionNode {
     std::set<StringData> multikeyFields;
 };
 
-struct ProjectionNode : public QuerySolutionNode {
-    /**
-     * We have a few implementations of the projection functionality.  The most general
-     * implementation 'DEFAULT' is much slower than the fast-path implementations
-     * below.  We only really have all the information available to choose a projection
-     * implementation at planning time.
-     */
-    enum ProjectionType {
-        // This is the most general implementation of the projection functionality.  It handles
-        // every case.
-        DEFAULT,
-
-        // This is a fast-path for when the projection is fully covered by one index.
-        COVERED_ONE_INDEX,
-
-        // This is a fast-path for when the projection only has inclusions on non-dotted fields.
-        SIMPLE_DOC,
-    };
-
-    ProjectionNode(ParsedProjection proj)
-        : _sorts(SimpleBSONObjComparator::kInstance.makeBSONObjSet()),
-          fullExpression(NULL),
-          projType(DEFAULT),
-          parsed(proj) {}
-
-    virtual ~ProjectionNode() {}
-
-    virtual StageType getType() const {
-        return STAGE_PROJECTION;
-    }
-
-    virtual void computeProperties();
+/**
+ * We have a few implementations of the projection functionality. They are chosen by constructing
+ * a type derived from this abstract struct. The most general implementation 'ProjectionNodeDefault'
+ * is much slower than the fast-path implementations. We only really have all the information
+ * available to choose a projection implementation at planning time.
+ */
+struct ProjectionNode : QuerySolutionNode {
+    ProjectionNode(std::unique_ptr<QuerySolutionNode> child,
+                   const MatchExpression& fullExpression,
+                   BSONObj projection,
+                   ParsedProjection parsed)
+        : QuerySolutionNode(std::move(child)),
+          _sorts(SimpleBSONObjComparator::kInstance.makeBSONObjSet()),
+          fullExpression(fullExpression),
+          projection(std::move(projection)),
+          parsed(parsed) {}
+
+    void computeProperties() final;
 
-    virtual void appendToString(mongoutils::str::stream* ss, int indent) const;
+    void appendToString(mongoutils::str::stream* ss, int indent) const final;
 
     /**
      * Data from the projection node is considered fetch iff the child provides fetched data.
@@ -584,29 +577,89 @@ struct ProjectionNode : public QuerySolutionNode {
         return _sorts;
     }
 
-    QuerySolutionNode* clone() const;
+protected:
+    void cloneProjectionData(ProjectionNode* copy) const;
+
+public:
+    /**
+     * Identify projectionImplementation type as a string.
+     */
+    virtual StringData projectionImplementationTypeToString() const = 0;
 
     BSONObjSet _sorts;
 
     // The full query tree.  Needed when we have positional operators.
     // Owned in the CanonicalQuery, not here.
-    MatchExpression* fullExpression;
+    const MatchExpression& fullExpression;
 
     // Given that we don't yet have a MatchExpression analogue for the expression language, we
     // use a BSONObj.
     BSONObj projection;
 
-    // What implementation of the projection algorithm should we use?
-    ProjectionType projType;
-
     ParsedProjection parsed;
+};
+
+/**
+ * This is the most general implementation of the projection functionality. It handles every case.
+ */
+struct ProjectionNodeDefault final : ProjectionNode {
+    using ProjectionNode::ProjectionNode;
+
+    StageType getType() const final {
+        return STAGE_PROJECTION_DEFAULT;
+    }
+
+    ProjectionNode* clone() const final;
+
+    StringData projectionImplementationTypeToString() const final {
+        return "DEFAULT"_sd;
+    }
+};
+
+/**
+ * This is a fast-path for when the projection is fully covered by one index.
+ */
+struct ProjectionNodeCovered final : ProjectionNode {
+    ProjectionNodeCovered(std::unique_ptr<QuerySolutionNode> child,
+                          const MatchExpression& fullExpression,
+                          BSONObj projection,
+                          ParsedProjection parsed,
+                          BSONObj coveredKeyObj)
+        : ProjectionNode(std::move(child), fullExpression, projection, parsed),
+          coveredKeyObj(std::move(coveredKeyObj)) {}
+
+    StageType getType() const final {
+        return STAGE_PROJECTION_COVERED;
+    }
+
+    ProjectionNode* clone() const final;
+
+    StringData projectionImplementationTypeToString() const final {
+        return "COVERED_ONE_INDEX"_sd;
+    }
 
-    // Only meaningful if projType == COVERED_ONE_INDEX.  This is the key pattern of the index
-    // supplying our covered data.  We can pre-compute which fields to include and cache that
-    // data for later if we know we only have one index.
+    // This is the key pattern of the index supplying our covered data. We can pre-compute which
+    // fields to include and cache that data for later if we know we only have one index.
     BSONObj coveredKeyObj;
 };
 
+/**
+ * This is a fast-path for when the projection only has inclusions on non-dotted fields.
+ */
+struct ProjectionNodeSimple final : ProjectionNode {
+    using ProjectionNode::ProjectionNode;
+
+    StageType getType() const final {
+        return STAGE_PROJECTION_SIMPLE;
+    }
+
+    ProjectionNode* clone() const final;
+
+    StringData projectionImplementationTypeToString() const final {
+        return "SIMPLE_DOC"_sd;
+    }
+};
+
 struct SortKeyGeneratorNode : public QuerySolutionNode {
     StageType getType() const final {
         return STAGE_SORT_KEY_GENERATOR;