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

Alexis Engelke via llvm-commits llvm-commits at lists.llvm.org
Sun Aug 4 02:55:08 PDT 2024


https://github.com/aengelke updated https://github.com/llvm/llvm-project/pull/101705

>From 5b907f0f5c34fcaf20ad1caa8e53b7d616e75c8d Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Wed, 31 Jul 2024 10:10:11 +0000
Subject: [PATCH 1/2] [Support] Store dominator tree nodes in a vector

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.
---
 llvm/include/llvm/ADT/GraphTraits.h           |  14 ++
 llvm/include/llvm/Support/GenericDomTree.h    | 149 ++++++++++++++----
 .../llvm/Support/GenericDomTreeConstruction.h |  20 ++-
 3 files changed, 149 insertions(+), 34 deletions(-)

diff --git a/llvm/include/llvm/ADT/GraphTraits.h b/llvm/include/llvm/ADT/GraphTraits.h
index 3a7773592af3d..69f9facebefe7 100644
--- a/llvm/include/llvm/ADT/GraphTraits.h
+++ b/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
diff --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h
index e05e5f0f842e3..4d3412b24399e 100644
--- a/llvm/include/llvm/Support/GenericDomTree.h
+++ b/llvm/include/llvm/Support/GenericDomTree.h
@@ -258,14 +258,17 @@ 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 block 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>;
 
@@ -273,22 +276,22 @@ class DominatorTreeBase {
   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,80 @@ 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())
+    size_t NumNodes = 0;
+    for (const auto &Node : DomTreeNodes) {
+      if (!Node)
+        continue;
+      if (Node->compare(Other.getNode(Node->getBlock())))
+        return true;
+      NumNodes++;
+    }
+
+    size_t NumOtherNodes = 0;
+    for (const auto &OtherNode : Other.DomTreeNodes)
+      if (OtherNode)
+        NumOtherNodes++;
+    if (NumNodes != NumOtherNodes)
       return true;
 
-    for (const auto &DomTreeNode : DomTreeNodes) {
-      NodeT *BB = DomTreeNode.first;
-      typename DomTreeNodeMapType::const_iterator OI =
-          OtherDomTreeNodes.find(BB);
-      if (OI == OtherDomTreeNodes.end())
-        return true;
+    return false;
+  }
 
-      DomTreeNodeBase<NodeT> &MyNd = *DomTreeNode.second;
-      DomTreeNodeBase<NodeT> &OtherNd = *OI->second;
+private:
+  template <typename T>
+  using has_number_t =
+      decltype(GraphTraits<T *>::getNumber(std::declval<T *>()));
+
+  template <class T_ = NodeT>
+  std::enable_if_t<is_detected<has_number_t, T_>::value,
+                   std::optional<unsigned>>
+  getNodeIndex(const NodeT *BB) const {
+    // 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;
+  }
 
-      if (MyNd.compare(&OtherNd))
-        return true;
+  template <class T_ = NodeT>
+  std::enable_if_t<!is_detected<has_number_t, T_>::value,
+                   std::optional<unsigned>>
+  getNodeIndex(const NodeT *BB) const {
+    if (auto It = NodeNumberMap.find(BB); It != NodeNumberMap.end())
+      return It->second;
+    return std::nullopt;
+  }
+
+  template <class T_ = NodeT>
+  std::enable_if_t<is_detected<has_number_t, T_>::value, unsigned>
+  getNodeIndexForInsert(const NodeT *BB) {
+    // 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;
+  }
 
-    return false;
+  template <class T_ = NodeT>
+  std::enable_if_t<!is_detected<has_number_t, T_>::value, unsigned>
+  getNodeIndexForInsert(const NodeT *BB) {
+    // 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 +726,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 +745,8 @@ class DominatorTreeBase {
       IDom->Children.pop_back();
     }
 
-    DomTreeNodes.erase(BB);
+    DomTreeNodes[*IdxOpt] = nullptr;
+    NodeNumberMap.erase(BB);
 
     if (!IsPostDom) return;
 
@@ -786,17 +837,54 @@ class DominatorTreeBase {
     DFSInfoValid = true;
   }
 
+private:
+  template <class T_ = NodeT>
+  std::enable_if_t<is_detected<has_number_t, T_>::value, void>
+  updateBlockNumberEpoch() {
+    BlockNumberEpoch = GraphTraits<ParentPtr>::getNumberEpoch(Parent);
+  }
+
+  template <class T_ = NodeT>
+  std::enable_if_t<!is_detected<has_number_t, T_>::value, void>
+  updateBlockNumberEpoch() {
+    // Nothing to do for graphs that don't number their blocks.
+  }
+
+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 necssarily 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 +905,7 @@ class DominatorTreeBase {
 
   void reset() {
     DomTreeNodes.clear();
+    NodeNumberMap.clear();
     Roots.clear();
     RootNode = nullptr;
     Parent = nullptr;
@@ -831,7 +920,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 +1005,7 @@ class DominatorTreeBase {
   /// assignable and destroyable state, but otherwise invalid.
   void wipe() {
     DomTreeNodes.clear();
+    NodeNumberMap.clear();
     RootNode = nullptr;
     Parent = nullptr;
   }
diff --git a/llvm/include/llvm/Support/GenericDomTreeConstruction.h b/llvm/include/llvm/Support/GenericDomTreeConstruction.h
index af7ac04a5ab29..21ab042251c25 100644
--- a/llvm/include/llvm/Support/GenericDomTreeConstruction.h
+++ b/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;

>From bf7fbacc6b80bf1dd64dd5d33ca0707da1ce3019 Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Sun, 4 Aug 2024 09:53:14 +0000
Subject: [PATCH 2/2] Address comments and add unit test

---
 llvm/include/llvm/Support/GenericDomTree.h    |  87 ++++++--------
 llvm/unittests/Support/CMakeLists.txt         |   1 +
 llvm/unittests/Support/GenericDomTreeTest.cpp | 109 ++++++++++++++++++
 3 files changed, 146 insertions(+), 51 deletions(-)
 create mode 100644 llvm/unittests/Support/GenericDomTreeTest.cpp

diff --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h
index 4d3412b24399e..f870be32dfe8f 100644
--- a/llvm/include/llvm/Support/GenericDomTree.h
+++ b/llvm/include/llvm/Support/GenericDomTree.h
@@ -261,7 +261,7 @@ class DominatorTreeBase {
   using DomTreeNodeStorageTy =
       SmallVector<std::unique_ptr<DomTreeNodeBase<NodeT>>>;
   DomTreeNodeStorageTy DomTreeNodes;
-  // For graphs where block don't have numbers, create a numbering here.
+  // For graphs where blocks don't have numbers, create a numbering here.
   DenseMap<const NodeT *, unsigned> NodeNumberMap;
   DomTreeNodeBase<NodeT> *RootNode = nullptr;
   ParentPtr Parent = nullptr;
@@ -337,6 +337,7 @@ class DominatorTreeBase {
       return true;
 
     size_t NumNodes = 0;
+    // All node we have must exist and be equal in the other tree
     for (const auto &Node : DomTreeNodes) {
       if (!Node)
         continue;
@@ -345,14 +346,12 @@ class DominatorTreeBase {
       NumNodes++;
     }
 
+    // If we 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++;
-    if (NumNodes != NumOtherNodes)
-      return true;
-
-    return false;
+    return NumNodes != NumOtherNodes;
   }
 
 private:
@@ -361,45 +360,37 @@ class DominatorTreeBase {
       decltype(GraphTraits<T *>::getNumber(std::declval<T *>()));
 
   template <class T_ = NodeT>
-  std::enable_if_t<is_detected<has_number_t, T_>::value,
-                   std::optional<unsigned>>
-  getNodeIndex(const NodeT *BB) const {
-    // 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;
-  }
-
-  template <class T_ = NodeT>
-  std::enable_if_t<!is_detected<has_number_t, T_>::value,
-                   std::optional<unsigned>>
-  getNodeIndex(const NodeT *BB) const {
-    if (auto It = NodeNumberMap.find(BB); It != NodeNumberMap.end())
-      return It->second;
-    return std::nullopt;
-  }
-
-  template <class T_ = NodeT>
-  std::enable_if_t<is_detected<has_number_t, T_>::value, unsigned>
-  getNodeIndexForInsert(const NodeT *BB) {
-    // 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);
+  std::optional<unsigned> getNodeIndex(const NodeT *BB) const {
+    if constexpr (is_detected<has_number_t, T_>::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 Idx;
   }
 
-  template <class T_ = NodeT>
-  std::enable_if_t<!is_detected<has_number_t, T_>::value, unsigned>
-  getNodeIndexForInsert(const NodeT *BB) {
-    // 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;
+  template <class T_ = NodeT> unsigned getNodeIndexForInsert(const NodeT *BB) {
+    if constexpr (is_detected<has_number_t, T_>::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:
@@ -838,16 +829,10 @@ class DominatorTreeBase {
   }
 
 private:
-  template <class T_ = NodeT>
-  std::enable_if_t<is_detected<has_number_t, T_>::value, void>
-  updateBlockNumberEpoch() {
-    BlockNumberEpoch = GraphTraits<ParentPtr>::getNumberEpoch(Parent);
-  }
-
-  template <class T_ = NodeT>
-  std::enable_if_t<!is_detected<has_number_t, T_>::value, void>
-  updateBlockNumberEpoch() {
+  template <class T_ = NodeT> void updateBlockNumberEpoch() {
     // Nothing to do for graphs that don't number their blocks.
+    if constexpr (is_detected<has_number_t, T_>::value)
+      BlockNumberEpoch = GraphTraits<ParentPtr>::getNumberEpoch(Parent);
   }
 
 public:
@@ -877,7 +862,7 @@ class DominatorTreeBase {
       if (!Node)
         continue;
       unsigned Idx = *getNodeIndex(Node->getBlock());
-      // getMaxNumber is not necssarily supported
+      // getMaxNumber is not necessarily supported
       if (Idx >= NewVector.size())
         NewVector.resize(Idx + 1);
       NewVector[Idx] = std::move(Node);
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index 631f2e6bf00df..db47a170e814a 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/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
diff --git a/llvm/unittests/Support/GenericDomTreeTest.cpp b/llvm/unittests/Support/GenericDomTreeTest.cpp
new file mode 100644
index 0000000000000..f0f87e3a98908
--- /dev/null
+++ b/llvm/unittests/Support/GenericDomTreeTest.cpp
@@ -0,0 +1,109 @@
+//===- unittests/Support/GenericDomTreeTest.cpp - GenericDomTree.h tests --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/GenericDomTree.h"
+#include "llvm/ADT/GraphTraits.h"
+#include "llvm/Support/DataTypes.h"
+#include "gtest/gtest.h"
+using namespace llvm;
+
+namespace {
+
+// Very simple (fake) graph structure to test dominator tree on.
+struct NumberedGraph;
+
+struct NumberedNode {
+  NumberedGraph *Parent;
+  unsigned Number;
+
+  NumberedNode(NumberedGraph *Parent, unsigned Number)
+      : Parent(Parent), Number(Number) {}
+
+  NumberedGraph *getParent() const { return Parent; }
+};
+
+struct NumberedGraph {
+  SmallVector<std::unique_ptr<NumberedNode>> Nodes;
+  unsigned NumberEpoch = 0;
+
+  NumberedNode *addNode() {
+    unsigned Num = Nodes.size();
+    return Nodes.emplace_back(std::make_unique<NumberedNode>(this, Num)).get();
+  }
+};
+} // namespace
+
+namespace llvm {
+template <> struct GraphTraits<NumberedNode *> {
+  using NodeRef = NumberedNode *;
+  static unsigned getNumber(NumberedNode *Node) { return Node->Number; }
+};
+
+template <> struct GraphTraits<const NumberedNode *> {
+  using NodeRef = NumberedNode *;
+  static unsigned getNumber(const NumberedNode *Node) { return Node->Number; }
+};
+
+template <> struct GraphTraits<NumberedGraph *> {
+  using NodeRef = NumberedNode *;
+  static unsigned getMaxNumber(NumberedGraph *G) { return G->Nodes.size(); }
+  static unsigned getNumberEpoch(NumberedGraph *G) { return G->NumberEpoch; }
+};
+
+namespace DomTreeBuilder {
+// Dummy specialization. Only needed so that we can call recalculate(), which
+// sets DT.Parent -- but we can't access DT.Parent here.
+template <> void Calculate(DomTreeBase<NumberedNode> &DT) {}
+} // end namespace DomTreeBuilder
+} // end namespace llvm
+
+namespace {
+
+TEST(GenericDomTree, BlockNumbers) {
+  NumberedGraph G;
+  NumberedNode *N1 = G.addNode();
+  NumberedNode *N2 = G.addNode();
+
+  DomTreeBase<NumberedNode> DT;
+  DT.recalculate(G); // only sets parent
+  // Construct fake domtree: node 0 dominates all other nodes
+  DT.setNewRoot(N1);
+  DT.addNewBlock(N2, N1);
+
+  // Roundtrip is correct
+  for (auto &N : G.Nodes)
+    EXPECT_EQ(DT.getNode(N.get())->getBlock(), N.get());
+  // If we manually change a number, we should get a different node.
+  ASSERT_EQ(N1->Number, 0u);
+  ASSERT_EQ(N2->Number, 1u);
+  N1->Number = 1;
+  EXPECT_EQ(DT.getNode(N1)->getBlock(), N2);
+  EXPECT_EQ(DT.getNode(N2)->getBlock(), N2);
+  N2->Number = 0;
+  EXPECT_EQ(DT.getNode(N2)->getBlock(), N1);
+
+  // Renumer blocks, should fix node domtree-internal node map
+  DT.updateBlockNumbers();
+  for (auto &N : G.Nodes)
+    EXPECT_EQ(DT.getNode(N.get())->getBlock(), N.get());
+
+  // Adding a new node with a higher number is no problem
+  NumberedNode *N3 = G.addNode();
+  EXPECT_EQ(DT.getNode(N3), nullptr);
+  // ... even if it exceeds getMaxNumber()
+  NumberedNode *N4 = G.addNode();
+  N4->Number = 1000;
+  EXPECT_EQ(DT.getNode(N4), nullptr);
+
+  DT.addNewBlock(N3, N1);
+  DT.addNewBlock(N4, N1);
+  for (auto &N : G.Nodes)
+    EXPECT_EQ(DT.getNode(N.get())->getBlock(), N.get());
+}
+
+} // namespace



More information about the llvm-commits mailing list