Merge branch 'master' into jenniferpeshansky/SERVER-70854jenniferpeshansky/SERVER-70854

1 files changed, 74 insertions, 13 deletions

diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.h b/src/mongo/db/query/sbe_stage_builder_helpers.h
index f3b5d9ca455..634005c504c 100644
--- a/src/mongo/db/query/sbe_stage_builder_helpers.h
+++ b/src/mongo/db/query/sbe_stage_builder_helpers.h
@@ -44,6 +44,10 @@
 #include "mongo/db/query/sbe_stage_builder_eval_frame.h"
 #include "mongo/db/query/stage_types.h"
 
+namespace mongo::projection_ast {
+class Projection;
+}
+
 namespace mongo::stage_builder {
 
 std::unique_ptr<sbe::EExpression> makeUnaryOp(sbe::EPrimUnary::Op unaryOp,
@@ -524,19 +528,20 @@ std::unique_ptr<sbe::EExpression> makeLocalBind(sbe::value::FrameIdGenerator* fr
     return sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(innerExpr));
 }
 
-
-std::tuple<sbe::value::SlotId, sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>>
-makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage,
-                     sbe::value::SlotId seekKeySlot,
-                     sbe::value::SlotId snapshotIdSlot,
-                     sbe::value::SlotId indexIdSlot,
-                     sbe::value::SlotId indexKeySlot,
-                     sbe::value::SlotId indexKeyPatternSlot,
-                     const CollectionPtr& collToFetch,
-                     StringMap<const IndexAccessMethod*> iamMap,
-                     PlanNodeId planNodeId,
-                     sbe::value::SlotVector slotsToForward,
-                     sbe::value::SlotIdGenerator& slotIdGenerator);
+std::unique_ptr<sbe::PlanStage> makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage,
+                                                     sbe::value::SlotId resultSlot,
+                                                     sbe::value::SlotId recordIdSlot,
+                                                     std::vector<std::string> fields,
+                                                     sbe::value::SlotVector fieldSlots,
+                                                     sbe::value::SlotId seekKeySlot,
+                                                     sbe::value::SlotId snapshotIdSlot,
+                                                     sbe::value::SlotId indexIdSlot,
+                                                     sbe::value::SlotId indexKeySlot,
+                                                     sbe::value::SlotId indexKeyPatternSlot,
+                                                     const CollectionPtr& collToFetch,
+                                                     StringMap<const IndexAccessMethod*> iamMap,
+                                                     PlanNodeId planNodeId,
+                                                     sbe::value::SlotVector slotsToForward);
 
 /**
  * Trees generated by 'generateFilter' maintain state during execution. There are two types of state
@@ -1010,4 +1015,60 @@ std::unique_ptr<sbe::PlanStage> rehydrateIndexKey(std::unique_ptr<sbe::PlanStage
                                                   const sbe::value::SlotVector& indexKeySlots,
                                                   sbe::value::SlotId resultSlot);
 
+IndexKeyPatternTreeNode buildPatternTree(const projection_ast::Projection& projection);
+
+/**
+ * This method retrieves the values of the specified top-level fields ('fields') from 'resultSlot'
+ * and stores the values into slots.
+ *
+ * This method returns a pair containing: (1) the updated SBE plan stage tree and; (2) a vector of
+ * the slots ('slots') containing the field values.
+ *
+ * The order of slots in 'slots' will match the order of fields in 'fields'.
+ */
+std::pair<std::unique_ptr<sbe::PlanStage>, sbe::value::SlotVector> projectTopLevelFields(
+    std::unique_ptr<sbe::PlanStage> stage,
+    const std::vector<std::string>& fields,
+    sbe::value::SlotId resultSlot,
+    PlanNodeId nodeId,
+    sbe::value::SlotIdGenerator* slotIdGenerator);
+
+/**
+ * This method projects the constant value Nothing to multiple slots (the specific number of slots
+ * being specified by parameter 'n').
+ *
+ * This method returns a pair containing: (1) the updated SBE plan stage tree and; (2) a vector of
+ * slots ('slots') containing Nothing.
+ *
+ * The number of slots in 'slots' will always be equal to parameter 'n'.
+ */
+std::pair<std::unique_ptr<sbe::PlanStage>, sbe::value::SlotVector> projectNothingToSlots(
+    std::unique_ptr<sbe::PlanStage> stage,
+    size_t n,
+    PlanNodeId nodeId,
+    sbe::value::SlotIdGenerator* slotIdGenerator);
+
+template <typename T, typename FuncT>
+std::vector<T> filterVector(std::vector<T> vec, FuncT fn) {
+    std::vector<T> result;
+    std::copy_if(std::make_move_iterator(vec.begin()),
+                 std::make_move_iterator(vec.end()),
+                 std::back_inserter(result),
+                 fn);
+    return result;
+}
+
+template <typename T, typename FuncT>
+std::pair<std::vector<T>, std::vector<T>> splitVector(std::vector<T> vec, FuncT fn) {
+    std::pair<std::vector<T>, std::vector<T>> result;
+    for (size_t i = 0; i < vec.size(); ++i) {
+        if (fn(vec[i])) {
+            result.first.emplace_back(std::move(vec[i]));
+        } else {
+            result.second.emplace_back(std::move(vec[i]));
+        }
+    }
+    return result;
+}
+
 }  // namespace mongo::stage_builder