[AArch64] Fix a compiler crash in MachineSink (#67705)

There were a couple of issues with maintaining register def/uses held
in `MachineRegisterInfo`:

* when an operand is changed from one register to another, the
corresponding instruction must already be inserted into the function,
or MRI won't be updated

* when traversing the set of all uses of a register, that set must not
change

1 files changed, 19 insertions, 17 deletions

diff --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index 480ac23..9d4e0c6 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -461,7 +461,7 @@ bool MachineSinking::PerformSinkAndFold(MachineInstr &MI,
       }
 
       if (UseInst.getParent() != MI.getParent()) {
-        // If the register class of the register we are replacingis a superset
+        // If the register class of the register we are replacing is a superset
         // of any of the register classes of the operands of the materialized
         // instruction don't consider that live range extended.
         const TargetRegisterClass *RCS = MRI->getRegClass(Reg);
@@ -527,8 +527,8 @@ bool MachineSinking::PerformSinkAndFold(MachineInstr &MI,
         if (U.getInt())
           continue;
         MachineInstr *NewDbgMI = SinkDst->getMF()->CloneMachineInstr(DbgMI);
-        NewDbgMI->getOperand(0).setReg(DstReg);
         SinkMBB.insertAfter(InsertPt, NewDbgMI);
+        NewDbgMI->getOperand(0).setReg(DstReg);
       }
     } else {
       // Fold instruction into the addressing mode of a memory instruction.
@@ -538,33 +538,35 @@ bool MachineSinking::PerformSinkAndFold(MachineInstr &MI,
     SinkDst->eraseFromParent();
   }
 
-  MI.eraseFromParent();
-
-  // Collect instructions that need to be deleted (COPYs). We cannot delete them
-  // while traversing register uses.
-  SmallVector<MachineInstr *> CleanupInstrs;
+  // Collect operands that need to be cleaned up because the registers no longer
+  // exist (in COPYs and debug instructions). We cannot delete instructions or
+  // clear operands while traversing register uses.
+  SmallVector<MachineOperand *> Cleanup;
   Worklist.push_back(DefReg);
   while (!Worklist.empty()) {
     Register Reg = Worklist.pop_back_val();
-
     for (MachineOperand &MO : MRI->use_operands(Reg)) {
       MachineInstr *U = MO.getParent();
       assert((U->isCopy() || U->isDebugInstr()) &&
              "Only debug uses and copies must remain");
-      if (U->isCopy()) {
+      if (U->isCopy())
         Worklist.push_back(U->getOperand(0).getReg());
-        CleanupInstrs.push_back(U);
-      } else {
-        MO.setReg(0);
-        MO.setSubReg(0);
-      }
+      Cleanup.push_back(&MO);
     }
   }
 
-  // Delete the dead COPYs.
-  for (MachineInstr *Del : CleanupInstrs)
-    Del->eraseFromParent();
+  // Delete the dead COPYs and clear operands in debug instructions
+  for (MachineOperand *MO : Cleanup) {
+    MachineInstr *I = MO->getParent();
+    if (I->isCopy()) {
+      I->eraseFromParent();
+    } else {
+      MO->setReg(0);
+      MO->setSubReg(0);
+    }
+  }
 
+  MI.eraseFromParent();
   return true;
 }