Merge pull request #67760 from atrick/fix-closure-moveonly-arg

Fix compiler crashes with consuming and borrowing keywords.

Author: swift-ci

include/swift/SIL/ApplySite.h

@@ -375,9 +375,11 @@ class ApplySite {
   ///
   /// For full applies, this is equivalent to `getArgumentConvention`. But for
   /// a partial_apply, the argument ownership convention at the partial_apply
-  /// instruction itself is different from the argument convention of the callee.
+  /// instruction itself is different from the argument convention of the
+  /// callee.
+  /// 
   /// For details see the partial_apply documentation in SIL.rst.
-  SILArgumentConvention getArgumentOperandConvention(const Operand &oper) const {
+  SILArgumentConvention getCaptureConvention(const Operand &oper) const {
     SILArgumentConvention conv = getArgumentConvention(oper);
     auto *pai = dyn_cast<PartialApplyInst>(Inst);
     if (!pai)
@@ -388,6 +390,9 @@ class ApplySite {
     case SILArgumentConvention::Pack_Inout:
       return conv;
     case SILArgumentConvention::Direct_Owned:
+      assert(!pai->isOnStack() &&
+             "on-stack closures do not support owned arguments");
+      return conv;
     case SILArgumentConvention::Direct_Unowned:
     case SILArgumentConvention::Direct_Guaranteed:
       return pai->isOnStack() ? SILArgumentConvention::Direct_Guaranteed

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -281,17 +281,24 @@ void releasePartialApplyCapturedArg(
     SILParameterInfo paramInfo,
     InstModCallbacks callbacks = InstModCallbacks());
 
-/// Insert destroys of captured arguments of partial_apply [stack].
+/// Insert destroys of captured arguments of partial_apply [stack]. \p builder
+/// indicates a position at which the closure's lifetime ends.
+///
+/// The \p getValueToDestroy callback allows the caller to handle some captured
+/// arguments specially. For example, ClosureLifetimeFixup generates borrow
+/// scopes for captured arguments; each getValueToDestroy callback then inserts
+/// the corresponding end_borrow and returns the owned operand of the borrow,
+/// which will then be destroyed as usual.
 void insertDestroyOfCapturedArguments(
     PartialApplyInst *pai, SILBuilder &builder,
-    llvm::function_ref<bool(SILValue)> shouldInsertDestroy =
-        [](SILValue arg) -> bool { return true; },
+    llvm::function_ref<SILValue(SILValue)> getValueToDestroy =
+        [](SILValue arg) -> SILValue { return arg; },
     SILLocation loc = RegularLocation::getAutoGeneratedLocation());
 
 void insertDeallocOfCapturedArguments(PartialApplyInst *pai,
     DominanceInfo *domInfo,
-    llvm::function_ref<bool(SILValue)> shouldInsertDestroy =
-      [](SILValue arg) -> bool { return true; });
+    llvm::function_ref<SILValue(SILValue)> getAddressToDealloc =
+      [](SILValue arg) -> SILValue { return arg; });
 
 /// This iterator 'looks through' one level of builtin expect users exposing all
 /// users of the looked through builtin expect instruction i.e it presents a

lib/SIL/Verifier/MemoryLifetimeVerifier.cpp

@@ -434,7 +434,7 @@ void MemoryLifetimeVerifier::initDataflowInBlock(SILBasicBlock *block,
         ApplySite AS(&I);
         for (Operand &op : I.getAllOperands()) {
           if (AS.isArgumentOperand(op)) {
-            setFuncOperandBits(state, op, AS.getArgumentOperandConvention(op),
+            setFuncOperandBits(state, op, AS.getCaptureConvention(op),
                               isa<TryApplyInst>(&I));
           }
         }
@@ -771,7 +771,7 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
         ApplySite AS(&I);
         for (Operand &op : I.getAllOperands()) {
           if (AS.isArgumentOperand(op))
-            checkFuncArgument(bits, op, AS.getArgumentOperandConvention(op), &I);
+            checkFuncArgument(bits, op, AS.getCaptureConvention(op), &I);
         }
         break;
       }

lib/SIL/Verifier/SILVerifier.cpp

@@ -1956,6 +1956,11 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     unsigned appliedArgStartIdx =
         substConv.getNumSILArguments() - PAI->getNumArguments();
     for (auto p : llvm::enumerate(PAI->getArguments())) {
+      if (PAI->isOnStack()) {
+        require(
+          substConv.getSILArgumentConvention(appliedArgStartIdx + p.index()),
+          "on-stack closures do not support owned arguments");
+      }
       requireSameType(
           p.value()->getType(),
           substConv.getSILArgumentType(appliedArgStartIdx + p.index(),

lib/SILOptimizer/Mandatory/CapturePromotion.cpp

@@ -1555,7 +1555,9 @@ processPartialApplyInst(SILOptFunctionBuilder &funcBuilder,
         builder.setInsertionPoint(std::next(SILBasicBlock::iterator(dsi)));
         insertDestroyOfCapturedArguments(
             newPAI, builder,
-            [&](SILValue arg) -> bool { return newCaptures.count(arg); });
+            [&](SILValue arg) -> SILValue {
+              return newCaptures.count(arg) ? arg : SILValue();
+            });
       }
     }
     // Map the mark dependence arguments.

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -889,12 +889,28 @@ static SILValue tryRewriteToPartialApplyStack(
     destroy->dump();
     */
     SILBuilderWithScope builder(std::next(destroy->getIterator()));
-    insertDestroyOfCapturedArguments(newPA, builder,
-                                     [&](SILValue arg) -> bool {
-                                       // Don't need to destroy if we borrowed
-                                       // in place .
-                                       return !borrowedOriginals.count(arg);
-                                     },
+    // This getCapturedArg hack attempts to perfectly compensate for all the
+    // other hacks involved in gathering new arguments above.
+    auto getArgToDestroy = [&](SILValue argValue) -> SILValue {
+      // A MoveOnlyWrapperToCopyableValueInst may produce a trivial value. Be
+      // careful not to emit an extra destroy of the original.
+      if (argValue->getType().isTrivial(argValue->getFunction()))
+        return SILValue();
+
+      // We may have inserted a new begin_borrow->moveonlywrapper_to_copyvalue
+      // when creating the new arguments. Now we need to end that borrow.
+      if (auto *m = dyn_cast<MoveOnlyWrapperToCopyableValueInst>(argValue))
+        if (m->hasGuaranteedInitialKind())
+          argValue = m->getOperand();
+      auto *argBorrow = dyn_cast<BeginBorrowInst>(argValue);
+      if (argBorrow) {
+        argValue = argBorrow->getOperand();
+        builder.createEndBorrow(newPA->getLoc(), argBorrow);
+      }
+      // Don't need to destroy if we borrowed in place .
+      return borrowedOriginals.count(argValue) ? SILValue() : argValue;
+    };
+    insertDestroyOfCapturedArguments(newPA, builder, getArgToDestroy,
                                      newPA->getLoc());
   }
   /* DEBUG
@@ -919,13 +935,17 @@ static SILValue tryRewriteToPartialApplyStack(
   if (unreachableBlocks.count(newPA->getParent()))
     return closureOp;
 
+  auto getAddressToDealloc = [&](SILValue argAddress) -> SILValue {
+    if (auto moveWrapper =
+        dyn_cast<MoveOnlyWrapperToCopyableAddrInst>(argAddress)) {
+      argAddress = moveWrapper->getOperand();
+    }
+    // Don't need to destroy if we borrowed in place .
+    return borrowedOriginals.count(argAddress) ? SILValue() : argAddress;
+  };
   insertDeallocOfCapturedArguments(
       newPA, dominanceAnalysis->get(closureUser->getFunction()),
-      [&](SILValue arg) -> bool {
-        // Don't need to destroy if we borrowed
-        // in place.
-        return !borrowedOriginals.count(arg);
-      });
+      getAddressToDealloc);
 
   return closureOp;
 }

lib/SILOptimizer/Mandatory/ConsumeOperatorCopyableAddressesChecker.cpp

@@ -1332,7 +1332,7 @@ bool GatherLexicalLifetimeUseVisitor::visitUse(Operand *op,
       return false;
     }
 
-    if (fas.getArgumentOperandConvention(*op) ==
+    if (fas.getCaptureConvention(*op) ==
         SILArgumentConvention::Indirect_InoutAliasable) {
       // If we don't find the function, we can't handle this, so bail.
       auto *func = fas.getCalleeFunction();

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -2378,7 +2378,7 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
     // either.
     if (auto *f = pas->getCalleeFunction()) {
       if (f->hasSemanticsAttr(semantics::NO_MOVEONLY_DIAGNOSTICS)) {
-        if (ApplySite(pas).getArgumentOperandConvention(*op).isInoutConvention()) {
+        if (ApplySite(pas).getCaptureConvention(*op).isInoutConvention()) {
           diagnosticEmitter.emitEarlierPassEmittedDiagnostic(markedValue);
           return false;
         }

lib/SILOptimizer/Mandatory/MoveOnlyUtils.cpp

@@ -259,7 +259,7 @@ bool noncopyable::memInstMustReinitialize(Operand *memOper) {
   case SILInstructionKind::TryApplyInst:
   case SILInstructionKind::ApplyInst: {
     FullApplySite applySite(memInst);
-    return applySite.getArgumentOperandConvention(*memOper).isInoutConvention();
+    return applySite.getCaptureConvention(*memOper).isInoutConvention();
   }
   case SILInstructionKind::StoreInst:
     return cast<StoreInst>(memInst)->getOwnershipQualifier() ==
@@ -296,7 +296,7 @@ bool noncopyable::memInstMustConsume(Operand *memOper) {
   case SILInstructionKind::TryApplyInst:
   case SILInstructionKind::ApplyInst: {
     FullApplySite applySite(memInst);
-    return applySite.getArgumentOperandConvention(*memOper).isOwnedConvention();
+    return applySite.getCaptureConvention(*memOper).isOwnedConvention();
   }
   case SILInstructionKind::PartialApplyInst: {
     // If we are on the stack or have an inout convention, we do not

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -902,13 +902,14 @@ void swift::emitDestroyOperation(SILBuilder &builder, SILLocation loc,
   callbacks.createdNewInst(u.get<ReleaseValueInst *>());
 }
 
-// *HEY YOU, YES YOU, PLEASE READ*. Even though a textual partial apply is
-// printed with the convention of the closed over function upon it, all
-// non-inout arguments to a partial_apply are passed at +1. This includes
-// arguments that will eventually be passed as guaranteed or in_guaranteed to
-// the closed over function. This is because the partial apply is building up a
-// boxed aggregate to send off to the closed over function. Of course when you
-// call the function, the proper conventions will be used.
+// NOTE: The ownership of the partial_apply argument does not match the
+// convention of the closed over function. All non-inout arguments to a
+// partial_apply are passed at +1 for regular escaping closures and +0 for
+// closures that have been promoted to partial_apply [on_stack]. An escaping
+// partial apply stores each capture in an owned box, even for guaranteed and
+// in_guaranteed argument convention. A non-escaping/on-stack closure either
+// borrows its arguments or takes an inout_aliasable address. Non-escaping
+// closures do not support owned arguments.
 void swift::releasePartialApplyCapturedArg(SILBuilder &builder, SILLocation loc,
                                            SILValue arg,
                                            SILParameterInfo paramInfo,
@@ -1479,7 +1480,7 @@ swift::findLocalApplySites(FunctionRefBaseInst *fri) {
 /// Insert destroys of captured arguments of partial_apply [stack].
 void swift::insertDestroyOfCapturedArguments(
     PartialApplyInst *pai, SILBuilder &builder,
-    llvm::function_ref<bool(SILValue)> shouldInsertDestroy,
+    llvm::function_ref<SILValue(SILValue)> getValueToDestroy,
     SILLocation origLoc) {
   assert(pai->isOnStack());
 
@@ -1488,27 +1489,25 @@ void swift::insertDestroyOfCapturedArguments(
                                     pai->getModule());
   auto loc = CleanupLocation(origLoc);
   for (auto &arg : pai->getArgumentOperands()) {
-    SILValue argValue = arg.get();
-    if (auto *m = dyn_cast<MoveOnlyWrapperToCopyableValueInst>(argValue))
-      if (m->hasGuaranteedInitialKind())
-        argValue = m->getOperand();
-    auto *argBorrow = dyn_cast<BeginBorrowInst>(argValue);
-    if (argBorrow) {
-      argValue = argBorrow->getOperand();
-      builder.createEndBorrow(loc, argBorrow);
-    }
-    if (!shouldInsertDestroy(argValue))
+    SILValue argValue = getValueToDestroy(arg.get());
+    if (!argValue)
       continue;
+
+    assert(!pai->getFunction()->hasOwnership()
+           || (argValue->getOwnershipKind().isCompatibleWith(
+                 OwnershipKind::Owned)));
+
     unsigned calleeArgumentIndex = site.getCalleeArgIndex(arg);
     assert(calleeArgumentIndex >= calleeConv.getSILArgIndexOfFirstParam());
     auto paramInfo = calleeConv.getParamInfoForSILArg(calleeArgumentIndex);
     releasePartialApplyCapturedArg(builder, loc, argValue, paramInfo);
   }
 }
 
-void swift::insertDeallocOfCapturedArguments(PartialApplyInst *pai,
-                         DominanceInfo *domInfo,
-                         llvm::function_ref<bool(SILValue)> shouldInsertDestroy)
+void swift::insertDeallocOfCapturedArguments(
+  PartialApplyInst *pai,
+  DominanceInfo *domInfo,
+  llvm::function_ref<SILValue(SILValue)> getAddressToDealloc)
 {
   assert(pai->isOnStack());
 
@@ -1522,13 +1521,10 @@ void swift::insertDeallocOfCapturedArguments(PartialApplyInst *pai,
     if (!paramInfo.isIndirectInGuaranteed())
       continue;
 
-    SILValue argValue = arg.get();
-    if (!shouldInsertDestroy(argValue)) {
+    SILValue argValue = getAddressToDealloc(arg.get());
+    if (!argValue) {
       continue;
     }
-    if (auto moveWrapper =
-            dyn_cast<MoveOnlyWrapperToCopyableAddrInst>(argValue))
-      argValue = moveWrapper->getOperand();
 
     SmallVector<SILBasicBlock *, 4> boundary;
     auto *asi = cast<AllocStackInst>(argValue);