AddressLowering: Handle coroutines (#58821)

Author: meg-gupta

include/swift/AST/Types.h

@@ -4460,6 +4460,32 @@ class SILFunctionType final
   SILType getDirectFormalResultsType(SILModule &M,
                                      TypeExpansionContext expansion);
 
+  unsigned getNumIndirectFormalYields() const {
+    return isCoroutine() ? NumAnyIndirectFormalResults : 0;
+  }
+  /// Does this function have any formally indirect yields?
+  bool hasIndirectFormalYields() const {
+    return getNumIndirectFormalYields() != 0;
+  }
+  unsigned getNumDirectFormalYields() const {
+    return isCoroutine() ? NumAnyResults - NumAnyIndirectFormalResults : 0;
+  }
+
+  struct IndirectFormalYieldFilter {
+    bool operator()(SILYieldInfo yield) const {
+      return !yield.isFormalIndirect();
+    }
+  };
+
+  using IndirectFormalYieldIter =
+      llvm::filter_iterator<const SILYieldInfo *, IndirectFormalYieldFilter>;
+  using IndirectFormalYieldRange = iterator_range<IndirectFormalYieldIter>;
+
+  /// A range of SILYieldInfo for all formally Indirect Yields.
+  IndirectFormalYieldRange getIndirectFormalYields() const {
+    return llvm::make_filter_range(getYields(), IndirectFormalYieldFilter());
+  }
+
   /// Get a single non-address SILType for all SIL results regardless of whether
   /// they are formally indirect. The actual SIL result type of an apply
   /// instruction that calls this function depends on the current SIL stage and

include/swift/SIL/SILValue.h

@@ -255,6 +255,9 @@ struct ValueOwnershipKind {
   explicit ValueOwnershipKind(unsigned newValue) : value(innerty(newValue)) {}
   ValueOwnershipKind(const SILFunction &f, SILType type,
                      SILArgumentConvention convention);
+  ValueOwnershipKind(const SILFunction &f, SILType type,
+                     SILArgumentConvention convention,
+                     SILModuleConventions moduleConventions);
 
   /// Parse Value into a ValueOwnershipKind.
   ///

lib/AST/ASTContext.cpp

@@ -3949,7 +3949,10 @@ SILFunctionType::SILFunctionType(
   } else {
     assert(normalResults.empty());    
     NumAnyResults = yields.size();
-    NumAnyIndirectFormalResults = 0; // unused
+    NumAnyIndirectFormalResults = std::count_if(
+        yields.begin(), yields.end(), [](const SILYieldInfo &yieldInfo) {
+          return yieldInfo.isFormalIndirect();
+        });
     memcpy(getMutableYields().data(), yields.data(),
            yields.size() * sizeof(SILYieldInfo));
   }

lib/SIL/IR/SILInstructions.cpp

@@ -536,10 +536,12 @@ BeginApplyInst::create(SILDebugLocation loc, SILValue callee,
 
   for (auto &yield : substCalleeType->getYields()) {
     auto yieldType = conv.getSILType(yield, F.getTypeExpansionContext());
-    auto convention = SILArgumentConvention(yield.getConvention());
+    auto argConvention = SILArgumentConvention(yield.getConvention());
     resultTypes.push_back(yieldType);
-    resultOwnerships.push_back(
-      ValueOwnershipKind(F, yieldType, convention));
+    resultOwnerships.push_back(ValueOwnershipKind(
+        F, yieldType, argConvention,
+        moduleConventions.hasValue() ? moduleConventions.getValue()
+                                     : SILModuleConventions(F.getModule())));
   }
 
   resultTypes.push_back(SILType::getSILTokenType(F.getASTContext()));

lib/SIL/IR/SILValue.cpp

@@ -186,9 +186,13 @@ StringRef OwnershipKind::asString() const {
 
 ValueOwnershipKind::ValueOwnershipKind(const SILFunction &F, SILType Type,
                                        SILArgumentConvention Convention)
-    : value(OwnershipKind::Any) {
-  auto &M = F.getModule();
+    : ValueOwnershipKind(F, Type, Convention,
+                         SILModuleConventions(F.getModule())) {}
 
+ValueOwnershipKind::ValueOwnershipKind(const SILFunction &F, SILType Type,
+                                       SILArgumentConvention Convention,
+                                       SILModuleConventions moduleConventions)
+    : value(OwnershipKind::Any) {
   // Trivial types can be passed using a variety of conventions. They always
   // have trivial ownership.
   if (Type.isTrivial(F)) {
@@ -199,14 +203,12 @@ ValueOwnershipKind::ValueOwnershipKind(const SILFunction &F, SILType Type,
   switch (Convention) {
   case SILArgumentConvention::Indirect_In:
   case SILArgumentConvention::Indirect_In_Constant:
-    value = SILModuleConventions(M).useLoweredAddresses()
-                ? OwnershipKind::None
-                : OwnershipKind::Owned;
+    value = moduleConventions.useLoweredAddresses() ? OwnershipKind::None
+                                                    : OwnershipKind::Owned;
     break;
   case SILArgumentConvention::Indirect_In_Guaranteed:
-    value = SILModuleConventions(M).useLoweredAddresses()
-                ? OwnershipKind::None
-                : OwnershipKind::Guaranteed;
+    value = moduleConventions.useLoweredAddresses() ? OwnershipKind::None
+                                                    : OwnershipKind::Guaranteed;
     break;
   case SILArgumentConvention::Indirect_Inout:
   case SILArgumentConvention::Indirect_InoutAliasable:

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -645,8 +645,10 @@ void OpaqueValueVisitor::checkForIndirectApply(FullApplySite applySite) {
     }
     ++calleeArgIdx;
   }
-  if (applySite.getSubstCalleeType()->hasIndirectFormalResults())
+
+  if (applySite.getSubstCalleeType()->hasIndirectFormalResults()) {
     pass.indirectApplies.insert(applySite);
+  }
 }
 
 /// If `value` is address-only, add it to the `valueStorageMap`.
@@ -883,6 +885,16 @@ static SILValue getProjectedUseValue(Operand *operand) {
   return SILValue();
 }
 
+static bool doesNotNeedStackAllocation(SILValue value) {
+  auto *defInst = value->getDefiningInstruction();
+  if (!defInst)
+    return false;
+
+  if (isa<LoadBorrowInst>(defInst) || isa<BeginApplyInst>(defInst))
+    return true;
+
+  return false;
+}
 //===----------------------------------------------------------------------===//
 //                          OpaqueStorageAllocation
 //
@@ -1042,11 +1054,15 @@ void OpaqueStorageAllocation::allocateValue(SILValue value) {
   if (getReusedStorageOperand(value))
     return;
 
+  if (doesNotNeedStackAllocation(value))
+    return;
+
   // Check for values that inherently project storage from their operand.
   if (auto *storageOper = getProjectedDefOperand(value)) {
     pass.valueStorageMap.recordDefProjection(storageOper, value);
     return;
   }
+
   if (value->getOwnershipKind() == OwnershipKind::Guaranteed) {
     value->dump();
     llvm::report_fatal_error("^^^ guaranteed values must reuse storage");
@@ -1183,7 +1199,6 @@ void OpaqueStorageAllocation::removeAllocation(SILValue value) {
 AllocStackInst *OpaqueStorageAllocation::createStackAllocation(SILValue value) {
   assert(value.getOwnershipKind() != OwnershipKind::Guaranteed &&
          "creating storage for a guaranteed value implies a copy");
-
   // Instructions that produce an opened type never reach here because they
   // have guaranteed ownership--they project their storage. We reach this
   // point after the opened value has been copied.
@@ -1822,10 +1837,11 @@ class ApplyRewriter {
         loweredCalleeConv(getLoweredCallConv(oldCall)) {}
 
   void convertApplyWithIndirectResults();
+  void convertBeginApplyWithOpaqueYield();
 
 protected:
   SILBasicBlock::iterator getCallResultInsertionPoint() {
-    if (isa<ApplyInst>(apply))
+    if (isa<ApplyInst>(apply) || isa<BeginApplyInst>(apply))
       return std::next(SILBasicBlock::iterator(apply.getInstruction()));
 
     auto *bb = cast<TryApplyInst>(apply)->getNormalBB();
@@ -2080,6 +2096,35 @@ void ApplyRewriter::rewriteApply(ArrayRef<SILValue> newCallArgs) {
   // will be deleted with its destructure_tuple.
 }
 
+void ApplyRewriter::convertBeginApplyWithOpaqueYield() {
+  auto *origCall = cast<BeginApplyInst>(apply.getInstruction());
+  SmallVector<SILValue, 4> opValues;
+
+  for (auto &oper : origCall->getArgumentOperands()) {
+    opValues.push_back(oper.get());
+  }
+
+  // Recreate the begin_apply so that the instruction results have the right
+  // ownership kind as per the lowered addresses convention.
+  auto *newCall = argBuilder.createBeginApply(
+      callLoc, apply.getCallee(), apply.getSubstitutionMap(), opValues,
+      origCall->getApplyOptions(), origCall->getSpecializationInfo());
+  this->apply = FullApplySite(newCall);
+
+  // Replace uses of orig begin_apply with the new begin_apply
+  auto oldResults = origCall->getAllResultsBuffer();
+  auto newResults = newCall->getAllResultsBuffer();
+  assert(oldResults.size() == newResults.size());
+  for (auto i : indices(oldResults)) {
+    if (oldResults[i].getType().isAddressOnly(*pass.function)) {
+      pass.valueStorageMap.setStorageAddress(&oldResults[i], &newResults[i]);
+      pass.valueStorageMap.getStorage(&oldResults[i]).markRewritten();
+    } else {
+      oldResults[i].replaceAllUsesWith(&newResults[i]);
+    }
+  }
+}
+
 // Replace \p tryApply with a new try_apply using \p newCallArgs.
 //
 // If the old result was a single opaque value, then create and return a
@@ -2545,6 +2590,16 @@ class UseRewriter : SILInstructionVisitor<UseRewriter> {
     CallArgRewriter(applyInst, pass).rewriteIndirectArgument(use);
   }
 
+  void visitBeginApplyInst(BeginApplyInst *bai) {
+    CallArgRewriter(bai, pass).rewriteIndirectArgument(use);
+  }
+
+  void visitYieldInst(YieldInst *yield) {
+    SILValue addr =
+        pass.valueStorageMap.getStorage(yield->getOperand(0)).storageAddress;
+    yield->setOperand(0, addr);
+  }
+
   void visitAssignInst(AssignInst *assignInst);
 
   void visitBeginBorrowInst(BeginBorrowInst *borrow);
@@ -3039,6 +3094,11 @@ class DefRewriter : SILInstructionVisitor<DefRewriter> {
     ApplyRewriter(applyInst, pass).convertApplyWithIndirectResults();
   }
 
+  void visitBeginApplyInst(BeginApplyInst *bai) {
+    CallArgRewriter(bai, pass).rewriteArguments();
+    ApplyRewriter(bai, pass).convertBeginApplyWithOpaqueYield();
+  }
+
   // Rewrite the apply for an indirect result.
   void visitDestructureTupleInst(DestructureTupleInst *destructure) {
     SILValue srcVal = destructure->getOperand();
@@ -3107,6 +3167,10 @@ class DefRewriter : SILInstructionVisitor<DefRewriter> {
     }
   }
 
+  void visitLoadBorrowInst(LoadBorrowInst *lbi) {
+    pass.valueStorageMap.setStorageAddress(lbi, lbi->getOperand());
+  }
+
   // Define an opaque struct.
   void visitStructInst(StructInst *structInst) {
     // For each element, initialize the operand's memory. Some struct elements
@@ -3159,8 +3223,9 @@ static void rewriteIndirectApply(FullApplySite apply,
   // If all indirect args were loadable, then they still need to be rewritten.
   CallArgRewriter(apply, pass).rewriteArguments();
 
-  if (!apply.getSubstCalleeType()->hasIndirectFormalResults())
+  if (!apply.getSubstCalleeType()->hasIndirectFormalResults()) {
     return;
+  }
 
   // If the call has indirect results and wasn't already rewritten, rewrite it
   // now. This handles try_apply, which is not rewritten when DefRewriter visits

test/SILOptimizer/address_lowering.sil

@@ -48,6 +48,11 @@ enum Mixed<T> {
   case o(AnyObject)
 };
 
+class TestGeneric<T> {
+  init()
+  @_borrowed @_hasStorage var borrowedGeneric: T { get set }
+}
+
 precedencegroup ComparisonPrecedence {
   assignment: true
   associativity: right
@@ -1171,6 +1176,87 @@ bb2(%9 : $Error):
   throw %9 : $Error
 }
 
+sil hidden [ossa] @testBeginApplyDeadYield : $@convention(thin) <T> (@guaranteed TestGeneric<T>) -> () {
+bb0(%0 : @guaranteed $TestGeneric<T>):
+  %2 = class_method %0 : $TestGeneric<T>, #TestGeneric.borrowedGeneric!read : <T> (TestGeneric<T>) -> () -> (), $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+  (%3, %4) = begin_apply %2<T>(%0) : $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+  end_apply %4
+  %10 = tuple ()
+  return %10 : $()
+}
+
+class Klass {}
+
+// CHECK-LABEL: sil hidden [ossa] @testBeginApplyLoadableYieldWithIndirectConv :
+// CHECK: bb0(%0 : @guaranteed $TestGeneric<Klass>):
+// CHECK:   [[STK:%.*]] = alloc_stack $Klass
+// CHECK:   [[METH:%.*]] = class_method %0 : $TestGeneric<Klass>, #TestGeneric.borrowedGeneric!read : <T> (TestGeneric<T>) -> () -> (), $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+// CHECK:   ([[Y:%.*]], [[TOK:%.*]]) = begin_apply [[METH]]<Klass>(%0) : $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+// CHECK:   copy_addr [[Y]] to [initialization] [[STK]] : $*Klass
+// CHECK:   end_apply [[TOK]]
+// CHECK:   destroy_addr [[STK]] : $*Klass
+// CHECK-LABEL: }
+sil hidden [ossa] @testBeginApplyLoadableYieldWithIndirectConv : $@convention(thin) <Klass> (@guaranteed TestGeneric<Klass>) -> () {
+bb0(%0 : @guaranteed $TestGeneric<Klass>):
+  %2 = class_method %0 : $TestGeneric<Klass>, #TestGeneric.borrowedGeneric!read : <Klass> (TestGeneric<Klass>) -> () -> (), $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+  (%3, %4) = begin_apply %2<Klass>(%0) : $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+  %5 = copy_value %3 : $Klass
+  end_apply %4
+  destroy_value %5 : $Klass
+  %10 = tuple ()
+  return %10 : $()
+}
+
+// CHECK-LABEL: sil hidden [ossa] @testBeginApplyOpaqueYield :
+// CHECK: bb0(%0 : @guaranteed $TestGeneric<T>):
+// CHECK:   [[STK:%.*]] = alloc_stack $T
+// CHECK:   [[M:%.*]] = class_method %0 : $TestGeneric<T>, #TestGeneric.borrowedGeneric!read : <T> (TestGeneric<T>) -> () -> (), $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+// CHECK:   ([[Y:%.*]], [[TOK:%.*]]) = begin_apply %2<T>(%0) : $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+// CHECK:   copy_addr [[Y]] to [initialization] [[STK]] : $*T
+// CHECK:   end_apply [[TOK]]
+// CHECK:   destroy_addr [[STK]] : $*T
+// CHECK-LABEL: } // end sil function 'testBeginApplyOpaqueYield'
+sil hidden [ossa] @testBeginApplyOpaqueYield : $@convention(thin) <T> (@guaranteed TestGeneric<T>) -> () {
+bb0(%0 : @guaranteed $TestGeneric<T>):
+  %2 = class_method %0 : $TestGeneric<T>, #TestGeneric.borrowedGeneric!read : <T> (TestGeneric<T>) -> () -> (), $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+  (%3, %4) = begin_apply %2<T>(%0) : $@yield_once @convention(method) <τ_0_0> (@guaranteed TestGeneric<τ_0_0>) -> @yields @in_guaranteed τ_0_0
+  %5 = copy_value %3 : $T
+  end_apply %4
+  destroy_value %5 : $T
+  %10 = tuple ()
+  return %10 : $()
+}
+
+// CHECK-LABEL: sil hidden [ossa] @testOpaqueYield :
+// CHECK: bb0(%0 : @guaranteed $TestGeneric<T>):
+// CHECK:   [[REF:%.*]] = ref_element_addr %0 : $TestGeneric<T>, #TestGeneric.borrowedGeneric
+// CHECK:   [[ACCESS:%.*]] = begin_access [read] [dynamic] [[REF]] : $*T
+// CHECK:   yield [[ACCESS]] : $*T, resume bb2, unwind bb1
+// CHECK: bb1:
+// CHECK:   end_access [[ACCESS]] : $*T
+// CHECK:   unwind
+// CHECK: bb2:
+// CHECK:   end_access [[ACCESS]] : $*T
+// CHECK-LABEL: } // end sil function 'testOpaqueYield'
+sil hidden [ossa] @testOpaqueYield : $@yield_once @convention(method) <T> (@guaranteed TestGeneric<T>) -> @yields @in_guaranteed T {
+bb0(%0 : @guaranteed $TestGeneric<T>):
+  %2 = ref_element_addr %0 : $TestGeneric<T>, #TestGeneric.borrowedGeneric
+  %3 = begin_access [read] [dynamic] %2 : $*T
+  %4 = load_borrow %3 : $*T
+  yield %4 : $T, resume bb1, unwind bb2
+
+bb1:
+  end_borrow %4 : $T
+  end_access %3 : $*T
+  %8 = tuple ()
+  return %8 : $()
+
+bb2:
+  end_borrow %4 : $T
+  end_access %3 : $*T
+  unwind
+}
+
 sil [ossa] @use_T : $@convention(thin) <T> (@in T) -> ()
 
 // CHECK-LABEL: sil [ossa] @test_checked_cast_br1 : $@convention(method) <T> (@in Any) -> () {