1 files changed, 31 insertions, 10 deletions

diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp
index 5dd7a5bd8c7..9e160342c8c 100644
--- a/llvm/lib/Transforms/Utils/CloneFunction.cpp
+++ b/llvm/lib/Transforms/Utils/CloneFunction.cpp
@@ -739,12 +739,12 @@ Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
                                    const Twine &NameSuffix, LoopInfo *LI,
                                    DominatorTree *DT,
                                    SmallVectorImpl<BasicBlock *> &Blocks) {
-  assert(OrigLoop->getSubLoops().empty() &&
-         "Loop to be cloned cannot have inner loop");
   Function *F = OrigLoop->getHeader()->getParent();
   Loop *ParentLoop = OrigLoop->getParentLoop();
+  DenseMap<Loop *, Loop *> LMap;
 
   Loop *NewLoop = LI->AllocateLoop();
+  LMap[OrigLoop] = NewLoop;
   if (ParentLoop)
     ParentLoop->addChildLoop(NewLoop);
   else
@@ -764,17 +764,38 @@ Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
   // Update DominatorTree.
   DT->addNewBlock(NewPH, LoopDomBB);
 
-  for (BasicBlock *BB : OrigLoop->getBlocks()) {
-    BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F);
-    VMap[BB] = NewBB;
+  for (Loop *CurLoop : OrigLoop->getLoopsInPreorder()) {
+    for (BasicBlock *BB : CurLoop->getBlocks()) {
+      if (CurLoop != LI->getLoopFor(BB))
+        continue;
+
+      Loop *&NewLoop = LMap[CurLoop];
+      if (!NewLoop) {
+        NewLoop = LI->AllocateLoop();
 
-    // Update LoopInfo.
-    NewLoop->addBasicBlockToLoop(NewBB, *LI);
+        // Establish the parent/child relationship.
+        Loop *OrigParent = CurLoop->getParentLoop();
+        assert(OrigParent && "Could not find the original parent loop");
+        Loop *NewParentLoop = LMap[OrigParent];
+        assert(NewParentLoop && "Could not find the new parent loop");
+
+        NewParentLoop->addChildLoop(NewLoop);
+      }
 
-    // Add DominatorTree node. After seeing all blocks, update to correct IDom.
-    DT->addNewBlock(NewBB, NewPH);
+      BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F);
+      VMap[BB] = NewBB;
 
-    Blocks.push_back(NewBB);
+      // Update LoopInfo.
+      NewLoop->addBasicBlockToLoop(NewBB, *LI);
+      if (BB == CurLoop->getHeader())
+        NewLoop->moveToHeader(NewBB);
+
+      // Add DominatorTree node. After seeing all blocks, update to correct
+      // IDom.
+      DT->addNewBlock(NewBB, NewPH);
+
+      Blocks.push_back(NewBB);
+    }
   }
 
   for (BasicBlock *BB : OrigLoop->getBlocks()) {