SERVER-54745 Simplify covered projection plans in SBE

4 files changed, 307 insertions, 53 deletions

diff --git a/src/mongo/db/query/projection_ast_visitor.h b/src/mongo/db/query/projection_ast_visitor.h
index 30dd8b9df43..b765973faf9 100644
--- a/src/mongo/db/query/projection_ast_visitor.h
+++ b/src/mongo/db/query/projection_ast_visitor.h
@@ -33,6 +33,7 @@
 
 namespace mongo {
 namespace projection_ast {
+class ASTNode;
 class MatchExpressionASTNode;
 class ProjectionPathASTNode;
 class ProjectionPositionalASTNode;
@@ -64,5 +65,35 @@ public:
 
 using ProjectionASTMutableVisitor = ProjectionASTVisitor<false>;
 using ProjectionASTConstVisitor = ProjectionASTVisitor<true>;
+
+template <bool IsConst = false>
+class ProjectionASTWalker {
+public:
+    using VisitorPtr = ProjectionASTVisitor<IsConst>*;
+    using ASTNodePtr = tree_walker::MaybeConstPtr<IsConst, ASTNode>;
+
+    ProjectionASTWalker(VisitorPtr preVisitor, VisitorPtr inVisitor, VisitorPtr postVisitor)
+        : _preVisitor{preVisitor}, _inVisitor{inVisitor}, _postVisitor{postVisitor} {}
+
+    void preVisit(ASTNodePtr node) {
+        node->acceptVisitor(_preVisitor);
+    }
+
+    void postVisit(ASTNodePtr node) {
+        node->acceptVisitor(_postVisitor);
+    }
+
+    void inVisit(long count, ASTNodePtr node) {
+        node->acceptVisitor(_inVisitor);
+    }
+
+private:
+    VisitorPtr _preVisitor;
+    VisitorPtr _inVisitor;
+    VisitorPtr _postVisitor;
+};
+
+using ProjectionASTMutableWalker = ProjectionASTWalker<false>;
+using ProjectionASTConstWalker = ProjectionASTWalker<true>;
 }  // namespace projection_ast
 }  // namespace mongo
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index 52dc361235f..b97d1c3a340 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -67,7 +67,7 @@
 namespace mongo::stage_builder {
 namespace {
 /**
- * Tree representation of an index key pattern.
+ * Tree representing index key pattern or a subset of it.
  *
  * For example, the key pattern {a.b: 1, x: 1, a.c: 1} would look like:
  *
@@ -77,7 +77,7 @@ namespace {
  *       / \
  *      b   c
  *
- * This tree is used for building SBE subtrees to re-hydrate index keys.
+ * This tree is used for building SBE subtrees to re-hydrate index keys and for covered projections.
  */
 struct IndexKeyPatternTreeNode {
     IndexKeyPatternTreeNode* emplace(StringData fieldComponent) {
@@ -89,6 +89,30 @@ struct IndexKeyPatternTreeNode {
         return newNodeRaw;
     }
 
+    /**
+     * Returns leaf node matching field path. If the field path provided resolves to a non-leaf
+     * node, null will be returned.
+     *
+     * For example, if tree was built for key pattern {a: 1, a.b: 1}, this method will return
+     * nullptr for field path "a". On the other hand, this method will return corresponding node for
+     * field path "a.b".
+     */
+    IndexKeyPatternTreeNode* findLeafNode(const FieldRef& fieldRef, size_t currentIndex = 0) {
+        if (currentIndex == fieldRef.numParts()) {
+            if (children.empty()) {
+                return this;
+            }
+            return nullptr;
+        }
+
+        auto currentPart = fieldRef.getPart(currentIndex);
+        if (auto it = children.find(currentPart); it != children.end()) {
+            return it->second->findLeafNode(fieldRef, currentIndex + 1);
+        } else {
+            return nullptr;
+        }
+    }
+
     StringMap<std::unique_ptr<IndexKeyPatternTreeNode>> children;
     std::vector<std::string> childrenOrder;
 
@@ -98,6 +122,112 @@ struct IndexKeyPatternTreeNode {
 };
 
 /**
+ * For covered projections, each of the projection field paths represent respective index key. To
+ * rehydrate index keys into the result object, we first need to convert projection AST into
+ * 'IndexKeyPatternTreeNode' structure. Context structure and visitors below are used for this
+ * purpose.
+ */
+struct IndexKeysBuilderContext {
+    // Contains resulting tree of index keys converted from projection AST.
+    IndexKeyPatternTreeNode root;
+
+    // Full field path of the currently visited projection node.
+    std::vector<StringData> currentFieldPath;
+
+    // Each projection node has a vector of field names. This stack contains indexes of the
+    // currently visited field names for each of the projection nodes.
+    std::vector<size_t> currentFieldIndex;
+};
+
+/**
+ * Covered projections are always inclusion-only, so we ban all other operators.
+ */
+class IndexKeysBuilder : public projection_ast::ProjectionASTConstVisitor {
+public:
+    using projection_ast::ProjectionASTConstVisitor::visit;
+
+    IndexKeysBuilder(IndexKeysBuilderContext* context) : _context{context} {}
+
+    void visit(const projection_ast::ProjectionPositionalASTNode* node) final {
+        tasserted(5474501, "Positional projection is not allowed in covered projection");
+    }
+
+    void visit(const projection_ast::ProjectionSliceASTNode* node) final {
+        tasserted(5474502, "$slice is not allowed in covered projection");
+    }
+
+    void visit(const projection_ast::ProjectionElemMatchASTNode* node) final {
+        tasserted(5474503, "$elemMatch is not allowed in covered projection");
+    }
+
+    void visit(const projection_ast::ExpressionASTNode* node) final {
+        tasserted(5474504, "Expressions are not allowed in covered projection");
+    }
+
+    void visit(const projection_ast::MatchExpressionASTNode* node) final {
+        tasserted(5474505,
+                  "$elemMatch and positional projection are not allowed in covered projection");
+    }
+
+    void visit(const projection_ast::BooleanConstantASTNode* node) override {}
+
+protected:
+    IndexKeysBuilderContext* _context;
+};
+
+class IndexKeysPreBuilder final : public IndexKeysBuilder {
+public:
+    using IndexKeysBuilder::IndexKeysBuilder;
+    using IndexKeysBuilder::visit;
+
+    void visit(const projection_ast::ProjectionPathASTNode* node) final {
+        _context->currentFieldIndex.push_back(0);
+        _context->currentFieldPath.emplace_back(node->fieldNames().front());
+    }
+};
+
+class IndexKeysInBuilder final : public IndexKeysBuilder {
+public:
+    using IndexKeysBuilder::IndexKeysBuilder;
+    using IndexKeysBuilder::visit;
+
+    void visit(const projection_ast::ProjectionPathASTNode* node) final {
+        auto& currentIndex = _context->currentFieldIndex.back();
+        currentIndex++;
+        _context->currentFieldPath.back() = node->fieldNames()[currentIndex];
+    }
+};
+
+class IndexKeysPostBuilder final : public IndexKeysBuilder {
+public:
+    using IndexKeysBuilder::IndexKeysBuilder;
+    using IndexKeysBuilder::visit;
+
+    void visit(const projection_ast::ProjectionPathASTNode* node) final {
+        _context->currentFieldIndex.pop_back();
+        _context->currentFieldPath.pop_back();
+    }
+
+    void visit(const projection_ast::BooleanConstantASTNode* constantNode) final {
+        if (!constantNode->value()) {
+            // Even though only inclusion is allowed in covered projection, there still can be
+            // {_id: 0} component. We do not need to generate any nodes for it.
+            return;
+        }
+
+        // Insert current field path into the index keys tree if it does not exist yet.
+        auto* node = &_context->root;
+        for (const auto& part : _context->currentFieldPath) {
+            if (auto it = node->children.find(part); it != node->children.end()) {
+                node = it->second.get();
+            } else {
+                node = node->emplace(part);
+            }
+        }
+    }
+};
+
+/**
  * Given a key pattern and an array of slots of equal size, builds an IndexKeyPatternTreeNode
  * representing the mapping between key pattern component and slot.
  *
@@ -268,18 +398,63 @@ std::string PlanStageData::debugString() const {
 }
 
 namespace {
-const QuerySolutionNode* getNodeByType(const QuerySolutionNode* root, StageType type) {
+void getAllNodesByTypeHelper(const QuerySolutionNode* root,
+                             StageType type,
+                             std::vector<const QuerySolutionNode*>& results) {
+    if (root->getType() == type) {
+        results.push_back(root);
+    }
+
+    for (auto&& child : root->children) {
+        getAllNodesByTypeHelper(child, type, results);
+    }
+}
+
+std::vector<const QuerySolutionNode*> getAllNodesByType(const QuerySolutionNode* root,
+                                                        StageType type) {
+    std::vector<const QuerySolutionNode*> results;
+    getAllNodesByTypeHelper(root, type, results);
+    return results;
+}
+
+/**
+ * Returns pair consisting of:
+ *  - First node of the specified type found by pre-order traversal. If node was not found, this
+ *    pair element is nullptr.
+ *  - Total number of nodes with the specified type in tree.
+ */
+std::pair<const QuerySolutionNode*, size_t> getFirstNodeByType(const QuerySolutionNode* root,
+                                                               StageType type) {
+    const QuerySolutionNode* result = nullptr;
+    size_t count = 0;
     if (root->getType() == type) {
-        return root;
+        result = root;
+        count++;
     }
 
     for (auto&& child : root->children) {
-        if (auto result = getNodeByType(child, type)) {
-            return result;
+        auto [subTreeResult, subTreeCount] = getFirstNodeByType(child, type);
+        if (!result) {
+            result = subTreeResult;
         }
+        count += subTreeCount;
     }
 
-    return nullptr;
+    return {result, count};
+}
+
+/**
+ * Returns node of the specified type found in tree. If there is no such node, returns null. If
+ * there are more than one nodes of the specified type, throws an exception.
+ */
+const QuerySolutionNode* getLoneNodeByType(const QuerySolutionNode* root, StageType type) {
+    auto [result, count] = getFirstNodeByType(root, type);
+    const auto msgCount = count;
+    tassert(5474506,
+            str::stream() << "Found " << msgCount << " nodes of type " << type
+                          << ", expected one or zero",
+            count < 2);
+    return result;
 }
 
 sbe::LockAcquisitionCallback makeLockAcquisitionCallback(bool checkNodeCanServeReads) {
@@ -484,16 +659,21 @@ SlotBasedStageBuilder::SlotBasedStageBuilder(OperationContext* opCtx,
     // SERVER-52803: In the future if we need to gather more information from the QuerySolutionNode
     // tree, rather than doing one-off scans for each piece of information, we should add a formal
     // analysis pass here.
-    if (auto node = getNodeByType(solution.root(), STAGE_COLLSCAN)) {
+    // NOTE: Currently, we assume that each query operates on at most one collection, so there can
+    // be only one STAGE_COLLSCAN node.
+    if (auto node = getLoneNodeByType(solution.root(), STAGE_COLLSCAN)) {
         auto csn = static_cast<const CollectionScanNode*>(node);
         _data.shouldTrackLatestOplogTimestamp = csn->shouldTrackLatestOplogTimestamp;
         _data.shouldTrackResumeToken = csn->requestResumeToken;
         _data.shouldUseTailableScan = csn->tailable;
     }
 
-    if (auto node = getNodeByType(solution.root(), STAGE_VIRTUAL_SCAN)) {
+    for (const auto& node : getAllNodesByType(solution.root(), STAGE_VIRTUAL_SCAN)) {
         auto vsn = static_cast<const VirtualScanNode*>(node);
-        _shouldProduceRecordIdSlot = vsn->hasRecordId;
+        if (!vsn->hasRecordId) {
+            _shouldProduceRecordIdSlot = false;
+            break;
+        }
     }
 }
 
@@ -1034,7 +1214,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     sbe::IndexKeysInclusionSet sortPatternKeyBitSet;
     if (isCoveredQuery) {
         // If query is covered, we need to request index key for each part of the sort pattern.
-        auto indexScan = static_cast<const IndexScanNode*>(getNodeByType(child, STAGE_IXSCAN));
+        auto indexScan = static_cast<const IndexScanNode*>(getLoneNodeByType(child, STAGE_IXSCAN));
         tassert(5601701, "Expected index scan below sort for covered query", indexScan);
         indexKeyPattern = indexScan->index.keyPattern;
 
@@ -1283,7 +1463,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
         // 'b' and 'a'),  we wish to insert them into 'inputKeys' in the order that they appear in
         // the sort pattern.
         StringMap<size_t> indexKeyPositionMap;
-        auto ixnNode = getNodeByType(child, STAGE_IXSCAN);
+        auto ixnNode = getLoneNodeByType(child, STAGE_IXSCAN);
         tassert(5184300,
                 str::stream() << "Can't build exec tree for node: " << child->toString(),
                 ixnNode);
@@ -1455,28 +1635,93 @@ SlotBasedStageBuilder::buildProjectionCovered(const QuerySolutionNode* root,
 std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots>
 SlotBasedStageBuilder::buildProjectionDefault(const QuerySolutionNode* root,
                                               const PlanStageReqs& reqs) {
-    using namespace std::literals;
     invariant(!reqs.getIndexKeyBitset());
 
+    auto childReqs = reqs.copy();
+
     auto pn = static_cast<const ProjectionNodeDefault*>(root);
+    const auto& projection = pn->proj;
+    const auto [indexScanNode, indexScanCount] = getFirstNodeByType(root, STAGE_IXSCAN);
+    // TODO SERVER-57533: Support multiple index scan nodes located below OR and SORT_MERGE stages.
+    const auto isCoveredProjection =
+        !pn->fetched() && indexScanCount == 1 && projection.isInclusionOnly();
+
+    boost::optional<IndexKeyPatternTreeNode> patternRoot;
+    std::vector<IndexKeyPatternTreeNode*> patternNodesForSlots;
+    if (isCoveredProjection) {
+        // Convert projection fieldpaths into the tree of 'IndexKeyPatternTreeNode'.
+        IndexKeysBuilderContext context;
+        IndexKeysPreBuilder preVisitor{&context};
+        IndexKeysInBuilder inVisitor{&context};
+        IndexKeysPostBuilder postVisitor{&context};
+        projection_ast::ProjectionASTConstWalker walker{&preVisitor, &inVisitor, &postVisitor};
+        tree_walker::walk<true, projection_ast::ASTNode>(projection.root(), &walker);
+        patternRoot = std::move(context.root);
+
+        // Construct a bitset requesting slots from the underlying index scan. These slots
+        // correspond to index keys for projection fieldpaths.
+        auto& indexKeyPattern = static_cast<const IndexScanNode*>(indexScanNode)->index.keyPattern;
+        size_t i = 0;
+        sbe::IndexKeysInclusionSet patternBitSet;
+        for (const auto& element : indexKeyPattern) {
+            FieldRef fieldRef{element.fieldNameStringData()};
+            // Projection field paths are always leaf nodes. In other words, projection like
+            // {a: 1, 'a.b': 1} would produce a path collision error.
+            if (auto node = patternRoot->findLeafNode(fieldRef); node) {
+                patternBitSet.set(i);
+                patternNodesForSlots.push_back(node);
+            }
+
+            ++i;
+        }
+
+        childReqs.getIndexKeyBitset() = patternBitSet;
+
+        // We do not need index scan to restore the entire object. Instead, we will restore only
+        // necessary parts of it below.
+        childReqs.clear(kResult);
+    } else {
+        // The child must produce all of the slots required by the parent of this
+        // ProjectionNodeDefault. In addition to that, the child must always produce a 'resultSlot'
+        // because it's needed by the projection logic below.
+        childReqs.set(kResult);
+    }
 
-    // The child must produce all of the slots required by the parent of this ProjectionNodeDefault.
-    // In addition to that, the child must always produce a 'resultSlot' because it's needed by the
-    // projection logic below.
-    auto childReqs = reqs.copy().set(kResult);
     auto [inputStage, outputs] = build(pn->children[0], childReqs);
 
-    auto relevantSlots = sbe::makeSV();
-    outputs.forEachSlot(childReqs, [&](auto&& slot) { relevantSlots.push_back(slot); });
+    sbe::value::SlotId resultSlot;
+    std::unique_ptr<sbe::PlanStage> resultStage;
+    if (isCoveredProjection) {
+        auto indexKeySlots = *outputs.extractIndexKeySlots();
+
+        // Extract slots corresponding to each of the projection fieldpaths.
+        invariant(indexKeySlots.size() == patternNodesForSlots.size());
+        for (size_t i = 0; i < indexKeySlots.size(); i++) {
+            patternNodesForSlots[i]->indexKeySlot = indexKeySlots[i];
+        }
 
-    auto [slot, stage] = generateProjection(_state,
-                                            &pn->proj,
-                                            {std::move(inputStage), std::move(relevantSlots)},
-                                            outputs.get(kResult),
-                                            root->nodeId());
-    outputs.set(kResult, slot);
+        // Finally, build the expression to create object with requested projection fieldpaths.
+        resultSlot = _slotIdGenerator.generate();
+        auto resultExpr = buildNewObjExpr(&*patternRoot);
+        resultStage = sbe::makeProjectStage(
+            std::move(inputStage), root->nodeId(), resultSlot, std::move(resultExpr));
+    } else {
+        auto relevantSlots = sbe::makeSV();
+        outputs.forEachSlot(childReqs, [&](auto&& slot) { relevantSlots.push_back(slot); });
+
+        EvalStage stage;
+        std::tie(resultSlot, stage) =
+            generateProjection(_state,
+                               &projection,
+                               {std::move(inputStage), std::move(relevantSlots)},
+                               outputs.get(kResult),
+                               root->nodeId());
+
+        resultStage = std::move(stage.stage);
+    }
 
-    return {std::move(stage.stage), std::move(outputs)};
+    outputs.set(kResult, resultSlot);
+    return {std::move(resultStage), std::move(outputs)};
 }
 
 std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder::buildOr(
diff --git a/src/mongo/db/query/sbe_stage_builder_projection.cpp b/src/mongo/db/query/sbe_stage_builder_projection.cpp
index 5a4dd5e211f..6afc9e6bf88 100644
--- a/src/mongo/db/query/sbe_stage_builder_projection.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_projection.cpp
@@ -872,31 +872,6 @@ private:
     ProjectionTraversalVisitorContext* _context;
 };
 
-class ProjectionTraversalWalker final {
-public:
-    ProjectionTraversalWalker(projection_ast::ProjectionASTConstVisitor* preVisitor,
-                              projection_ast::ProjectionASTConstVisitor* inVisitor,
-                              projection_ast::ProjectionASTConstVisitor* postVisitor)
-        : _preVisitor{preVisitor}, _inVisitor{inVisitor}, _postVisitor{postVisitor} {}
-
-    void preVisit(const projection_ast::ASTNode* node) {
-        node->acceptVisitor(_preVisitor);
-    }
-
-    void postVisit(const projection_ast::ASTNode* node) {
-        node->acceptVisitor(_postVisitor);
-    }
-
-    void inVisit(long count, const projection_ast::ASTNode* node) {
-        node->acceptVisitor(_inVisitor);
-    }
-
-private:
-    projection_ast::ProjectionASTConstVisitor* _preVisitor;
-    projection_ast::ProjectionASTConstVisitor* _inVisitor;
-    projection_ast::ProjectionASTConstVisitor* _postVisitor;
-};
-
 /**
  * Generates expression that applies positional projection operator to the array stored in the
  * 'inputSlot' using optional index from 'maybeIndexSlot'.
@@ -1177,7 +1152,7 @@ std::pair<sbe::value::SlotId, EvalStage> generateProjection(
     ProjectionTraversalPreVisitor preVisitor{&context};
     ProjectionTraversalInVisitor inVisitor{&context};
     ProjectionTraversalPostVisitor postVisitor{&context};
-    ProjectionTraversalWalker walker{&preVisitor, &inVisitor, &postVisitor};
+    projection_ast::ProjectionASTConstWalker walker{&preVisitor, &inVisitor, &postVisitor};
     tree_walker::walk<true, projection_ast::ASTNode>(projection->root(), &walker);
     auto [resultSlot, resultStage] = context.done();
 
@@ -1192,7 +1167,8 @@ std::pair<sbe::value::SlotId, EvalStage> generateProjection(
         ProjectionTraversalPreVisitor slicePreVisitor{&sliceContext};
         ProjectionTraversalInVisitor sliceInVisitor{&sliceContext};
         SliceProjectionTraversalPostVisitor slicePostVisitor{&sliceContext};
-        ProjectionTraversalWalker sliceWalker{&slicePreVisitor, &sliceInVisitor, &slicePostVisitor};
+        projection_ast::ProjectionASTConstWalker sliceWalker{
+            &slicePreVisitor, &sliceInVisitor, &slicePostVisitor};
         tree_walker::walk<true, projection_ast::ASTNode>(projection->root(), &sliceWalker);
         std::tie(resultSlot, resultStage) = sliceContext.done();
     }
diff --git a/src/mongo/db/query/tree_walker.h b/src/mongo/db/query/tree_walker.h
index eb36a27e6ca..4452ee09516 100644
--- a/src/mongo/db/query/tree_walker.h
+++ b/src/mongo/db/query/tree_walker.h
@@ -29,6 +29,8 @@
 
 #pragma once
 
+#include <type_traits>
+
 namespace mongo::tree_walker {
 /**
  * A template type which resolves to 'const T*' if 'IsConst' argument is 'true', and to 'T*'