start of some new simplification code, not thoroughly tested, use at your own
[oota-llvm.git] / lib / Transforms / Scalar / LoopUnswitch.cpp
index 09f4976f100ba416643fe7465570513bc87b8a62..9f493b5912c40e3d02763a37076ff1d04dba1bda 100644 (file)
 using namespace llvm;
 
 namespace {
-  Statistic<> NumUnswitched("loop-unswitch", "Number of loops unswitched");
+  Statistic<> NumBranches("loop-unswitch", "Number of branches unswitched");
+  Statistic<> NumSwitches("loop-unswitch", "Number of switches unswitched");
+  Statistic<> NumSelects ("loop-unswitch", "Number of selects unswitched");
+  Statistic<> NumTrivial ("loop-unswitch",
+                          "Number of unswitches that are trivial");
+  Statistic<> NumSimplify("loop-unswitch", 
+                          "Number of simplifications of unswitched code");
   cl::opt<unsigned>
   Threshold("loop-unswitch-threshold", cl::desc("Max loop size to unswitch"),
             cl::init(10), cl::Hidden);
@@ -66,12 +72,16 @@ namespace {
     }
 
   private:
+    bool UnswitchIfProfitable(Value *LoopCond, Constant *Val,Loop *L);
     unsigned getLoopUnswitchCost(Loop *L, Value *LIC);
-    void VersionLoop(Value *LIC, Loop *L, Loop *&Out1, Loop *&Out2);
-    BasicBlock *SplitBlock(BasicBlock *BB, bool SplitAtTop);
-    void RewriteLoopBodyWithConditionConstant(Loop *L, Value *LIC, bool Val);
-    void UnswitchTrivialCondition(Loop *L, Value *Cond, bool EntersLoopOnCond,
-                                  BasicBlock *ExitBlock);
+    void VersionLoop(Value *LIC, Constant *OnVal,
+                     Loop *L, Loop *&Out1, Loop *&Out2);
+    BasicBlock *SplitEdge(BasicBlock *From, BasicBlock *To);
+    BasicBlock *SplitBlock(BasicBlock *Old, Instruction *SplitPt);
+    void RewriteLoopBodyWithConditionConstant(Loop *L, Value *LIC,Constant *Val,
+                                              bool isEqual);
+    void UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
+                                  bool EntersWhenTrue, BasicBlock *ExitBlock);
   };
   RegisterOpt<LoopUnswitch> X("loop-unswitch", "Unswitch loops");
 }
@@ -116,26 +126,50 @@ static bool LoopValuesUsedOutsideLoop(Loop *L) {
   return false;
 }
 
-/// FindTrivialLoopExitBlock - We know that we have a branch from the loop
-/// header to the specified latch block.   See if one of the successors of the
-/// latch block is an exit, and if so what block it is.
-static BasicBlock *FindTrivialLoopExitBlock(Loop *L, BasicBlock *Latch) {
+/// isTrivialLoopExitBlock - Check to see if all paths from BB either:
+///   1. Exit the loop with no side effects.
+///   2. Branch to the latch block with no side-effects.
+///
+/// If these conditions are true, we return true and set ExitBB to the block we
+/// exit through.
+///
+static bool isTrivialLoopExitBlockHelper(Loop *L, BasicBlock *BB,
+                                         BasicBlock *&ExitBB,
+                                         std::set<BasicBlock*> &Visited) {
   BasicBlock *Header = L->getHeader();
-  BranchInst *LatchBranch = dyn_cast<BranchInst>(Latch->getTerminator());
-  if (!LatchBranch || !LatchBranch->isConditional()) return 0;
-  
-  // Simple case, the latch block is a conditional branch.  The target that
-  // doesn't go to the loop header is our block if it is not in the loop.
-  if (LatchBranch->getSuccessor(0) == Header) {
-    if (L->contains(LatchBranch->getSuccessor(1))) return false;
-    return LatchBranch->getSuccessor(1);
-  } else {
-    assert(LatchBranch->getSuccessor(1) == Header);
-    if (L->contains(LatchBranch->getSuccessor(0))) return false;
-    return LatchBranch->getSuccessor(0);
+  for (succ_iterator SI = succ_begin(BB), E = succ_end(BB); SI != E; ++SI) {
+    if (!Visited.insert(*SI).second) {
+      // Already visited and Ok, end of recursion.
+    } else if (L->contains(*SI)) {
+      // Check to see if the successor is a trivial loop exit.
+      if (!isTrivialLoopExitBlockHelper(L, *SI, ExitBB, Visited))
+        return false;
+    } else {
+      // Otherwise, this is a loop exit, this is fine so long as this is the
+      // first exit.
+      if (ExitBB != 0) return false;
+      ExitBB = *SI;
+    }
   }
+
+  // Okay, everything after this looks good, check to make sure that this block
+  // doesn't include any side effects.
+  for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
+    if (I->mayWriteToMemory())
+      return false;
+  
+  return true;
 }
 
+static BasicBlock *isTrivialLoopExitBlock(Loop *L, BasicBlock *BB) {
+  std::set<BasicBlock*> Visited;
+  Visited.insert(L->getHeader());  // Branches to header are ok.
+  Visited.insert(BB);              // Don't revisit BB after we do.
+  BasicBlock *ExitBB = 0;
+  if (isTrivialLoopExitBlockHelper(L, BB, ExitBB, Visited))
+    return ExitBB;
+  return 0;
+}
 
 /// IsTrivialUnswitchCondition - Check to see if this unswitch condition is
 /// trivial: that is, that the condition controls whether or not the loop does
@@ -147,44 +181,58 @@ static BasicBlock *FindTrivialLoopExitBlock(Loop *L, BasicBlock *Latch) {
 /// runs when the condition is true, False if the loop body executes when the
 /// condition is false.  Otherwise, return null to indicate a complex condition.
 static bool IsTrivialUnswitchCondition(Loop *L, Value *Cond,
-                                       bool *CondEntersLoop = 0,
+                                       Constant **Val = 0,
+                                       bool *EntersWhenTrue = 0,
                                        BasicBlock **LoopExit = 0) {
   BasicBlock *Header = L->getHeader();
-  BranchInst *HeaderTerm = dyn_cast<BranchInst>(Header->getTerminator());
-  
-  // If the header block doesn't end with a conditional branch on Cond, we can't
-  // handle it.
-  if (!HeaderTerm || !HeaderTerm->isConditional() ||
-      HeaderTerm->getCondition() != Cond)
-    return false;
+  TerminatorInst *HeaderTerm = Header->getTerminator();
+
+  BasicBlock *LoopExitBB = 0;
+  if (BranchInst *BI = dyn_cast<BranchInst>(HeaderTerm)) {
+    // If the header block doesn't end with a conditional branch on Cond, we
+    // can't handle it.
+    if (!BI->isConditional() || BI->getCondition() != Cond)
+      return false;
   
-  // Check to see if the conditional branch goes to the latch block.  If not,
-  // it's not trivial.  This also determines the value of Cond that will execute
-  // the loop.
-  BasicBlock *Latch = L->getLoopLatch();
-  if (HeaderTerm->getSuccessor(1) == Latch) {
-    if (CondEntersLoop) *CondEntersLoop = true;
-  } else if (HeaderTerm->getSuccessor(0) == Latch)
-    if (CondEntersLoop) *CondEntersLoop = false;
-  else
-    return false;  // Doesn't branch to latch block.
+    // Check to see if a successor of the branch is guaranteed to go to the
+    // latch block or exit through a one exit block without having any 
+    // side-effects.  If so, determine the value of Cond that causes it to do
+    // this.
+    if ((LoopExitBB = isTrivialLoopExitBlock(L, BI->getSuccessor(0)))) {
+      if (Val) *Val = ConstantBool::False;
+    } else if ((LoopExitBB = isTrivialLoopExitBlock(L, BI->getSuccessor(1)))) {
+      if (Val) *Val = ConstantBool::True;
+    }
+  } else if (SwitchInst *SI = dyn_cast<SwitchInst>(HeaderTerm)) {
+    // If this isn't a switch on Cond, we can't handle it.
+    if (SI->getCondition() != Cond) return false;
+    
+    // Check to see if a successor of the switch is guaranteed to go to the
+    // latch block or exit through a one exit block without having any 
+    // side-effects.  If so, determine the value of Cond that causes it to do
+    // this.  Note that we can't trivially unswitch on the default case.
+    for (unsigned i = 1, e = SI->getNumSuccessors(); i != e; ++i)
+      if ((LoopExitBB = isTrivialLoopExitBlock(L, SI->getSuccessor(i)))) {
+        // Okay, we found a trivial case, remember the value that is trivial.
+        if (Val) *Val = SI->getCaseValue(i);
+        if (EntersWhenTrue) *EntersWhenTrue = false;
+        break;
+      }
+  }
+
+  if (!LoopExitBB)
+    return false;   // Can't handle this.
   
-  // The latch block must end with a conditional branch where one edge goes to
-  // the header (this much we know) and one edge goes OUT of the loop.
-  BasicBlock *LoopExitBlock = FindTrivialLoopExitBlock(L, Latch);
-  if (!LoopExitBlock) return 0;
-  if (LoopExit) *LoopExit = LoopExitBlock;
+  if (LoopExit) *LoopExit = LoopExitBB;
   
   // We already know that nothing uses any scalar values defined inside of this
   // loop.  As such, we just have to check to see if this loop will execute any
   // side-effecting instructions (e.g. stores, calls, volatile loads) in the
-  // part of the loop that the code *would* execute.
+  // part of the loop that the code *would* execute.  We already checked the
+  // tail, check the header now.
   for (BasicBlock::iterator I = Header->begin(), E = Header->end(); I != E; ++I)
     if (I->mayWriteToMemory())
       return false;
-  for (BasicBlock::iterator I = Latch->begin(), E = Latch->end(); I != E; ++I)
-    if (I->mayWriteToMemory())
-      return false;
   return true;
 }
 
@@ -256,106 +304,148 @@ bool LoopUnswitch::visitLoop(Loop *L) {
   // loop.
   for (Loop::block_iterator I = L->block_begin(), E = L->block_end();
        I != E; ++I) {
-    for (BasicBlock::iterator BBI = (*I)->begin(), E = (*I)->end(); 
-         BBI != E; ++BBI)
     TerminatorInst *TI = (*I)->getTerminator();
-    // FIXME: Handle invariant select instructions.
-    
-    if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
-      if (!isa<Constant>(SI) && L->isLoopInvariant(SI->getCondition()))
-        DEBUG(std::cerr << "TODO: Implement unswitching 'switch' loop %"
-              << L->getHeader()->getName() << ", cost = "
-              << L->getBlocks().size() << "\n" << **I);
-      continue;
+    if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
+      // If this isn't branching on an invariant condition, we can't unswitch
+      // it.
+      if (BI->isConditional()) {
+        // See if this, or some part of it, is loop invariant.  If so, we can
+        // unswitch on it if we desire.
+        Value *LoopCond = FindLIVLoopCondition(BI->getCondition(), L, Changed);
+        if (LoopCond && UnswitchIfProfitable(LoopCond, ConstantBool::True, L)) {
+          ++NumBranches;
+          return true;
+        }
+      }      
+    } else if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
+      Value *LoopCond = FindLIVLoopCondition(SI->getCondition(), L, Changed);
+      if (LoopCond && SI->getNumCases() > 1) {
+        // Find a value to unswitch on:
+        // FIXME: this should chose the most expensive case!
+        Constant *UnswitchVal = SI->getCaseValue(1);
+        if (UnswitchIfProfitable(LoopCond, UnswitchVal, L)) {
+          ++NumSwitches;
+          return true;
+        }
+      }
     }
     
-    BranchInst *BI = dyn_cast<BranchInst>(TI);
-    if (!BI) continue;
-    
-    // If this isn't branching on an invariant condition, we can't unswitch it.
-    if (!BI->isConditional())
-      continue;
-    
-    // See if this, or some part of it, is loop invariant.  If so, we can
-    // unswitch on it if we desire.
-    Value *LoopCond = FindLIVLoopCondition(BI->getCondition(), L, Changed);
-    if (LoopCond == 0) continue;
+    // Scan the instructions to check for unswitchable values.
+    for (BasicBlock::iterator BBI = (*I)->begin(), E = (*I)->end(); 
+         BBI != E; ++BBI)
+      if (SelectInst *SI = dyn_cast<SelectInst>(BBI)) {
+        Value *LoopCond = FindLIVLoopCondition(SI->getCondition(), L, Changed);
+        if (LoopCond && UnswitchIfProfitable(LoopCond, ConstantBool::True, L)) {
+          ++NumSelects;
+          return true;
+        }
+      }
+  }
     
-    // Check to see if it would be profitable to unswitch this loop.
-    if (getLoopUnswitchCost(L, LoopCond) > Threshold) {
-      // FIXME: this should estimate growth by the amount of code shared by the
-      // resultant unswitched loops.  This should have no code growth:
-      //    for () { if (iv) {...} }
-      // as one copy of the loop will be empty.
-      //
-      DEBUG(std::cerr << "NOT unswitching loop %"
-            << L->getHeader()->getName() << ", cost too high: "
-            << L->getBlocks().size() << "\n");
-      continue;
-    }
+  return Changed;
+}
+
+/// UnswitchIfProfitable - We have found that we can unswitch L when
+/// LoopCond == Val to simplify the loop.  If we decide that this is profitable,
+/// unswitch the loop, reprocess the pieces, then return true.
+bool LoopUnswitch::UnswitchIfProfitable(Value *LoopCond, Constant *Val,Loop *L){
+  // Check to see if it would be profitable to unswitch this loop.
+  if (getLoopUnswitchCost(L, LoopCond) > Threshold) {
+    // FIXME: this should estimate growth by the amount of code shared by the
+    // resultant unswitched loops.
+    //
+    DEBUG(std::cerr << "NOT unswitching loop %"
+                    << L->getHeader()->getName() << ", cost too high: "
+                    << L->getBlocks().size() << "\n");
+    return false;
+  }
     
-    // If this loop has live-out values, we can't unswitch it. We need something
-    // like loop-closed SSA form in order to know how to insert PHI nodes for
-    // these values.
-    if (LoopValuesUsedOutsideLoop(L)) {
-      DEBUG(std::cerr << "NOT unswitching loop %"
-                      << L->getHeader()->getName()
-                      << ", a loop value is used outside loop!\n");
-      continue;
-    }
+  // If this loop has live-out values, we can't unswitch it. We need something
+  // like loop-closed SSA form in order to know how to insert PHI nodes for
+  // these values.
+  if (LoopValuesUsedOutsideLoop(L)) {
+    DEBUG(std::cerr << "NOT unswitching loop %" << L->getHeader()->getName()
+                    << ", a loop value is used outside loop!\n");
+    return false;
+  }
       
-    //std::cerr << "BEFORE:\n"; LI->dump();
-    Loop *NewLoop1 = 0, *NewLoop2 = 0;
+  //std::cerr << "BEFORE:\n"; LI->dump();
+  Loop *NewLoop1 = 0, *NewLoop2 = 0;
  
-    // If this is a trivial condition to unswitch (which results in no code
-    // duplication), do it now.
-    bool EntersLoopOnCond;
-    BasicBlock *ExitBlock;
-    if (IsTrivialUnswitchCondition(L, LoopCond, &EntersLoopOnCond, &ExitBlock)){
-      UnswitchTrivialCondition(L, LoopCond, EntersLoopOnCond, ExitBlock);
-      NewLoop1 = L;
-    } else {
-      VersionLoop(LoopCond, L, NewLoop1, NewLoop2);
-    }
-    
-    //std::cerr << "AFTER:\n"; LI->dump();
-    
-    // Try to unswitch each of our new loops now!
-    if (NewLoop1) visitLoop(NewLoop1);
-    if (NewLoop2) visitLoop(NewLoop2);
-    return true;
+  // If this is a trivial condition to unswitch (which results in no code
+  // duplication), do it now.
+  Constant *CondVal;
+  bool EntersWhenTrue = true;
+  BasicBlock *ExitBlock;
+  if (IsTrivialUnswitchCondition(L, LoopCond, &CondVal,
+                                 &EntersWhenTrue, &ExitBlock)) {
+    UnswitchTrivialCondition(L, LoopCond, CondVal, EntersWhenTrue, ExitBlock);
+    NewLoop1 = L;
+  } else {
+    VersionLoop(LoopCond, Val, L, NewLoop1, NewLoop2);
   }
-
-  return Changed;
+  
+  //std::cerr << "AFTER:\n"; LI->dump();
+  
+  // Try to unswitch each of our new loops now!
+  if (NewLoop1) visitLoop(NewLoop1);
+  if (NewLoop2) visitLoop(NewLoop2);
+  return true;
 }
 
-/// SplitBlock - Split the specified basic block into two pieces.  If SplitAtTop
-/// is false, this splits the block so the second half only has an unconditional
-/// branch.  If SplitAtTop is true, it makes it so the first half of the block
-/// only has an unconditional branch in it.
+/// SplitBlock - Split the specified block at the specified instruction - every
+/// thing before SplitPt stays in Old and everything starting with SplitPt moves
+/// to a new block.  The two blocks are joined by an unconditional branch and
+/// the loop info is updated.
 ///
-/// This method updates the LoopInfo for this function to correctly reflect the
-/// CFG changes made.
-///
-/// This routine returns the new basic block that was inserted, which is always
-/// the later part of the block.
-BasicBlock *LoopUnswitch::SplitBlock(BasicBlock *BB, bool SplitAtTop) {
-  BasicBlock::iterator SplitPoint;
-  if (!SplitAtTop)
-    SplitPoint = BB->getTerminator();
-  else {
-    SplitPoint = BB->begin();
-    while (isa<PHINode>(SplitPoint)) ++SplitPoint;
-  }
-  
-  BasicBlock *New = BB->splitBasicBlock(SplitPoint, BB->getName()+".tail");
-  // New now lives in whichever loop that BB used to.
-  if (Loop *L = LI->getLoopFor(BB))
+BasicBlock *LoopUnswitch::SplitBlock(BasicBlock *Old, Instruction *SplitPt) {
+  BasicBlock::iterator SplitIt = SplitPt;
+  while (isa<PHINode>(SplitIt))
+    ++SplitIt;
+  BasicBlock *New = Old->splitBasicBlock(SplitIt, Old->getName()+".split");
+
+  // The new block lives in whichever loop the old one did.
+  if (Loop *L = LI->getLoopFor(Old))
     L->addBasicBlockToLoop(New, *LI);
+  
   return New;
 }
 
 
+BasicBlock *LoopUnswitch::SplitEdge(BasicBlock *BB, BasicBlock *Succ) {
+  TerminatorInst *LatchTerm = BB->getTerminator();
+  unsigned SuccNum = 0;
+  for (unsigned i = 0, e = LatchTerm->getNumSuccessors(); ; ++i) {
+    assert(i != e && "Didn't find edge?");
+    if (LatchTerm->getSuccessor(i) == Succ) {
+      SuccNum = i;
+      break;
+    }
+  }
+  
+  // If this is a critical edge, let SplitCriticalEdge do it.
+  if (SplitCriticalEdge(BB->getTerminator(), SuccNum, this))
+    return LatchTerm->getSuccessor(SuccNum);
+
+  // If the edge isn't critical, then BB has a single successor or Succ has a
+  // single pred.  Split the block.
+  BasicBlock::iterator SplitPoint;
+  if (BasicBlock *SP = Succ->getSinglePredecessor()) {
+    // If the successor only has a single pred, split the top of the successor
+    // block.
+    assert(SP == BB && "CFG broken");
+    return SplitBlock(Succ, Succ->begin());
+  } else {
+    // Otherwise, if BB has a single successor, split it at the bottom of the
+    // block.
+    assert(BB->getTerminator()->getNumSuccessors() == 1 &&
+           "Should have a single succ!"); 
+    return SplitBlock(BB, BB->getTerminator());
+  }
+}
+  
+
+
 // RemapInstruction - Convert the instruction operands from referencing the
 // current values into those specified by ValueMap.
 //
@@ -393,86 +483,98 @@ static Loop *CloneLoop(Loop *L, Loop *PL, std::map<const Value*, Value*> &VM,
   return New;
 }
 
+/// EmitPreheaderBranchOnCondition - Emit a conditional branch on two values
+/// if LIC == Val, branch to TrueDst, otherwise branch to FalseDest.  Insert the
+/// code immediately before InsertPt.
+static void EmitPreheaderBranchOnCondition(Value *LIC, Constant *Val,
+                                           BasicBlock *TrueDest,
+                                           BasicBlock *FalseDest,
+                                           Instruction *InsertPt) {
+  // Insert a conditional branch on LIC to the two preheaders.  The original
+  // code is the true version and the new code is the false version.
+  Value *BranchVal = LIC;
+  if (!isa<ConstantBool>(Val)) {
+    BranchVal = BinaryOperator::createSetEQ(LIC, Val, "tmp", InsertPt);
+  } else if (Val != ConstantBool::True) {
+    // We want to enter the new loop when the condition is true.
+    std::swap(TrueDest, FalseDest);
+  }
+
+  // Insert the new branch.
+  new BranchInst(TrueDest, FalseDest, BranchVal, InsertPt);
+}
+
+
 /// UnswitchTrivialCondition - Given a loop that has a trivial unswitchable
 /// condition in it (a cond branch from its header block to its latch block,
 /// where the path through the loop that doesn't execute its body has no 
 /// side-effects), unswitch it.  This doesn't involve any code duplication, just
 /// moving the conditional branch outside of the loop and updating loop info.
 void LoopUnswitch::UnswitchTrivialCondition(Loop *L, Value *Cond, 
-                                            bool EnterOnCond,
+                                            Constant *Val, bool EntersWhenTrue,
                                             BasicBlock *ExitBlock) {
   DEBUG(std::cerr << "loop-unswitch: Trivial-Unswitch loop %"
         << L->getHeader()->getName() << " [" << L->getBlocks().size()
         << " blocks] in Function " << L->getHeader()->getParent()->getName()
-        << " on cond:" << *Cond << "\n");
+        << " on cond: " << *Val << (EntersWhenTrue ? " == " : " != ") << 
+        *Cond << "\n");
   
-  // First step, split the preahder, so that we know that there is a safe place
+  // First step, split the preheader, so that we know that there is a safe place
   // to insert the conditional branch.  We will change 'OrigPH' to have a
   // conditional branch on Cond.
   BasicBlock *OrigPH = L->getLoopPreheader();
-  BasicBlock *NewPH = SplitBlock(OrigPH, false);
+  BasicBlock *NewPH = SplitEdge(OrigPH, L->getHeader());
 
   // Now that we have a place to insert the conditional branch, create a place
   // to branch to: this is the exit block out of the loop that we should
   // short-circuit to.
   
-  // Split this block now, so that the loop maintains its exit block.
+  // Split this block now, so that the loop maintains its exit block, and so
+  // that the jump from the preheader can execute the contents of the exit block
+  // without actually branching to it (the exit block should be dominated by the
+  // loop header, not the preheader).
   assert(!L->contains(ExitBlock) && "Exit block is in the loop?");
-  BasicBlock *NewExit;
-  if (BasicBlock *SinglePred = ExitBlock->getSinglePredecessor()) {
-    assert(SinglePred == L->getLoopLatch() && "Unexpected case");
-    NewExit = SplitBlock(ExitBlock, true);
-  } else {
-    // Otherwise, this is a critical edge.  Split block would split the wrong
-    // edge here, so we use SplitCriticalEdge, which allows us to specify the
-    // edge to split, not just the block.
-    TerminatorInst *LatchTerm = L->getLoopLatch()->getTerminator();
-    unsigned SuccNum = 0;
-    for (unsigned i = 0, e = LatchTerm->getNumSuccessors(); ; ++i) {
-      assert(i != e && "Didn't find edge?");
-      if (LatchTerm->getSuccessor(i) == ExitBlock) {
-        SuccNum = i;
-        break;
-      }
-    }
-    SplitCriticalEdge(LatchTerm, SuccNum, this);
-    NewExit = LatchTerm->getSuccessor(SuccNum);
-    assert(NewExit != ExitBlock && "Edge not split!");
-  }
-  
+  BasicBlock *NewExit = SplitBlock(ExitBlock, ExitBlock->begin());
+    
   // Okay, now we have a position to branch from and a position to branch to, 
   // insert the new conditional branch.
-  new BranchInst(EnterOnCond ? NewPH : NewExit, EnterOnCond ? NewExit : NewPH,
-                 Cond, OrigPH->getTerminator());
+  {
+    BasicBlock *TrueDest = NewPH, *FalseDest = NewExit;
+    if (!EntersWhenTrue) std::swap(TrueDest, FalseDest);
+    EmitPreheaderBranchOnCondition(Cond, Val, TrueDest, FalseDest, 
+                                   OrigPH->getTerminator());
+  }
   OrigPH->getTerminator()->eraseFromParent();
 
   // Now that we know that the loop is never entered when this condition is a
   // particular value, rewrite the loop with this info.  We know that this will
   // at least eliminate the old branch.
-  RewriteLoopBodyWithConditionConstant(L, Cond, EnterOnCond);
-  
-  ++NumUnswitched;
+  RewriteLoopBodyWithConditionConstant(L, Cond, Val, EntersWhenTrue);
+  ++NumTrivial;
 }
 
 
-/// VersionLoop - We determined that the loop is profitable to unswitch and
-/// contains a branch on a loop invariant condition.  Split it into loop
-/// versions and test the condition outside of either loop.  Return the loops
-/// created as Out1/Out2.
-void LoopUnswitch::VersionLoop(Value *LIC, Loop *L, Loop *&Out1, Loop *&Out2) {
+/// VersionLoop - We determined that the loop is profitable to unswitch when LIC
+/// equal Val.  Split it into loop versions and test the condition outside of
+/// either loop.  Return the loops created as Out1/Out2.
+void LoopUnswitch::VersionLoop(Value *LIC, Constant *Val, Loop *L,
+                               Loop *&Out1, Loop *&Out2) {
   Function *F = L->getHeader()->getParent();
   
   DEBUG(std::cerr << "loop-unswitch: Unswitching loop %"
-        << L->getHeader()->getName() << " [" << L->getBlocks().size()
-        << " blocks] in Function " << F->getName()
-        << " on cond:" << *LIC << "\n");
+                  << L->getHeader()->getName() << " [" << L->getBlocks().size()
+                  << " blocks] in Function " << F->getName()
+                  << " when '" << *Val << "' == " << *LIC << "\n");
 
+  // LoopBlocks contains all of the basic blocks of the loop, including the
+  // preheader of the loop, the body of the loop, and the exit blocks of the 
+  // loop, in that order.
   std::vector<BasicBlock*> LoopBlocks;
 
   // First step, split the preheader and exit blocks, and add these blocks to
   // the LoopBlocks list.
   BasicBlock *OrigPreheader = L->getLoopPreheader();
-  LoopBlocks.push_back(SplitBlock(OrigPreheader, false));
+  LoopBlocks.push_back(SplitEdge(OrigPreheader, L->getHeader()));
 
   // We want the loop to come after the preheader, but before the exit blocks.
   LoopBlocks.insert(LoopBlocks.end(), L->block_begin(), L->block_end());
@@ -482,10 +584,30 @@ void LoopUnswitch::VersionLoop(Value *LIC, Loop *L, Loop *&Out1, Loop *&Out2) {
   std::sort(ExitBlocks.begin(), ExitBlocks.end());
   ExitBlocks.erase(std::unique(ExitBlocks.begin(), ExitBlocks.end()),
                    ExitBlocks.end());
+  
+  // Split all of the edges from inside the loop to their exit blocks.  This
+  // unswitching trivial: no phi nodes to update.
+  unsigned NumBlocks = L->getBlocks().size();
   for (unsigned i = 0, e = ExitBlocks.size(); i != e; ++i) {
-    SplitBlock(ExitBlocks[i], true);
-    LoopBlocks.push_back(ExitBlocks[i]);
+    BasicBlock *ExitBlock = ExitBlocks[i];
+    std::vector<BasicBlock*> Preds(pred_begin(ExitBlock), pred_end(ExitBlock));
+
+    for (unsigned j = 0, e = Preds.size(); j != e; ++j) {
+      assert(L->contains(Preds[j]) &&
+             "All preds of loop exit blocks must be the same loop!");
+      SplitEdge(Preds[j], ExitBlock);
+    }
   }
+  
+  // The exit blocks may have been changed due to edge splitting, recompute.
+  ExitBlocks.clear();
+  L->getExitBlocks(ExitBlocks);
+  std::sort(ExitBlocks.begin(), ExitBlocks.end());
+  ExitBlocks.erase(std::unique(ExitBlocks.begin(), ExitBlocks.end()),
+                   ExitBlocks.end());
+  
+  // Add exit blocks to the loop blocks.
+  LoopBlocks.insert(LoopBlocks.end(), ExitBlocks.begin(), ExitBlocks.end());
 
   // Next step, clone all of the basic blocks that make up the loop (including
   // the loop preheader and exit blocks), keeping track of the mapping between
@@ -494,8 +616,9 @@ void LoopUnswitch::VersionLoop(Value *LIC, Loop *L, Loop *&Out1, Loop *&Out2) {
   NewBlocks.reserve(LoopBlocks.size());
   std::map<const Value*, Value*> ValueMap;
   for (unsigned i = 0, e = LoopBlocks.size(); i != e; ++i) {
-    NewBlocks.push_back(CloneBasicBlock(LoopBlocks[i], ValueMap, ".us", F));
-    ValueMap[LoopBlocks[i]] = NewBlocks.back();  // Keep the BB mapping.
+    BasicBlock *New = CloneBasicBlock(LoopBlocks[i], ValueMap, ".us", F);
+    NewBlocks.push_back(New);
+    ValueMap[LoopBlocks[i]] = New;  // Keep the BB mapping.
   }
 
   // Splice the newly inserted blocks into the function right before the
@@ -505,13 +628,32 @@ void LoopUnswitch::VersionLoop(Value *LIC, Loop *L, Loop *&Out1, Loop *&Out2) {
 
   // Now we create the new Loop object for the versioned loop.
   Loop *NewLoop = CloneLoop(L, L->getParentLoop(), ValueMap, LI);
-  if (Loop *Parent = L->getParentLoop()) {
+  Loop *ParentLoop = L->getParentLoop();
+  if (ParentLoop) {
     // Make sure to add the cloned preheader and exit blocks to the parent loop
     // as well.
-    Parent->addBasicBlockToLoop(NewBlocks[0], *LI);
-    for (unsigned i = 0, e = ExitBlocks.size(); i != e; ++i)
-      Parent->addBasicBlockToLoop(cast<BasicBlock>(ValueMap[ExitBlocks[i]]),
-                                  *LI);
+    ParentLoop->addBasicBlockToLoop(NewBlocks[0], *LI);
+  }
+  
+  for (unsigned i = 0, e = ExitBlocks.size(); i != e; ++i) {
+    BasicBlock *NewExit = cast<BasicBlock>(ValueMap[ExitBlocks[i]]);
+    if (ParentLoop)
+      ParentLoop->addBasicBlockToLoop(cast<BasicBlock>(NewExit), *LI);
+    
+    assert(NewExit->getTerminator()->getNumSuccessors() == 1 &&
+           "Exit block should have been split to have one successor!");
+    BasicBlock *ExitSucc = NewExit->getTerminator()->getSuccessor(0);
+    
+    // If the successor of the exit block had PHI nodes, add an entry for
+    // NewExit.
+    PHINode *PN;
+    for (BasicBlock::iterator I = ExitSucc->begin();
+         (PN = dyn_cast<PHINode>(I)); ++I) {
+      Value *V = PN->getIncomingValueForBlock(ExitBlocks[i]);
+      std::map<const Value *, Value*>::iterator It = ValueMap.find(V);
+      if (It != ValueMap.end()) V = It->second;
+      PN->addIncoming(V, NewExit);
+    }
   }
 
   // Rewrite the code to refer to itself.
@@ -521,45 +663,218 @@ void LoopUnswitch::VersionLoop(Value *LIC, Loop *L, Loop *&Out1, Loop *&Out2) {
       RemapInstruction(I, ValueMap);
   
   // Rewrite the original preheader to select between versions of the loop.
-  assert(isa<BranchInst>(OrigPreheader->getTerminator()) &&
-         cast<BranchInst>(OrigPreheader->getTerminator())->isUnconditional() &&
-         OrigPreheader->getTerminator()->getSuccessor(0) == LoopBlocks[0] &&
+  BranchInst *OldBR = cast<BranchInst>(OrigPreheader->getTerminator());
+  assert(OldBR->isUnconditional() && OldBR->getSuccessor(0) == LoopBlocks[0] &&
          "Preheader splitting did not work correctly!");
-  // Remove the unconditional branch to LoopBlocks[0].
-  OrigPreheader->getInstList().pop_back();
 
-  // Insert a conditional branch on LIC to the two preheaders.  The original
-  // code is the true version and the new code is the false version.
-  new BranchInst(LoopBlocks[0], NewBlocks[0], LIC, OrigPreheader);
+  // Emit the new branch that selects between the two versions of this loop.
+  EmitPreheaderBranchOnCondition(LIC, Val, NewBlocks[0], LoopBlocks[0], OldBR);
+  OldBR->eraseFromParent();
 
   // Now we rewrite the original code to know that the condition is true and the
   // new code to know that the condition is false.
-  RewriteLoopBodyWithConditionConstant(L, LIC, true);
-  RewriteLoopBodyWithConditionConstant(NewLoop, LIC, false);
-  ++NumUnswitched;
+  RewriteLoopBodyWithConditionConstant(L, LIC, Val, false);
+  RewriteLoopBodyWithConditionConstant(NewLoop, LIC, Val, true);
   Out1 = L;
   Out2 = NewLoop;
 }
 
-// RewriteLoopBodyWithConditionConstant - We know that the boolean value LIC has
-// the value specified by Val in the specified loop.  Rewrite any uses of LIC or
-// of properties correlated to it.
+/// RemoveFromWorklist - Remove all instances of I from the worklist vector
+/// specified.
+static void RemoveFromWorklist(Instruction *I, 
+                               std::vector<Instruction*> &Worklist) {
+  std::vector<Instruction*>::iterator WI = std::find(Worklist.begin(),
+                                                     Worklist.end(), I);
+  while (WI != Worklist.end()) {
+    unsigned Offset = WI-Worklist.begin();
+    Worklist.erase(WI);
+    WI = std::find(Worklist.begin()+Offset, Worklist.end(), I);
+  }
+}
+
+/// ReplaceUsesOfWith - When we find that I really equals V, remove I from the
+/// program, replacing all uses with V and update the worklist.
+static void ReplaceUsesOfWith(Instruction *I, Value *V, 
+                              std::vector<Instruction*> &Worklist) {
+  DEBUG(std::cerr << "Replace with '" << *V << "': " << *I);
+
+  // Add uses to the worklist, which may be dead now.
+  for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i)
+    if (Instruction *Use = dyn_cast<Instruction>(I->getOperand(i)))
+      Worklist.push_back(Use);
+
+  // Add users to the worklist which may be simplified now.
+  for (Value::use_iterator UI = I->use_begin(), E = I->use_end();
+       UI != E; ++UI)
+    Worklist.push_back(cast<Instruction>(*UI));
+  I->replaceAllUsesWith(V);
+  I->eraseFromParent();
+  RemoveFromWorklist(I, Worklist);
+  ++NumSimplify;
+}
+
+
+
+// RewriteLoopBodyWithConditionConstant - We know either that the value LIC has
+// the value specified by Val in the specified loop, or we know it does NOT have
+// that value.  Rewrite any uses of LIC or of properties correlated to it.
 void LoopUnswitch::RewriteLoopBodyWithConditionConstant(Loop *L, Value *LIC,
-                                                        bool Val) {
+                                                        Constant *Val,
+                                                        bool IsEqual) {
   assert(!isa<Constant>(LIC) && "Why are we unswitching on a constant?");
+  
   // FIXME: Support correlated properties, like:
   //  for (...)
   //    if (li1 < li2)
   //      ...
   //    if (li1 > li2)
   //      ...
-  ConstantBool *BoolVal = ConstantBool::get(Val);
 
+  // NotVal - If Val is a bool, this contains its inverse.
+  Constant *NotVal = 0;
+  if (ConstantBool *CB = dyn_cast<ConstantBool>(Val))
+    NotVal = ConstantBool::get(!CB->getValue());
+  
   // FOLD boolean conditions (X|LIC), (X&LIC).  Fold conditional branches,
   // selects, switches.
   std::vector<User*> Users(LIC->use_begin(), LIC->use_end());
-  for (unsigned i = 0, e = Users.size(); i != e; ++i)
-    if (Instruction *U = cast<Instruction>(Users[i]))
-      if (L->contains(U->getParent()))
-        U->replaceUsesOfWith(LIC, BoolVal);
+
+  std::vector<Instruction*> Worklist;
+  
+  // If we know that LIC == Val, or that LIC == NotVal, just replace uses of LIC
+  // in the loop with the appropriate one directly.
+  if (IsEqual || NotVal) {
+    Value *Replacement = NotVal ? NotVal : Val;
+    
+    for (unsigned i = 0, e = Users.size(); i != e; ++i)
+      if (Instruction *U = cast<Instruction>(Users[i])) {
+        if (!L->contains(U->getParent()))
+          continue;
+        U->replaceUsesOfWith(LIC, Replacement);
+        Worklist.push_back(U);
+      }
+  } else {
+    // Otherwise, we don't know the precise value of LIC, but we do know that it
+    // is certainly NOT "Val".  As such, simplify any uses in the loop that we
+    // can.  This case occurs when we unswitch switch statements.
+    for (unsigned i = 0, e = Users.size(); i != e; ++i)
+      if (Instruction *U = cast<Instruction>(Users[i])) {
+        if (!L->contains(U->getParent()))
+          continue;
+
+        Worklist.push_back(U);
+
+        // If we know that LIC is not Val, use this info to simplify code.
+        if (SwitchInst *SI = dyn_cast<SwitchInst>(U)) {
+          for (unsigned i = 1, e = SI->getNumCases(); i != e; ++i) {
+            if (SI->getCaseValue(i) == Val) {
+              // Found a dead case value.  Don't remove PHI nodes in the 
+              // successor if they become single-entry, those PHI nodes may
+              // be in the Users list.
+              SI->getSuccessor(i)->removePredecessor(SI->getParent(), true);
+              SI->removeCase(i);
+              break;
+            }
+          }
+        }
+        
+        // TODO: We could do other simplifications, for example, turning 
+        // LIC == Val -> false.
+      }
+  }
+    
+  // Okay, now that we have simplified some instructions in the loop, walk over
+  // it and constant prop, dce, and fold control flow where possible.  Note that
+  // this is effectively a very simple loop-structure-aware optimizer.
+  while (!Worklist.empty()) {
+    Instruction *I = Worklist.back();
+    Worklist.pop_back();
+    
+    // Simple constant folding.
+    if (Constant *C = ConstantFoldInstruction(I)) {
+      ReplaceUsesOfWith(I, C, Worklist);
+      continue;
+    }
+    
+    // Simple DCE.
+    if (isInstructionTriviallyDead(I)) {
+      DEBUG(std::cerr << "Remove dead instruction '" << *I);
+      
+      // Add uses to the worklist, which may be dead now.
+      for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i)
+        if (Instruction *Use = dyn_cast<Instruction>(I->getOperand(i)))
+          Worklist.push_back(Use);
+      I->eraseFromParent();
+      RemoveFromWorklist(I, Worklist);
+      ++NumSimplify;
+      continue;
+    }
+    
+    // Special case hacks that appear commonly in unswitched code.
+    switch (I->getOpcode()) {
+    case Instruction::Select:
+      if (ConstantBool *CB = dyn_cast<ConstantBool>(I->getOperand(0))) {
+        ReplaceUsesOfWith(I, I->getOperand(!CB->getValue()+1), Worklist);
+        continue;
+      }
+      break;
+    case Instruction::And:
+      if (isa<ConstantBool>(I->getOperand(0)))   // constant -> RHS
+        cast<BinaryOperator>(I)->swapOperands();
+      if (ConstantBool *CB = dyn_cast<ConstantBool>(I->getOperand(1))) {
+        if (CB->getValue())   // X & 1 -> X
+          ReplaceUsesOfWith(I, I->getOperand(0), Worklist);
+        else                  // X & 0 -> 0
+          ReplaceUsesOfWith(I, I->getOperand(1), Worklist);
+        continue;
+      }
+      break;
+    case Instruction::Or:
+      if (isa<ConstantBool>(I->getOperand(0)))   // constant -> RHS
+        cast<BinaryOperator>(I)->swapOperands();
+      if (ConstantBool *CB = dyn_cast<ConstantBool>(I->getOperand(1))) {
+        if (CB->getValue())   // X | 1 -> 1
+          ReplaceUsesOfWith(I, I->getOperand(1), Worklist);
+        else                  // X | 0 -> X
+          ReplaceUsesOfWith(I, I->getOperand(0), Worklist);
+        continue;
+      }
+      break;
+    case Instruction::Br: {
+      BranchInst *BI = cast<BranchInst>(I);
+      if (BI->isUnconditional()) {
+        // If BI's parent is the only pred of the successor, fold the two blocks
+        // together.
+        BasicBlock *Pred = BI->getParent();
+        BasicBlock *Succ = BI->getSuccessor(0);
+        BasicBlock *SinglePred = Succ->getSinglePredecessor();
+        if (!SinglePred) continue;  // Nothing to do.
+        assert(SinglePred == Pred && "CFG broken");
+
+        DEBUG(std::cerr << "Merging blocks: " << Pred->getName() << " <- " 
+                        << Succ->getName() << "\n");
+        
+        // Resolve any single entry PHI nodes in Succ.
+        while (PHINode *PN = dyn_cast<PHINode>(Succ->begin()))
+          ReplaceUsesOfWith(PN, PN->getIncomingValue(0), Worklist);
+        
+        // Move all of the successor contents from Succ to Pred.
+        Pred->getInstList().splice(BI, Succ->getInstList(), Succ->begin(),
+                                   Succ->end());
+        BI->eraseFromParent();
+        RemoveFromWorklist(BI, Worklist);
+        
+        // If Succ has any successors with PHI nodes, update them to have
+        // entries coming from Pred instead of Succ.
+        Succ->replaceAllUsesWith(Pred);
+        
+        // Remove Succ from the loop tree.
+        LI->removeBlock(Succ);
+        Succ->eraseFromParent();
+        break;
+      }
+      break;
+    }
+    }
+  }
 }