[MemorySSAUpdater] Fix iterator invalidation bug in applyInsertUpdates (#139370)

This patch defers resetting optimized accesses until all uses are
replaced, to avoid invalidating the iterator.

Closes #139103.
Closes #139289.
Closes #139308.

Author: dtcxzyw

llvm/lib/Analysis/MemorySSAUpdater.cpp

@@ -1119,6 +1119,9 @@ void MemorySSAUpdater::applyInsertUpdates(ArrayRef<CFGUpdate> Updates,
     if (auto DefsList = MSSA->getWritableBlockDefs(BlockWithDefsToReplace)) {
       for (auto &DefToReplaceUses : *DefsList) {
         BasicBlock *DominatingBlock = DefToReplaceUses.getBlock();
+        // We defer resetting optimized accesses until all uses are replaced, to
+        // avoid invalidating the iterator.
+        SmallVector<MemoryUseOrDef *, 4> ResetOptimized;
         for (Use &U : llvm::make_early_inc_range(DefToReplaceUses.uses())) {
           MemoryAccess *Usr = cast<MemoryAccess>(U.getUser());
           if (MemoryPhi *UsrPhi = dyn_cast<MemoryPhi>(Usr)) {
@@ -1135,10 +1138,13 @@ void MemorySSAUpdater::applyInsertUpdates(ArrayRef<CFGUpdate> Updates,
                 assert(IDom && "Block must have a valid IDom.");
                 U.set(GetLastDef(IDom->getBlock()));
               }
-              cast<MemoryUseOrDef>(Usr)->resetOptimized();
+              ResetOptimized.push_back(cast<MemoryUseOrDef>(Usr));
             }
           }
         }
+
+        for (auto *Usr : ResetOptimized)
+          Usr->resetOptimized();
       }
     }
   }

llvm/test/Analysis/MemorySSA/pr139103.ll

@@ -0,0 +1,47 @@
+; RUN: opt -disable-output -passes="loop-mssa(licm,loop-rotate,licm,simple-loop-unswitch<nontrivial>),print<memoryssa>" -verify-memoryssa < %s 2>&1 | FileCheck %s
+
+; Make sure that we update MSSA correctly in this case.
+
+; CHECK-LABEL: MemorySSA for function: test
+; CHECK: for.header2.preheader:
+; CHECK-NEXT: 11 = MemoryPhi({entry.split,liveOnEntry},{for.header,9})
+; CHECK: for.body.us:
+; CHECK-NEXT: 7 = MemoryPhi({for.header2.preheader.split.us,11},{for.header2.us,9})
+; CHECK-NEXT: 8 = MemoryDef(7)->7
+; CHECK-NEXT: store i32 0, ptr %p, align 4
+; CHECK-NEXT: 9 = MemoryDef(8)->8
+; CHECK-NEXT: store i8 0, ptr %p, align 1
+
+define void @test(ptr %p, i1 %cond) {
+entry:
+  br label %for.header
+
+for.header:
+  br i1 false, label %exit.loopexit1, label %for.header2.preheader
+
+for.header2.preheader:
+  br label %for.body
+
+for.header2:
+  br i1 false, label %for.latch, label %for.body
+
+for.body:
+  store i32 0, ptr %p, align 4
+  store i8 0, ptr %p, align 1
+  br i1 %cond, label %for.header2, label %exit.loopexit
+
+for.latch:
+  br i1 false, label %for.inc, label %exit.loopexit1
+
+for.inc:
+  br label %for.header
+
+exit.loopexit:
+  br label %exit
+
+exit.loopexit1:
+  br label %exit
+
+exit:
+  ret void
+}