Add forEachDescendant to DominatorTree.

[nicholas]

Summary

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

[kuhar]

Would it make sense to use one of the llvm graph iterators and have descendants_begin() and descendants_end() functions instead?

[kuhar]

OK, I can see that this is exactly what you want to avoid in the related revision. In this case, isn't it possible to provide a dummy storage to DepthFirstIterator and use it to efficiently traverse trees?

llvm/include/llvm/Support/GenericDomTree.h

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CFGUpdate.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
+#include <functional>
 #include <iterator>
 #include <memory>
 #include <type_traits>
 #include <utility>
 #include <vector>
 
 namespace llvm {
 
public:
   /// this root may be a node with the block == NULL.  This is the case when
   /// there are multiple exit nodes from a particular function.  Consumers of
   /// post-dominance information must be capable of dealing with this
   /// possibility.
   ///
   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.
   /// Note that this is not a constant time operation!
   ///
   bool properlyDominates(const DomTreeNodeBase<NodeT> *A,
                          const DomTreeNodeBase<NodeT> *B) const {
     if (!A || !B)

llvm/unittests/IR/DominatorTreeTest.cpp

runWithDomTree(
         EXPECT_TRUE(PDT->dominates(BB0, BB0));
         EXPECT_FALSE(PDT->dominates(BB1, BB0));
         EXPECT_FALSE(PDT->dominates(BB2, BB0));
         EXPECT_FALSE(PDT->dominates(BB3, BB0));
         EXPECT_TRUE(PDT->dominates(BB4, BB1));
 
         // 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.
         DominatedBBs.clear();
         PostDominatedBBs.clear();
         DT->getDescendants(BB3, DominatedBBs);
         DT->getDescendants(BB3, PostDominatedBBs);
         EXPECT_EQ(DominatedBBs.size(), 0UL);