2 files changed, 14 insertions, 6 deletions
diff --git a/lib/Transforms/Scalar/LoopUnroll.cpp b/lib/Transforms/Scalar/LoopUnroll.cpp
index 2140b22f2f..ff8ece2aa8 100644
--- a/lib/Transforms/Scalar/LoopUnroll.cpp
+++ b/lib/Transforms/Scalar/LoopUnroll.cpp
@@ -396,10 +396,14 @@ bool LoopUnroll::unrollLoop(Loop *L, unsigned Count, unsigned Threshold) {
     }
     
     // Remap all instructions in the most recent iteration
-    for (unsigned i = 0; i < NewBlocks.size(); ++i)
-      for (BasicBlock::iterator I = NewBlocks[i]->begin(),
-           E = NewBlocks[i]->end(); I != E; ++I)
+    for (unsigned i = 0; i < NewBlocks.size(); ++i) {
+      BasicBlock *NB = NewBlocks[i];
+      if (BasicBlock *UnwindDest = NB->getUnwindDest())
+        NB->setUnwindDest(cast<BasicBlock>(LastValueMap[UnwindDest]));
+
+      for (BasicBlock::iterator I = NB->begin(), E = NB->end(); I != E; ++I)
         RemapInstruction(I, LastValueMap);
+    }
   }
   
   // The latch block exits the loop.  If there are any PHI nodes in the
diff --git a/lib/Transforms/Scalar/LoopUnswitch.cpp b/lib/Transforms/Scalar/LoopUnswitch.cpp
index 8d39a4dece..966038ddff 100644
--- a/lib/Transforms/Scalar/LoopUnswitch.cpp
+++ b/lib/Transforms/Scalar/LoopUnswitch.cpp
@@ -819,10 +819,14 @@ void LoopUnswitch::UnswitchNontrivialCondition(Value *LIC, Constant *Val,
   }
 
   // Rewrite the code to refer to itself.
-  for (unsigned i = 0, e = NewBlocks.size(); i != e; ++i)
-    for (BasicBlock::iterator I = NewBlocks[i]->begin(),
-           E = NewBlocks[i]->end(); I != E; ++I)
+  for (unsigned i = 0, e = NewBlocks.size(); i != e; ++i) {
+    BasicBlock *NB = NewBlocks[i];
+    if (BasicBlock *UnwindDest = NB->getUnwindDest())
+      NB->setUnwindDest(cast<BasicBlock>(ValueMap[UnwindDest]));
+
+    for (BasicBlock::iterator I = NB->begin(), E = NB->end(); I != E; ++I)
       RemapInstruction(I, ValueMap);
+  }
   
   // Rewrite the original preheader to select between versions of the loop.
   BranchInst *OldBR = cast<BranchInst>(OrigPreheader->getTerminator());