[PATCH] D60754: Add forEachDescendant to DominatorTree.

[Nick Lewycky via Phabricator via llvm-commits]

nicholas created this revision.
Herald added subscribers: llvm-commits, kuhar.
Herald added a project: LLVM.

DominatorTree already has getDescendants, but that always builds up a SmallVector. This is more efficient when the caller doesn't need a vector.


https://reviews.llvm.org/D60754

Files:
  llvm/include/llvm/Support/GenericDomTree.h
  llvm/unittests/IR/DominatorTreeTest.cpp


Index: llvm/unittests/IR/DominatorTreeTest.cpp
===================================================================
--- llvm/unittests/IR/DominatorTreeTest.cpp
+++ llvm/unittests/IR/DominatorTreeTest.cpp
@@ -205,10 +205,17 @@
 
         // Dominance descendants.
         SmallVector<BasicBlock *, 8> DominatedBBs, PostDominatedBBs;
+        int CallCount = 0, PostCallCount = 0;
 
+        DT->forEachDescendant(DT->getNode(BB0),
+                              [&](DomTreeNode *) { ++CallCount; });
         DT->getDescendants(BB0, DominatedBBs);
+        PDT->forEachDescendant(DT->getNode(BB0),
+                               [&](DomTreeNode *) { ++PostCallCount; });
         PDT->getDescendants(BB0, PostDominatedBBs);
+        EXPECT_EQ(CallCount, 4);
         EXPECT_EQ(DominatedBBs.size(), 4UL);
+        EXPECT_EQ(PostCallCount, 4);
         EXPECT_EQ(PostDominatedBBs.size(), 1UL);
 
         // BB3 is unreachable. It should have no dominators nor postdominators.
Index: llvm/include/llvm/Support/GenericDomTree.h
===================================================================
--- llvm/include/llvm/Support/GenericDomTree.h
+++ llvm/include/llvm/Support/GenericDomTree.h
@@ -34,6 +34,7 @@
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
+#include <functional>
 #include <iterator>
 #include <memory>
 #include <type_traits>
@@ -347,20 +348,29 @@
   DomTreeNodeBase<NodeT> *getRootNode() { return RootNode; }
   const DomTreeNodeBase<NodeT> *getRootNode() const { return RootNode; }
 
+  /// Calls Callback on all nodes dominated by R, including R itself, in no
+  /// particular order.
+  void forEachDescendant(
+      DomTreeNodeBase<NodeT> *R,
+      std::function<void(DomTreeNodeBase<NodeT> *)> Callback) const {
+    SmallVector<DomTreeNodeBase<NodeT> *, 8> WL;
+    WL.push_back(R);
+    while (!WL.empty()) {
+      DomTreeNodeBase<NodeT> *N = WL.pop_back_val();
+      Callback(N);
+      WL.append(N->begin(), N->end());
+    }
+  }
+
   /// Get all nodes dominated by R, including R itself.
   void getDescendants(NodeT *R, SmallVectorImpl<NodeT *> &Result) const {
     Result.clear();
-    const DomTreeNodeBase<NodeT> *RN = getNode(R);
+    DomTreeNodeBase<NodeT> *RN = getNode(R);
     if (!RN)
       return; // If R is unreachable, it will not be present in the DOM tree.
-    SmallVector<const DomTreeNodeBase<NodeT> *, 8> WL;
-    WL.push_back(RN);
-
-    while (!WL.empty()) {
-      const DomTreeNodeBase<NodeT> *N = WL.pop_back_val();
-      Result.push_back(N->getBlock());
-      WL.append(N->begin(), N->end());
-    }
+    forEachDescendant(RN, [&](DomTreeNodeBase<NodeT> *N) {
+      Result.emplace_back(N->getBlock());
+    });
   }
 
   /// properlyDominates - Returns true iff A dominates B and A != B.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D60754.195312.patch
Type: text/x-patch
Size: 2763 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20190416/75fb1a91/attachment.bin>