Fix compiler crashes with consuming and borrowing keywords.

Without this fix, the new 'consuming' and 'borrowing' keywords cannot
be used with trivial types. Which means, for example, they can't be
used in macro expansions that work on various types.

Fixes patterns like:

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

// Sadly, test4 is still incorrectly diagnosed.
public func test4(i: borrowing Int) -> Int {
  takeClosure { i }
}

Fixes rdar://112795074 (Crash compiling function that has a macro annotation and uses `consuming`)

(cherry picked from commit 9c24933)

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 }
+// }