Merge pull request #59160 from meg-gupta/closurefixupbug

ClosureLifetimeFixup : Fix bug in lifetime extension of non escaping closure

Author: meg-gupta

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -201,7 +201,12 @@ cleanupDeadTrivialPhiArgs(SILValue initialValue,
 
 /// Extend the lifetime of the convert_escape_to_noescape's operand to the end
 /// of the function.
-///
+/// Create a copy of the escaping closure operand and end its lifetime at
+/// function exits. Since, the cvt may not be dominating function exits, we
+/// need to create an optional and use the SSAUpdater to extend the lifetime. In
+/// order to prevent the optional being optimized away, create a borrow scope
+/// and insert a mark_dependence of the non escaping closure on the borrow.
+
 /// NOTE: Since we are lifetime extending a copy that we have introduced, we do
 /// not need to consider destroy_value emitted by SILGen unlike
 /// copy_block_without_escaping which consumes its sentinel parameter. Unlike
@@ -215,91 +220,108 @@ static void extendLifetimeToEndOfFunction(SILFunction &fn,
   auto optionalEscapingClosureTy = SILType::getOptionalType(escapingClosureTy);
   auto loc = RegularLocation::getAutoGeneratedLocation();
 
-  // If our Cvt is in the initial block, we do not need to use the SSA updater
-  // since we know Cvt can not be in a loop and must dominate all exits
-  // (*). Just insert a copy of the escaping closure at the Cvt and destroys at
+  SmallVector<SILBasicBlock *, 4> exitingBlocks;
+  fn.findExitingBlocks(exitingBlocks);
+
+  auto createLifetimeEnd = [](SILLocation loc, SILInstruction *insertPt,
+                              SILValue value) {
+    SILBuilderWithScope builder(insertPt);
+    if (value.getOwnershipKind() == OwnershipKind::Owned) {
+      builder.emitDestroyOperation(loc, value);
+      return;
+    }
+    builder.emitEndBorrowOperation(loc, value);
+  };
+  auto createLifetimeEndAtFunctionExits =
+      [&](std::function<SILValue(SILBasicBlock *)> getValue) {
+        for (auto *block : exitingBlocks) {
+          auto *safeDestructionPoint =
+              getDeinitSafeClosureDestructionPoint(block);
+          createLifetimeEnd(loc, safeDestructionPoint, getValue(block));
+        }
+      };
+
+  // If our cvt is in the initial block, we do not need to use the SSA updater
+  // since we know cvt cannot be in a loop and must dominate all exits
+  // (*). Just insert a copy of the escaping closure at the cvt and destroys at
   // the exit blocks of the function.
   //
   // (*) In fact we can't use the SILSSAUpdater::GetValueInMiddleOfBlock.
   if (cvt->getParent() == cvt->getFunction()->getEntryBlock()) {
-    auto *innerCVI =
-        SILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
+    auto *copy = SILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
     cvt->setLifetimeGuaranteed();
-    cvt->setOperand(innerCVI);
-    SmallVector<SILBasicBlock *, 4> exitingBlocks;
-    fn.findExitingBlocks(exitingBlocks);
-
-    for (auto *block : exitingBlocks) {
-      auto *safeDestructPoint = getDeinitSafeClosureDestructionPoint(block);
-      SILBuilderWithScope(safeDestructPoint).createDestroyValue(loc, innerCVI);
-    }
+    cvt->setOperand(copy);
+    createLifetimeEndAtFunctionExits([&copy](SILBasicBlock *) { return copy; });
     return;
   }
 
-  // Ok. At this point we know that Cvt is not in the entry block... so we can
-  // use SILSSAUpdater::GetValueInMiddleOfBlock() to extend the object's
-  // lifetime respecting loops.
-  SmallVector<SILPhiArgument *, 8> insertedPhis;
-  updater.setInsertedPhis(&insertedPhis);
-  updater.initialize(optionalEscapingClosureTy, fn.hasOwnership()
-                                                    ? OwnershipKind::Owned
-                                                    : OwnershipKind::None);
-
-  // Create an Optional<() -> ()>.none in the entry block of the function and
-  // add it as an available value to the SSAUpdater.
-  //
-  // Since we know that Cvt is not in the entry block and this must be, we know
-  // that it is safe to use the SSAUpdater's getValueInMiddleOfBlock with this
-  // value.
-  SILValue entryBlockOptionalNone = [&]() -> SILValue {
-    SILBuilderWithScope b(fn.getEntryBlock()->begin());
-    return b.createOptionalNone(loc, optionalEscapingClosureTy);
-  }();
-  updater.addAvailableValue(fn.getEntryBlock(), entryBlockOptionalNone);
-
-  // Create a copy of the convert_escape_to_no_escape and add it as an available
-  // value to the SSA updater.
-  //
+  // Create a copy of the convert_escape_to_no_escape.
   // NOTE: The SSAUpdater does not support providing multiple values in the same
-  // block without extra work. So the fact that Cvt is not in the entry block
+  // block without extra work. So the fact that cvt is not in the entry block
   // means that we don't have to worry about overwriting the .none value.
-  auto *cvi = [&]() -> CopyValueInst * {
-    auto *innerCVI =
-        SILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
-    cvt->setLifetimeGuaranteed();
-    cvt->setOperand(innerCVI);
-    SILBuilderWithScope b(std::next(cvt->getIterator()));
-    updater.addAvailableValue(
-        cvt->getParent(),
-        b.createOptionalSome(loc, innerCVI, optionalEscapingClosureTy));
-    return innerCVI;
-  }();
-
-  // Then we use the SSA updater to insert a destroy_value before the cvt and at
-  // the reachable exit blocks.
-  SmallVector<SILBasicBlock *, 4> exitingBlocks;
-  findReachableExitBlocks(cvt, exitingBlocks);
-
-  {
-    // Before the copy value, insert an extra destroy_value to handle
-    // loops. Since we used our enum value this is safe.
-    SILValue v = updater.getValueInMiddleOfBlock(cvi->getParent());
-    SILBuilderWithScope(cvi).createDestroyValue(loc, v);
+  auto *copy = SILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
+  cvt->setLifetimeGuaranteed();
+  cvt->setOperand(copy);
+
+  // Create an optional some to extend the lifetime of copy until function
+  // exits.
+  SILBuilderWithScope lifetimeExtendBuilder(std::next(cvt->getIterator()));
+  auto *optionalSome = lifetimeExtendBuilder.createOptionalSome(
+      loc, copy, optionalEscapingClosureTy);
+
+  // Create a borrow scope and a mark_dependence to prevent the enum being
+  // optimized away.
+  auto *borrow = lifetimeExtendBuilder.createBeginBorrow(loc, optionalSome);
+  auto *mdi = lifetimeExtendBuilder.createMarkDependence(loc, cvt, borrow);
+
+  // Replace all uses of the non escaping closure with mark_dependence
+  SmallVector<Operand *, 4> convertUses;
+  for (auto *cvtUse : cvt->getUses()) {
+    convertUses.push_back(cvtUse);
   }
-
-  for (auto *block : exitingBlocks) {
-    auto *safeDestructionPt = getDeinitSafeClosureDestructionPoint(block);
-    SILValue v = updater.getValueAtEndOfBlock(block);
-    SILBuilderWithScope(safeDestructionPt).createDestroyValue(loc, v);
+  for (auto *cvtUse : convertUses) {
+    auto *cvtUser = cvtUse->getUser();
+    if (cvtUser == mdi)
+      continue;
+    cvtUser->setOperand(cvtUse->getOperandNumber(), mdi);
   }
 
-  // Finally, we need to prune phis inserted by the SSA updater that only take
-  // the .none from the entry block.
-  //
-  // TODO: Should we sort inserted phis before or after we initialize
-  // the worklist or maybe backwards? We should investigate how the
-  // SSA updater adds phi nodes to this list to resolve this question.
-  cleanupDeadTrivialPhiArgs(entryBlockOptionalNone, insertedPhis);
+  auto fixupSILForLifetimeExtension = [&](SILValue value) {
+    // Use SSAUpdater to find insertion points for lifetime ends.
+    updater.initialize(optionalEscapingClosureTy, value.getOwnershipKind());
+    SmallVector<SILPhiArgument *, 8> insertedPhis;
+    updater.setInsertedPhis(&insertedPhis);
+
+    // Create an optional none at the function entry.
+    auto *optionalNone =
+        SILBuilderWithScope(fn.getEntryBlock()->begin())
+            .createOptionalNone(loc, optionalEscapingClosureTy);
+    updater.addAvailableValue(fn.getEntryBlock(), optionalNone);
+    updater.addAvailableValue(value->getParentBlock(), value);
+    {
+      // Since value maybe in a loop, insert an extra lifetime end. Since we
+      // used our enum value, this is safe.
+      SILValue midValue =
+          updater.getValueInMiddleOfBlock(value->getParentBlock());
+      createLifetimeEnd(loc, cvt, midValue);
+    }
+
+    // Insert lifetime ends.
+    createLifetimeEndAtFunctionExits([&updater](SILBasicBlock *block) {
+      return updater.getValueAtEndOfBlock(block);
+    });
+
+    // Prune the phis inserted by the SSA updater that only take
+    // the .none from the entry block.
+    // TODO: Should we sort inserted phis before or after we initialize
+    // the worklist or maybe backwards? We should investigate how the
+    // SSA updater adds phi nodes to this list to resolve this question.
+    cleanupDeadTrivialPhiArgs(optionalNone, insertedPhis);
+  };
+
+  // Use the SSAUpdater to create lifetime ends for the copy and the borrow.
+  fixupSILForLifetimeExtension(borrow);
+  fixupSILForLifetimeExtension(optionalSome);
 }
 
 static SILInstruction *lookThroughRebastractionUsers(
@@ -979,8 +1001,8 @@ static bool fixupClosureLifetimes(SILFunction &fn, bool &checkStackNesting,
 
   for (auto &block : fn) {
     SILSSAUpdater updater;
-    for (SILInstruction *inst : updater.getDeleter().updatingRange(&block)) {
 
+    for (SILInstruction *inst : updater.getDeleter().updatingRange(&block)) {
       // Handle, copy_block_without_escaping instructions.
       if (auto *cb = dyn_cast<CopyBlockWithoutEscapingInst>(inst)) {
         if (fixupCopyBlockWithoutEscaping(cb, updater.getDeleter(), modifiedCFG)) {

lib/SILOptimizer/Utils/SILSSAUpdater.cpp

@@ -229,7 +229,7 @@ SILValue SILSSAUpdater::getValueInMiddleOfBlock(SILBasicBlock *block) {
   }
 
   // Create a new phi node.
-  SILPhiArgument *phiArg = block->createPhiArgument(type, OwnershipKind::Owned);
+  SILPhiArgument *phiArg = block->createPhiArgument(type, ownershipKind);
   for (auto &pair : predVals) {
     addNewEdgeValueToBranch(pair.first->getTerminator(), block, pair.second,
                             deleter);
@@ -341,6 +341,7 @@ class SSAUpdaterTraits<SILSSAUpdater> {
     auto *phiBlock = phi->getParent();
     size_t phiArgIndex = phi->getIndex();
     auto *ti = predBlock->getTerminator();
+
     changeEdgeValue(ti, phiBlock, phiArgIndex, value);
   }
 

test/SILOptimizer/closure-lifetime-fixup.sil

@@ -208,3 +208,50 @@ bb3(%13 : $@noescape @callee_guaranteed () -> ()):
   return %16 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @test_lifetime_extension_until_function_exit :
+// CHECK: [[NONE1:%.*]] = enum $Optional<@callee_guaranteed () -> ()>, #Optional.none!enumelt
+// CHECK: [[NONE2:%.*]] = enum $Optional<@callee_guaranteed () -> ()>, #Optional.none!enumelt
+// CHECK: cond_br undef, bb1, bb2
+// CHECK: bb1:
+// CHECK:  [[ORIG:%.*]] = function_ref @originalClosure : $@convention(thin) () -> ()
+// CHECK:  [[PA1:%.*]] = partial_apply [callee_guaranteed] [[ORIG]]() : $@convention(thin) () -> ()
+// CHECK:  [[COPY:%.*]] = copy_value [[PA1]] : $@callee_guaranteed () -> ()
+// CHECK:  [[NOESC:%.*]] = convert_escape_to_noescape [[COPY]] : $@callee_guaranteed () -> () to $@noescape @callee_guaranteed () -> ()
+// CHECK:  [[SOME:%.*]] = enum $Optional<@callee_guaranteed () -> ()>, #Optional.some!enumelt, [[COPY]] : $@callee_guaranteed () -> ()
+// CHECK:  [[BORROW:%.*]] = begin_borrow [[SOME]] : $Optional<@callee_guaranteed () -> ()>
+// CHECK:  [[MDI:%.*]] = mark_dependence [[NOESC]] : $@noescape @callee_guaranteed () -> () on [[BORROW]] : $Optional<@callee_guaranteed () -> ()>
+// CHECK:  br bb3([[BORROW]] : $Optional<@callee_guaranteed () -> ()>, [[SOME]] : $Optional<@callee_guaranteed () -> ()>, {{.*}})
+// CHECK-LABEL: } // end sil function 'test_lifetime_extension_until_function_exit'
+sil [ossa] @test_lifetime_extension_until_function_exit : $@convention(thin) (@guaranteed Klass, @guaranteed @callee_guaranteed () -> (), @guaranteed Klass, @guaranteed @callee_guaranteed () -> ()) -> () {
+bb0(%0 : @guaranteed $Klass, %1 : @guaranteed $@callee_guaranteed () -> (), %2 : @guaranteed $Klass, %3 : @guaranteed $@callee_guaranteed () -> ()):
+  cond_br undef, bb1, bb2
+
+bb1:
+  %5 = function_ref @originalClosure : $@convention(thin) () -> ()
+  %6 = partial_apply [callee_guaranteed] %5() : $@convention(thin) () -> ()
+  %7 = convert_escape_to_noescape [not_guaranteed] %6 : $@callee_guaranteed () -> () to $@noescape @callee_guaranteed () -> ()
+  %8 = function_ref @noEscapeThunk : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> ()
+  %9 = partial_apply [callee_guaranteed] %8(%7) : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> ()
+  %10 = mark_dependence %9 : $@callee_guaranteed () -> () on %7 : $@noescape @callee_guaranteed () -> ()
+  %11 = copy_value %10 : $@callee_guaranteed () -> ()
+  %12 = alloc_stack $@block_storage @callee_guaranteed () -> ()
+  %13 = project_block_storage %12 : $*@block_storage @callee_guaranteed () -> ()
+  store %11 to [init] %13 : $*@callee_guaranteed () -> ()
+  %15 = function_ref @blockThunk : $@convention(c) (@inout_aliasable @block_storage @callee_guaranteed () -> ()) -> ()
+  %16 = init_block_storage_header %12 : $*@block_storage @callee_guaranteed () -> (), invoke %15 : $@convention(c) (@inout_aliasable @block_storage @callee_guaranteed () -> ()) -> (), type $@convention(block) @noescape () -> ()
+  %17 = copy_block_without_escaping %16 : $@convention(block) @noescape () -> () withoutEscaping %10 : $@callee_guaranteed () -> ()
+  %18 = enum $Optional<@convention(block) @noescape () -> ()>, #Optional.some!enumelt, %17 : $@convention(block) @noescape () -> ()
+  destroy_addr %13 : $*@callee_guaranteed () -> ()
+  dealloc_stack %12 : $*@block_storage @callee_guaranteed () -> ()
+  destroy_value %6 : $@callee_guaranteed () -> ()
+  destroy_value %18 : $Optional<@convention(block) @noescape () -> ()>
+  br bb3
+
+bb2:
+  br bb3
+
+bb3:
+  %25 = tuple ()
+  return %25 : $()
+}
+

test/SILOptimizer/definite-init-convert-to-escape.swift

@@ -93,13 +93,13 @@ public func returnOptionalEscape() -> (() ->())?
 // NOPEEPHOLE:  switch_enum [[V1]] : $Optional<{{.*}}>, case #Optional.some!enumelt: [[SOME_BB:bb[0-9]+]], case #Optional.none!enumelt: [[NONE_BB:bb[0-9]+]]
 //
 // NOPEEPHOLE: [[SOME_BB]]([[V2:%.*]]: $@callee_guaranteed () -> ()):
-// NOPEEPHOLE-NEXT:  release_value [[NONE_2]]
 // NOPEEPHOLE-NEXT:  [[CVT:%.*]] = convert_escape_to_noescape [[V2]]
 // NOPEEPHOLE-NEXT:  [[SOME:%.*]] = enum $Optional<{{.*}}>, #Optional.some!enumelt, [[V2]]
-// NOPEEPHOLE-NEXT:  [[NOESCAPE_SOME:%.*]] = enum $Optional<{{.*}}>, #Optional.some!enumelt, [[CVT]]
-// NOPEEPHOLE-NEXT:  br bb2([[NOESCAPE_SOME]] : $Optional<{{.*}}>, [[SOME]] :
+// NOPEEPHOLE-NEXT:  [[MDI:%.*]] = mark_dependence [[CVT]] : $@noescape @callee_guaranteed () -> () on [[SOME]] : $Optional<@callee_guaranteed () -> ()>
+// NOPEEPHOLE-NEXT:  [[NOESCAPE_SOME:%.*]] = enum $Optional<{{.*}}>, #Optional.some!enumelt, [[MDI]]
+// NOPEEPHOLE-NEXT:  br bb2([[NOESCAPE_SOME]] : $Optional<{{.*}}>, [[SOME]] : $Optional<{{.*}}>, [[SOME]] : $Optional<{{.*}}>)
 //
-// NOPEEPHOLE: bb2([[NOESCAPE_SOME:%.*]] : $Optional<{{.*}}>, [[SOME:%.*]] :
+// NOPEEPHOLE: bb2([[NOESCAPE_SOME:%.*]] : $Optional<{{.*}}>, [[SOME1:%.*]] : $Optional<{{.*}}>, [[SOME:%.*]] : $Optional<{{.*}}>):
 // NOPEEPHOLE:   switch_enum [[NOESCAPE_SOME]] : $Optional<{{.*}}>, case #Optional.some!enumelt: [[SOME_BB_2:bb[0-9]+]], case #Optional.none!enumelt: [[NONE_BB_2:bb[0-9]+]]
 //
 // NOPEEPHOLE: [[SOME_BB_2]](