[llvm] [msan] Don't modify CFG iterating it (PR #90691)

[Vitaly Buka via llvm-commits]

https://github.com/vitalybuka created https://github.com/llvm/llvm-project/pull/90691

In rare cases `SplitBlockAndInsertSimpleForLoop` in `paintOrigin` crashes outsize iterators.

Somehow existing `SplitBlockAndInsertIfThen` do not invalidate iterators.


>From 7cfab955232245a158ba7eeb32d83b2c0e61628b Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka at google.com>
Date: Tue, 30 Apr 2024 17:20:42 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.4
---
 .../Instrumentation/MemorySanitizer.cpp       | 27 +++++++++++++------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index cc2295c44023c4..4a5f4a40726574 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -1135,6 +1135,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   std::unique_ptr<VarArgHelper> VAHelper;
   const TargetLibraryInfo *TLI;
   Instruction *FnPrologueEnd;
+  SmallVector<Instruction *, 128> Instructions;
 
   // The following flags disable parts of MSan instrumentation based on
   // exclusion list contents and command-line options.
@@ -1520,6 +1521,11 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     for (BasicBlock *BB : depth_first(FnPrologueEnd->getParent()))
       visit(*BB);
 
+    // `visit` above only collects instructions. Process them after iterating
+    // CFG to avoid requirement on CFG transformations.
+    for (Instruction *I : Instructions)
+      instrument(*I);
+
     // Finalize PHI nodes.
     for (PHINode *PN : ShadowPHINodes) {
       PHINode *PNS = cast<PHINode>(getShadow(PN));
@@ -2181,14 +2187,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     return ConstantDataVector::get(IRB.getContext(), OrderingTable);
   }
 
-  // ------------------- Visitors.
-  using InstVisitor<MemorySanitizerVisitor>::visit;
-  void visit(Instruction &I) {
-    if (I.getMetadata(LLVMContext::MD_nosanitize))
-      return;
-    // Don't want to visit if we're in the prologue
-    if (isInPrologue(I))
-      return;
+  void instrument(Instruction &I) {
     if (!DebugCounter::shouldExecute(DebugInstrumentInstruction)) {
       LLVM_DEBUG(dbgs() << "Skipping instruction: " << I << "\n");
       // We still need to set the shadow and origin to clean values.
@@ -2199,6 +2198,18 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     InstVisitor<MemorySanitizerVisitor>::visit(I);
   }
 
+  // ------------------- Visitors.
+  using InstVisitor<MemorySanitizerVisitor>::visit;
+  void visit(Instruction &I) {
+    if (I.getMetadata(LLVMContext::MD_nosanitize))
+      return;
+    // Don't want to visit if we're in the prologue
+    if (isInPrologue(I))
+      return;
+
+    Instructions.push_back(&I);
+  }
+
   /// Instrument LoadInst
   ///
   /// Loads the corresponding shadow and (optionally) origin.