1 files changed, 20 insertions, 2 deletions
diff --git a/llvm/lib/Transforms/Instrumentation/CFGMST.h b/llvm/lib/Transforms/Instrumentation/CFGMST.h
index 16e2e6b4e73..a1c2670c659 100644
--- a/llvm/lib/Transforms/Instrumentation/CFGMST.h
+++ b/llvm/lib/Transforms/Instrumentation/CFGMST.h
@@ -20,6 +20,7 @@
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -136,6 +137,21 @@ public:
                      << " w = " << BBWeight << "\n");
       }
     }
+    // check if there is any infinite loop. If yes, add a fake edge from
+    // the header block to the fake node:
+    for (auto *L : *LI) {
+      SmallVector<BasicBlock *, 2> ExitingBlocks;
+      L->getExitingBlocks(ExitingBlocks);
+      if (!ExitingBlocks.empty())
+        continue;
+      auto *HB = L->getHeader();
+      if (!HB)
+        continue;
+      addEdge(HB, nullptr, UINT64_MAX);
+      DEBUG(dbgs() << "  Edge: from infinite loop header " 
+                   << HB->getName() << " to exit"
+                   << " w = " << UINT64_MAX << "\n");
+    }
   }
 
   // Sort CFG edges based on its weight.
@@ -212,11 +228,13 @@ public:
 
   BranchProbabilityInfo *BPI;
   BlockFrequencyInfo *BFI;
+  LoopInfo *LI;
 
 public:
   CFGMST(Function &Func, BranchProbabilityInfo *BPI_ = nullptr,
-         BlockFrequencyInfo *BFI_ = nullptr)
-      : F(Func), BPI(BPI_), BFI(BFI_) {
+         BlockFrequencyInfo *BFI_ = nullptr,
+         LoopInfo *LI_ = nullptr)
+      : F(Func), BPI(BPI_), BFI(BFI_), LI(LI_) {
     buildEdges();
     sortEdgesByWeight();
     computeMinimumSpanningTree();