Index: llvm/include/llvm/Support/GenericDomTree.h =================================================================== --- llvm/include/llvm/Support/GenericDomTree.h +++ llvm/include/llvm/Support/GenericDomTree.h @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -347,20 +348,29 @@ DomTreeNodeBase *getRootNode() { return RootNode; } const DomTreeNodeBase *getRootNode() const { return RootNode; } + /// Calls Callback on all nodes dominated by R, including R itself, in no + /// particular order. + void forEachDescendant( + DomTreeNodeBase *R, + std::function *)> Callback) const { + SmallVector *, 8> WL; + WL.push_back(R); + while (!WL.empty()) { + DomTreeNodeBase *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 &Result) const { Result.clear(); - const DomTreeNodeBase *RN = getNode(R); + DomTreeNodeBase *RN = getNode(R); if (!RN) return; // If R is unreachable, it will not be present in the DOM tree. - SmallVector *, 8> WL; - WL.push_back(RN); - - while (!WL.empty()) { - const DomTreeNodeBase *N = WL.pop_back_val(); - Result.push_back(N->getBlock()); - WL.append(N->begin(), N->end()); - } + forEachDescendant(RN, [&](DomTreeNodeBase *N) { + Result.emplace_back(N->getBlock()); + }); } /// properlyDominates - Returns true iff A dominates B and A != B. Index: llvm/unittests/IR/DominatorTreeTest.cpp =================================================================== --- llvm/unittests/IR/DominatorTreeTest.cpp +++ llvm/unittests/IR/DominatorTreeTest.cpp @@ -205,10 +205,17 @@ // Dominance descendants. SmallVector 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.