1 files changed, 11 insertions, 24 deletions
diff --git a/llvm/lib/CodeGen/AtomicExpandPass.cpp b/llvm/lib/CodeGen/AtomicExpandPass.cpp
index 2da723a0..39a7055 100644
--- a/llvm/lib/CodeGen/AtomicExpandPass.cpp
+++ b/llvm/lib/CodeGen/AtomicExpandPass.cpp
@@ -351,30 +351,17 @@ bool AtomicExpandImpl::run(Function &F, const TargetMachine *TM) {
 
   bool MadeChange = false;
 
-  for (Function::iterator BBI = F.begin(), BBE = F.end(); BBI != BBE;) {
-    BasicBlock *BB = &*BBI;
-    ++BBI;
-
-    BasicBlock::iterator Next;
-
-    for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E;
-         I = Next) {
-      Instruction &Inst = *I;
-      Next = std::next(I);
-
-      if (processAtomicInstr(&Inst)) {
-        MadeChange = true;
-
-        // Detect control flow change and resume iteration from the original
-        // block to inspect any newly inserted blocks. This allows incremental
-        // legalization of atomicrmw and cmpxchg.
-        if (Next == E || BB != Next->getParent()) {
-          BBI = BB->getIterator();
-          BBE = F.end();
-          break;
-        }
-      }
-    }
+  SmallVector<Instruction *, 1> AtomicInsts;
+
+  // Changing control-flow while iterating through it is a bad idea, so gather a
+  // list of all atomic instructions before we start.
+  for (Instruction &I : instructions(F))
+    if (I.isAtomic() && !isa<FenceInst>(&I))
+      AtomicInsts.push_back(&I);
+
+  for (auto *I : AtomicInsts) {
+    if (processAtomicInstr(I))
+      MadeChange = true;
   }
 
   return MadeChange;