ClosureLifetimeFixup: fix a crash due to an iterator invalidation problem.

include/swift/SILOptimizer/Utils/CFGOptUtils.h

@@ -47,21 +47,7 @@ struct InstModCallbacks;
 /// \return The created branch. The old branch is deleted.
 /// The argument is appended at the end of the argument tuple.
 TermInst *addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
-                                  SILValue val, InstModCallbacks &callbacks);
-
-/// Adds a new argument to an edge between a branch and a destination
-/// block.
-///
-/// \param branch The terminator to add the argument to.
-/// \param dest The destination block of the edge.
-/// \param val The value to the arguments of the branch.
-/// \return The created branch. The old branch is deleted.
-/// The argument is appended at the end of the argument tuple.
-inline TermInst *addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
-                                         SILValue val) {
-  InstModCallbacks callbacks;
-  return addNewEdgeValueToBranch(branch, dest, val, callbacks);
-}
+                                  SILValue val, InstructionDeleter &deleter);
 
 /// Changes the edge value between a branch and destination basic block
 /// at the specified index. Changes all edges from \p Branch to \p Dest to carry

include/swift/SILOptimizer/Utils/SILSSAUpdater.h

@@ -14,6 +14,7 @@
 #define SWIFT_SIL_SILSSAUPDATER_H
 
 #include "llvm/Support/Allocator.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILValue.h"
 
@@ -58,6 +59,10 @@ class SILSSAUpdater {
   // If not null updated with inserted 'phi' nodes (SILArgument).
   SmallVectorImpl<SILPhiArgument *> *insertedPhis;
 
+  // Used to delete branch instructions when they are replaced for adding
+  // phi arguments.
+  InstructionDeleter deleter;
+
   // Not copyable.
   void operator=(const SILSSAUpdater &) = delete;
   SILSSAUpdater(const SILSSAUpdater &) = delete;
@@ -67,6 +72,8 @@ class SILSSAUpdater {
       SmallVectorImpl<SILPhiArgument *> *insertedPhis = nullptr);
   ~SILSSAUpdater();
 
+  InstructionDeleter &getDeleter() { return deleter; }
+
   void setInsertedPhis(SmallVectorImpl<SILPhiArgument *> *inputInsertedPhis) {
     insertedPhis = inputInsertedPhis;
   }

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -208,7 +208,8 @@ cleanupDeadTrivialPhiArgs(SILValue initialValue,
 /// that case where we have to consider that destroy_value, we have a simpler
 /// time here.
 static void extendLifetimeToEndOfFunction(SILFunction &fn,
-                                          ConvertEscapeToNoEscapeInst *cvt) {
+                                          ConvertEscapeToNoEscapeInst *cvt,
+                                          SILSSAUpdater &updater) {
   auto escapingClosure = cvt->getOperand();
   auto escapingClosureTy = escapingClosure->getType();
   auto optionalEscapingClosureTy = SILType::getOptionalType(escapingClosureTy);
@@ -239,7 +240,7 @@ static void extendLifetimeToEndOfFunction(SILFunction &fn,
   // use SILSSAUpdater::GetValueInMiddleOfBlock() to extend the object's
   // lifetime respecting loops.
   SmallVector<SILPhiArgument *, 8> insertedPhis;
-  SILSSAUpdater updater(&insertedPhis);
+  updater.setInsertedPhis(&insertedPhis);
   updater.initialize(optionalEscapingClosureTy, fn.hasOwnership()
                                                     ? OwnershipKind::Owned
                                                     : OwnershipKind::None);
@@ -443,7 +444,7 @@ static SILValue skipConvert(SILValue v) {
 /// nesting.
 static SILValue tryRewriteToPartialApplyStack(
     ConvertEscapeToNoEscapeInst *cvt,
-    SILInstruction *closureUser, SILBasicBlock::iterator &advanceIfDelete,
+    SILInstruction *closureUser, InstructionDeleter &deleter,
     llvm::DenseMap<SILInstruction *, SILInstruction *> &memoized) {
 
   auto *origPA = dyn_cast<PartialApplyInst>(skipConvert(cvt->getOperand()));
@@ -457,10 +458,8 @@ static SILValue tryRewriteToPartialApplyStack(
   // Whenever we delete an instruction advance the iterator and remove the
   // instruction from the memoized map.
   auto saveDeleteInst = [&](SILInstruction *i) {
-    if (&*advanceIfDelete == i)
-      ++advanceIfDelete;
     memoized.erase(i);
-    i->eraseFromParent();
+    deleter.forceDelete(i);
   };
 
   // Look for a single non ref count user of the partial_apply.
@@ -534,7 +533,7 @@ static SILValue tryRewriteToPartialApplyStack(
       SmallVector<Operand*, 16> Uses(arg.get()->getUses());
       for (auto use : Uses)
         if (auto *deallocInst = dyn_cast<DeallocStackInst>(use->getUser()))
-          deallocInst->eraseFromParent();
+          deleter.forceDelete(deallocInst);
     }
   }
 
@@ -552,7 +551,7 @@ static SILValue tryRewriteToPartialApplyStack(
 static bool tryExtendLifetimeToLastUse(
     ConvertEscapeToNoEscapeInst *cvt,
     llvm::DenseMap<SILInstruction *, SILInstruction *> &memoized,
-    SILBasicBlock::iterator &advanceIfDelete) {
+    InstructionDeleter &deleter) {
   // If there is a single user that is an apply this is simple: extend the
   // lifetime of the operand until after the apply.
   auto *singleUser = lookThroughRebastractionUsers(cvt, memoized);
@@ -574,7 +573,7 @@ static bool tryExtendLifetimeToLastUse(
   }
 
   if (SILValue closure = tryRewriteToPartialApplyStack(cvt, singleUser,
-                                                   advanceIfDelete, memoized)) {
+                                                       deleter, memoized)) {
     if (auto *cfi = dyn_cast<ConvertFunctionInst>(closure))
       closure = cfi->getOperand();
     if (endAsyncLet && isa<MarkDependenceInst>(closure)) {
@@ -587,7 +586,7 @@ static bool tryExtendLifetimeToLastUse(
       builder.createBuiltin(endAsyncLet->getLoc(), endAsyncLet->getName(),
         endAsyncLet->getType(), endAsyncLet->getSubstitutions(),
         {endAsyncLet->getOperand(0), closure});
-      endAsyncLet->eraseFromParent();
+      deleter.forceDelete(endAsyncLet);
     }
     return true;
   }
@@ -797,6 +796,7 @@ static SILInstruction *getOnlyDestroy(CopyBlockWithoutEscapingInst *cb) {
 ///      cond_fail %e
 ///      destroy_value %closure
 static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *cb,
+                                          InstructionDeleter &deleter,
                                           bool &modifiedCFG) {
   // Find the end of the lifetime of the copy_block_without_escaping
   // instruction.
@@ -837,7 +837,7 @@ static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *cb,
   SILBuilderWithScope b(cb);
   auto *newCB = b.createCopyBlock(loc, cb->getBlock());
   cb->replaceAllUsesWith(newCB);
-  cb->eraseFromParent();
+  deleter.forceDelete(cb);
 
   auto autoGenLoc = RegularLocation::getAutoGeneratedLocation();
 
@@ -977,14 +977,12 @@ static bool fixupClosureLifetimes(SILFunction &fn, bool &checkStackNesting,
   llvm::DenseMap<SILInstruction *, SILInstruction *> memoizedQueries;
 
   for (auto &block : fn) {
-    auto i = block.begin();
-    while (i != block.end()) {
-      SILInstruction *inst = &*i;
-      ++i;
+    SILSSAUpdater updater;
+    for (SILInstruction *inst : updater.getDeleter().updatingRange(&block)) {
 
       // Handle, copy_block_without_escaping instructions.
       if (auto *cb = dyn_cast<CopyBlockWithoutEscapingInst>(inst)) {
-        if (fixupCopyBlockWithoutEscaping(cb, modifiedCFG)) {
+        if (fixupCopyBlockWithoutEscaping(cb, updater.getDeleter(), modifiedCFG)) {
           changed = true;
         }
         continue;
@@ -1004,15 +1002,15 @@ static bool fixupClosureLifetimes(SILFunction &fn, bool &checkStackNesting,
         }
       }
 
-      if (tryExtendLifetimeToLastUse(cvt, memoizedQueries, i)) {
+      if (tryExtendLifetimeToLastUse(cvt, memoizedQueries, updater.getDeleter())) {
         changed = true;
         checkStackNesting = true;
         continue;
       }
 
       // Otherwise, extend the lifetime of the operand to the end of the
       // function.
-      extendLifetimeToEndOfFunction(fn, cvt);
+      extendLifetimeToEndOfFunction(fn, cvt, updater);
       changed = true;
     }
   }

lib/SILOptimizer/Utils/CFGOptUtils.cpp

@@ -26,7 +26,7 @@ using namespace swift;
 
 TermInst *swift::addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
                                          SILValue val,
-                                         InstModCallbacks &callbacks) {
+                                         InstructionDeleter &deleter) {
   SILBuilderWithScope builder(branch);
   TermInst *newBr = nullptr;
 
@@ -53,7 +53,7 @@ TermInst *swift::addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
         cbi->getLoc(), cbi->getCondition(), cbi->getTrueBB(), trueArgs,
         cbi->getFalseBB(), falseArgs, cbi->getTrueBBCount(),
         cbi->getFalseBBCount());
-    callbacks.createdNewInst(newBr);
+    deleter.getCallbacks().createdNewInst(newBr);
   } else if (auto *bi = dyn_cast<BranchInst>(branch)) {
     SmallVector<SILValue, 8> args;
 
@@ -63,13 +63,13 @@ TermInst *swift::addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
     args.push_back(val);
     assert(args.size() == dest->getNumArguments());
     newBr = builder.createBranch(bi->getLoc(), bi->getDestBB(), args);
-    callbacks.createdNewInst(newBr);
+    deleter.getCallbacks().createdNewInst(newBr);
   } else {
     // At the moment we can only add arguments to br and cond_br.
     llvm_unreachable("Can't add argument to terminator");
   }
 
-  callbacks.deleteInst(branch);
+  deleter.forceDelete(branch);
 
   return newBr;
 }

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -78,7 +78,8 @@ insertOwnedBaseValueAlongBranchEdge(BranchInst *bi, SILValue innerCopy,
   // argument.
   auto *phiArg =
       destBB->createPhiArgument(innerCopy->getType(), OwnershipKind::Owned);
-  addNewEdgeValueToBranch(bi, destBB, innerCopy, callbacks);
+  InstructionDeleter deleter(callbacks);
+  addNewEdgeValueToBranch(bi, destBB, innerCopy, deleter);
 
   // Grab our predecessor blocks, ignoring us, add to the branch edge an
   // undef corresponding to our value.
@@ -93,7 +94,7 @@ insertOwnedBaseValueAlongBranchEdge(BranchInst *bi, SILValue innerCopy,
       continue;
     addNewEdgeValueToBranch(
         predBlock->getTerminator(), destBB,
-        SILUndef::get(innerCopy->getType(), *destBB->getParent()), callbacks);
+        SILUndef::get(innerCopy->getType(), *destBB->getParent()), deleter);
   }
 
   return phiArg;

lib/SILOptimizer/Utils/SILSSAUpdater.cpp

@@ -230,9 +230,10 @@ SILValue SILSSAUpdater::getValueInMiddleOfBlock(SILBasicBlock *block) {
 
   // Create a new phi node.
   SILPhiArgument *phiArg = block->createPhiArgument(type, OwnershipKind::Owned);
-  for (auto &pair : predVals)
-    addNewEdgeValueToBranch(pair.first->getTerminator(), block, pair.second);
-
+  for (auto &pair : predVals) {
+    addNewEdgeValueToBranch(pair.first->getTerminator(), block, pair.second,
+                            deleter);
+  }
   if (insertedPhis)
     insertedPhis->push_back(phiArg);
 
@@ -326,7 +327,8 @@ class SSAUpdaterTraits<SILSSAUpdater> {
 
     for (auto *predBlock : predBlockList) {
       TermInst *ti = predBlock->getTerminator();
-      addNewEdgeValueToBranch(ti, block, ssaUpdater->phiSentinel.get());
+      addNewEdgeValueToBranch(ti, block, ssaUpdater->phiSentinel.get(),
+                              ssaUpdater->deleter);
     }
 
     return phi;

test/SILOptimizer/closure-lifetime-fixup.sil

@@ -178,3 +178,33 @@ bb0(%0 : @guaranteed $Klass, %1 : @guaranteed $Klass):
   %42 = tuple ()
   return %42 : $()
 }
+
+sil @simpleClosure : $@convention(thin) () -> ()
+
+// Don't crash.
+// CHECK-LABEL: sil [ossa] @testIteratorInvalidation
+// CHECK:    [[C:%.*]] = thin_to_thick_function
+// CHECK:    [[CC:%.*]] = copy_value [[C]]
+// CHECK:    [[NE:%.*]] = convert_escape_to_noescape [[CC]]
+// CHECK:    br bb3([[NE:%.*]]
+// CHECK:  } // end sil function 'testIteratorInvalidation'
+sil [ossa] @testIteratorInvalidation : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> () {
+bb0(%0 : $@noescape @callee_guaranteed () -> ()):
+  %2 = function_ref @simpleClosure : $@convention(thin) () -> ()
+  %3 = thin_to_thick_function %2 : $@convention(thin) () -> () to $@callee_guaranteed () -> ()
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb3(%0 : $@noescape @callee_guaranteed () -> ())
+
+bb2:
+  %11 = convert_escape_to_noescape [not_guaranteed] %3 : $@callee_guaranteed () -> () to $@noescape @callee_guaranteed () -> ()
+  br bb3(%11 : $@noescape @callee_guaranteed () -> ())
+
+
+bb3(%13 : $@noescape @callee_guaranteed () -> ()):
+  %15 = apply %13() : $@noescape @callee_guaranteed () -> ()
+  %16 = tuple ()
+  return %16 : $()
+}
+