[flang-commits] [flang] [flang][OpenMP] Fix use-after-free in OMPFunctionFiltering (PR #84373)

[Krzysztof Parzyszek via flang-commits]

https://github.com/kparzysz updated https://github.com/llvm/llvm-project/pull/84373

>From 4bf5b953eff90ca1198a5feca8c3d5dcee52b4af Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Thu, 7 Mar 2024 14:11:33 -0600
Subject: [PATCH] [flang][OpenMP] Fix use-after-free in OMPFunctionFiltering

Erasing the element of the range that it being iterated on can
invalidate the range. Instead of erasing function as we see them,
store them in a separate list, then erase them after the range
has been traversed.

This was detected by address sanitizer.
---
 flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
index 466bf53e8dbd60..f8bede88293ff5 100644
--- a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
+++ b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
@@ -40,6 +40,8 @@ class OMPFunctionFilteringPass
     if (!op || !op.getIsTargetDevice())
       return;
 
+    llvm::SmallVector<func::FuncOp> removedFuncs;
+
     op->walk<WalkOrder::PreOrder>([&](func::FuncOp funcOp) {
       // Do not filter functions with target regions inside, because they have
       // to be available for both host and device so that regular and reverse
@@ -80,12 +82,14 @@ class OMPFunctionFilteringPass
           callOp->erase();
         }
         if (!hasTargetRegion)
-          funcOp.erase();
+          removedFuncs.push_back(funcOp);
         else if (declareTargetOp)
           declareTargetOp.setDeclareTarget(declareType,
                                            omp::DeclareTargetCaptureClause::to);
       }
     });
+    for (func::FuncOp f : removedFuncs)
+      f.erase();
   }
 };
 } // namespace