1 files changed, 17 insertions, 6 deletions
diff --git a/include/llvm/CodeGen/BreakCriticalMachineEdge.h b/include/llvm/CodeGen/BreakCriticalMachineEdge.h
index 743cff3f2b..b8bb2d209c 100644
--- a/include/llvm/CodeGen/BreakCriticalMachineEdge.h
+++ b/include/llvm/CodeGen/BreakCriticalMachineEdge.h
@@ -52,16 +52,27 @@ MachineBasicBlock* SplitCriticalMachineEdge(MachineBasicBlock* src,
       break;
     }
     
-    // Scan the operands of this branch, replacing any uses of dst with
-    // crit_mbb.
+    // Scan the operands of this branch, finding all uses of this MBB
+    std::vector<unsigned> toRemove;
+    unsigned reg = 0;
     for (unsigned i = 0, e = mii->getNumOperands(); i != e; ++i) {
       MachineOperand & mo = mii->getOperand(i);
       if (mo.isMachineBasicBlock() &&
-          mo.getMachineBasicBlock() == dst) {
-        found_branch = true;
-        mo.setMachineBasicBlock(crit_mbb);
-      }
+          mo.getMachineBasicBlock() == dst)
+        reg = mii->getOperand(i-1).getReg();
+        toRemove.push_back(i-1);
     }
+    
+    // Remove all uses of this MBB
+    for (std::vector<unsigned>::reverse_iterator I = toRemove.rbegin(),
+         E = toRemove.rend(); I != E; ++I) {
+      mii->RemoveOperand(*I+1);
+      mii->RemoveOperand(*I);
+    }
+    
+    // Add a single use corresponding to the new MBB
+    mii->addRegOperand(reg, false);
+    mii->addMachineBasicBlockOperand(crit_mbb);
   }
 
   // TODO: This is tentative. It may be necessary to fix this code. Maybe