arangodb · jsteemann · May 18, 2020 · May 13, 2020 · May 13, 2020 · May 13, 2020
diff --git a/arangod/Aql/ClusterNodes.cpp b/arangod/Aql/ClusterNodes.cpp
@@ -142,8 +142,8 @@ RemoteNode::RemoteNode(ExecutionPlan* plan, arangodb::velocypack::Slice const& b
 /// @brief creates corresponding ExecutionBlock
 std::unique_ptr<ExecutionBlock> RemoteNode::createBlock(
     ExecutionEngine& engine, std::unordered_map<ExecutionNode*, ExecutionBlock*> const&) const {
-  RegisterId const nrOutRegs = getRegisterPlan()->nrRegs[getDepth()];
-  RegisterId const nrInRegs = nrOutRegs;
+  auto const nrOutRegs = getRegisterPlan()->nrRegs[getDepth()];
+  auto const nrInRegs = nrOutRegs;
 
   auto regsToKeep = getRegsToKeepStack();
   auto regsToClear = getRegsToClear();

diff --git a/arangod/Aql/ExecutionNode.cpp b/arangod/Aql/ExecutionNode.cpp
@@ -1315,6 +1315,9 @@ bool ExecutionNode::isIncreaseDepth(ExecutionNode::NodeType type) {
     case REMOTESINGLE:
     case ENUMERATE_IRESEARCH_VIEW:
     case MATERIALIZE:
+
+    case SUBQUERY_START:
+    case SUBQUERY_END:
       return true;
 
     default:

diff --git a/arangod/Aql/RegisterInfos.cpp b/arangod/Aql/RegisterInfos.cpp
@@ -61,11 +61,6 @@ RegisterInfos::RegisterInfos(RegIdSet readableInputRegisters, RegIdSet writeable
       _numOutRegs(nrOutputRegisters),
       _registersToKeep(std::move(registersToKeep)),
       _registersToClear(std::move(registersToClear)) {
-  // The second assert part is from ReturnExecutor special case, we shrink all
-  // results into a single Register column.
-  TRI_ASSERT((nrInputRegisters <= nrOutputRegisters) ||
-             (nrOutputRegisters == 1 && _registersToKeep.back().empty() &&
-              _registersToClear.empty()));
 
 #ifdef ARANGODB_ENABLE_MAINTAINER_MODE
   for (RegisterId const inReg : _inRegs) {
@@ -76,13 +71,14 @@ RegisterInfos::RegisterInfos(RegIdSet readableInputRegisters, RegIdSet writeable
   }
   for (RegisterId const regToClear : _registersToClear) {
     // sic: It's possible that a current output register is immediately cleared!
-    TRI_ASSERT(regToClear < nrOutputRegisters);
+    TRI_ASSERT(regToClear < nrInputRegisters);
     TRI_ASSERT(_registersToKeep.back().find(regToClear) ==
                _registersToKeep.back().end());
   }
   TRI_ASSERT(!_registersToKeep.empty());
   for (RegisterId const regToKeep : _registersToKeep.back()) {
     TRI_ASSERT(regToKeep < nrInputRegisters);
+    TRI_ASSERT(regToKeep < nrOutputRegisters);
     TRI_ASSERT(_registersToClear.find(regToKeep) == _registersToClear.end());
   }
 #endif

diff --git a/arangod/Aql/RegisterPlan.cpp b/arangod/Aql/RegisterPlan.cpp
@@ -153,6 +153,7 @@ void RegisterPlanWalkerT<T>::after(T* en) {
 
   switch (en->getType()) {
     case ExecutionNode::SUBQUERY_START: {
+      previousSubqueryNrRegs.emplace(plan->nrRegs.back());
       auto topUnused = unusedRegisters.back();
       unusedRegisters.emplace_back(std::move(topUnused));
 
@@ -168,6 +169,14 @@ void RegisterPlanWalkerT<T>::after(T* en) {
     case ExecutionNode::SUBQUERY_END: {
       unusedRegisters.pop_back();
       regsToKeepStack.pop_back();
+      // This must have added a section, otherwise we would decrease the
+      // number of registers available inside the subquery.
+      TRI_ASSERT(en->isIncreaseDepth());
+      // We must plan the registers after this, so newly added registers are
+      // based upon this nrRegs.
+      TRI_ASSERT(mayReuseRegisterImmediately);
+      plan->nrRegs.back() = previousSubqueryNrRegs.top();
+      previousSubqueryNrRegs.pop();
     } break;
     default: {
       auto regsToClear = calculateRegistersToClear(en);
@@ -197,7 +206,7 @@ void RegisterPlanWalkerT<T>::after(T* en) {
 }
 
 template <typename T>
-RegisterPlanT<T>::RegisterPlanT() : depth(0), totalNrRegs(0) {
+RegisterPlanT<T>::RegisterPlanT() : depth(0) {
   nrRegs.reserve(8);
   nrRegs.emplace_back(0);
 }
@@ -208,8 +217,7 @@ RegisterPlanT<T>::RegisterPlanT(RegisterPlan const& v, unsigned int newdepth)
     : varInfo(v.varInfo),
       nrRegs(v.nrRegs),
       subQueryNodes(),
-      depth(newdepth + 1),
-      totalNrRegs(v.nrRegs[newdepth]) {
+      depth(newdepth + 1) {
   if (depth + 1 < 8) {
     // do a minium initial allocation to avoid frequent reallocations
     nrRegs.reserve(8);
@@ -218,14 +226,13 @@ RegisterPlanT<T>::RegisterPlanT(RegisterPlan const& v, unsigned int newdepth)
   // this is required because back returns a reference and emplace/push_back may
   // invalidate all references
   nrRegs.resize(depth);
-  RegisterId registerId = nrRegs.back();
-  nrRegs.emplace_back(registerId);
+  auto regCount = nrRegs.back();
+  nrRegs.emplace_back(regCount);
 }
 
 template <typename T>
 RegisterPlanT<T>::RegisterPlanT(VPackSlice slice, unsigned int depth)
-    : depth(depth),
-      totalNrRegs(slice.get("totalNrRegs").getNumericValue<unsigned int>()) {
+    : depth(depth) {
   VPackSlice varInfoList = slice.get("varInfoList");
   if (!varInfoList.isArray()) {
     THROW_ARANGO_EXCEPTION_MESSAGE(
@@ -266,7 +273,6 @@ auto RegisterPlanT<T>::clone() -> std::shared_ptr<RegisterPlanT> {
 
   other->nrRegs = nrRegs;
   other->depth = depth;
-  other->totalNrRegs = totalNrRegs;
   other->varInfo = varInfo;
 
   // No need to clone subQueryNodes because this was only used during
@@ -281,14 +287,13 @@ void RegisterPlanT<T>::increaseDepth() {
   // create a copy of the last value here
   // this is required because back returns a reference and emplace/push_back
   // may invalidate all references
-  RegisterId registerId = nrRegs.back();
-  nrRegs.emplace_back(registerId);
+  auto regCount = nrRegs.back();
+  nrRegs.emplace_back(regCount);
 }
 
 template <typename T>
 RegisterId RegisterPlanT<T>::addRegister() {
-  nrRegs[depth]++;
-  return totalNrRegs++;
+  return static_cast<RegisterId>(nrRegs[depth]++);
 }
 
 template <typename T>
@@ -355,19 +360,16 @@ void RegisterPlanT<T>::toVelocyPack(VPackBuilder& builder) const {
   builder.add(VPackValue("nrRegsHere"));
   { VPackArrayBuilder guard(&builder); }
 
-  builder.add("totalNrRegs", VPackValue(totalNrRegs));
+  // totalNrRegs is not used anymore and is intentionally left empty
+  // can be removed in ArangoDB 3.8
+  builder.add("totalNrRegs", VPackSlice::noneSlice());
 }
 
 template <typename T>
 void RegisterPlanT<T>::addSubqueryNode(T* subquery) {
   subQueryNodes.emplace_back(subquery);
 }
 
-template <typename T>
-auto RegisterPlanT<T>::getTotalNrRegs() -> unsigned int {
-  return totalNrRegs;
-}
-
 template <typename T>
 auto RegisterPlanT<T>::calcRegsToKeep(VarSetStack const& varsUsedLaterStack,
                                       VarSetStack const& varsValidStack,

diff --git a/arangod/Aql/RegisterPlan.h b/arangod/Aql/RegisterPlan.h
@@ -31,6 +31,7 @@
 #include "Basics/Common.h"
 
 #include <memory>
+#include <stack>
 #include <unordered_map>
 #include <vector>
 
@@ -70,9 +71,12 @@ struct RegisterPlanWalkerT final : public WalkerWorker<T> {
     return false;  // do not walk into subquery
   }
 
+  using RegCountStack = std::stack<RegisterCount>;
+
   RegIdOrderedSetStack unusedRegisters{{}};
   RegIdSetStack regsToKeepStack{{}};
   std::shared_ptr<RegisterPlanT<T>> plan;
+  RegCountStack previousSubqueryNrRegs{};
 };
 
 template <typename T>
@@ -88,7 +92,7 @@ struct RegisterPlanT final : public std::enable_shared_from_this<RegisterPlanT<T
   // the entry with index i here is always the sum of all values
   // in nrRegsHere from index 0 to i (inclusively) and the two
   // have the same length:
-  std::vector<RegisterId> nrRegs;
+  std::vector<RegisterCount> nrRegs;
 
   // We collect the subquery nodes to deal with them at the end:
   std::vector<T*> subQueryNodes;
@@ -110,7 +114,6 @@ struct RegisterPlanT final : public std::enable_shared_from_this<RegisterPlanT<T
   void increaseDepth();
   auto addRegister() -> RegisterId;
   void addSubqueryNode(T* subquery);
-  auto getTotalNrRegs() -> unsigned int;
 
   void toVelocyPack(arangodb::velocypack::Builder& builder) const;
   static void toVelocyPackEmpty(arangodb::velocypack::Builder& builder);
@@ -123,7 +126,6 @@ struct RegisterPlanT final : public std::enable_shared_from_this<RegisterPlanT<T
 
  private:
   unsigned int depth;
-  unsigned int totalNrRegs;
 };
 
 template <typename T>