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

The fix is the use the InstructionDeleter to iterate over (the not yet deleted) instructions in a block.

rdar://80093482

Author: eeckstein

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;
     }
   }

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 : $()
+}
+