[inlining] Fix two bugs around inlining of coroutines.

include/swift/SIL/SILBuilder.h

@@ -174,6 +174,9 @@ class SILBuilder {
     setInsertionPoint(I);
   }
 
+  SILBuilder(SILBasicBlock *BB, const SILDebugScope *DS, SILBuilder &B)
+      : SILBuilder(BB, DS, B.getBuilderContext()) {}
+
   /// Build instructions before the given insertion point, inheriting the debug
   /// location.
   ///
@@ -2264,6 +2267,11 @@ class SILBuilderWithScope : public SILBuilder {
     inheritScopeFrom(InheritScopeFrom);
   }
 
+  explicit SILBuilderWithScope(SILBasicBlock *BB, SILBuilder &B,
+                               SILInstruction *InheritScopeFrom)
+      : SILBuilder(BB, InheritScopeFrom->getDebugScope(),
+                   B.getBuilderContext()) {}
+
   /// Creates a new SILBuilder with an insertion point at the
   /// beginning of BB and the debug scope from the first
   /// non-metainstruction in the BB.

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -571,10 +571,11 @@ replaceBeginApplyInst(SILBuilder &builder, SILLocation loc,
   SILValue token = newBAI->getTokenResult();
 
   // The token will only be used by end_apply and abort_apply. Use that to
-  // insert the end_borrows we need.
+  // insert the end_borrows we need /after/ those uses.
   for (auto *use : token->getUses()) {
-    SILBuilderWithScope borrowBuilder(use->getUser(),
-                                      builder.getBuilderContext());
+    SILBuilderWithScope borrowBuilder(
+        &*std::next(use->getUser()->getIterator()),
+        builder.getBuilderContext());
     for (SILValue borrow : newArgBorrows) {
       borrowBuilder.createEndBorrow(loc, borrow);
     }

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -102,16 +102,20 @@ class BeginApplySite {
     return BeginApplySite(BeginApply, Loc, Builder);
   }
 
-  void preprocess(SILBasicBlock *returnToBB) {
+  void preprocess(SILBasicBlock *returnToBB,
+                  SmallVectorImpl<SILInstruction *> &endBorrowInsertPts) {
     SmallVector<EndApplyInst *, 1> endApplyInsts;
     SmallVector<AbortApplyInst *, 1> abortApplyInsts;
     BeginApply->getCoroutineEndPoints(endApplyInsts, abortApplyInsts);
     while (!endApplyInsts.empty()) {
       auto *endApply = endApplyInsts.pop_back_val();
       collectEndApply(endApply);
+      endBorrowInsertPts.push_back(&*std::next(endApply->getIterator()));
     }
     while (!abortApplyInsts.empty()) {
-      collectAbortApply(abortApplyInsts.pop_back_val());
+      auto *abortApply = abortApplyInsts.pop_back_val();
+      collectAbortApply(abortApply);
+      endBorrowInsertPts.push_back(&*std::next(abortApply->getIterator()));
     }
   }
 
@@ -428,29 +432,38 @@ SILInlineCloner::cloneInline(ArrayRef<SILValue> AppliedArgs) {
 
   SmallVector<SILValue, 4> entryArgs;
   entryArgs.reserve(AppliedArgs.size());
+  SmallBitVector borrowedArgs(AppliedArgs.size());
+
   auto calleeConv = getCalleeFunction()->getConventions();
-  for (unsigned argIdx = 0, endIdx = AppliedArgs.size(); argIdx < endIdx;
-       ++argIdx) {
-    SILValue callArg = AppliedArgs[argIdx];
+  for (auto p : llvm::enumerate(AppliedArgs)) {
+    SILValue callArg = p.value();
+    unsigned idx = p.index();
     // Insert begin/end borrow for guaranteed arguments.
-    if (argIdx >= calleeConv.getSILArgIndexOfFirstParam()
-        && calleeConv.getParamInfoForSILArg(argIdx).isGuaranteed()) {
-      callArg = borrowFunctionArgument(callArg, Apply);
+    if (idx >= calleeConv.getSILArgIndexOfFirstParam() &&
+        calleeConv.getParamInfoForSILArg(idx).isGuaranteed()) {
+      if (SILValue newValue = borrowFunctionArgument(callArg, Apply)) {
+        callArg = newValue;
+        borrowedArgs[idx] = true;
+      }
     }
     entryArgs.push_back(callArg);
   }
 
   // Create the return block and set ReturnToBB for use in visitTerminator
   // callbacks.
-  SILBasicBlock *callerBB = Apply.getParent();
+  SILBasicBlock *callerBlock = Apply.getParent();
+  SILBasicBlock *throwBlock = nullptr;
+  SmallVector<SILInstruction *, 1> endBorrowInsertPts;
+
   switch (Apply.getKind()) {
   case FullApplySiteKind::ApplyInst: {
     auto *AI = dyn_cast<ApplyInst>(Apply);
 
     // Split the BB and do NOT create a branch between the old and new
     // BBs; we will create the appropriate terminator manually later.
     ReturnToBB =
-      callerBB->split(std::next(Apply.getInstruction()->getIterator()));
+        callerBlock->split(std::next(Apply.getInstruction()->getIterator()));
+    endBorrowInsertPts.push_back(&*ReturnToBB->begin());
 
     // Create an argument on the return-to BB representing the returned value.
     auto *retArg =
@@ -459,22 +472,49 @@ SILInlineCloner::cloneInline(ArrayRef<SILValue> AppliedArgs) {
     AI->replaceAllUsesWith(retArg);
     break;
   }
-  case FullApplySiteKind::BeginApplyInst:
+  case FullApplySiteKind::BeginApplyInst: {
     ReturnToBB =
-      callerBB->split(std::next(Apply.getInstruction()->getIterator()));
-    BeginApply->preprocess(ReturnToBB);
+        callerBlock->split(std::next(Apply.getInstruction()->getIterator()));
+    // For begin_apply, we insert the end_borrow in the end_apply, abort_apply
+    // blocks to ensure that our borrowed values live over both the body and
+    // resume block of our coroutine.
+    BeginApply->preprocess(ReturnToBB, endBorrowInsertPts);
     break;
-
-  case FullApplySiteKind::TryApplyInst:
-    ReturnToBB = cast<TryApplyInst>(Apply)->getNormalBB();
+  }
+  case FullApplySiteKind::TryApplyInst: {
+    auto *tai = cast<TryApplyInst>(Apply);
+    ReturnToBB = tai->getNormalBB();
+    endBorrowInsertPts.push_back(&*ReturnToBB->begin());
+    throwBlock = tai->getErrorBB();
     break;
   }
+  }
+
+  // Then insert end_borrow in our end borrow block and in the throw
+  // block if we have one.
+  if (borrowedArgs.any()) {
+    for (unsigned i : indices(AppliedArgs)) {
+      if (!borrowedArgs.test(i)) {
+        continue;
+      }
+
+      for (auto *insertPt : endBorrowInsertPts) {
+        SILBuilderWithScope returnBuilder(insertPt, getBuilder());
+        returnBuilder.createEndBorrow(Apply.getLoc(), entryArgs[i]);
+      }
+
+      if (throwBlock) {
+        SILBuilderWithScope throwBuilder(throwBlock->begin(), getBuilder());
+        throwBuilder.createEndBorrow(Apply.getLoc(), entryArgs[i]);
+      }
+    }
+  }
 
   // Visit original BBs in depth-first preorder, starting with the
   // entry block, cloning all instructions and terminators.
   //
   // NextIter is initialized during `fixUp`.
-  cloneFunctionBody(getCalleeFunction(), callerBB, entryArgs);
+  cloneFunctionBody(getCalleeFunction(), callerBlock, entryArgs);
 
   // For non-throwing applies, the inlined body now unconditionally branches to
   // the returned-to-code, which was previously part of the call site's basic
@@ -560,25 +600,11 @@ SILValue SILInlineCloner::borrowFunctionArgument(SILValue callArg,
                                                  FullApplySite AI) {
   if (!AI.getFunction()->hasOwnership()
       || callArg.getOwnershipKind() != ValueOwnershipKind::Owned) {
-    return callArg;
+    return SILValue();
   }
 
   SILBuilderWithScope beginBuilder(AI.getInstruction(), getBuilder());
-  auto *borrow = beginBuilder.createBeginBorrow(AI.getLoc(), callArg);
-  if (auto *tryAI = dyn_cast<TryApplyInst>(AI)) {
-    SILBuilderWithScope returnBuilder(tryAI->getNormalBB()->begin(),
-                                      getBuilder());
-    returnBuilder.createEndBorrow(AI.getLoc(), borrow, callArg);
-
-    SILBuilderWithScope throwBuilder(tryAI->getErrorBB()->begin(),
-                                     getBuilder());
-    throwBuilder.createEndBorrow(AI.getLoc(), borrow, callArg);
-  } else {
-    SILBuilderWithScope returnBuilder(
-        std::next(AI.getInstruction()->getIterator()), getBuilder());
-    returnBuilder.createEndBorrow(AI.getLoc(), borrow, callArg);
-  }
-  return borrow;
+  return beginBuilder.createBeginBorrow(AI.getLoc(), callArg);
 }
 
 void SILInlineCloner::visitDebugValueInst(DebugValueInst *Inst) {

test/SILOptimizer/mandatory_inlining_ownership.sil

@@ -10,6 +10,11 @@ class C {
   init(i: Builtin.Int64)
 }
 
+private class C2 {
+  var i: Builtin.Int64 { get set }
+  init(i: Builtin.Int64)
+}
+
 class Klass {}
 
 sil [transparent] [ossa] @calleeWithGuaranteed : $@convention(thin) (@guaranteed C) -> Builtin.Int64 {
@@ -414,3 +419,168 @@ bb0(%0 : @owned $Builtin.NativeObject):
   %9999 = tuple()
   return %9999 : $()
 }
+
+///////////////////////
+// Begin Apply Tests //
+///////////////////////
+
+// Make sure that we do not violate any ownership invariants after inlining this
+// code.
+
+sil @get_hidden_int_field_of_klass : $@convention(method) (@guaranteed Klass) -> Builtin.Int32
+sil @int_klass_pair_user : $@convention(method) (Builtin.Int32, @guaranteed Klass) -> ()
+
+sil [transparent] [ossa] @begin_apply_callee : $@yield_once @convention(method) (@guaranteed Klass) -> @yields @inout Builtin.Int32 {
+bb0(%0 : @guaranteed $Klass):
+  %2 = alloc_stack $Builtin.Int32
+  %3 = function_ref @get_hidden_int_field_of_klass : $@convention(method) (@guaranteed Klass) -> Builtin.Int32
+  %4 = apply %3(%0) : $@convention(method) (@guaranteed Klass) -> Builtin.Int32
+  store %4 to [trivial] %2 : $*Builtin.Int32
+  yield %2 : $*Builtin.Int32, resume bb1, unwind bb2
+
+bb1:
+  %7 = load [trivial] %2 : $*Builtin.Int32
+  %8 = function_ref @int_klass_pair_user : $@convention(method) (Builtin.Int32, @guaranteed Klass) -> ()
+  %9 = apply %8(%7, %0) : $@convention(method) (Builtin.Int32, @guaranteed Klass) -> ()
+  dealloc_stack %2 : $*Builtin.Int32
+  %11 = tuple ()
+  return %11 : $()
+
+bb2:
+  %13 = load [trivial] %2 : $*Builtin.Int32
+  %14 = function_ref @int_klass_pair_user : $@convention(method) (Builtin.Int32, @guaranteed Klass) -> ()
+  %15 = apply %14(%13, %0) : $@convention(method) (Builtin.Int32, @guaranteed Klass) -> ()
+  dealloc_stack %2 : $*Builtin.Int32
+  unwind
+}
+
+// CHECK-LABEL: sil [ossa] @begin_apply_caller : $@convention(method) (@guaranteed Klass) -> @error Error {
+// CHECK-NOT: begin_apply
+// CHECK: } // end sil function 'begin_apply_caller'
+sil [ossa] @begin_apply_caller : $@convention(method) (@guaranteed Klass) -> @error Error {
+bb0(%0 : @guaranteed $Klass):
+  %6 = copy_value %0 : $Klass
+  %12 = function_ref @begin_apply_callee : $@yield_once @convention(method) (@guaranteed Klass) -> @yields @inout Builtin.Int32
+  (%13, %14) = begin_apply %12(%6) : $@yield_once @convention(method) (@guaranteed Klass) -> @yields @inout Builtin.Int32
+  end_apply %14
+  destroy_value %6 : $Klass
+  %19 = tuple ()
+  return %19 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @begin_apply_caller_2 : $@convention(method) (@guaranteed Klass) -> @error Error {
+// CHECK-NOT: begin_apply
+// CHECK: } // end sil function 'begin_apply_caller_2'
+sil [ossa] @begin_apply_caller_2 : $@convention(method) (@guaranteed Klass) -> @error Error {
+bb0(%0 : @guaranteed $Klass):
+  %6 = copy_value %0 : $Klass
+  %12 = function_ref @begin_apply_callee : $@yield_once @convention(method) (@guaranteed Klass) -> @yields @inout Builtin.Int32
+  (%13, %14) = begin_apply %12(%6) : $@yield_once @convention(method) (@guaranteed Klass) -> @yields @inout Builtin.Int32
+  abort_apply %14
+  destroy_value %6 : $Klass
+  %19 = tuple ()
+  return %19 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @begin_apply_caller_3 : $@convention(method) (@guaranteed Klass) -> @error Error {
+// CHECK-NOT: begin_apply
+// CHECK: } // end sil function 'begin_apply_caller_3'
+sil [ossa] @begin_apply_caller_3 : $@convention(method) (@guaranteed Klass) -> @error Error {
+bb0(%0 : @guaranteed $Klass):
+  %6 = copy_value %0 : $Klass
+  %12 = function_ref @begin_apply_callee : $@yield_once @convention(method) (@guaranteed Klass) -> @yields @inout Builtin.Int32
+  (%13, %14) = begin_apply %12(%6) : $@yield_once @convention(method) (@guaranteed Klass) -> @yields @inout Builtin.Int32
+  cond_br undef, bb1, bb2
+
+bb1:
+  end_apply %14
+  br bb3
+
+bb2:
+  abort_apply %14
+  br bb3
+
+bb3:
+  destroy_value %6 : $Klass
+  %19 = tuple ()
+  return %19 : $()
+}
+
+sil [ossa] [transparent] @devirt_callee : $@yield_once @convention(method) (@guaranteed C) -> @yields @inout Builtin.Int64 {
+bb0(%0 : @guaranteed $C):
+  %1 = alloc_stack $Builtin.Int64
+  %1a = integer_literal $Builtin.Int64, 0
+  store %1a to [trivial] %1 : $*Builtin.Int64
+  yield %1 : $*Builtin.Int64, resume bb1, unwind bb2
+
+bb1:
+  dealloc_stack %1 : $*Builtin.Int64
+  %6 = tuple ()
+  return %6 : $()
+
+bb2:
+  dealloc_stack %1 : $*Builtin.Int64
+  unwind
+}
+
+// Just make sure we actually inlined the begin_apply. We just want to make sure
+// we are not breaking ownership invariants by not properly borrowing %0.
+// CHECK-LABEL: sil [ossa] @begin_apply_devirt_caller : $@convention(method) (@owned C2) -> @error Error {
+// CHECK-NOT: begin_apply
+// CHECK: } // end sil function 'begin_apply_devirt_caller'
+sil [ossa] @begin_apply_devirt_caller : $@convention(method) (@owned C2) -> @error Error {
+bb0(%0 : @owned $C2):
+  %1 = class_method %0 : $C2, #C2.i!modify.1 : (C2) -> () -> (), $@yield_once @convention(method) (@guaranteed C2) -> @yields @inout Builtin.Int64
+  (%mem, %tok) = begin_apply %1(%0) : $@yield_once @convention(method) (@guaranteed C2) -> @yields @inout Builtin.Int64
+  br bb1
+
+bb1:
+  end_apply %tok
+  destroy_value %0 : $C2
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @begin_apply_devirt_caller_2 : $@convention(method) (@owned C2) -> @error Error {
+// CHECK-NOT: begin_apply
+// CHECK: } // end sil function 'begin_apply_devirt_caller_2'
+sil [ossa] @begin_apply_devirt_caller_2 : $@convention(method) (@owned C2) -> @error Error {
+bb0(%0 : @owned $C2):
+  %1 = class_method %0 : $C2, #C2.i!modify.1 : (C2) -> () -> (), $@yield_once @convention(method) (@guaranteed C2) -> @yields @inout Builtin.Int64
+  (%mem, %tok) = begin_apply %1(%0) : $@yield_once @convention(method) (@guaranteed C2) -> @yields @inout Builtin.Int64
+  br bb1
+
+bb1:
+  abort_apply %tok
+  destroy_value %0 : $C2
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @begin_apply_devirt_caller_3 : $@convention(method) (@owned C2) -> @error Error {
+// CHECK-NOT: begin_apply
+// CHECK: } // end sil function 'begin_apply_devirt_caller_3'
+sil [ossa] @begin_apply_devirt_caller_3 : $@convention(method) (@owned C2) -> @error Error {
+bb0(%0 : @owned $C2):
+  %1 = class_method %0 : $C2, #C2.i!modify.1 : (C2) -> () -> (), $@yield_once @convention(method) (@guaranteed C2) -> @yields @inout Builtin.Int64
+  (%mem, %tok) = begin_apply %1(%0) : $@yield_once @convention(method) (@guaranteed C2) -> @yields @inout Builtin.Int64
+  cond_br undef, bb1, bb2
+
+bb1:
+  end_apply %tok
+  br bb3
+
+bb2:
+  abort_apply %tok
+  br bb3
+
+bb3:
+  destroy_value %0 : $C2
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+
+sil_vtable C2 {
+  #C2.i!modify.1: (C2) -> () -> () : @devirt_callee
+}