commit 046a5301acfdd2e084f039377e80ef8482a9e6f7
Author: Timour Katchaounov <timour.katchaounov@mongodb.com>
Date:   Fri Jun 23 12:33:05 2023 +0300

    SERVER-78213 Replace the unordered map 'useMap' via bitset
    
    With the goal of improving merge speed, the unordered map CollectedInfo::useMap
    is replaced by a bitset of Vairable IDs, and a global storage of resolved and
    enumeratd Variables.

diff --git a/src/mongo/db/query/optimizer/reference_tracker.cpp b/src/mongo/db/query/optimizer/reference_tracker.cpp
index 85ff03c4c40..ddba3e44889 100644
--- a/src/mongo/db/query/optimizer/reference_tracker.cpp
+++ b/src/mongo/db/query/optimizer/reference_tracker.cpp
@@ -27,20 +27,33 @@
  *    it in the license file.
  */
 
+#include <algorithm>
+#include <boost/dynamic_bitset.hpp>
+#include <cstddef>
 #include <iostream>
+#include <memory>
+#include <ostream>
+#include <unordered_map>
+#include <utility>
+#include <vector>
 
 #include "mongo/db/query/optimizer/reference_tracker.h"
 
 
 namespace mongo::optimizer {
 
+struct Collector;
 struct CollectedInfo {
     using VarRefsMap = ProjectionNameMap<opt::unordered_map<const Variable*, bool>>;
 
+    CollectedInfo(Collector* coll) : collector(coll) {}
+
+    Collector* collector;
+
     /**
      * All resolved variables so far, regardless of visibility in the ABT.
      */
-    opt::unordered_map<const Variable*, Definition> useMap;
+    boost::dynamic_bitset<size_t> resolvedVariables;
 
     /**
      * Current definitions available for use in ancestor nodes (projections).
@@ -66,6 +79,8 @@ struct CollectedInfo {
     VarRefsMap varLastRefs;
     opt::unordered_set<const Variable*> lastRefs;
 
+    void mergeResolvedVars(boost::dynamic_bitset<size_t>& otherResolvedVars);
+
     /**
      * This is a destructive merge, the 'other' will be siphoned out.
      */
@@ -90,8 +105,7 @@ struct CollectedInfo {
 
         // It should be impossible to have duplicate Variable pointer so every Variable should be
         // moved from other.
-        useMap.merge(other.useMap);
-        tassert(6624024, "Found a duplicate Variable pointer", other.useMap.empty());
+        mergeResolvedVars(other.resolvedVariables);
 
         // There should not be two projections of the same name propagated up by a single operator,
         // so every definition should be moved from other.
@@ -221,14 +235,7 @@ struct CollectedInfo {
     /**
      * Resolve any free Variables matching the given the name with the corresponding definition.
      */
-    void resolveFreeVars(const ProjectionName& name, const Definition& def) {
-        if (auto it = freeVars.find(name); it != freeVars.end()) {
-            for (const Variable& var : it->second) {
-                useMap.emplace(&var, def);
-            }
-            freeVars.erase(it);
-        }
-    }
+    void resolveFreeVars(const ProjectionName& name, const Definition& def);
 
     void assertEmptyDefs() {
         tassert(6624028, "Definitions must be empty", defs.empty());
@@ -271,7 +278,8 @@ private:
 
 struct Collector {
     explicit Collector(const cascades::MemoGroupBinderInterface* memoInterface)
-        : _memoInterface(memoInterface) {}
+        : allResolvedVariables(std::make_unique<ResolvedVarsStorage>()),
+          _memoInterface(memoInterface) {}
 
     template <typename T, typename... Ts>
     CollectedInfo transport(const ABT&, const T& op, Ts&&... ts) {
@@ -279,14 +287,14 @@ struct Collector {
 
         // The default behavior resolves free variables, merges known definitions and propagates
         // them up unmodified.
-        CollectedInfo result{};
+        CollectedInfo result{this};
         (result.merge(std::forward<Ts>(ts)), ...);
 
         return result;
     }
 
     CollectedInfo transport(const ABT& n, const Variable& variable) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // Every variable starts as a free variable until it is resolved.
         result.freeVars[variable.name()].push_back(variable);
@@ -301,7 +309,7 @@ struct Collector {
                             const Let& let,
                             CollectedInfo bindResult,
                             CollectedInfo inResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // The 'in' portion of the Let will execute after the bind, so its last refs should be kept
         // over the last refs from the bind. Then, its safe to finalize the last ref for the
@@ -320,7 +328,7 @@ struct Collector {
     }
 
     CollectedInfo transport(const ABT& n, const LambdaAbstraction& lam, CollectedInfo inResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // As in the Let case, we can finalize the last ref for the local variable.
         inResult.finalizeLastRefs(lam.varName());
@@ -338,7 +346,7 @@ struct Collector {
                             CollectedInfo condResult,
                             CollectedInfo thenResult,
                             CollectedInfo elseResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // Only one of the 'then' or 'else' will be executed, so it's safe to union the last refs.
         // Since the condition will be executed before either of the then/else, its last refs should
@@ -358,7 +366,7 @@ struct Collector {
                                         const Node& node,
                                         const ExpressionBinder& binder,
                                         CollectedInfo refs) {
-        CollectedInfo result{};
+        CollectedInfo result{refs.collector};
 
         // 'refs' should just track references to projections from any children of a Scan/Seek.
         result.mergeNoDefs(std::move(refs));
@@ -372,21 +380,21 @@ struct Collector {
     }
 
     CollectedInfo transport(const ABT& n, const ScanNode& node, CollectedInfo /*bindResult*/) {
-        return collectForScan(n, node, node.binder(), {});
+        return collectForScan(n, node, node.binder(), {this});
     }
 
     CollectedInfo transport(const ABT& n, const ValueScanNode& node, CollectedInfo /*bindResult*/) {
-        return collectForScan(n, node, node.binder(), {});
+        return collectForScan(n, node, node.binder(), {this});
     }
 
     CollectedInfo transport(const ABT& n,
                             const PhysicalScanNode& node,
                             CollectedInfo /*bindResult*/) {
-        return collectForScan(n, node, node.binder(), {});
+        return collectForScan(n, node, node.binder(), {this});
     }
 
     CollectedInfo transport(const ABT& n, const IndexScanNode& node, CollectedInfo /*bindResult*/) {
-        return collectForScan(n, node, node.binder(), {});
+        return collectForScan(n, node, node.binder(), {this});
     }
 
     CollectedInfo transport(const ABT& n,
@@ -397,14 +405,14 @@ struct Collector {
     }
 
     CollectedInfo transport(const ABT& n, const CoScanNode& node) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
         result.nodeDefs[&node] = result.defs;
         return result;
     }
 
     CollectedInfo transport(const ABT& n,
                             const MemoLogicalDelegatorNode& memoLogicalDelegatorNode) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         tassert(6624029, "Uninitialized memo interface", _memoInterface);
         const auto& binder =
@@ -429,7 +437,7 @@ struct Collector {
                             const FilterNode& filterNode,
                             CollectedInfo childResult,
                             CollectedInfo exprResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
         result.merge(std::move(childResult));
         result.mergeNoDefs(std::move(exprResult));
         result.nodeDefs[&filterNode] = result.defs;
@@ -440,7 +448,7 @@ struct Collector {
                             const EvaluationNode& evaluationNode,
                             CollectedInfo childResult,
                             CollectedInfo exprResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         tassert(6624030,
                 str::stream() << "Cannot overwrite project " << evaluationNode.getProjectionName(),
@@ -463,7 +471,7 @@ struct Collector {
                             CollectedInfo childResult,
                             CollectedInfo bindResult,
                             CollectedInfo /*refResult*/) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.merge(std::move(childResult));
         result.mergeNoDefs(std::move(bindResult));
@@ -483,7 +491,7 @@ struct Collector {
                             const RIDIntersectNode& node,
                             CollectedInfo leftChildResult,
                             CollectedInfo rightChildResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // This is a special case where both children of 'node' have a definition for the scan
         // projection. Remove the definition from one side to avoid running into the conflict of two
@@ -502,7 +510,7 @@ struct Collector {
                             const RIDUnionNode& node,
                             CollectedInfo leftChildResult,
                             CollectedInfo rightChildResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // For simplicity, don't provide any projections. In principle we could preserve any
         // projections that are common to both the left and right children, but for now we don't
@@ -528,7 +536,7 @@ struct Collector {
                                                 CollectedInfo leftChildResult,
                                                 CollectedInfo rightChildResult,
                                                 CollectedInfo filterResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // Note correlated projections might be coming either from the left child or from the
         // parent.
@@ -584,7 +592,7 @@ struct Collector {
                             CollectedInfo leftChildResult,
                             CollectedInfo rightChildResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.merge(std::move(leftChildResult));
         // Do not resolve further free variables.
@@ -601,7 +609,7 @@ struct Collector {
                             CollectedInfo leftChildResult,
                             CollectedInfo rightChildResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.merge(std::move(leftChildResult));
         // Do not resolve further free variables.
@@ -618,7 +626,7 @@ struct Collector {
                             std::vector<CollectedInfo> childResults,
                             CollectedInfo bindResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         const auto& names = node.binder().names();
 
@@ -636,7 +644,8 @@ struct Collector {
                 tassert(7063706,
                         str::stream() << "SortedMerge projection does not exist: " << name,
                         u.defs.count(name) != 0);
-                u.useMap.emplace(&refsResult.freeVars[name][counter].get(), u.defs[name]);
+                resolveVariable(
+                    refsResult.freeVars[name][counter].get(), u.defs[name], u.resolvedVariables);
             }
             u.defs.clear();
             result.merge(std::move(u));
@@ -663,7 +672,7 @@ struct Collector {
                             std::vector<CollectedInfo> childResults,
                             CollectedInfo bindResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         const auto& names = unionNode.binder().names();
 
@@ -680,7 +689,8 @@ struct Collector {
                 tassert(6624031,
                         str::stream() << "Union projection does not exist:  " << name,
                         u.defs.count(name) != 0);
-                u.useMap.emplace(&refsResult.freeVars[name][counter].get(), u.defs[name]);
+                resolveVariable(
+                    refsResult.freeVars[name][counter].get(), u.defs[name], u.resolvedVariables);
             }
             u.defs.clear();
             result.merge(std::move(u));
@@ -708,7 +718,7 @@ struct Collector {
                             CollectedInfo refsAggResult,
                             CollectedInfo bindGbResult,
                             CollectedInfo refsGbResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // First resolve all variables from the inside point of view; i.e. agg expressions and group
         // by expressions reference variables from the input child.
@@ -752,7 +762,7 @@ struct Collector {
                             CollectedInfo childResult,
                             CollectedInfo bindResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         // First resolve all variables from the inside point of view.
         result.mergeNoDefs(std::move(refsResult));
@@ -780,7 +790,7 @@ struct Collector {
                             const UniqueNode& uniqueNode,
                             CollectedInfo childResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.merge(std::move(refsResult));
         result.merge(std::move(childResult));
@@ -800,7 +810,7 @@ struct Collector {
                             const CollationNode& collationNode,
                             CollectedInfo childResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.mergeNoDefs(std::move(refsResult));
         result.merge(std::move(childResult));
@@ -819,7 +829,7 @@ struct Collector {
     CollectedInfo transport(const ABT& n,
                             const LimitSkipNode& limitSkipNode,
                             CollectedInfo childResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
         result.merge(std::move(childResult));
         result.nodeDefs[&limitSkipNode] = result.defs;
         return result;
@@ -829,7 +839,7 @@ struct Collector {
                             const ExchangeNode& exchangeNode,
                             CollectedInfo childResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.mergeNoDefs(std::move(refsResult));
         result.merge(std::move(childResult));
@@ -848,7 +858,7 @@ struct Collector {
                             const RootNode& rootNode,
                             CollectedInfo childResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.mergeNoDefs(std::move(refsResult));
         result.merge(std::move(childResult));
@@ -870,7 +880,7 @@ struct Collector {
                             CollectedInfo filterResult,
                             CollectedInfo bindResult,
                             CollectedInfo refsResult) {
-        CollectedInfo result{};
+        CollectedInfo result{this};
 
         result.merge(std::move(refsResult));
         result.merge(std::move(childResult));
@@ -895,42 +905,128 @@ struct Collector {
     }
 
     CollectedInfo transport(const ABT& n, const SpoolConsumerNode& node, CollectedInfo bindResult) {
-        return collectForScan(n, node, node.binder(), {});
+        return collectForScan(n, node, node.binder(), {this});
     }
 
     CollectedInfo collect(const ABT& n) {
         return algebra::transport<true>(n, *this);
     }
 
+    void resolveVariable(const Variable& var,
+                         const Definition& def,
+                         boost::dynamic_bitset<size_t>& varSet) {
+        allResolvedVariables->emplace(&var, std::make_pair(def, _lastVarId));
+        if (varSet.size() < _lastVarId + 1) {
+            varSet.resize(_lastVarId + 1);
+        }
+        varSet.set(_lastVarId, true);
+        ++_lastVarId;
+    }
+
+    std::unique_ptr<ResolvedVarsStorage> allResolvedVariables;
+    size_t _lastVarId = 0;
+
 private:
     const cascades::MemoGroupBinderInterface* _memoInterface;
 };
 
+void CollectedInfo::resolveFreeVars(const ProjectionName& name, const Definition& def) {
+    if (auto it = freeVars.find(name); it != freeVars.end()) {
+        for (const Variable& var : it->second) {
+            collector->resolveVariable(var, def, resolvedVariables);
+        }
+        freeVars.erase(it);
+    }
+}
+
+/*
+void CollectedInfo::mergeResolvedVars(boost::dynamic_bitset<size_t>& otherResolvedVars) {
+    if (otherResolvedVars.size() < resolvedVariables.size()) {
+        otherResolvedVars.resize(resolvedVariables.size());
+    }
+    if (resolvedVariables.size() < otherResolvedVars.size()) {
+        resolvedVariables.resize(otherResolvedVars.size());
+    }
+
+    auto newResolvedVars = resolvedVariables | otherResolvedVars;
+    auto newOtherResolvedVars = resolvedVariables & otherResolvedVars;
+
+    if (newOtherResolvedVars.any()) {
+        std::cout << "Duplicate variables:\n"
+                  << "resolvedVariables:" << resolvedVariables << std::endl
+                  << "otherResolvedVars:" << otherResolvedVars << std::endl
+                  << "newResolvedVars:" << newResolvedVars << std::endl
+                  << "newOtherResolvedVars:" << newOtherResolvedVars << std::endl
+                  << "collector->_lastVarId:" << collector->_lastVarId << std::endl;
+        std::cout << "allResolvedVariables:" << std::endl;
+        for (auto const& elem : *collector->allResolvedVariables) {
+            // std::cout << "{" << elem.first->name() << ": "
+            //           << "[" << elem.second.first << ", " << elem.second.second << "]}\n";
+            std::cout << "{" << elem.first->name() << ": " << elem.second.second << "}\n";
+        }
+    }
+
+    // otherResolvedVars &= resolvedVariables;
+    otherResolvedVars = newOtherResolvedVars;
+    resolvedVariables = newResolvedVars;
+}
+*/
+void CollectedInfo::mergeResolvedVars(boost::dynamic_bitset<size_t>& otherResolvedVars) {
+    if (otherResolvedVars.size() < resolvedVariables.size()) {
+        otherResolvedVars.resize(resolvedVariables.size());
+    }
+    if (resolvedVariables.size() < otherResolvedVars.size()) {
+        resolvedVariables.resize(otherResolvedVars.size());
+    }
+
+    if (false) {
+        // This version computes both the new this, and new other in order to be able to check that
+        // new other is empty. However, this approach involves constructing one more bitmap, and
+        // shows on the benchmark as a regression.
+        auto newResolvedVars = resolvedVariables | otherResolvedVars;
+        otherResolvedVars &= resolvedVariables;
+        resolvedVariables = newResolvedVars;
+        tassert(6624024, "Found a duplicate Variable pointer", otherResolvedVars.none());
+    } else {
+        resolvedVariables |= otherResolvedVars;
+    }
+}
+
+
 VariableEnvironment VariableEnvironment::build(
     const ABT& root, const cascades::MemoGroupBinderInterface* memoInterface) {
     Collector c(memoInterface);
     auto info = std::make_unique<CollectedInfo>(c.collect(root));
 
-    return VariableEnvironment{std::move(info), memoInterface};
+    return VariableEnvironment{std::move(info), memoInterface, std::move(c.allResolvedVariables)};
 }
 
 void VariableEnvironment::rebuild(const ABT& root) {
-    _info = std::make_unique<CollectedInfo>(Collector{_memoInterface}.collect(root));
+    Collector c(_memoInterface);
+    _info = std::make_unique<CollectedInfo>(c.collect(root));
+    _allResolvedVariables = std::move(c.allResolvedVariables);
 }
 
 VariableEnvironment::VariableEnvironment(std::unique_ptr<CollectedInfo> info,
-                                         const cascades::MemoGroupBinderInterface* memoInterface)
-    : _info(std::move(info)), _memoInterface(memoInterface) {}
+                                         const cascades::MemoGroupBinderInterface* memoInterface,
+                                         std::unique_ptr<ResolvedVarsStorage> allResolvedVars)
+    : _info(std::move(info)),
+      _memoInterface(memoInterface),
+      _allResolvedVariables(std::move(allResolvedVars)) {}
 
 VariableEnvironment::~VariableEnvironment() {}
 
 Definition VariableEnvironment::getDefinition(const Variable& var) const {
-    auto it = _info->useMap.find(&var);
-    if (it == _info->useMap.end()) {
+    auto it = _allResolvedVariables->find(&var);
+    if (it == _allResolvedVariables->end()) {
         return Definition();
     }
 
-    return it->second;
+    size_t varId = it->second.second;
+    if (_info->resolvedVariables.test(varId)) {
+        return it->second.first;
+    }
+    return Definition();
 }
 
 const DefinitionsMap& VariableEnvironment::getDefinitions(const Node& node) const {
diff --git a/src/mongo/db/query/optimizer/reference_tracker.h b/src/mongo/db/query/optimizer/reference_tracker.h
index c11ddce41f4..13dfd58e9b9 100644
--- a/src/mongo/db/query/optimizer/reference_tracker.h
+++ b/src/mongo/db/query/optimizer/reference_tracker.h
@@ -54,13 +54,15 @@ struct Definition {
 
 struct CollectedInfo;
 using DefinitionsMap = ProjectionNameMap<Definition>;
+using ResolvedVarsStorage = std::unordered_map<const Variable*, std::pair<Definition, size_t>>;
 
 /**
  * Helps enforce scoping and validity rules for definitions and Variable references.
  */
 class VariableEnvironment {
     VariableEnvironment(std::unique_ptr<CollectedInfo> info,
-                        const cascades::MemoGroupBinderInterface* memoInterface);
+                        const cascades::MemoGroupBinderInterface* memoInterface,
+                        std::unique_ptr<ResolvedVarsStorage> allResolvedVars);
 
 public:
     /**
@@ -132,6 +134,7 @@ private:
     // '_memoInterface' is required to track references in an ABT containing
     // MemoLogicalDelegatorNodes.
     const cascades::MemoGroupBinderInterface* _memoInterface{nullptr};
+    std::unique_ptr<ResolvedVarsStorage> _allResolvedVariables;
 };
 
 }  // namespace mongo::optimizer
