DomTree: Remove getRoots() accessor

Summary:
Avoid exposing details about how roots are stored. This enables subsequent
type-erasure changes.

v5:
- cleanup a unit test by using EXPECT_EQ instead of EXPECT_TRUE

Change-Id: I532b774cc71f2224e543bc7d79131d97f63f093d

Reviewers: arsenm, RKSimon, mehdi_amini, courbet

Subscribers: jvesely, wdng, hiraditya, kuhar, kerbowa, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D83085

8 files changed, 26 insertions, 24 deletions

diff --git a/llvm/include/llvm/Analysis/DominanceFrontier.h b/llvm/include/llvm/Analysis/DominanceFrontier.h
index c0bf30e..f67929c 100644
--- a/llvm/include/llvm/Analysis/DominanceFrontier.h
+++ b/llvm/include/llvm/Analysis/DominanceFrontier.h
@@ -130,7 +130,7 @@ public:
   using DomSetType = typename DominanceFrontierBase<BlockT, false>::DomSetType;
 
   void analyze(DomTreeT &DT) {
-    assert(DT.getRoots().size() == 1 &&
+    assert(DT.root_size() == 1 &&
            "Only one entry block for forward domfronts!");
     this->Roots = {DT.getRoot()};
     calculate(DT, DT[this->Roots[0]]);
diff --git a/llvm/include/llvm/CodeGen/MachineDominators.h b/llvm/include/llvm/CodeGen/MachineDominators.h
index 9d31232..2d26163 100644
--- a/llvm/include/llvm/CodeGen/MachineDominators.h
+++ b/llvm/include/llvm/CodeGen/MachineDominators.h
@@ -93,15 +93,6 @@ public:
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
-  /// getRoots -  Return the root blocks of the current CFG.  This may include
-  /// multiple blocks if we are computing post dominators.  For forward
-  /// dominators, this will always be a single block (the entry node).
-  ///
-  const SmallVectorImpl<MachineBasicBlock*> &getRoots() const {
-    applySplitCriticalEdges();
-    return DT->getRoots();
-  }
-
   MachineBasicBlock *getRoot() const {
     applySplitCriticalEdges();
     return DT->getRoot();
diff --git a/llvm/include/llvm/CodeGen/MachinePostDominators.h b/llvm/include/llvm/CodeGen/MachinePostDominators.h
index 597bb40..cee4294 100644
--- a/llvm/include/llvm/CodeGen/MachinePostDominators.h
+++ b/llvm/include/llvm/CodeGen/MachinePostDominators.h
@@ -41,10 +41,6 @@ public:
 
   FunctionPass *createMachinePostDominatorTreePass();
 
-  const SmallVectorImpl<MachineBasicBlock *> &getRoots() const {
-    return PDT->getRoots();
-  }
-
   MachineDomTreeNode *getRootNode() const { return PDT->getRootNode(); }
 
   MachineDomTreeNode *operator[](MachineBasicBlock *BB) const {
diff --git a/llvm/include/llvm/Support/GenericDomTree.h b/llvm/include/llvm/Support/GenericDomTree.h
index e83e7aa..407a060 100644
--- a/llvm/include/llvm/Support/GenericDomTree.h
+++ b/llvm/include/llvm/Support/GenericDomTree.h
@@ -283,11 +283,27 @@ protected:
   DominatorTreeBase(const DominatorTreeBase &) = delete;
   DominatorTreeBase &operator=(const DominatorTreeBase &) = delete;
 
-  /// getRoots - Return the root blocks of the current CFG.  This may include
-  /// multiple blocks if we are computing post dominators.  For forward
-  /// dominators, this will always be a single block (the entry node).
+  /// Iteration over roots.
   ///
-  const SmallVectorImpl<NodeT *> &getRoots() const { return Roots; }
+  /// This may include multiple blocks if we are computing post dominators.
+  /// For forward dominators, this will always be a single block (the entry
+  /// block).
+  using root_iterator = typename SmallVectorImpl<NodeT *>::iterator;
+  using const_root_iterator = typename SmallVectorImpl<NodeT *>::const_iterator;
+
+  root_iterator root_begin() { return Roots.begin(); }
+  const_root_iterator root_begin() const { return Roots.begin(); }
+  root_iterator root_end() { return Roots.end(); }
+  const_root_iterator root_end() const { return Roots.end(); }
+
+  size_t root_size() const { return Roots.size(); }
+
+  iterator_range<root_iterator> roots() {
+    return make_range(root_begin(), root_end());
+  }
+  iterator_range<const_root_iterator> roots() const {
+    return make_range(root_begin(), root_end());
+  }
 
   /// isPostDominator - Returns true if analysis based of postdoms
   ///
diff --git a/llvm/include/llvm/Support/GenericDomTreeConstruction.h b/llvm/include/llvm/Support/GenericDomTreeConstruction.h
index 1e9b0f2..bde59ff 100644
--- a/llvm/include/llvm/Support/GenericDomTreeConstruction.h
+++ b/llvm/include/llvm/Support/GenericDomTreeConstruction.h
@@ -1372,7 +1372,7 @@ struct SemiNCAInfo {
     if (!DT.DFSInfoValid || !DT.Parent)
       return true;
 
-    const NodePtr RootBB = IsPostDom ? nullptr : DT.getRoots()[0];
+    const NodePtr RootBB = IsPostDom ? nullptr : *DT.root_begin();
     const TreeNodePtr Root = DT.getNode(RootBB);
 
     auto PrintNodeAndDFSNums = [](const TreeNodePtr TN) {
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp b/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp
index ab0d216..4182966 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp
@@ -199,8 +199,7 @@ bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
   // If there's only one exit, we don't need to do anything, unless this is a
   // pixel shader and that exit is an infinite loop, since we still have to
   // insert an export in that case.
-  if (PDT.getRoots().size() <= 1 &&
-      F.getCallingConv() != CallingConv::AMDGPU_PS)
+  if (PDT.root_size() <= 1 && F.getCallingConv() != CallingConv::AMDGPU_PS)
     return false;
 
   LegacyDivergenceAnalysis &DA = getAnalysis<LegacyDivergenceAnalysis>();
@@ -217,7 +216,7 @@ bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
   bool InsertExport = false;
 
   bool Changed = false;
-  for (BasicBlock *BB : PDT.getRoots()) {
+  for (BasicBlock *BB : PDT.roots()) {
     if (isa<ReturnInst>(BB->getTerminator())) {
       if (!isUniformlyReached(DA, *BB))
         ReturningBlocks.push_back(BB);
diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 7e17254..dd8dc84 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -1853,7 +1853,7 @@ struct DSEState {
         if (CommonPred)
           WorkList.insert(CommonPred);
         else
-          for (BasicBlock *R : PDT.getRoots())
+          for (BasicBlock *R : PDT.roots())
             WorkList.insert(R);
 
         NumCFGTries++;
diff --git a/llvm/unittests/IR/DominatorTreeTest.cpp b/llvm/unittests/IR/DominatorTreeTest.cpp
index 66e1227..16c12b2 100644
--- a/llvm/unittests/IR/DominatorTreeTest.cpp
+++ b/llvm/unittests/IR/DominatorTreeTest.cpp
@@ -805,7 +805,7 @@ TEST(DominatorTree, InsertFromUnreachable) {
   BasicBlock *To = B.getOrAddBlock(LastUpdate->Edge.To);
   PDT.insertEdge(From, To);
   EXPECT_TRUE(PDT.verify());
-  EXPECT_TRUE(PDT.getRoots().size() == 2);
+  EXPECT_EQ(PDT.root_size(), 2);
   // Make sure we can use a const pointer with getNode.
   const BasicBlock *BB5 = B.getOrAddBlock("5");
   EXPECT_NE(PDT.getNode(BB5), nullptr);