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

[5.9] Fix compiler crashes with consuming and borrowing keywords.

Author: atrick

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/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/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,
@@ -1468,7 +1469,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());
 
@@ -1477,27 +1478,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());
 
@@ -1511,13 +1510,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);

test/SILOptimizer/closure-lifetime-fixup.sil

@@ -310,3 +310,41 @@ bb3(%25 : @owned $any Error):
 bb4(%29 : @owned $any Error):
   unreachable
 }
+
+sil @testTrivialBorrowClosure : $@convention(thin) (Int) -> Int
+sil @testTrivialBorrowTakeClosure : $@convention(thin) (@guaranteed @noescape @callee_guaranteed () -> Int) -> Int
+
+// rdar://112795074 (Crash compiling function that has a macro annotation and uses `consuming`)
+//
+// CHECK-LABEL: sil [ossa] @testTrivialBorrow : $@convention(thin) (Int) -> Int {
+// CHECK: [[ARG:%.*]] = moveonlywrapper_to_copyable [guaranteed]
+// CHECK: [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] %{{.*}}([[ARG]]) : $@convention(thin) (Int) -> Int
+// CHECK-NOT: destroy 
+// CHECK: destroy_value [[PA]] : $@noescape @callee_guaranteed () -> Int
+// CHECK-NOT: destroy
+// CHECK: destroy_addr
+// CHECK-NEXT: dealloc_stack
+// CHECK-NEXT: return
+// CHECK-LABEL: } // end sil function 'testTrivialBorrow'
+sil [ossa] @testTrivialBorrow : $@convention(thin) (Int) -> Int {
+bb0(%0 : @noImplicitCopy @_eagerMove $Int):
+  %1 = alloc_stack $@moveOnly Int, var, name "i"
+  %2 = mark_must_check [consumable_and_assignable] %1 : $*@moveOnly Int
+  %3 = moveonlywrapper_to_copyable_addr %2 : $*@moveOnly Int
+  store %0 to [trivial] %3 : $*Int
+  %5 = begin_access [read] [static] %2 : $*@moveOnly Int
+  %6 = load_borrow %5 : $*@moveOnly Int
+  %7 = moveonlywrapper_to_copyable [guaranteed] %6 : $@moveOnly Int
+  end_borrow %6 : $@moveOnly Int
+  end_access %5 : $*@moveOnly Int
+  %10 = function_ref @testTrivialBorrowClosure : $@convention(thin) (Int) -> Int
+  %11 = partial_apply [callee_guaranteed] %10(%7) : $@convention(thin) (Int) -> Int
+  %12 = convert_escape_to_noescape [not_guaranteed] %11 : $@callee_guaranteed () -> Int to $@noescape @callee_guaranteed () -> Int
+  %13 = function_ref @testTrivialBorrowTakeClosure : $@convention(thin) (@guaranteed @noescape @callee_guaranteed () -> Int) -> Int
+  %14 = apply %13(%12) : $@convention(thin) (@guaranteed @noescape @callee_guaranteed () -> Int) -> Int
+  destroy_value %12 : $@noescape @callee_guaranteed () -> Int
+  destroy_value %11 : $@callee_guaranteed () -> Int
+  destroy_addr %2 : $*@moveOnly Int
+  dealloc_stack %1 : $*@moveOnly Int
+  return %14 : $Int
+}

test/SILOptimizer/noimplicitcopy_trivial.swift

@@ -1078,3 +1078,23 @@ public func noImplicitCopyReturnUse(_ x: Int) -> Int {
     let _ = z
     return y // expected-note {{consumed again here}}
 }
+
+func takeClosure(_ f: ()->Int) -> Int { f() }
+
+public func test1(i: consuming Int) -> Int {
+  takeClosure { [i = copy i] in i }
+}
+
+public func test2(i: borrowing Int) -> Int {
+  takeClosure { [i = copy i] in i }
+}
+
+public func test3(i: consuming Int) -> Int {
+  takeClosure { i }
+}
+
+// TODO: incorrect diagnostic:
+//       error: 'i' cannot be captured by an escaping closure since it is a borrowed parameter
+// public func test4(i: borrowing Int) -> Int {
+//  takeClosure { i }
+// }