1 files changed, 16 insertions, 5 deletions
diff --git a/llvm/lib/Analysis/IteratedDominanceFrontier.cpp b/llvm/lib/Analysis/IteratedDominanceFrontier.cpp
index e7751d32aab..000fe5ddad5 100644
--- a/llvm/lib/Analysis/IteratedDominanceFrontier.cpp
+++ b/llvm/lib/Analysis/IteratedDominanceFrontier.cpp
@@ -17,6 +17,7 @@
 #include <queue>
 
 namespace llvm {
+
 template <class NodeTy, bool IsPostDom>
 void IDFCalculator<NodeTy, IsPostDom>::calculate(
     SmallVectorImpl<BasicBlock *> &PHIBlocks) {
@@ -61,29 +62,39 @@ void IDFCalculator<NodeTy, IsPostDom>::calculate(
       BasicBlock *BB = Node->getBlock();
       // Succ is the successor in the direction we are calculating IDF, so it is
       // successor for IDF, and predecessor for Reverse IDF.
-      for (auto *Succ : children<NodeTy>(BB)) {
+      auto DoWork = [&](BasicBlock *Succ) {
         DomTreeNode *SuccNode = DT.getNode(Succ);
 
         // Quickly skip all CFG edges that are also dominator tree edges instead
         // of catching them below.
         if (SuccNode->getIDom() == Node)
-          continue;
+          return;
 
         const unsigned SuccLevel = SuccNode->getLevel();
         if (SuccLevel > RootLevel)
-          continue;
+          return;
 
         if (!VisitedPQ.insert(SuccNode).second)
-          continue;
+          return;
 
         BasicBlock *SuccBB = SuccNode->getBlock();
         if (useLiveIn && !LiveInBlocks->count(SuccBB))
-          continue;
+          return;
 
         PHIBlocks.emplace_back(SuccBB);
         if (!DefBlocks->count(SuccBB))
           PQ.push(std::make_pair(
               SuccNode, std::make_pair(SuccLevel, SuccNode->getDFSNumIn())));
+      };
+
+      if (GD) {
+        for (auto Pair : children<
+                 std::pair<const GraphDiff<BasicBlock *, IsPostDom> *, NodeTy>>(
+                 {GD, BB}))
+          DoWork(Pair.second);
+      } else {
+        for (auto *Succ : children<NodeTy>(BB))
+          DoWork(Succ);
       }
 
       for (auto DomChild : *Node) {