SERVER-71530: Add VariableTransporter for walking variables with callback

6 files changed, 53 insertions, 60 deletions

diff --git a/src/mongo/db/query/optimizer/cascades/implementers.cpp b/src/mongo/db/query/optimizer/cascades/implementers.cpp
index 08aa9811069..b315461560b 100644
--- a/src/mongo/db/query/optimizer/cascades/implementers.cpp
+++ b/src/mongo/db/query/optimizer/cascades/implementers.cpp
@@ -1140,10 +1140,9 @@ public:
                 const ABT& aggExpr = node.getAggregationExpressions().at(aggIndex);
                 aggregationProjections.push_back(aggExpr);
 
-                for (const Variable& var : VariableEnvironment::getVariables(aggExpr)._variables) {
-                    // Add all references this expression requires.
-                    projectionsToAdd.insert(var.name());
-                }
+                // Add all references this expression requires.
+                VariableEnvironment::walkVariables(
+                    aggExpr, [&](const Variable& var) { projectionsToAdd.insert(var.name()); });
             }
         }
 
diff --git a/src/mongo/db/query/optimizer/optimizer_test.cpp b/src/mongo/db/query/optimizer/optimizer_test.cpp
index 5a08bafd135..a3f68bc83c7 100644
--- a/src/mongo/db/query/optimizer/optimizer_test.cpp
+++ b/src/mongo/db/query/optimizer/optimizer_test.cpp
@@ -207,12 +207,14 @@ TEST(Optimizer, Tracker4) {
     ASSERT(!env.hasFreeVariables());
 
     // Get all variables from the expression
-    auto vars = VariableEnvironment::getVariables(filterNode.cast<FilterNode>()->getFilter());
-    ASSERT(vars._variables.size() == 1);
+    std::vector<std::reference_wrapper<const Variable>> vars;
+    VariableEnvironment::walkVariables(filterNode.cast<FilterNode>()->getFilter(),
+                                       [&](const Variable& var) { vars.push_back(var); });
+    ASSERT(vars.size() == 1);
     // Get all definitions from the scan and below (even though there is nothing below the scan).
     auto defs = env.getDefinitions(scanNodeRef);
     // Make sure that variables are defined by the scan (and not by Eval).
-    for (const Variable& v : vars._variables) {
+    for (const Variable& v : vars) {
         auto it = defs.find(v.name());
         ASSERT(it != defs.end());
         ASSERT(it->second.definedBy == env.getDefinition(v).definedBy);
diff --git a/src/mongo/db/query/optimizer/reference_tracker.cpp b/src/mongo/db/query/optimizer/reference_tracker.cpp
index ec242879425..2ba8fd22bd0 100644
--- a/src/mongo/db/query/optimizer/reference_tracker.cpp
+++ b/src/mongo/db/query/optimizer/reference_tracker.cpp
@@ -236,37 +236,37 @@ struct CollectedInfo {
 };
 
 /**
- * Collect all Variables into a set.
+ * Walks over all variables in the ABT and calls a callback for each variable.
  */
-class VariableCollector {
+class VariableTransporter {
 public:
+    VariableTransporter(
+        const std::function<void(const Variable&)>& variableCallback,
+        const std::function<void(const ProjectionName&)>& variableDefinitionCallback)
+        : _variableCallback(variableCallback),
+          _variableDefinitionCallback(variableDefinitionCallback) {}
+
     template <typename T, typename... Ts>
     void transport(const T& /*op*/, Ts&&... /*ts*/) {}
 
     void transport(const Variable& op) {
-        _result._variables.push_back(op);
+        _variableCallback(op);
     }
 
     void transport(const LambdaAbstraction& op, const ABT& /*bind*/) {
-        _result._definedVars.insert(op.varName());
+        _variableDefinitionCallback(op.varName());
     }
 
     void transport(const Let& op, const ABT& /*bind*/, const ABT& /*expr*/) {
-        _result._definedVars.insert(op.varName());
-    }
-
-    static VariableCollectorResult collect(const ABT& n) {
-        VariableCollector collector;
-        collector.collectInternal(n);
-        return std::move(collector._result);
+        _variableDefinitionCallback(op.varName());
     }
 
 private:
-    void collectInternal(const ABT& n) {
-        algebra::transport<false>(n, *this);
-    }
+    // Callback used on each Variable in the ABT.
+    const std::function<void(const Variable&)>& _variableCallback;
 
-    VariableCollectorResult _result;
+    // Callback used on any defined variable name (via a Let or Lambda) in the ABT.
+    const std::function<void(const ProjectionName&)>& _variableDefinitionCallback;
 };
 
 struct Collector {
@@ -991,17 +991,12 @@ bool VariableEnvironment::isLastRef(const Variable& var) const {
     return _info->lastRefs.count(&var) > 0;
 }
 
-VariableCollectorResult VariableEnvironment::getVariables(const ABT& n) {
-    return VariableCollector::collect(n);
-}
-
-void VariableEnvironment::walkVariables(const ABT& n, std::function<void(const Variable&)> func) {
-    // TODO SERVER-71530 consider passing the lambda into the transport, to avoid building a vector.
-    VariableCollectorResult collected = VariableEnvironment::getVariables(n);
-
-    for (auto&& var : collected._variables) {
-        func(var.get());
-    }
+void VariableEnvironment::walkVariables(
+    const ABT& n,
+    const std::function<void(const Variable&)>& variableCallback,
+    const std::function<void(const ProjectionName&)>& variableDefinitionCallback) {
+    VariableTransporter transporter(variableCallback, variableDefinitionCallback);
+    algebra::transport<false>(n, transporter);
 }
 
 
diff --git a/src/mongo/db/query/optimizer/reference_tracker.h b/src/mongo/db/query/optimizer/reference_tracker.h
index 1c0bc87705a..c11ddce41f4 100644
--- a/src/mongo/db/query/optimizer/reference_tracker.h
+++ b/src/mongo/db/query/optimizer/reference_tracker.h
@@ -55,16 +55,6 @@ struct Definition {
 struct CollectedInfo;
 using DefinitionsMap = ProjectionNameMap<Definition>;
 
-struct VariableCollectorResult {
-    // The Variables referenced by the subtree.
-    std::vector<std::reference_wrapper<const Variable>> _variables;
-    // The names of locally-defined Variables. These aren't propagated up the tree during normal
-    // variable resolution. Tracking these separately allows us to easily check, for example, which
-    // variables are referenced in but not defined by the subtree (i.e. variables which should be
-    // defined elsewhere in the ABT).
-    ProjectionNameSet _definedVars;
-};
-
 /**
  * Helps enforce scoping and validity rules for definitions and Variable references.
  */
@@ -83,15 +73,14 @@ public:
     void rebuild(const ABT& root);
 
     /**
-     * Get information about Variables in the subtree rooted at 'n', including Variables referenced
-     * by the subtree and locally-defined Variables.
-     */
-    static VariableCollectorResult getVariables(const ABT& n);
-
-    /**
-     * Call 'func' on each Variable in the subtree.
+     * Calls 'variableCallback' on each Variable and `variableDefinitionCallback` on each
+     * variable name defined via a Let or Lambda in the ABT.
      */
-    static void walkVariables(const ABT& n, std::function<void(const Variable&)> func);
+    static void walkVariables(
+        const ABT& n,
+        const std::function<void(const Variable&)>& variableCallback,
+        const std::function<void(const ProjectionName&)>& variableDefinitionCallback =
+            [](const ProjectionName&) {});
 
     ~VariableEnvironment();
 
diff --git a/src/mongo/db/query/optimizer/reference_tracker_test.cpp b/src/mongo/db/query/optimizer/reference_tracker_test.cpp
index f3efa486be4..6ce0c385165 100644
--- a/src/mongo/db/query/optimizer/reference_tracker_test.cpp
+++ b/src/mongo/db/query/optimizer/reference_tracker_test.cpp
@@ -175,10 +175,12 @@ TEST(ReferenceTrackerTest, GetDefinitionsForLet) {
 
     // But, the environment keeps the info about the definitions for all variables in the ABT. Check
     // that the local variable is defined by the Let.
-    auto variables = VariableEnvironment::getVariables(evalNode);
-    auto xVar = std::find_if(variables._variables.begin(),
-                             variables._variables.end(),
-                             [&](const Variable& var) { return var.name() == "x"; });
+    const Variable* xVar = nullptr;
+    VariableEnvironment::walkVariables(evalNode, [&](const Variable& var) {
+        if (var.name() == "x") {
+            xVar = &var;
+        }
+    });
     ASSERT(env.isLastRef(*xVar));
     ASSERT(env.getDefinition(*xVar).definedBy == letRef);
 }
diff --git a/src/mongo/db/query/optimizer/utils/reftracker_utils.cpp b/src/mongo/db/query/optimizer/utils/reftracker_utils.cpp
index d148306b818..436e666d2f4 100644
--- a/src/mongo/db/query/optimizer/utils/reftracker_utils.cpp
+++ b/src/mongo/db/query/optimizer/utils/reftracker_utils.cpp
@@ -207,11 +207,17 @@ public:
 
 private:
     void extractFromABT(ProjectionNameSet& vars, const ABT& v) {
-        const auto& result = VariableEnvironment::getVariables(v);
-        for (const Variable& var : result._variables) {
-            if (result._definedVars.count(var.name()) == 0) {
-                // We are interested in either free variables, or variables defined on other nodes.
-                vars.insert(var.name());
+        // Mark variables as defined or not in this subtree.
+        ProjectionNameMap<bool> varHasDefinitionMap;
+        VariableEnvironment::walkVariables(
+            v,
+            [&](const Variable& var) { varHasDefinitionMap.emplace(var.name(), false); },
+            [&](const ProjectionName& definedVar) { varHasDefinitionMap[definedVar] = true; });
+        for (const auto& varHasDefinition : varHasDefinitionMap) {
+            if (!varHasDefinition.second) {
+                // We are interested in when the variable has no definition in this subtree of the
+                // ABT: either free variables, or variables defined on other nodes.
+                vars.insert(varHasDefinition.first);
             }
         }
     }