2 files changed, 53 insertions, 32 deletions

diff --git a/lib/Transforms/Scalar/LoopRotation.cpp b/lib/Transforms/Scalar/LoopRotation.cpp
index 199c42009d..d1c975f2e9 100644
--- a/lib/Transforms/Scalar/LoopRotation.cpp
+++ b/lib/Transforms/Scalar/LoopRotation.cpp
@@ -282,30 +282,57 @@ bool LoopRotate::rotateLoop(Loop *L) {
   L->moveToHeader(NewHeader);
   assert(L->getHeader() == NewHeader && "Latch block is our new header");
 
-  // Update DominatorTree to reflect the CFG change we just made.  Then split
-  // edges as necessary to preserve LoopSimplify form.
-  if (DominatorTree *DT = getAnalysisIfAvailable<DominatorTree>()) {
-    // Since OrigPreheader now has the conditional branch to Exit block, it is
-    // the dominator of Exit.
-    DT->changeImmediateDominator(Exit, OrigPreheader);
-    DT->changeImmediateDominator(NewHeader, OrigPreheader);
+  
+  // At this point, we've finished our major CFG changes.  As part of cloning
+  // the loop into the preheader we've simplified instructions and the
+  // duplicated conditional branch may now be branching on a constant.  If it is
+  // branching on a constant and if that constant means that we enter the loop,
+  // then we fold away the cond branch to an uncond branch.  This simplifies the
+  // loop in cases important for nested loops, and it also means we don't have
+  // to split as many edges.
+  BranchInst *PHBI = cast<BranchInst>(OrigPreheader->getTerminator());
+  assert(PHBI->isConditional() && "Should be clone of BI condbr!");
+  if (!isa<ConstantInt>(PHBI->getCondition()) ||
+      PHBI->getSuccessor(cast<ConstantInt>(PHBI->getCondition())->isZero())
+          != NewHeader) {
+    // The conditional branch can't be folded, handle the general case.
+    // Update DominatorTree to reflect the CFG change we just made.  Then split
+    // edges as necessary to preserve LoopSimplify form.
+    if (DominatorTree *DT = getAnalysisIfAvailable<DominatorTree>()) {
+      // Since OrigPreheader now has the conditional branch to Exit block, it is
+      // the dominator of Exit.
+      DT->changeImmediateDominator(Exit, OrigPreheader);
+      DT->changeImmediateDominator(NewHeader, OrigPreheader);
+      
+      // Update OrigHeader to be dominated by the new header block.
+      DT->changeImmediateDominator(OrigHeader, OrigLatch);
+    }
+    
+    // Right now OrigPreHeader has two successors, NewHeader and ExitBlock, and
+    // thus is not a preheader anymore.  Split the edge to form a real preheader.
+    BasicBlock *NewPH = SplitCriticalEdge(OrigPreheader, NewHeader, this);
+    NewPH->setName(NewHeader->getName() + ".lr.ph");
     
-    // Update OrigHeader to be dominated by the new header block.
-    DT->changeImmediateDominator(OrigHeader, OrigLatch);
+    // Preserve canonical loop form, which means that 'Exit' should have only one
+    // predecessor.
+    BasicBlock *ExitSplit = SplitCriticalEdge(L->getLoopLatch(), Exit, this);
+    ExitSplit->moveBefore(Exit);
+  } else {
+    // We can fold the conditional branch in the preheader, this makes things
+    // simpler. The first step is to remove the extra edge to the Exit block.
+    Exit->removePredecessor(OrigPreheader, true /*preserve LCSSA*/);
+    BranchInst::Create(NewHeader, PHBI);
+    PHBI->eraseFromParent();
+    
+    // With our CFG finalized, update DomTree if it is available.
+    if (DominatorTree *DT = getAnalysisIfAvailable<DominatorTree>()) {
+      // Update OrigHeader to be dominated by the new header block.
+      DT->changeImmediateDominator(NewHeader, OrigPreheader);
+      DT->changeImmediateDominator(OrigHeader, OrigLatch);
+    }
   }
   
-  // Right now OrigPreHeader has two successors, NewHeader and ExitBlock, and
-  // thus is not a preheader anymore.  Split the edge to form a real preheader.
-  BasicBlock *NewPH = SplitCriticalEdge(OrigPreheader, NewHeader, this);
-  NewPH->setName(NewHeader->getName() + ".lr.ph");
-  
-  // Preserve canonical loop form, which means that 'Exit' should have only one
-  // predecessor.
-  BasicBlock *ExitSplit = SplitCriticalEdge(L->getLoopLatch(), Exit, this);
-  ExitSplit->moveBefore(Exit);
-  
-  assert(L->getLoopPreheader() == NewPH &&
-         "Invalid loop preheader after loop rotation");
+  assert(L->getLoopPreheader() && "Invalid loop preheader after loop rotation");
   assert(L->getLoopLatch() && "Invalid loop latch after loop rotation");
 
   // Now that the CFG and DomTree are in a consistent state again, merge the
diff --git a/test/Transforms/LoopRotate/phi-duplicate.ll b/test/Transforms/LoopRotate/phi-duplicate.ll
index 96bec82c19..7372830922 100644
--- a/test/Transforms/LoopRotate/phi-duplicate.ll
+++ b/test/Transforms/LoopRotate/phi-duplicate.ll
@@ -28,19 +28,13 @@ for.end:                                          ; preds = %for.cond
   ret void
 }
 
-; Should only end up with one phi. Also, the original for.cond block should
-; be moved to the end of the loop so that the new loop header pleasantly
-; ends up at the top.
-
+; Should only end up with one phi.
 ; CHECK:      define void @test
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   br i1 true, label %for.body.lr.ph, label %for.end
-; CHECK-NOT:  :
-; CHECK:      for.body.lr.ph:
 ; CHECK-NEXT:   br label %for.body
-; CHECK-NOT:  :
 ; CHECK:      for.body:
 ; CHECK-NEXT:   %j.01 = phi i64
-; CHECK-NOT:    phi
-; CHECK:        ret void
-; CHECK-NEXT: }
+; CHECK-NOT:  br
+; CHECK:   br i1 %cmp, label %for.body, label %for.end
+; CHECK:      for.end:
+; CHECK-NEXT:        ret void