Merge pull request #80701 from nate-chandler/rdar148941214

[CoroutineAccessors] Use yield_once_2 on Darwin and Linux.

Author: nate-chandler

include/swift/AST/SILOptions.h

@@ -338,6 +338,10 @@ class SILOptions {
   // Whether to allow merging traps and cond_fails.
   bool MergeableTraps = false;
 
+  /// Whether the @yield_once_2 convention is used by accessors added with the
+  /// CoroutineAccessors feature (i.e. read2/modify2).
+  bool CoroutineAccessorsUseYieldOnce2 = false;
+
   SILOptions() {}
 
   /// Return a hash code of any components from these options that should

include/swift/Option/FrontendOptions.td

@@ -1175,6 +1175,12 @@ def enable_arm64_corocc : Flag<["-"], "enable-arm64-corocc">,
 def disable_arm64_corocc : Flag<["-"], "disable-arm64-corocc">,
   HelpText<"Don't use swiftcorocc for yield_once_2 routines on arm64 variants.">;
 
+def enable_callee_allocated_coro_abi : Flag<["-"], "enable-callee-allocated-coro-abi">,
+  HelpText<"Override per-platform settings and use yield_once_2.">;
+
+def disable_callee_allocated_coro_abi : Flag<["-"], "disable-callee-allocated-coro-abi">,
+  HelpText<"Override per-platform settings and don't use yield_once_2.">;
+
 def enable_cond_fail_message_annotation : Flag<["-"], "enable-cond-fail-message-annotation">,
   HelpText<"Enable cond_fail message annotation. Will serialize a .o.yaml file per .o file.">;
 

lib/DriverTool/sil_opt_main.cpp

@@ -597,6 +597,16 @@ struct SILOptOptions {
       "enable-address-dependencies",
       llvm::cl::desc("Enable enforcement of lifetime dependencies on addressable values."));
 
+  llvm::cl::opt<bool> EnableCalleeAllocatedCoroAbi = llvm::cl::opt<bool>(
+      "enable-callee-allocated-coro-abi",
+      llvm::cl::desc("Override per-platform settings and use yield_once_2."),
+      llvm::cl::init(false));
+  llvm::cl::opt<bool> DisableCalleeAllocatedCoroAbi = llvm::cl::opt<bool>(
+      "disable-callee-allocated-coro-abi",
+      llvm::cl::desc(
+          "Override per-platform settings and don't use yield_once_2."),
+      llvm::cl::init(false));
+
   llvm::cl::opt<bool> MergeableTraps = llvm::cl::opt<bool>(
       "mergeable-traps",
       llvm::cl::desc("Enable cond_fail merging."));
@@ -918,6 +928,10 @@ int sil_opt_main(ArrayRef<const char *> argv, void *MainAddr) {
       options.EnablePackMetadataStackPromotion;
 
   SILOpts.EnableAddressDependencies = options.EnableAddressDependencies;
+  if (options.EnableCalleeAllocatedCoroAbi)
+    SILOpts.CoroutineAccessorsUseYieldOnce2 = true;
+  if (options.DisableCalleeAllocatedCoroAbi)
+    SILOpts.CoroutineAccessorsUseYieldOnce2 = false;
   SILOpts.MergeableTraps = options.MergeableTraps;
 
   if (options.OptModeFlag == OptimizationMode::NotSet) {

lib/Frontend/CompilerInvocation.cpp

@@ -3130,6 +3130,16 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
   Opts.EnableAddressDependencies = Args.hasFlag(
       OPT_enable_address_dependencies, OPT_disable_address_dependencies,
       Opts.EnableAddressDependencies);
+
+  if (LangOpts.Target.isOSDarwin() || LangOpts.Target.isOSLinux()) {
+    // On Darwin and Linux, use yield_once_2 by default.
+    Opts.CoroutineAccessorsUseYieldOnce2 = true;
+  }
+  Opts.CoroutineAccessorsUseYieldOnce2 =
+      Args.hasFlag(OPT_enable_callee_allocated_coro_abi,
+                   OPT_disable_callee_allocated_coro_abi,
+                   Opts.CoroutineAccessorsUseYieldOnce2);
+
   Opts.MergeableTraps = Args.hasArg(OPT_mergeable_traps);
 
   return false;

lib/SIL/IR/SILDeclRef.cpp

@@ -1855,5 +1855,8 @@ bool SILDeclRef::isCalleeAllocatedCoroutine() const {
   if (!accessor)
     return false;
 
-  return requiresFeatureCoroutineAccessors(accessor->getAccessorKind());
+  if (!requiresFeatureCoroutineAccessors(accessor->getAccessorKind()))
+    return false;
+
+  return getASTContext().SILOpts.CoroutineAccessorsUseYieldOnce2;
 }

lib/SIL/IR/SILFunctionType.cpp

@@ -2481,8 +2481,10 @@ static CanSILFunctionType getSILFunctionType(
 
   if (auto accessor = getAsCoroutineAccessor(constant)) {
     auto origAccessor = cast<AccessorDecl>(origConstant->getDecl());
+    auto &ctx = origAccessor->getASTContext();
     coroutineKind =
-        requiresFeatureCoroutineAccessors(accessor->getAccessorKind())
+        (requiresFeatureCoroutineAccessors(accessor->getAccessorKind()) &&
+         ctx.SILOpts.CoroutineAccessorsUseYieldOnce2)
             ? SILCoroutineKind::YieldOnce2
             : SILCoroutineKind::YieldOnce;
 

lib/SILGen/SILGenApply.cpp

@@ -4999,14 +4999,16 @@ class CallEmission {
   FormalEvaluationScope initialWritebackScope;
   std::optional<ActorIsolation> implicitActorHopTarget;
   bool implicitlyThrows;
+  bool canUnwind;
 
 public:
   /// Create an emission for a call of the given callee.
   CallEmission(SILGenFunction &SGF, Callee &&callee,
                FormalEvaluationScope &&writebackScope)
       : SGF(SGF), callee(std::move(callee)),
         initialWritebackScope(std::move(writebackScope)),
-        implicitActorHopTarget(std::nullopt), implicitlyThrows(false) {}
+        implicitActorHopTarget(std::nullopt), implicitlyThrows(false),
+        canUnwind(false) {}
 
   /// A factory method for decomposing the apply expr \p e into a call
   /// emission.
@@ -5057,6 +5059,8 @@ class CallEmission {
   /// function can throw or not.
   void setImplicitlyThrows(bool flag) { implicitlyThrows = flag; }
 
+  void setCanUnwind(bool flag) { canUnwind = flag; }
+
   CleanupHandle applyCoroutine(SmallVectorImpl<ManagedValue> &yields);
 
   RValue apply(SGFContext C = SGFContext()) {
@@ -5114,9 +5118,11 @@ namespace {
 /// Cleanup to end a coroutine application.
 class EndCoroutineApply : public Cleanup {
   SILValue ApplyToken;
+  bool CanUnwind;
   std::vector<BeginBorrowInst *> BorrowedMoveOnlyValues;
 public:
-  EndCoroutineApply(SILValue applyToken) : ApplyToken(applyToken) {}
+  EndCoroutineApply(SILValue applyToken, bool CanUnwind)
+      : ApplyToken(applyToken), CanUnwind(CanUnwind) {}
 
   void setBorrowedMoveOnlyValues(ArrayRef<BeginBorrowInst *> values) {
     BorrowedMoveOnlyValues.insert(BorrowedMoveOnlyValues.end(),
@@ -5129,14 +5135,7 @@ class EndCoroutineApply : public Cleanup {
       SGF.B.createEndBorrow(l, *i);
       SGF.B.createDestroyValue(l, (*i)->getOperand());
     }
-    auto *beginApply =
-        cast<BeginApplyInst>(ApplyToken->getDefiningInstruction());
-    auto isCalleeAllocated = beginApply->isCalleeAllocated();
-    auto unwindOnCallerError =
-        !isCalleeAllocated ||
-        SGF.SGM.getASTContext().LangOpts.hasFeature(
-            Feature::CoroutineAccessorsUnwindOnCallerError);
-    if (forUnwind && unwindOnCallerError) {
+    if (forUnwind && CanUnwind) {
       SGF.B.createAbortApply(l, ApplyToken);
     } else {
       SGF.B.createEndApply(l, ApplyToken,
@@ -5179,18 +5178,15 @@ CallEmission::applyCoroutine(SmallVectorImpl<ManagedValue> &yields) {
 
   auto fnValue = callee.getFnValue(SGF, borrowedSelf);
 
-  return SGF.emitBeginApply(uncurriedLoc.value(), fnValue,
+  return SGF.emitBeginApply(uncurriedLoc.value(), fnValue, canUnwind,
                             callee.getSubstitutions(), uncurriedArgs,
                             calleeTypeInfo.substFnType, options, yields);
 }
 
-CleanupHandle
-SILGenFunction::emitBeginApply(SILLocation loc, ManagedValue fn,
-                               SubstitutionMap subs,
-                               ArrayRef<ManagedValue> args,
-                               CanSILFunctionType substFnType,
-                               ApplyOptions options,
-                               SmallVectorImpl<ManagedValue> &yields) {
+CleanupHandle SILGenFunction::emitBeginApply(
+    SILLocation loc, ManagedValue fn, bool canUnwind, SubstitutionMap subs,
+    ArrayRef<ManagedValue> args, CanSILFunctionType substFnType,
+    ApplyOptions options, SmallVectorImpl<ManagedValue> &yields) {
   // Emit the call.
   SmallVector<SILValue, 4> rawResults;
   emitRawApply(*this, loc, fn, subs, args, substFnType, options,
@@ -5206,7 +5202,7 @@ SILGenFunction::emitBeginApply(SILLocation loc, ManagedValue fn,
 
   // Push a cleanup to end the application.
   // TODO: destroy all the arguments at exactly this point?
-  Cleanups.pushCleanup<EndCoroutineApply>(token);
+  Cleanups.pushCleanup<EndCoroutineApply>(token, canUnwind);
   auto endApplyHandle = getTopCleanup();
 
   // Manage all the yielded values.
@@ -6183,7 +6179,7 @@ SILValue SILGenFunction::emitApplyWithRethrow(SILLocation loc, SILValue fn,
 std::tuple<MultipleValueInstructionResult *, CleanupHandle, SILValue,
            CleanupHandle>
 SILGenFunction::emitBeginApplyWithRethrow(SILLocation loc, SILValue fn,
-                                          SILType substFnType,
+                                          SILType substFnType, bool canUnwind,
                                           SubstitutionMap subs,
                                           ArrayRef<SILValue> args,
                                           SmallVectorImpl<SILValue> &yields) {
@@ -6208,7 +6204,7 @@ SILGenFunction::emitBeginApplyWithRethrow(SILLocation loc, SILValue fn,
     deallocCleanup = enterDeallocStackCleanup(allocation);
   }
 
-  Cleanups.pushCleanup<EndCoroutineApply>(token);
+  Cleanups.pushCleanup<EndCoroutineApply>(token, canUnwind);
   auto abortCleanup = Cleanups.getTopCleanup();
 
   return {token, abortCleanup, allocation, deallocCleanup};
@@ -7561,6 +7557,21 @@ ManagedValue SILGenFunction::emitAddressorAccessor(
   return ManagedValue::forLValue(address);
 }
 
+bool SILGenFunction::canUnwindAccessorDeclRef(SILDeclRef accessorRef) {
+  auto *accessor =
+      dyn_cast_or_null<AccessorDecl>(accessorRef.getAbstractFunctionDecl());
+  ASSERT(accessor && "only accessors can unwind");
+  auto kind = accessor->getAccessorKind();
+  ASSERT(isYieldingAccessor(kind) && "only yielding accessors can unwind");
+  if (!requiresFeatureCoroutineAccessors(kind)) {
+    // _read and _modify can unwind
+    return true;
+  }
+  // Coroutine accessors can only unwind with the experimental feature.
+  return getASTContext().LangOpts.hasFeature(
+      Feature::CoroutineAccessorsUnwindOnCallerError);
+}
+
 CleanupHandle
 SILGenFunction::emitCoroutineAccessor(SILLocation loc, SILDeclRef accessor,
                                       SubstitutionMap substitutions,
@@ -7596,6 +7607,8 @@ SILGenFunction::emitCoroutineAccessor(SILLocation loc, SILDeclRef accessor,
 
   emission.addCallSite(loc, std::move(subscriptIndices));
 
+  emission.setCanUnwind(canUnwindAccessorDeclRef(accessor));
+
   auto endApplyHandle = emission.applyCoroutine(yields);
 
   return endApplyHandle;

lib/SILGen/SILGenFunction.h

@@ -2034,6 +2034,7 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
       PreparedArguments &&optionalSubscripts,
       SILType addressType, bool isOnSelfParameter);
 
+  bool canUnwindAccessorDeclRef(SILDeclRef accessorRef);
   CleanupHandle emitCoroutineAccessor(SILLocation loc, SILDeclRef accessor,
                                       SubstitutionMap substitutions,
                                       ArgumentSource &&optionalSelfValue,
@@ -2301,8 +2302,9 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                                         PreparedArguments &&args, Type overriddenSelfType,
                                         SGFContext ctx);
 
-  CleanupHandle emitBeginApply(SILLocation loc, ManagedValue fn,
-                               SubstitutionMap subs, ArrayRef<ManagedValue> args,
+  CleanupHandle emitBeginApply(SILLocation loc, ManagedValue fn, bool canUnwind,
+                               SubstitutionMap subs,
+                               ArrayRef<ManagedValue> args,
                                CanSILFunctionType substFnType,
                                ApplyOptions options,
                                SmallVectorImpl<ManagedValue> &yields);
@@ -2315,7 +2317,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   std::tuple<MultipleValueInstructionResult *, CleanupHandle, SILValue,
              CleanupHandle>
   emitBeginApplyWithRethrow(SILLocation loc, SILValue fn, SILType substFnType,
-                            SubstitutionMap subs, ArrayRef<SILValue> args,
+                            bool canUnwind, SubstitutionMap subs,
+                            ArrayRef<SILValue> args,
                             SmallVectorImpl<SILValue> &yields);
   void emitEndApplyWithRethrow(SILLocation loc,
                                MultipleValueInstructionResult *token,

lib/SILGen/SILGenLValue.cpp

@@ -2535,9 +2535,9 @@ namespace {
 
       // Perform the begin_apply.
       SmallVector<ManagedValue, 1> yields;
-      auto cleanup =
-        SGF.emitBeginApply(loc, projectFnRef, subs, { base, keyPathValue },
-                           substFnType, ApplyOptions(), yields);
+      auto cleanup = SGF.emitBeginApply(loc, projectFnRef, /*canUnwind=*/true,
+                                        subs, {base, keyPathValue}, substFnType,
+                                        ApplyOptions(), yields);
 
       // Push an operation to do the end_apply.
       pushEndApplyWriteback(SGF, loc, cleanup, getTypeData());

lib/SILGen/SILGenPoly.cpp

@@ -7045,8 +7045,8 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
   case SILCoroutineKind::YieldOnce2: {
     SmallVector<SILValue, 4> derivedYields;
     auto tokenAndCleanups = emitBeginApplyWithRethrow(
-        loc, derivedRef, SILType::getPrimitiveObjectType(derivedFTy), subs,
-        args, derivedYields);
+        loc, derivedRef, SILType::getPrimitiveObjectType(derivedFTy),
+        canUnwindAccessorDeclRef(base), subs, args, derivedYields);
     auto token = std::get<0>(tokenAndCleanups);
     auto abortCleanup = std::get<1>(tokenAndCleanups);
     auto allocation = std::get<2>(tokenAndCleanups);
@@ -7513,7 +7513,8 @@ void SILGenFunction::emitProtocolWitness(
   case SILCoroutineKind::YieldOnce2: {
     SmallVector<SILValue, 4> witnessYields;
     auto tokenAndCleanups = emitBeginApplyWithRethrow(
-        loc, witnessFnRef, witnessSILTy, witnessSubs, args, witnessYields);
+        loc, witnessFnRef, witnessSILTy, canUnwindAccessorDeclRef(requirement),
+        witnessSubs, args, witnessYields);
     auto token = std::get<0>(tokenAndCleanups);
     auto abortCleanup = std::get<1>(tokenAndCleanups);
     auto allocation = std::get<2>(tokenAndCleanups);