[pseudo] WIP: GSS Node refcount (smart pointers)

[sammccall]

Summary

Here is the version where ownership is managed by smart-pointers.
It uses some thread_local magic but it's not hairy, and it seems much easier to
reason about to me.

clang-tools-extra/pseudo/include/clang-pseudo/GLR.h

 // The structure is a DAG rather than a linear stack:
 // - GLR allows multiple actions (conflicts) on the same head, producing forks
 //   where several nodes have the same parent
 // - The parser merges nodes with the same (state, ForestNode), producing joins
 //   where one node has multiple parents
 //
 // The parser is responsible for creating nodes and keeping track of the set of
 // heads. The GSS class is mostly an arena for them.
+//
+// GSS is bound to a single thread, and there can only be one GSS per thread.
+// (Thread-local state is used to implement NodePtr without extra storage).
 struct GSS {
+private:
+  static thread_local GSS *Current;
+
+public:
+  GSS();
   ~GSS();
+
+  struct Node;
+  // A smart-pointer for GSS nodes, which are refcounted.
+  class NodePtr {
+  public:
+    using element_type = const Node;
+    ~NodePtr() {
+      if (P)
+        Current->unref(P);
+    }
+    NodePtr() {}
+    explicit NodePtr(Node *N) : P(N) { Current->ref(P); }
+    NodePtr(const NodePtr &O) : P(O.P) { Current->ref(P); }
+    NodePtr(NodePtr &&O) { std::swap(O.P, P); }
+    NodePtr &operator=(const NodePtr &O) {
+      P = O.P;
+      Current->ref(P);
+      return *this;
+    }
+    NodePtr &operator=(NodePtr &&O) {
+      std::swap(O.P, P);
+      return *this;
+    }
+
+    const Node *get() const { return P; }
+    const Node &operator*() const { return *P; }
+    const Node *operator->() const { return P; }
+    friend bool operator==(const NodePtr &A, const NodePtr &B) {
+      return A.P == B.P;
+    }
+    friend bool operator<(const NodePtr &A, const NodePtr &B) {
+      return A.P < B.P;
+    }
+    friend llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
+                                         const NodePtr &N) {
+      return OS << N.P;
+    }
+
+  private:
+    Node *P = nullptr;
+  };
+
   // A node represents a partial parse of the input up to some point.
   //
   // It is the equivalent of a frame in an LR parse stack.
   // Like such a frame, it has an LR parse state and a syntax-tree node
   // representing the last parsed symbol (a ForestNode in our case).
   // Unlike a regular LR stack frame, it may have multiple parents.
   //
   // Nodes are not exactly pushed and popped on the stack: pushing is just
   // allocating a new head node with a parent pointer to the old head. Popping
   // is just forgetting about a node and remembering its parent instead.
-  struct alignas(struct Node *) Node {
+  struct alignas(class NodePtr) Node {
     // LR state describing how parsing should continue from this head.
     LRTable::StateID State;
     // Number of the parents of this node.
     // The parents hold previous parsed symbols, and may resume control after
     // this node is reduced.
     uint16_t ParentCount;
     // Number of references to this node.
     // Typically this is 1 if the node is a head, plus the number of children.
     uint16_t RefCount;
     // The parse node for the last parsed symbol.
     // This symbol appears on the left of the dot in the parse state's items.
     // (In the literature, the node is attached to the *edge* to the parent).
     const ForestNode *Payload = nullptr;
 
-    llvm::ArrayRef<const Node *> parents() const {
-      return llvm::makeArrayRef(reinterpret_cast<const Node *const *>(this + 1),
-                                ParentCount);
+    llvm::ArrayRef<NodePtr> parents() const {
+      return llvm::makeArrayRef(
+          reinterpret_cast<NodePtr *>(const_cast<Node *>(this) + 1),
     };
-    // Parents are stored as a trailing array of Node*.
+    // Parents are stored as a trailing array of NodePtr.
   };
 
   // Allocates a new node in the graph.
   // It has a refcount of 1, and should be pased to deref() when done.
-  const Node *addNode(LRTable::StateID State, const ForestNode *Symbol,
-                      llvm::ArrayRef<const Node *> Parents);
-  void ref(const Node *N);
-  void unref(const Node *N);
-
+  NodePtr addNode(LRTable::StateID State, const ForestNode *Symbol,
+                  llvm::ArrayRef<NodePtr> Parents);
   size_t bytes() const { return Arena.getTotalMemory() + sizeof(*this); }
   size_t nodeCount() const { return NodesCreated; }
 
 private:
   // Nodes are recycled using freelists.
   // They are variable size, so use one free-list per distinct #parents.
   std::vector<std::vector<Node *>> FreeList;
 
+  void ref(Node *N);
+  void unref(Node *N);
   Node *allocate(unsigned Parents);
   void destroy(Node *N);
 
   llvm::BumpPtrAllocator Arena;
   unsigned NodesCreated = 0;
   unsigned NodesDestroyed = 0;
   unsigned NodesAllocated = 0;
 };

 const ForestNode &glrParse(const TokenStream &Code, const ParseParams &Params,
                            SymbolID StartSymbol);
 
 // An active stack head can have multiple available actions (reduce/reduce
 // actions, reduce/shift actions).
 // A step is any one action applied to any one stack head.
 struct ParseStep {
   // A specific stack head.
-  const GSS::Node *Head = nullptr;
+  GSS::NodePtr Head;
   // An action associated with the head.
   LRTable::Action Action = LRTable::Action::sentinel();
 };
 // A callback is invoked whenever a new GSS head is created during the GLR
 // parsing process (glrShift, or glrReduce).
-using NewHeadCallback = std::function<void(const GSS::Node *)>;
+using NewHeadCallback = std::function<void(GSS::NodePtr)>;
 // Apply all PendingShift actions on a given GSS state, newly-created heads are
 // passed to the callback.
 //
 // When this function returns, PendingShift is empty.
 //
 // Exposed for testing only.
 void glrShift(std::vector<ParseStep> &PendingShift, const ForestNode &NextTok,
               const ParseParams &Params, NewHeadCallback NewHeadCB);

clang-tools-extra/pseudo/lib/GLR.cpp

 
 namespace clang {
 namespace pseudo {
 
 using StateID = LRTable::StateID;
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const GSS::Node &N) {
   std::vector<std::string> ParentStates;
-  for (const auto *Parent : N.parents())
+  for (const auto &Parent : N.parents())
     ParentStates.push_back(llvm::formatv("{0}", Parent->State));
   OS << llvm::formatv("state {0}, parsed symbol {1}, parents {{{2}}, refcount={3}",
                       N.State, N.Payload ? N.Payload->symbol() : 0,
                       llvm::join(ParentStates, ", "), N.RefCount);
   return OS;
 }
 
 #ifndef NDEBUG
 static void checkRefcounts(std::vector<const GSS::Node *> Queue) {
   llvm::DenseSet<const GSS::Node *> Seen;
   llvm::DenseMap<const GSS::Node *, unsigned> Expected;
   for (const auto *N : Queue)
     ++Expected[N];
   while (!Queue.empty()) {
     const GSS::Node *N = Queue.back();
     Queue.pop_back();
     if (!Seen.insert(N).second)
       continue;
-    for (const auto *P : N->parents()) {
-      ++Expected[P];
-      Queue.push_back(P);
+    for (const auto &P : N->parents()) {
+      ++Expected[P.get()];
+      Queue.push_back(P.get());
     }
   }
   bool AnyWrong = false;
   llvm::dbgs() << "Checking refcounts:\n";
   for (const auto& E : Expected) {
     llvm::dbgs() << E.first << " " << *E.first << "\n";
     if (E.second != E.first->RefCount) {
       llvm::dbgs() << "Wrong refcount: actual=" << E.second << " for "
const ForestNode &glrParse(const TokenStream &Tokens, const ParseParams &Params,
                            SymbolID StartSymbol) {
   llvm::ArrayRef<ForestNode> Terminals = Params.Forest.createTerminals(Tokens);
   auto &G = Params.G;
   (void)G;
   auto &GSS = Params.GSStack;
 
   // Lists of active shift, reduce, accept actions.
   std::vector<ParseStep> PendingShift, PendingReduce, PendingAccept;
-  auto AddSteps = [&](const GSS::Node *Head, SymbolID NextTok) {
+  auto AddSteps = [&](const GSS::NodePtr &Head, SymbolID NextTok) {
     for (const auto &Action : Params.Table.getActions(Head->State, NextTok)) {
       switch (Action.kind()) {
       case LRTable::Action::Shift:
-        GSS.ref(Head);
         PendingShift.push_back({Head, Action});
         break;
       case LRTable::Action::Reduce:
-        GSS.ref(Head);
         PendingReduce.push_back({Head, Action});
         break;
       case LRTable::Action::Accept:
-        GSS.ref(Head);
         PendingAccept.push_back({Head, Action});
         break;
       default:
         llvm_unreachable("unexpected action kind!");
       }
     }
   };
-  std::vector<const GSS::Node *> NewHeads = {
+  std::vector<GSS::NodePtr> NewHeads = {
       GSS.addNode(/*State=*/Params.Table.getStartState(StartSymbol),
                   /*ForestNode=*/nullptr, {})};
   auto CheckRefcounts = [&] {
     DEBUG_WITH_TYPE("GSSNode", {
-      std::vector<const GSS::Node *> Roots = NewHeads;
+      std::vector<const GSS::Node *> Roots;
+      for (const auto &H : NewHeads)
+        Roots.push_back(H.get());
       for (const auto *Pending :
            {&PendingShift, &PendingReduce, &PendingAccept})
         for (const ParseStep &Step : *Pending)
-          Roots.push_back(Step.Head);
+          Roots.push_back(Step.Head.get());
       checkRefcounts(Roots);
     });
   };
   for (const ForestNode &Terminal : Terminals) {
     LLVM_DEBUG(llvm::dbgs() << llvm::formatv("Next token {0} (id={1})\n",
                                              G.symbolName(Terminal.symbol()),
                                              Terminal.symbol()));
     CheckRefcounts();
-    for (const auto *Head : NewHeads) {
+    for (const auto &Head : NewHeads)
       AddSteps(Head, Terminal.symbol());
-      GSS.unref(Head);
-    }
     NewHeads.clear();
     CheckRefcounts();
     LLVM_DEBUG(llvm::dbgs() << "  Reduce\n");
-    glrReduce(PendingReduce, Params,
-              [&](const GSS::Node * NewHead) {
-                // A reduce will enable more steps.
-                AddSteps(NewHead, Terminal.symbol());
-              });
+    glrReduce(PendingReduce, Params, [&](GSS::NodePtr NewHead) {
     CheckRefcounts();
 
-    glrShift(PendingShift, Terminal, Params, [&](const GSS::Node *NewHead) {
-      GSS.ref(NewHead);
-      NewHeads.push_back(NewHead);
+    glrShift(PendingShift, Terminal, Params, [&](GSS::NodePtr NewHead) {
+      NewHeads.push_back(std::move(NewHead));
     });
   }
   LLVM_DEBUG(llvm::dbgs() << llvm::formatv("Next is eof\n"));
   CheckRefcounts();
-  for (const auto *Head : NewHeads) {
+  for (const auto &Head : NewHeads)
     AddSteps(Head, tokenSymbol(tok::eof));
-    GSS.unref(Head);
-  }
   NewHeads.clear();
   CheckRefcounts();
-  glrReduce(PendingReduce, Params,
-            [&](const GSS::Node * NewHead) {
-              // A reduce will enable more steps.
-              AddSteps(NewHead, tokenSymbol(tok::eof));
-            });
+  glrReduce(PendingReduce, Params, [&](GSS::NodePtr NewHead) {
   CheckRefcounts();
 
   assert(PendingShift.empty());
   assert(PendingReduce.empty());
   if (!PendingAccept.empty()) {
     LLVM_DEBUG({
       llvm::dbgs() << llvm::formatv("Accept: {0} accepted result:\n",
                                              PendingAccept.size());
       for (const auto &Accept : PendingAccept)
         llvm::dbgs() << "  - " << G.symbolName(Accept.Head->Payload->symbol())
                      << "\n";
     });
     assert(PendingAccept.size() == 1);
-    auto *Result = PendingAccept.front().Head->Payload;
-    GSS.unref(PendingAccept.front().Head);
-    return *Result;
+    return *PendingAccept.front().Head->Payload;
   }
   // We failed to parse the input, returning an opaque forest node for recovery.
   return Params.Forest.createOpaque(StartSymbol, /*Token::Index=*/0);
 }
 
 // Apply all pending shift actions.
 // In theory, LR parsing doesn't have shift/shift conflicts on a single head.
 // But we may have multiple active heads, and each head has a shift action.
 //
 // We merge the stack -- if multiple heads will reach the same state after
 // shifting a token, we shift only once by combining these heads.
 //
 // E.g. we have two heads (2, 3) in the GSS, and will shift both to reach 4:
 //   0---1---2
 //       └---3
 // After the shift action, the GSS is:
 //   0---1---2---4
 //       └---3---┘
 void glrShift(std::vector<ParseStep> &PendingShift, const ForestNode &NewTok,
               const ParseParams &Params, NewHeadCallback NewHeadCB) {
-  auto &GSS = Params.GSStack;
   assert(NewTok.kind() == ForestNode::Terminal);
   assert(llvm::all_of(PendingShift,
                       [](const ParseStep &Step) {
                         return Step.Action.kind() == LRTable::Action::Shift;
                       }) &&
          "Pending shift actions must be shift actions");
   LLVM_DEBUG(llvm::dbgs() << llvm::formatv("  Shift {0} ({1} active heads):\n",
                                            Params.G.symbolName(NewTok.symbol()),
                                            PendingShift.size()));
 
   // We group pending shifts by their target state so we can merge them.
   llvm::stable_sort(PendingShift, [](const ParseStep &L, const ParseStep &R) {
     return L.Action.getShiftState() < R.Action.getShiftState();
   });
   auto Rest = llvm::makeArrayRef(PendingShift);
-  llvm::SmallVector<const GSS::Node *> Parents;
+  llvm::SmallVector<GSS::NodePtr> Parents;
   while (!Rest.empty()) {
     // Collect the batch of PendingShift that have compatible shift states.
     // Their heads become TempParents, the parents of the new GSS node.
     StateID NextState = Rest.front().Action.getShiftState();
 
     Parents.clear();
     for (const auto &Base : Rest) {
       if (Base.Action.getShiftState() != NextState)
         break;
-      Parents.push_back(Base.Head); // safe borrow
+      Parents.push_back(Base.Head);
     }
     Rest = Rest.drop_front(Parents.size());
 
     LLVM_DEBUG(llvm::dbgs() << llvm::formatv("    --> S{0} ({1} heads)\n",
                                              NextState, Parents.size()));
-    const auto *NewHead = Params.GSStack.addNode(NextState, &NewTok, Parents);
-    NewHeadCB(NewHead);
-    GSS.unref(NewHead);
+    NewHeadCB(Params.GSStack.addNode(NextState, &NewTok, Parents));
   }
-  for (const auto & PS : PendingShift)
-    GSS.unref(PS.Head);
   PendingShift.clear();
 }
 
 namespace {
 // A KeyedQueue yields pairs of keys and values in order of the keys.
 template <typename Key, typename Value>
 using KeyedQueue =
     std::priority_queue<std::pair<Key, Value>,

 //   Before (joining due to same goto state, the same base):
 //     0--1(class-name)--3(STAR)
 //     └--2(enum-name)--4(STAR)
 //   After reducing 3 by `pointer := class-name STAR` and
 //                  2 by`enum-name := class-name STAR`:
 //     0--5(pointer)       // 5 is goto(0, pointer)
 void glrReduce(std::vector<ParseStep> &PendingReduce, const ParseParams &Params,
                NewHeadCallback NewHeadCB) {
-  auto &GSS = Params.GSStack;
   // There are two interacting complications:
   // 1.  Performing one reduce can unlock new reduces on the newly-created head.
   // 2a. The ambiguous ForestNodes must be complete (have all sequence nodes).
   //     This means we must have unlocked all the reduces that contribute to it.
   // 2b. Similarly, the new GSS nodes must be complete (have all parents).
   //
   // We define a "family" of reduces as those that produce the same symbol and
   // cover the same range of tokens. These are exactly the set of reductions
struct Family {
     }
   };
 
   // The base nodes are the heads after popping the GSS nodes we are reducing.
   // We don't care which rule yielded each base. If Family.Symbol is S, the
   // base includes an item X := ... • S ... and since the grammar is
   // context-free, *all* parses of S are valid here.
   // FIXME: reuse the queues across calls instead of reallocating.
-  KeyedQueue<Family, const GSS::Node *> Bases;
+  KeyedQueue<Family, GSS::NodePtr> Bases;
 
   // A sequence is the ForestNode payloads of the GSS nodes we are reducing.
   // These are the RHS of the rule, the RuleID is stored in the Family.
   // They specify a sequence ForestNode we may build (but we dedup first).
   using Sequence = llvm::SmallVector<const ForestNode *, Rule::MaxElements>;
   KeyedQueue<Family, Sequence> Sequences;
 
   Sequence TempSequence;
   // Pop walks up the parent chain(s) for a reduction from Head by to Rule.
   // Once we reach the end, record the bases and sequences.
-  auto Pop = [&](const GSS::Node *Head, RuleID RID) {
+  auto Pop = [&](const GSS::NodePtr &Head, RuleID RID) {
     LLVM_DEBUG(llvm::dbgs() << "  Pop " << Params.G.dumpRule(RID) << "\n");
     const auto &Rule = Params.G.lookupRule(RID);
     Family F{/*Start=*/0, /*Symbol=*/Rule.Target, /*Rule=*/RID};
     TempSequence.resize_for_overwrite(Rule.Size);
-    auto DFS = [&](const GSS::Node *N, unsigned I, auto &DFS) {
+    auto DFS = [&](const GSS::NodePtr &N, unsigned I, auto &DFS) {
       if (I == Rule.Size) {
         F.Start = TempSequence.front()->startTokenIndex();
-        GSS.ref(N);
         Bases.emplace(F, N);
         LLVM_DEBUG(llvm::dbgs() << "    --> base at S" << N->State << "\n");
         Sequences.emplace(F, TempSequence);
         return;
       }
       TempSequence[Rule.Size - 1 - I] = N->Payload;
-      for (const GSS::Node *Parent : N->parents())
+      for (const auto &Parent : N->parents())
         DFS(Parent, I + 1, DFS);
     };
     DFS(Head, 0, DFS);
   };
   auto PopPending = [&] {
-    for (const ParseStep &Pending : PendingReduce) {
+    for (const ParseStep &Pending : PendingReduce)
       Pop(Pending.Head, Pending.Action.getReduceRule());
-      GSS.unref(Pending.Head);
-    }
     PendingReduce.clear();
   };
 
-  std::vector<std::pair</*Goto*/ StateID, const GSS::Node *>> FamilyBases;
+  std::vector<std::pair</*Goto*/ StateID, GSS::NodePtr>> FamilyBases;
   std::vector<std::pair<RuleID, Sequence>> FamilySequences;
 
-  std::vector<const GSS::Node *> TempGSSNodes;
+  std::vector<GSS::NodePtr> TempGSSNodes;
   std::vector<const ForestNode *> TempForestNodes;
 
   // Main reduction loop:
   //  - pop as much as we can
   //  - process one family at a time, forming a forest node
   //  - produces new GSS heads which may enable more pops
   PopPending();
   while (!Bases.empty()) {
const ForestNode *Parsed =
         SequenceNodes.size() == 1
             ? SequenceNodes.front()
             : &Params.Forest.createAmbiguous(F.Symbol, SequenceNodes);
     LLVM_DEBUG(llvm::dbgs() << "    --> " << Parsed->dump(Params.G) << "\n");
 
     // Grab the bases for this family.
     // As well as deduplicating them, we'll group by the goto state.
     do {
-      FamilyBases.emplace_back(
-          Params.Table.getGoToState(Bases.top().second->State, F.Symbol),
-          Bases.top().second); // transfer ownership
+      StateID State =
+          Params.Table.getGoToState(Bases.top().second->State, F.Symbol);
+      FamilyBases.emplace_back(State, std::move(Bases.top().second));
       Bases.pop();
     } while (!Bases.empty() && Bases.top().first == F);
-    llvm::sort(FamilyBases);
-    // XXX sortAndUnique is not aware of ref/unref, copy for simplicity.
-    std::vector<decltype(FamilyBases)::value_type> UniqueBases;
-    std::unique_copy(FamilyBases.begin(), FamilyBases.end(),
-                     std::back_inserter(UniqueBases));
+    sortAndUnique(FamilyBases);
     // Create a GSS node for each unique goto state.
-    llvm::ArrayRef<decltype(FamilyBases)::value_type> BasesLeft = UniqueBases;
+    llvm::ArrayRef<decltype(FamilyBases)::value_type> BasesLeft = FamilyBases;
     while (!BasesLeft.empty()) {
       StateID NextState = BasesLeft.front().first;
       auto &Parents = TempGSSNodes;
       Parents.clear();
       for (const auto &Base : BasesLeft) {
         if (Base.first != NextState)
           break;
-        Parents.push_back(Base.second); // safe borrow
+        Parents.push_back(Base.second);
       }
       BasesLeft = BasesLeft.drop_front(Parents.size());
 
       // Invoking the callback for new heads, a real GLR parser may add new
       // reduces to the PendingReduce queue!
-      const auto *NewNode = Params.GSStack.addNode(NextState, Parsed, Parents);
-      NewHeadCB(NewNode);
-      GSS.unref(NewNode);
+      NewHeadCB(Params.GSStack.addNode(NextState, Parsed, Parents));
     }
-    for (auto &FB : FamilyBases)
-      GSS.unref(FB.second);
     FamilyBases.clear();
     PopPending();
   }
   assert(Sequences.empty());
 }
 
 static void logLifetime(llvm::StringRef Sigil, const GSS::Node* N) {
   DEBUG_WITH_TYPE("GSSNode", llvm::dbgs()
                                  << Sigil << " " << N << " " << *N << "\n");
 }
 
-const GSS::Node *GSS::addNode(LRTable::StateID State, const ForestNode *Symbol,
-                              llvm::ArrayRef<const Node *> Parents) {
+GSS::NodePtr GSS::addNode(LRTable::StateID State, const ForestNode *Symbol,
+                          llvm::ArrayRef<NodePtr> Parents) {
   ++NodesCreated;
   Node *Result = new (allocate(Parents.size()))
       Node({State, static_cast<uint16_t>(Parents.size()), /*RefCount=*/0});
   Result->Payload = Symbol;
-  auto *ParentArray = reinterpret_cast<const Node **>(Result + 1);
-  for (const auto *Parent : Parents) {
-    ref(Parent);
-    *ParentArray++ = Parent;
-  }
+  auto *ParentArray = reinterpret_cast<NodePtr *>(Result + 1);
+  for (const auto &Parent : Parents)
+    new (ParentArray++) NodePtr(Parent);
   logLifetime("*", Result);
-  ref(Result);
-  return Result;
+  return NodePtr(Result);
 }
 
-void GSS::ref(const Node *N) {
-  const_cast<Node *>(N)->RefCount++;
+void GSS::ref(Node *N) {
+  N->RefCount++;
   logLifetime("+", N);
 }
 
-void GSS::unref(const Node *N) {
+void GSS::unref(Node *N) {
   assert(N->RefCount != 0);
-  Node *NN = const_cast<Node *>(N); // I created you, and I can destroy you!
-  --NN->RefCount;
+  --N->RefCount;
   logLifetime("-", N);
-  // Node is const from the caller's perspective, but we created it
-  if (0 == NN->RefCount)
-    destroy(NN);
+  if (0 == N->RefCount)
+    destroy(N);
 }
 
 GSS::Node *GSS::allocate(unsigned Parents) {
   if (FreeList.size() <= Parents)
     FreeList.resize(Parents + 1);
   auto &SizedList = FreeList[Parents];
   if (!SizedList.empty()) {
     auto *Result = SizedList.back();
     SizedList.pop_back();
     return Result;
   }
   ++NodesAllocated;
   return static_cast<Node *>(
-      Arena.Allocate(sizeof(Node) + Parents * sizeof(Node *), alignof(Node)));
+      Arena.Allocate(sizeof(Node) + Parents * sizeof(NodePtr), alignof(Node)));
 }
 
 void GSS::destroy(Node *N) {
   logLifetime("~", N);
   ++NodesDestroyed;
-  for (auto *Parent : N->parents())
-    unref(Parent);
   unsigned ParentCount = N->ParentCount;
+  for (auto &Parent : N->parents())
+    Parent.~NodePtr();
   N->~Node();
   assert(FreeList.size() > ParentCount && "established on construction!");
   FreeList[ParentCount].push_back(N);
 }
 
-GSS::~GSS() { assert(NodesCreated == NodesDestroyed); }
+GSS::GSS() {
+  assert(Current == nullptr);
+  Current = this;
+}
+
+GSS::~GSS() {
+  assert(Current == this);
+  assert(NodesCreated == NodesDestroyed);
+  Current = nullptr;
+}
+
+thread_local GSS *GSS::Current = nullptr;
 
 } // namespace pseudo
 } // namespace clang

clang-tools-extra/pseudo/unittests/GLRTest.cpp

 
 using Action = LRTable::Action;
 using testing::AllOf;
 
 MATCHER_P(state, StateID, "") { return arg->State == StateID; }
 MATCHER_P(parsedSymbol, FNode, "") { return arg->Payload == FNode; }
 MATCHER_P(parsedSymbolID, SID, "") { return arg->Payload->symbol() == SID; }
 
-testing::Matcher<const GSS::Node *>
-parents(llvm::ArrayRef<const GSS::Node *> Parents) {
-  return testing::Property(&GSS::Node::parents,
-                           testing::UnorderedElementsAreArray(Parents));
+testing::Matcher<const GSS::NodePtr &>
+parents(llvm::ArrayRef<GSS::NodePtr> Parents) {
+  return testing::Pointee(testing::Property(
+      &GSS::Node::parents, testing::UnorderedElementsAreArray(Parents)));
 }
 
 class GLRTest : public ::testing::Test {
 public:
   void build(llvm::StringRef GrammarBNF) {
     std::vector<std::string> Diags;
     G = Grammar::parseBNF(GrammarBNF, Diags);
   }
if (RuleRange.End - RuleRange.Start == 1)
       return G->table().Nonterminals[id(NonterminalName)].RuleRange.Start;
     ADD_FAILURE() << "Expected a single rule for " << NonterminalName
                   << ", but it has " << RuleRange.End - RuleRange.Start
                   << " rule!\n";
     return 0;
   }
 
   NewHeadCallback captureNewHeads() {
-    return [this](const GSS::Node *NewHead) {
-      GSStack.ref(NewHead);
-      NewHeadResults.push_back(NewHead);
-    };
+    return [this](GSS::NodePtr P) { NewHeadResults.push_back(std::move(P)); };
   };
 
 protected:
-  ~GLRTest() {
-    for (auto *N : NewHeadResults)
-      GSStack.unref(N);
-  }
-
   std::unique_ptr<Grammar> G;
   ForestArena Arena;
   GSS GSStack;
-  std::vector<const GSS::Node*> NewHeadResults;
+  std::vector<GSS::NodePtr> NewHeadResults;
 };
 
 TEST_F(GLRTest, ShiftMergingHeads) {
   // Given a test case where we have two heads 1, 2, 3 in the GSS, the heads 1,
   // 2 have shift actions to reach state 4, and the head 3 has a shift action to
   // reach state 5:
   //   0--1
   //   └--2
   //   └--3
   // After the shift action, the GSS (with new heads 4, 5) is:
   //   0---1---4
   //   └---2---┘
   //   └---3---5
-  auto *GSSNode0 =
+  auto GSSNode0 =
       GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr, /*Parents=*/{});
-  auto *GSSNode1 = GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr,
+  auto GSSNode1 = GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr,
                                    /*Parents=*/{GSSNode0});
-  auto *GSSNode2 = GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr,
+  auto GSSNode2 = GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr,
                                    /*Parents=*/{GSSNode0});
-  auto *GSSNode3 = GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr,
-                                   /*Parents=*/{GSSNode0});
-  GSStack.unref(GSSNode0);
+  auto GSSNode3 = GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr,
 
   buildGrammar({}, {}); // Create a fake empty grammar.
   LRTable T = LRTable::buildForTests(G->table(), /*Entries=*/{});
 
   ForestNode &SemiTerminal = Arena.createTerminal(tok::semi, 0);
   std::vector<ParseStep> PendingShift = {
       {GSSNode1, Action::shift(4)},
       {GSSNode3, Action::shift(5)},
       {GSSNode2, Action::shift(4)},
   };
   glrShift(PendingShift, SemiTerminal, {*G, T, Arena, GSStack},
            captureNewHeads());
 
   EXPECT_THAT(NewHeadResults, testing::UnorderedElementsAre(
                                   AllOf(state(4), parsedSymbol(&SemiTerminal),
                                         parents({GSSNode1, GSSNode2})),
                                   AllOf(state(5), parsedSymbol(&SemiTerminal),
-                                        parents({GSSNode3}))))
-      << NewHeadResults;
+                                        parents({GSSNode3}))));
 }
 
 TEST_F(GLRTest, ReduceConflictsSplitting) {
   //  Before (splitting due to R/R conflict):
   //    0--1(IDENTIFIER)
   //  After reducing 1 by `class-name := IDENTIFIER` and
   //                      `enum-name := IDENTIFIER`:
   //    0--2(class-name)    // 2 is goto(0, class-name)
   //    └--3(enum-name)     // 3 is goto(0, enum-name)
   buildGrammar({"class-name", "enum-name"},
                {"class-name := IDENTIFIER", "enum-name := IDENTIFIER"});
 
   LRTable Table = LRTable::buildForTests(
       G->table(), {{/*State=*/0, id("class-name"), Action::goTo(2)},
                    {/*State=*/0, id("enum-name"), Action::goTo(3)}});
 
-  const auto *GSSNode0 =
+  auto GSSNode0 =
       GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr, /*Parents=*/{});
-  const auto *GSSNode1 =
+  auto GSSNode1 =
       GSStack.addNode(3, &Arena.createTerminal(tok::identifier, 0), {GSSNode0});
-  GSStack.unref(GSSNode0);
-  GSStack.ref(GSSNode1);
 
   std::vector<ParseStep> PendingReduce = {
       {GSSNode1, Action::reduce(ruleFor("class-name"))},
       {GSSNode1, Action::reduce(ruleFor("enum-name"))}};
-  glrReduce(PendingReduce, {*G, Table, Arena, GSStack},
-            captureNewHeads());
+  glrReduce(PendingReduce, {*G, Table, Arena, GSStack}, captureNewHeads());
   EXPECT_THAT(NewHeadResults,
               testing::UnorderedElementsAre(
                   AllOf(state(2), parsedSymbolID(id("class-name")),
                         parents({GSSNode0})),
                   AllOf(state(3), parsedSymbolID(id("enum-name")),
-                        parents({GSSNode0})))) << NewHeadResults;
+                        parents({GSSNode0}))));
 }
 
 TEST_F(GLRTest, ReduceSplittingDueToMultipleBases) {
   //  Before (splitting due to multiple bases):
   //    2(class-name)--4(*)
   //    3(enum-name)---┘
   //  After reducing 4 by `ptr-operator := *`:
   //    2(class-name)--5(ptr-operator)    // 5 is goto(2, ptr-operator)
   //    3(enum-name)---6(ptr-operator)    // 6 is goto(3, ptr-operator)
   buildGrammar({"ptr-operator", "class-name", "enum-name"},
                {"ptr-operator := *"});
 
   auto *ClassNameNode = &Arena.createOpaque(id("class-name"), /*TokenIndex=*/0);
   auto *EnumNameNode = &Arena.createOpaque(id("enum-name"), /*TokenIndex=*/0);
 
-  const auto *GSSNode2 =
+  auto GSSNode2 =
       GSStack.addNode(/*State=*/2, /*ForestNode=*/ClassNameNode, /*Parents=*/{});
-  const auto *GSSNode3 =
+  auto GSSNode3 =
       GSStack.addNode(/*State=*/3, /*ForestNode=*/EnumNameNode, /*Parents=*/{});
-  const auto *GSSNode4 = GSStack.addNode(
+  auto GSSNode4 = GSStack.addNode(
       /*State=*/4, &Arena.createTerminal(tok::star, /*TokenIndex=*/1),
       /*Parents=*/{GSSNode2, GSSNode3});
-  GSStack.unref(GSSNode2);
-  GSStack.unref(GSSNode3);
 
   LRTable Table = LRTable::buildForTests(
       G->table(),
       {{/*State=*/2, id("ptr-operator"), Action::goTo(/*NextState=*/5)},
        {/*State=*/3, id("ptr-operator"), Action::goTo(/*NextState=*/6)}});
   std::vector<ParseStep> PendingReduce = {
       {GSSNode4, Action::reduce(ruleFor("ptr-operator"))}};
   glrReduce(PendingReduce, {*G, Table, Arena, GSStack},
             captureNewHeads());
 
   EXPECT_THAT(NewHeadResults,
               testing::UnorderedElementsAre(
                   AllOf(state(5), parsedSymbolID(id("ptr-operator")),
                         parents({GSSNode2})),
                   AllOf(state(6), parsedSymbolID(id("ptr-operator")),
-                        parents({GSSNode3})))) << NewHeadResults;
+                        parents({GSSNode3}))));
   // Verify that the payload of the two new heads is shared, only a single
   // ptr-operator node is created in the forest.
   EXPECT_EQ(NewHeadResults[0]->Payload, NewHeadResults[1]->Payload);
 }
 
 TEST_F(GLRTest, ReduceJoiningWithMultipleBases) {
   //  Before (joining due to same goto state, multiple bases):
   //    0--1(cv-qualifier)--3(class-name)
   //    └--2(cv-qualifier)--4(enum-name)
   //  After reducing 3 by `type-name := class-name` and
   //                 4 by `type-name := enum-name`:
   //    0--1(cv-qualifier)--5(type-name)  // 5 is goto(1, type-name) and
   //    └--2(cv-qualifier)--┘             // goto(2, type-name)
   buildGrammar({"type-name", "class-name", "enum-name", "cv-qualifier"},
                {"type-name := class-name", "type-name := enum-name"});
 
   auto *CVQualifierNode =
       &Arena.createOpaque(id("cv-qualifier"), /*TokenIndex=*/0);
   auto *ClassNameNode = &Arena.createOpaque(id("class-name"), /*TokenIndex=*/1);
   auto *EnumNameNode = &Arena.createOpaque(id("enum-name"), /*TokenIndex=*/1);
 
-  const auto *GSSNode0 =
+  auto GSSNode0 =
       GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr, /*Parents=*/{});
-  const auto *GSSNode1 = GSStack.addNode(
+  auto GSSNode1 = GSStack.addNode(
       /*State=*/1, /*ForestNode=*/CVQualifierNode, /*Parents=*/{GSSNode0});
-  const auto *GSSNode2 = GSStack.addNode(
+  auto GSSNode2 = GSStack.addNode(
       /*State=*/2, /*ForestNode=*/CVQualifierNode, /*Parents=*/{GSSNode0});
-  const auto *GSSNode3 =
-      GSStack.addNode(/*State=*/3, /*ForestNode=*/ClassNameNode,
+  auto GSSNode3 = GSStack.addNode(/*State=*/3, /*ForestNode=*/ClassNameNode,
                                   /*Parents=*/{GSSNode1});
-  const auto *GSSNode4 =
-      GSStack.addNode(/*State=*/4, /*ForestNode=*/EnumNameNode,
+  auto GSSNode4 = GSStack.addNode(/*State=*/4, /*ForestNode=*/EnumNameNode,
                                   /*Parents=*/{GSSNode2});
-  GSStack.unref(GSSNode0);
-  GSStack.unref(GSSNode1);
-  GSStack.unref(GSSNode2);
 
   LRTable Table = LRTable::buildForTests(
       G->table(),
       {{/*State=*/1, id("type-name"), Action::goTo(/*NextState=*/5)},
        {/*State=*/2, id("type-name"), Action::goTo(/*NextState=*/5)}});
   // FIXME: figure out a way to get rid of the hard-coded reduce RuleID!
   std::vector<ParseStep> PendingReduce = {
       {
           GSSNode3, Action::reduce(/*RuleID=*/0) // type-name := class-name
       },
       {
           GSSNode4, Action::reduce(/*RuleID=*/1) // type-name := enum-name
       }};
   glrReduce(PendingReduce, {*G, Table, Arena, GSStack},
             captureNewHeads());
 
   // Verify that the stack heads are joint at state 5 after reduces.
   EXPECT_THAT(NewHeadResults, testing::UnorderedElementsAre(AllOf(
                                   state(5), parsedSymbolID(id("type-name")),
-                                  parents({GSSNode1, GSSNode2}))))
-      << NewHeadResults;
+                                  parents({GSSNode1, GSSNode2}))));
   // Verify that we create an ambiguous ForestNode of two parses of `type-name`.
   EXPECT_EQ(NewHeadResults.front()->Payload->dumpRecursive(*G),
             "[  1, end) type-name := <ambiguous>\n"
             "[  1, end) ├─type-name := class-name\n"
             "[  1, end) │ └─class-name := <opaque>\n"
             "[  1, end) └─type-name := enum-name\n"
             "[  1, end)   └─enum-name := <opaque>\n");
 }
 
 TEST_F(GLRTest, ReduceJoiningWithSameBase) {
   //  Before (joining due to same goto state, the same base):
   //    0--1(class-name)--3(*)
   //    └--2(enum-name)--4(*)
   //  After reducing 3 by `pointer := class-name *` and
   //                 2 by `pointer := enum-name *`:
   //    0--5(pointer)  // 5 is goto(0, pointer)
   buildGrammar({"pointer", "class-name", "enum-name"},
                {"pointer := class-name *", "pointer := enum-name *"});
 
   auto *ClassNameNode = &Arena.createOpaque(id("class-name"), /*TokenIndex=*/0);
   auto *EnumNameNode = &Arena.createOpaque(id("enum-name"), /*TokenIndex=*/0);
   auto *StartTerminal = &Arena.createTerminal(tok::star, /*TokenIndex=*/1);
 
-  const auto *GSSNode0 =
+  auto GSSNode0 =
       GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr, /*Parents=*/{});
-  const auto *GSSNode1 =
+  auto GSSNode1 =
       GSStack.addNode(/*State=*/1, /*ForestNode=*/ClassNameNode,
                       /*Parents=*/{GSSNode0});
-  const auto *GSSNode2 =
+  auto GSSNode2 =
       GSStack.addNode(/*State=*/2, /*ForestNode=*/EnumNameNode,
                       /*Parents=*/{GSSNode0});
-  const auto *GSSNode3 =
+  auto GSSNode3 =
       GSStack.addNode(/*State=*/3, /*ForestNode=*/StartTerminal,
                       /*Parents=*/{GSSNode1});
-  const auto *GSSNode4 =
+  auto GSSNode4 =
       GSStack.addNode(/*State=*/4, /*ForestNode=*/StartTerminal,
                       /*Parents=*/{GSSNode2});
-  GSStack.unref(GSSNode0);
-  GSStack.unref(GSSNode1);
-  GSStack.unref(GSSNode2);
 
   LRTable Table = LRTable::buildForTests(
       G->table(), {{/*State=*/0, id("pointer"), Action::goTo(5)}});
   // FIXME: figure out a way to get rid of the hard-coded reduce RuleID!
   std::vector<ParseStep> PendingReduce = {
       {
           GSSNode3, Action::reduce(/*RuleID=*/0) // pointer := class-name *
       },
       {
           GSSNode4, Action::reduce(/*RuleID=*/1) // pointer := enum-name *
       }};
   glrReduce(PendingReduce, {*G, Table, Arena, GSStack},
             captureNewHeads());
 
   EXPECT_THAT(NewHeadResults, testing::UnorderedElementsAre(
                                   AllOf(state(5), parsedSymbolID(id("pointer")),
-                                        parents({GSSNode0}))))
-      << NewHeadResults;
+                                        parents({GSSNode0}))));
   EXPECT_EQ(NewHeadResults.front()->Payload->dumpRecursive(*G),
             "[  0, end) pointer := <ambiguous>\n"
             "[  0, end) ├─pointer := class-name *\n"
             "[  0,   1) │ ├─class-name := <opaque>\n"
             "[  1, end) │ └─* := tok[1]\n"
             "[  0, end) └─pointer := enum-name *\n"
             "[  0,   1)   ├─enum-name := <opaque>\n"
             "[  1, end)   └─* := tok[1]\n");