[flang-commits] [flang] 10b0156 - [flang][HLFIR] Fix use-after-free when rewriting users in `canonicalize` (#84371)

[via flang-commits]

Author: Krzysztof Parzyszek
Date: 2024-03-08T07:20:20-06:00
New Revision: 10b01563da1d7c98d5fff6aa8b4a78be56a8099d

URL: https://github.com/llvm/llvm-project/commit/10b01563da1d7c98d5fff6aa8b4a78be56a8099d
DIFF: https://github.com/llvm/llvm-project/commit/10b01563da1d7c98d5fff6aa8b4a78be56a8099d.diff

LOG: [flang][HLFIR] Fix use-after-free when rewriting users in `canonicalize` (#84371)

Rewriting an op can invalidate the operator range being iterated on.
Store the users in a separate list, and iterate over the list instead.

This was detected by address sanitizer.

Added: 
    

Modified: 
    flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index 8bc92a991a69cf..3568fe202caf14 100644
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
+
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
@@ -1152,7 +1153,9 @@ hlfir::MatmulOp::canonicalize(MatmulOp matmulOp,
 
       // but we do need to get rid of the hlfir.destroy for the hlfir.transpose
       // result (which is entirely removed)
-      for (mlir::Operation *user : transposeOp->getResult(0).getUsers())
+      llvm::SmallVector<mlir::Operation *> users(
+          transposeOp->getResult(0).getUsers());
+      for (mlir::Operation *user : users)
         if (auto destroyOp = mlir::dyn_cast_or_null<hlfir::DestroyOp>(user))
           rewriter.eraseOp(destroyOp);
       rewriter.eraseOp(transposeOp);
@@ -1864,7 +1867,8 @@ hlfir::ForallIndexOp::canonicalize(hlfir::ForallIndexOp indexOp,
       return mlir::failure();
 
   auto insertPt = rewriter.saveInsertionPoint();
-  for (mlir::Operation *user : indexOp->getResult(0).getUsers())
+  llvm::SmallVector<mlir::Operation *> users(indexOp->getResult(0).getUsers());
+  for (mlir::Operation *user : users)
     if (auto loadOp = mlir::dyn_cast<fir::LoadOp>(user)) {
       rewriter.setInsertionPoint(loadOp);
       rewriter.replaceOpWithNewOp<fir::ConvertOp>(