4 files changed, 22 insertions, 12 deletions

diff --git a/include/llvm/Target/TargetInstrInfo.h b/include/llvm/Target/TargetInstrInfo.h
index 844b96566a..f4652011e9 100644
--- a/include/llvm/Target/TargetInstrInfo.h
+++ b/include/llvm/Target/TargetInstrInfo.h
@@ -586,9 +586,11 @@ public:
   }
 
   /// ConvertToSetZeroFlag - Convert the instruction to set the zero flag so
-  /// that we can remove a "comparison with zero".
-  virtual bool ConvertToSetZeroFlag(MachineInstr *Instr,
-                                    MachineInstr *CmpInstr) const {
+  /// that we can remove a "comparison with zero".  Update the iterator *only*
+  /// if a transformation took place.
+  virtual bool ConvertToSetZeroFlag(MachineInstr * /*Instr*/,
+                                    MachineInstr * /*CmpInstr*/,
+                                    MachineBasicBlock::iterator &) const {
     return false;
   }
 
diff --git a/lib/CodeGen/PeepholeOptimizer.cpp b/lib/CodeGen/PeepholeOptimizer.cpp
index ea14d276d5..a4ff04bec1 100644
--- a/lib/CodeGen/PeepholeOptimizer.cpp
+++ b/lib/CodeGen/PeepholeOptimizer.cpp
@@ -76,7 +76,8 @@ namespace {
     }
 
   private:
-    bool OptimizeCmpInstr(MachineInstr *MI, MachineBasicBlock *MBB);
+    bool OptimizeCmpInstr(MachineInstr *MI, MachineBasicBlock *MBB,
+                          MachineBasicBlock::iterator &MII);
     bool OptimizeExtInstr(MachineInstr *MI, MachineBasicBlock *MBB,
                           SmallPtrSet<MachineInstr*, 8> &LocalMIs);
   };
@@ -232,7 +233,8 @@ OptimizeExtInstr(MachineInstr *MI, MachineBasicBlock *MBB,
 /// set) the same flag as the compare, then we can remove the comparison and use
 /// the flag from the previous instruction.
 bool PeepholeOptimizer::OptimizeCmpInstr(MachineInstr *MI,
-                                         MachineBasicBlock *MBB) {
+                                         MachineBasicBlock *MBB,
+                                         MachineBasicBlock::iterator &NextIter){
   // If this instruction is a comparison against zero and isn't comparing a
   // physical register, we can try to optimize it.
   unsigned SrcReg;
@@ -247,7 +249,7 @@ bool PeepholeOptimizer::OptimizeCmpInstr(MachineInstr *MI,
     return false;
 
   // Attempt to convert the defining instruction to set the "zero" flag.
-  if (TII->ConvertToSetZeroFlag(&*DI, MI)) {
+  if (TII->ConvertToSetZeroFlag(&*DI, MI, NextIter)) {
     ++NumEliminated;
     return true;
   }
@@ -269,12 +271,14 @@ bool PeepholeOptimizer::runOnMachineFunction(MachineFunction &MF) {
     LocalMIs.clear();
 
     for (MachineBasicBlock::iterator
-           MII = I->begin(), ME = I->end(); MII != ME; ) {
+           MII = I->begin(), MIE = I->end(); MII != MIE; ) {
       MachineInstr *MI = &*MII;
 
       if (MI->getDesc().isCompare()) {
-        ++MII; // The iterator may become invalid if the compare is deleted.
-        Changed |= OptimizeCmpInstr(MI, MBB);
+        if (OptimizeCmpInstr(MI, MBB, MII))
+          Changed = true;
+        else
+          ++MII;
       } else {
         Changed |= OptimizeExtInstr(MI, MBB, LocalMIs);
         ++MII;
diff --git a/lib/Target/ARM/ARMBaseInstrInfo.cpp b/lib/Target/ARM/ARMBaseInstrInfo.cpp
index e7024d6da7..142633c667 100644
--- a/lib/Target/ARM/ARMBaseInstrInfo.cpp
+++ b/lib/Target/ARM/ARMBaseInstrInfo.cpp
@@ -1378,9 +1378,11 @@ AnalyzeCompare(const MachineInstr *MI, unsigned &SrcReg, int &CmpValue) const {
 }
 
 /// ConvertToSetZeroFlag - Convert the instruction to set the "zero" flag so
-/// that we can remove a "comparison with zero".
+/// that we can remove a "comparison with zero". Update the iterator *only* if a
+/// transformation took place.
 bool ARMBaseInstrInfo::
-ConvertToSetZeroFlag(MachineInstr *MI, MachineInstr *CmpInstr) const {
+ConvertToSetZeroFlag(MachineInstr *MI, MachineInstr *CmpInstr,
+                     MachineBasicBlock::iterator &MII) const {
   // Conservatively refuse to convert an instruction which isn't in the same BB
   // as the comparison.
   if (MI->getParent() != CmpInstr->getParent())
@@ -1414,6 +1416,7 @@ ConvertToSetZeroFlag(MachineInstr *MI, MachineInstr *CmpInstr) const {
     MI->RemoveOperand(5);
     MachineInstrBuilder(MI)
       .addReg(ARM::CPSR, RegState::Define | RegState::Implicit);
+    MII = llvm::next(MachineBasicBlock::iterator(CmpInstr));
     CmpInstr->eraseFromParent();
     return true;
   }
diff --git a/lib/Target/ARM/ARMBaseInstrInfo.h b/lib/Target/ARM/ARMBaseInstrInfo.h
index f471b6772e..622402eba3 100644
--- a/lib/Target/ARM/ARMBaseInstrInfo.h
+++ b/lib/Target/ARM/ARMBaseInstrInfo.h
@@ -347,7 +347,8 @@ public:
   /// ConvertToSetZeroFlag - Convert the instruction to set the zero flag so
   /// that we can remove a "comparison with zero".
   virtual bool ConvertToSetZeroFlag(MachineInstr *Instr,
-                                    MachineInstr *CmpInstr) const;
+                                    MachineInstr *CmpInstr,
+                                    MachineBasicBlock::iterator &MII) const;
 
   virtual unsigned getNumMicroOps(const MachineInstr *MI,
                                   const InstrItineraryData *ItinData) const;