Expand cloneLoopWithPreheader() to support cloning loop nest

[Whitney]

Summary

cloneLoopWithPreheader() currently only support innermost loop, and assert otherwise.

[Meinersbur]

llvm/lib/Transforms/Utils/CloneFunction.cpp

744

Consider using DenseMap.

772

[style] I like getting a reference that can be updated instead of another map lookup. I.e.

Loop *&NewLoop = LMap[CurLoop];
...
NewLoop = LI->AllocateLoop();

777–778

[serious] Since the iteration order of OrigLoop->getBlocks() is undefined, we might get a block that is nested in two loops without having visited the parent first, i.e. NewParentLoop would be null in this case.

Consider iterating over LoopInfo's loop tree before iterating over all block to ensure that the new loop structure has been created.

[Whitney]

Addressed all review comments from Michael.

[Whitney]

llvm/lib/Transforms/Utils/CloneFunction.cpp

777–778

Good catch! Modified as suggested.

[Whitney]

@Meinersbur Thanks for the review! What do you think of the updated changeset?

[Meinersbur]

LGTM

llvm/lib/Transforms/Utils/CloneFunction.cpp

 ///
 /// Updates LoopInfo and DominatorTree assuming the loop is dominated by block
 /// \p LoopDomBB.  Insert the new blocks before block specified in \p Before.
 Loop *llvm::cloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
                                    Loop *OrigLoop, ValueToValueMapTy &VMap,
                                    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();
+  std::map<Loop *, Loop *> LMap;
 
   Loop *NewLoop = LI->AllocateLoop();
+  LMap[OrigLoop] = NewLoop;
   if (ParentLoop)
     ParentLoop->addChildLoop(NewLoop);
   else
     LI->addTopLevelLoop(NewLoop);
 
   BasicBlock *OrigPH = OrigLoop->getLoopPreheader();
   assert(OrigPH && "No preheader");
   BasicBlock *NewPH = CloneBasicBlock(OrigPH, VMap, NameSuffix, F);
   // To rename the loop PHIs.
   VMap[OrigPH] = NewPH;
   Blocks.push_back(NewPH);
 
   // Update LoopInfo.
   if (ParentLoop)
     ParentLoop->addBasicBlockToLoop(NewPH, *LI);
 
   // Update DominatorTree.
   DT->addNewBlock(NewPH, LoopDomBB);
 
   for (BasicBlock *BB : OrigLoop->getBlocks()) {
+    Loop *CurLoop = LI->getLoopFor(BB);
+    Loop *NewLoop = LMap[CurLoop];
+    if (!NewLoop) {
+      NewLoop = LI->AllocateLoop();
+      LMap[CurLoop] = NewLoop;
+
+      // 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);
+    }
+
     BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F);
     VMap[BB] = 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()) {