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

[via llvm-commits]

Author: Vitaly Buka
Date: 2024-05-01T14:47:00-07:00
New Revision: a2be1b8d03320908c4d4b42e5bd8b7569b1f9aae

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

LOG: [msan] Don't modify CFG iterating it (#90691)

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

Somehow existing `SplitBlockAndInsertIfThen` do not invalidate
iterators.

Added: 
    

Modified: 
    llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index cc2295c44023c4..2b504b893ddb0d 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 *, 16> 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)
+      InstVisitor<MemorySanitizerVisitor>::visit(*I);
+
     // Finalize PHI nodes.
     for (PHINode *PN : ShadowPHINodes) {
       PHINode *PNS = cast<PHINode>(getShadow(PN));
@@ -2196,7 +2202,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       setOrigin(&I, getCleanOrigin());
       return;
     }
-    InstVisitor<MemorySanitizerVisitor>::visit(I);
+
+    Instructions.push_back(&I);
   }
 
   /// Instrument LoadInst