[Mlir-commits] [mlir] [mlir][bufferization] Fix use-after-free in ownership-based buffer deallocation (PR #184118)

[Mehdi Amini]

================
@@ -717,8 +717,18 @@ BufferDeallocation::handleInterface(RegionBranchOpInterface op) {
   int counter = op->getNumResults();
   unsigned numMemrefResults = llvm::count_if(op->getResults(), isMemref);
   SmallVector<Type> ownershipResults(numMemrefResults, builder.getI1Type());
+  // Save the old result values before appendOpResults erases the op. The
+  // liveness analysis holds references to these values and they may be queried
+  // later (e.g., from handleInterface(BranchOpInterface) in the same block).
+  SmallVector<Value> oldResults(op->getResults());
   RegionBranchOpInterface newOp = appendOpResults(op, ownershipResults);
 
+  // Register the replacement of each old result with the corresponding new
+  // result so that stale liveness entries can be translated on demand.
+  for (auto [oldResult, newResult] :
+       llvm::zip(oldResults, newOp->getResults().take_front(oldResults.size())))
+    state.mapValue(oldResult, newResult);
----------------
joker-eph wrote:

I moved the logic to appendOpResult, I'm not sure if this mapping is right, I'm not fond of it either, but other alternatives (like the set you mentioned seem involved, likely just as intrusive?). I don't know...

https://github.com/llvm/llvm-project/pull/184118