[PATCH] D118591: [Function Specialisation] Fix use after free

[Alexandros Lamprineas via Phabricator via llvm-commits]

labrinea updated this revision to Diff 404727.
labrinea edited the summary of this revision.
labrinea added a comment.
Herald added subscribers: snehasish, ormris.

As suggested by Florian, instead of using a `WeakVH` I am lazily removing the replaced instructions after the Solver has run. None of the existing tests actually covers this code path I am afraid. Examining the debug output of the reproducer I found two `PhiNode` instructions being replaced with `null`.


CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D118591/new/

https://reviews.llvm.org/D118591

Files:
  llvm/lib/Transforms/IPO/FunctionSpecialization.cpp


Index: llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
===================================================================
--- llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
+++ llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
@@ -276,6 +276,7 @@
   std::function<TargetLibraryInfo &(Function &)> GetTLI;
 
   SmallPtrSet<Function *, 2> SpecializedFuncs;
+  SmallVector<Instruction *> ReplacedWithConstant;
 
 public:
   FunctionSpecializer(SCCPSolver &Solver,
@@ -320,6 +321,12 @@
     return Changed;
   }
 
+  void removeDeadInstructions() {
+    for (auto *I : ReplacedWithConstant)
+      I->eraseFromParent();
+    ReplacedWithConstant.clear();
+  }
+
   bool tryToReplaceWithConstant(Value *V) {
     if (!V->getType()->isSingleValueType() || isa<CallBase>(V) ||
         V->user_empty())
@@ -340,7 +347,7 @@
     // Remove the instruction from Block and Solver.
     if (auto *I = dyn_cast<Instruction>(V)) {
       if (I->isSafeToRemove()) {
-        I->eraseFromParent();
+        ReplacedWithConstant.push_back(I);
         Solver.removeLatticeValueFor(I);
       }
     }
@@ -886,7 +893,8 @@
     Changed = true;
   }
 
-  // Clean up the IR by removing ssa_copy intrinsics.
+  // Clean up the IR by removing dead instructions and ssa_copy intrinsics.
+  FS.removeDeadInstructions();
   removeSSACopy(M);
   return Changed;
 }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D118591.404727.patch
Type: text/x-patch
Size: 1348 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220131/5fd8dfa2/attachment.bin>