[Opt] Fix use after free (erase) in mem2reg (#628)

Co-authored-by: Alex McCaskey <amccaskey@nvidia.com>
NVIDIA · Sep 11, 2023 · 4e63416 · 4e63416
1 parent 77cbcd1
commit 4e63416
Showing 1 changed file with 10 additions and 9 deletions.
diff --git a/lib/Optimizer/Transforms/MemToReg.cpp b/lib/Optimizer/Transforms/MemToReg.cpp
@@ -302,20 +302,21 @@ class MemToRegPass : public cudaq::opt::impl::MemToRegBase<MemToRegPass> {
     processOpWithRegions(func, memAnalysis, cleanUps);
 
     // 3) Cleanup the dead ops.
+    SmallVector<quake::WrapOp> wrapOps;
     for (auto *op : cleanUps) {
-      if (auto wrap = dyn_cast<quake::WrapOp>(op))
+      if (auto wrap = dyn_cast<quake::WrapOp>(op)) {
+        wrapOps.push_back(wrap);
         continue;
+      }
       op->dropAllUses();
       op->erase();
     }
-    for (auto *op : cleanUps) {
-      if (auto wrap = dyn_cast<quake::WrapOp>(op)) {
-        auto ref = wrap.getRefValue();
-        auto wire = wrap.getWireValue();
-        if (!ref || (++wire.getUses().begin() != wire.getUses().end())) {
-          op->dropAllUses();
-          op->erase();
-        }
+    for (auto wrap : wrapOps) {
+      auto ref = wrap.getRefValue();
+      auto wire = wrap.getWireValue();
+      if (!ref || wire.getUses().empty()) {
+        wrap->dropAllUses();
+        wrap->erase();
       }
     }