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

[Vitaly Buka via llvm-commits]

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

>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 1/3] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20in?=
 =?UTF-8?q?itial=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.

>From 2fdcc90e1c444e76682aebb22b35a517609150ba Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka at google.com>
Date: Tue, 30 Apr 2024 17:23:49 -0700
Subject: [PATCH 2/3] move debug counter back

Created using spr 1.3.4
---
 .../Transforms/Instrumentation/MemorySanitizer.cpp | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 4a5f4a40726574..a903863e0265b2 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -2188,13 +2188,6 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   }
 
   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.
-      setShadow(&I, getCleanShadow(&I));
-      setOrigin(&I, getCleanOrigin());
-      return;
-    }
     InstVisitor<MemorySanitizerVisitor>::visit(I);
   }
 
@@ -2206,6 +2199,13 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     // Don't want to visit if we're in the prologue
     if (isInPrologue(I))
       return;
+    if (!DebugCounter::shouldExecute(DebugInstrumentInstruction)) {
+      LLVM_DEBUG(dbgs() << "Skipping instruction: " << I << "\n");
+      // We still need to set the shadow and origin to clean values.
+      setShadow(&I, getCleanShadow(&I));
+      setOrigin(&I, getCleanOrigin());
+      return;
+    }
 
     Instructions.push_back(&I);
   }

>From e31d66f23c960697486cb71b98060df503a7dff7 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka at google.com>
Date: Tue, 30 Apr 2024 23:53:01 -0700
Subject: [PATCH 3/3] remove "instrument"

Created using spr 1.3.4
---
 llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index a903863e0265b2..2b504b893ddb0d 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -1135,7 +1135,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   std::unique_ptr<VarArgHelper> VAHelper;
   const TargetLibraryInfo *TLI;
   Instruction *FnPrologueEnd;
-  SmallVector<Instruction *, 128> Instructions;
+  SmallVector<Instruction *, 16> Instructions;
 
   // The following flags disable parts of MSan instrumentation based on
   // exclusion list contents and command-line options.
@@ -1524,7 +1524,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     // `visit` above only collects instructions. Process them after iterating
     // CFG to avoid requirement on CFG transformations.
     for (Instruction *I : Instructions)
-      instrument(*I);
+      InstVisitor<MemorySanitizerVisitor>::visit(*I);
 
     // Finalize PHI nodes.
     for (PHINode *PN : ShadowPHINodes) {
@@ -2187,10 +2187,6 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     return ConstantDataVector::get(IRB.getContext(), OrderingTable);
   }
 
-  void instrument(Instruction &I) {
-    InstVisitor<MemorySanitizerVisitor>::visit(I);
-  }
-
   // ------------------- Visitors.
   using InstVisitor<MemorySanitizerVisitor>::visit;
   void visit(Instruction &I) {