[Support] Store dominator tree nodes in a vector

llvm/include/llvm/ADT/GraphTraits.h

@@ -71,6 +71,20 @@ struct GraphTraits {
   // static unsigned       size       (GraphType *G)
   //    Return total number of nodes in the graph
 
+  // Optionally implement the following:
+  // static unsigned getNumber(NodeRef)
+  //    Return a unique number of a node. Numbers are ideally dense, these are
+  //    used to store nodes in a vector.
+  // static unsigned getMaxNumber(GraphType *G)
+  //    Return the maximum number that getNumber() will return, or 0 if this is
+  //    unknown. Intended for reserving large enough buffers.
+  // static unsigned getNumberEpoch(GraphType *G)
+  //    Return the "epoch" of the node numbers. Should return a different
+  //    number after renumbering, so users can assert that the epoch didn't
+  //    change => numbers are still valid. If renumberings are not tracked, it
+  //    is always valid to return a constant value. This is solely for to ease
+  //    debugging by having a way to detect use of outdated numbers.
+
   // If anyone tries to use this class without having an appropriate
   // specialization, make an error.  If you get this error, it's because you
   // need to include the appropriate specialization of GraphTraits<> for your

llvm/include/llvm/Support/GenericDomTree.h

@@ -258,37 +258,40 @@ class DominatorTreeBase {
   // Dominators always have a single root, postdominators can have more.
   SmallVector<NodeT *, IsPostDom ? 4 : 1> Roots;
 
-  using DomTreeNodeMapType =
-     DenseMap<NodeT *, std::unique_ptr<DomTreeNodeBase<NodeT>>>;
-  DomTreeNodeMapType DomTreeNodes;
+  using DomTreeNodeStorageTy =
+      SmallVector<std::unique_ptr<DomTreeNodeBase<NodeT>>>;
+  DomTreeNodeStorageTy DomTreeNodes;
+  // For graphs where blocks don't have numbers, create a numbering here.
+  DenseMap<const NodeT *, unsigned> NodeNumberMap;
   DomTreeNodeBase<NodeT> *RootNode = nullptr;
   ParentPtr Parent = nullptr;
 
   mutable bool DFSInfoValid = false;
   mutable unsigned int SlowQueries = 0;
+  unsigned BlockNumberEpoch = 0;
 
   friend struct DomTreeBuilder::SemiNCAInfo<DominatorTreeBase>;
 
  public:
   DominatorTreeBase() = default;
 
   DominatorTreeBase(DominatorTreeBase &&Arg)
-      : Roots(std::move(Arg.Roots)),
-        DomTreeNodes(std::move(Arg.DomTreeNodes)),
-        RootNode(Arg.RootNode),
-        Parent(Arg.Parent),
-        DFSInfoValid(Arg.DFSInfoValid),
-        SlowQueries(Arg.SlowQueries) {
+      : Roots(std::move(Arg.Roots)), DomTreeNodes(std::move(Arg.DomTreeNodes)),
+        NodeNumberMap(std::move(Arg.NodeNumberMap)), RootNode(Arg.RootNode),
+        Parent(Arg.Parent), DFSInfoValid(Arg.DFSInfoValid),
+        SlowQueries(Arg.SlowQueries), BlockNumberEpoch(Arg.BlockNumberEpoch) {
     Arg.wipe();
   }
 
   DominatorTreeBase &operator=(DominatorTreeBase &&RHS) {
     Roots = std::move(RHS.Roots);
     DomTreeNodes = std::move(RHS.DomTreeNodes);
+    NodeNumberMap = std::move(RHS.NodeNumberMap);
     RootNode = RHS.RootNode;
     Parent = RHS.Parent;
     DFSInfoValid = RHS.DFSInfoValid;
     SlowQueries = RHS.SlowQueries;
+    BlockNumberEpoch = RHS.BlockNumberEpoch;
     RHS.wipe();
     return *this;
   }
@@ -333,35 +336,70 @@ class DominatorTreeBase {
     if (!std::is_permutation(Roots.begin(), Roots.end(), Other.Roots.begin()))
       return true;
 
-    const DomTreeNodeMapType &OtherDomTreeNodes = Other.DomTreeNodes;
-    if (DomTreeNodes.size() != OtherDomTreeNodes.size())
-      return true;
-
-    for (const auto &DomTreeNode : DomTreeNodes) {
-      NodeT *BB = DomTreeNode.first;
-      typename DomTreeNodeMapType::const_iterator OI =
-          OtherDomTreeNodes.find(BB);
-      if (OI == OtherDomTreeNodes.end())
+    size_t NumNodes = 0;
+    // All nodes we have must exist and be equal in the other tree.
+    for (const auto &Node : DomTreeNodes) {
+      if (!Node)
+        continue;
+      if (Node->compare(Other.getNode(Node->getBlock())))
         return true;
+      NumNodes++;
+    }
 
-      DomTreeNodeBase<NodeT> &MyNd = *DomTreeNode.second;
-      DomTreeNodeBase<NodeT> &OtherNd = *OI->second;
+    // If the other tree has more nodes than we have, they're not equal.
+    size_t NumOtherNodes = 0;
+    for (const auto &OtherNode : Other.DomTreeNodes)
+      if (OtherNode)
+        NumOtherNodes++;
+    return NumNodes != NumOtherNodes;
+  }
 
-      if (MyNd.compare(&OtherNd))
-        return true;
+private:
+  template <typename T>
+  using has_number_t =
+      decltype(GraphTraits<T *>::getNumber(std::declval<T *>()));
+
+  std::optional<unsigned> getNodeIndex(const NodeT *BB) const {
+    if constexpr (is_detected<has_number_t, NodeT>::value) {
+      // BB can be nullptr, map nullptr to index 0.
+      assert(BlockNumberEpoch ==
+                 GraphTraits<ParentPtr>::getNumberEpoch(Parent) &&
+             "dominator tree used with outdated block numbers");
+      return BB ? GraphTraits<const NodeT *>::getNumber(BB) + 1 : 0;
+    } else {
+      if (auto It = NodeNumberMap.find(BB); It != NodeNumberMap.end())
+        return It->second;
+      return std::nullopt;
     }
+  }
 
-    return false;
+  unsigned getNodeIndexForInsert(const NodeT *BB) {
+    if constexpr (is_detected<has_number_t, NodeT>::value) {
+      // getNodeIndex will never fail if nodes have getNumber().
+      unsigned Idx = *getNodeIndex(BB);
+      if (Idx >= DomTreeNodes.size()) {
+        unsigned Max = GraphTraits<ParentPtr>::getMaxNumber(Parent);
+        DomTreeNodes.resize(Max > Idx + 1 ? Max : Idx + 1);
+      }
+      return Idx;
+    } else {
+      // We might already have a number stored for BB.
+      unsigned Idx =
+          NodeNumberMap.try_emplace(BB, DomTreeNodes.size()).first->second;
+      if (Idx >= DomTreeNodes.size())
+        DomTreeNodes.resize(Idx + 1);
+      return Idx;
+    }
   }
 
+public:
   /// getNode - return the (Post)DominatorTree node for the specified basic
   /// block.  This is the same as using operator[] on this class.  The result
   /// may (but is not required to) be null for a forward (backwards)
   /// statically unreachable block.
   DomTreeNodeBase<NodeT> *getNode(const NodeT *BB) const {
-    auto I = DomTreeNodes.find(BB);
-    if (I != DomTreeNodes.end())
-      return I->second.get();
+    if (auto Idx = getNodeIndex(BB); Idx && *Idx < DomTreeNodes.size())
+      return DomTreeNodes[*Idx].get();
     return nullptr;
   }
 
@@ -678,8 +716,10 @@ class DominatorTreeBase {
   /// dominate any other blocks. Removes node from its immediate dominator's
   /// children list. Deletes dominator node associated with basic block BB.
   void eraseNode(NodeT *BB) {
-    DomTreeNodeBase<NodeT> *Node = getNode(BB);
-    assert(Node && "Removing node that isn't in dominator tree.");
+    std::optional<unsigned> IdxOpt = getNodeIndex(BB);
+    assert(IdxOpt && DomTreeNodes[*IdxOpt] &&
+           "Removing node that isn't in dominator tree.");
+    DomTreeNodeBase<NodeT> *Node = DomTreeNodes[*IdxOpt].get();
     assert(Node->isLeaf() && "Node is not a leaf node.");
 
     DFSInfoValid = false;
@@ -695,7 +735,8 @@ class DominatorTreeBase {
       IDom->Children.pop_back();
     }
 
-    DomTreeNodes.erase(BB);
+    DomTreeNodes[*IdxOpt] = nullptr;
+    NodeNumberMap.erase(BB);
 
     if (!IsPostDom) return;
 
@@ -786,17 +827,48 @@ class DominatorTreeBase {
     DFSInfoValid = true;
   }
 
+private:
+  void updateBlockNumberEpoch() {
+    // Nothing to do for graphs that don't number their blocks.
+    if constexpr (is_detected<has_number_t, NodeT>::value)
+      BlockNumberEpoch = GraphTraits<ParentPtr>::getNumberEpoch(Parent);
+  }
+
+public:
   /// recalculate - compute a dominator tree for the given function
   void recalculate(ParentType &Func) {
     Parent = &Func;
+    updateBlockNumberEpoch();
     DomTreeBuilder::Calculate(*this);
   }
 
   void recalculate(ParentType &Func, ArrayRef<UpdateType> Updates) {
     Parent = &Func;
+    updateBlockNumberEpoch();
     DomTreeBuilder::CalculateWithUpdates(*this, Updates);
   }
 
+  /// Update dominator tree after renumbering blocks.
+  template <class T_ = NodeT>
+  std::enable_if_t<is_detected<has_number_t, T_>::value, void>
+  updateBlockNumbers() {
+    updateBlockNumberEpoch();
+
+    unsigned MaxNumber = GraphTraits<ParentPtr>::getMaxNumber(Parent);
+    DomTreeNodeStorageTy NewVector;
+    NewVector.resize(MaxNumber + 1); // +1, because index 0 is for nullptr
+    for (auto &Node : DomTreeNodes) {
+      if (!Node)
+        continue;
+      unsigned Idx = *getNodeIndex(Node->getBlock());
+      // getMaxNumber is not necessarily supported
+      if (Idx >= NewVector.size())
+        NewVector.resize(Idx + 1);
+      NewVector[Idx] = std::move(Node);
+    }
+    DomTreeNodes = std::move(NewVector);
+  }
+
   /// verify - checks if the tree is correct. There are 3 level of verification:
   ///  - Full --  verifies if the tree is correct by making sure all the
   ///             properties (including the parent and the sibling property)
@@ -817,6 +889,7 @@ class DominatorTreeBase {
 
   void reset() {
     DomTreeNodes.clear();
+    NodeNumberMap.clear();
     Roots.clear();
     RootNode = nullptr;
     Parent = nullptr;
@@ -831,7 +904,8 @@ class DominatorTreeBase {
                                      DomTreeNodeBase<NodeT> *IDom = nullptr) {
     auto Node = std::make_unique<DomTreeNodeBase<NodeT>>(BB, IDom);
     auto *NodePtr = Node.get();
-    DomTreeNodes[BB] = std::move(Node);
+    unsigned NodeIdx = getNodeIndexForInsert(BB);
+    DomTreeNodes[NodeIdx] = std::move(Node);
     if (IDom)
       IDom->addChild(NodePtr);
     return NodePtr;
@@ -915,6 +989,7 @@ class DominatorTreeBase {
   /// assignable and destroyable state, but otherwise invalid.
   void wipe() {
     DomTreeNodes.clear();
+    NodeNumberMap.clear();
     RootNode = nullptr;
     Parent = nullptr;
   }

llvm/include/llvm/Support/GenericDomTreeConstruction.h

@@ -1231,7 +1231,9 @@ struct SemiNCAInfo {
     doFullDFSWalk(DT, AlwaysDescend);
 
     for (auto &NodeToTN : DT.DomTreeNodes) {
-      const TreeNodePtr TN = NodeToTN.second.get();
+      const TreeNodePtr TN = NodeToTN.get();
+      if (!TN)
+        continue;
       const NodePtr BB = TN->getBlock();
 
       // Virtual root has a corresponding virtual CFG node.
@@ -1264,7 +1266,9 @@ struct SemiNCAInfo {
   // Running time: O(N).
   static bool VerifyLevels(const DomTreeT &DT) {
     for (auto &NodeToTN : DT.DomTreeNodes) {
-      const TreeNodePtr TN = NodeToTN.second.get();
+      const TreeNodePtr TN = NodeToTN.get();
+      if (!TN)
+        continue;
       const NodePtr BB = TN->getBlock();
       if (!BB) continue;
 
@@ -1319,7 +1323,9 @@ struct SemiNCAInfo {
     // For each tree node verify if children's DFS numbers cover their parent's
     // DFS numbers with no gaps.
     for (const auto &NodeToTN : DT.DomTreeNodes) {
-      const TreeNodePtr Node = NodeToTN.second.get();
+      const TreeNodePtr Node = NodeToTN.get();
+      if (!Node)
+        continue;
 
       // Handle tree leaves.
       if (Node->isLeaf()) {
@@ -1432,7 +1438,9 @@ struct SemiNCAInfo {
   // the nodes it dominated previously will now become unreachable.
   bool verifyParentProperty(const DomTreeT &DT) {
     for (auto &NodeToTN : DT.DomTreeNodes) {
-      const TreeNodePtr TN = NodeToTN.second.get();
+      const TreeNodePtr TN = NodeToTN.get();
+      if (!TN)
+        continue;
       const NodePtr BB = TN->getBlock();
       if (!BB || TN->isLeaf())
         continue;
@@ -1466,7 +1474,9 @@ struct SemiNCAInfo {
   // siblings will now still be reachable.
   bool verifySiblingProperty(const DomTreeT &DT) {
     for (auto &NodeToTN : DT.DomTreeNodes) {
-      const TreeNodePtr TN = NodeToTN.second.get();
+      const TreeNodePtr TN = NodeToTN.get();
+      if (!TN)
+        continue;
       const NodePtr BB = TN->getBlock();
       if (!BB || TN->isLeaf())
         continue;

llvm/unittests/Support/CMakeLists.txt

@@ -44,6 +44,7 @@ add_llvm_unittest(SupportTests
   FileOutputBufferTest.cpp
   FormatVariadicTest.cpp
   FSUniqueIDTest.cpp
+  GenericDomTreeTest.cpp
   GlobPatternTest.cpp
   HashBuilderTest.cpp
   IndexedAccessorTest.cpp