From df807e0ae8c8123f2503678dde0584d2a6114c57 Mon Sep 17 00:00:00 2001
From: Vladimir Vuksanovic <Vladimir.Vuksanovic@Syrmia.com>
Date: Wed, 15 Nov 2023 15:37:27 +0100
Subject: [PATCH] [RegisterCoalescer] Fix reuse of instruction pointers

If a pointer gets recycled while it is already in a work list, it can
cause duplicates in the work lists. In that case the same instruction
may be processed twice. If it is coalesced the first time, processing
it again would read deallocated memory. This is fixed by removing the
first occurrence from the work lists when we detect a pointer was
recycled.

Fixes #71178
---
 llvm/lib/CodeGen/RegisterCoalescer.cpp | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/RegisterCoalescer.cpp b/llvm/lib/CodeGen/RegisterCoalescer.cpp
index 9858482cd51b4..a68ebaa9b6828 100644
--- a/llvm/lib/CodeGen/RegisterCoalescer.cpp
+++ b/llvm/lib/CodeGen/RegisterCoalescer.cpp
@@ -1031,6 +1031,19 @@ RegisterCoalescer::removeCopyByCommutingDef(const CoalescerPair &CP,
   return { true, ShrinkB };
 }
 
+// Replace the given MachineInstr * with nullptr inside a work list. Used to
+// remove an instruction when its pointer is reused. Returns true if replaced,
+// false otherwise.
+static bool removeFromWorkList(MutableArrayRef<MachineInstr *> WorkList,
+                               MachineInstr *MI) {
+  auto It = std::find(WorkList.begin(), WorkList.end(), MI);
+  if (It != WorkList.end()) {
+    *It = nullptr;
+    return true;
+  }
+  return false;
+}
+
 /// For copy B = A in BB2, if A is defined by A = B in BB0 which is a
 /// predecessor of BB2, and if B is not redefined on the way from A = B
 /// in BB0 to B = A in BB2, B = A in BB2 is partially redundant if the
@@ -1190,7 +1203,15 @@ bool RegisterCoalescer::removePartialRedundancy(const CoalescerPair &CP,
     // If the newly created Instruction has an address of an instruction that was
     // deleted before (object recycled by the allocator) it needs to be removed from
     // the deleted list.
-    ErasedInstrs.erase(NewCopyMI);
+    bool WasErased = ErasedInstrs.erase(NewCopyMI);
+    // Also remove the deleted instruction from work lists. There shouldn't be
+    // duplicate instructions there.
+    if (WasErased) {
+      // Attempt to remove from WorkList. If not found, it could be in
+      // LocalWorkList.
+      if (!removeFromWorkList(WorkList, NewCopyMI))
+        removeFromWorkList(LocalWorkList, NewCopyMI);
+    }
   } else {
     LLVM_DEBUG(dbgs() << "\tremovePartialRedundancy: Remove the copy from "
                       << printMBBReference(MBB) << '\t' << CopyMI);