[Support] Store dominator tree nodes in a vector (#101705)

Use basic block numbers to store dominator tree nodes in a vector. This
avoids frequent map lookups. Use block number epochs to validate that no
renumbering occured since the tree was created; after a renumbering, the
dominator tree can be updated with updateBlockNumbers().

Block numbers, block number epoch, and max block number are fetched via
newly added GraphTraits methods. Graphs which do not implement
getNumber() for blocks will use a DenseMap for an ad-hoc numbering.

Author: aengelke

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