swiftlang · DougGregor · May 4, 2022 · Apr 30, 2022
diff --git a/lib/SILGen/ExecutorBreadcrumb.h b/lib/SILGen/ExecutorBreadcrumb.h
@@ -10,6 +10,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#ifndef SWIFT_SILGEN_EXECUTORBREADCRUMB_H
+#define SWIFT_SILGEN_EXECUTORBREADCRUMB_H
+
 #include "swift/SIL/SILValue.h"
 
 namespace swift {
@@ -46,3 +49,5 @@ class ExecutorBreadcrumb {
 
 } // namespace Lowering
 } // namespace swift
+
+#endif
diff --git a/lib/SILGen/ResultPlan.cpp b/lib/SILGen/ResultPlan.cpp
@@ -498,6 +498,7 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
   SILType opaqueResumeType;
   SILValue resumeBuf;
   SILValue continuation;
+  ExecutorBreadcrumb breadcrumb;
 
 public:
   ForeignAsyncInitializationPlan(SILGenFunction &SGF, SILLocation loc,
@@ -596,6 +597,10 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
     return ManagedValue::forUnmanaged(block);
   }
 
+  void deferExecutorBreadcrumb(ExecutorBreadcrumb &&breadcrumb) override {
+    this->breadcrumb = breadcrumb;
+  }
+
   RValue finish(SILGenFunction &SGF, SILLocation loc, CanType substType,
                 ArrayRef<ManagedValue> &directResults,
                 SILValue bridgedForeignError) override {
@@ -691,6 +696,7 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
     // Propagate an error if we have one.
     if (errorBlock) {
       SGF.B.emitBlock(errorBlock);
+      breadcrumb.emit(SGF, loc);
 
       Scope errorScope(SGF, loc);
 
@@ -702,6 +708,7 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
     }
 
     SGF.B.emitBlock(resumeBlock);
+    breadcrumb.emit(SGF, loc);
 
     // The incoming value is the maximally-abstracted result type of the
     // continuation. Move it out of the resume buffer and reabstract it if
@@ -771,6 +778,10 @@ class ForeignErrorInitializationPlan final : public ResultPlan {
                                 AbstractionPattern(errorType), errorType);
   }
 
+  void deferExecutorBreadcrumb(ExecutorBreadcrumb &&breadcrumb) override {
+    subPlan->deferExecutorBreadcrumb(std::move(breadcrumb));
+  }
+
   RValue finish(SILGenFunction &SGF, SILLocation loc, CanType substType,
                 ArrayRef<ManagedValue> &directResults,
                 SILValue bridgedForeignError) override {

diff --git a/lib/SILGen/ResultPlan.h b/lib/SILGen/ResultPlan.h
@@ -14,6 +14,7 @@
 #define SWIFT_SILGEN_RESULTPLAN_H
 
 #include "Callee.h"
+#include "ExecutorBreadcrumb.h"
 #include "ManagedValue.h"
 #include "swift/AST/Types.h"
 #include "swift/Basic/LLVM.h"
@@ -42,6 +43,11 @@ class ResultPlan {
                         SILValue bridgedForeignError) = 0;
   virtual ~ResultPlan() = default;
 
+  /// Defers the emission of the given breadcrumb until \p finish is invoked.
+  virtual void deferExecutorBreadcrumb(ExecutorBreadcrumb &&breadcrumb) {
+    llvm_unreachable("this ResultPlan does not handle deferred breadcrumbs!");
+  }
+
   virtual void
   gatherIndirectResultAddrs(SILGenFunction &SGF, SILLocation loc,
                             SmallVectorImpl<SILValue> &outList) const = 0;

diff --git a/lib/SILGen/SILGenApply.cpp b/lib/SILGen/SILGenApply.cpp
@@ -4427,27 +4427,6 @@ class FixLifetimeDestroyCleanup : public Cleanup {
 #endif
   }
 };
-
-class EmitBreadcrumbCleanup : public Cleanup {
-  ExecutorBreadcrumb breadcrumb;
-
-public:
-  EmitBreadcrumbCleanup(ExecutorBreadcrumb &&breadcrumb)
-    : breadcrumb(std::move(breadcrumb)) {}
-
-  void emit(SILGenFunction &SGF, CleanupLocation l,
-            ForUnwind_t forUnwind) override {
-    breadcrumb.emit(SGF, l);
-  }
-
-  void dump(SILGenFunction &SGF) const override {
-#ifndef NDEBUG
-    llvm::errs() << "EmitBreadcrumbCleanup "
-                 << "State:" << getState()
-                 << "NeedsEmit:" << breadcrumb.needsEmit();
-#endif
-  }
-};
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
@@ -4598,11 +4577,6 @@ RValue SILGenFunction::emitApply(
     rawDirectResult = rawDirectResults[0];
   }
 
-  if (!calleeTypeInfo.foreign.async) {
-    // hop back to the current executor
-    breadcrumb.emit(*this, loc);
-  }
-
   // For objc async calls, lifetime extend the args until the result plan which
   // generates `await_async_continuation`.
   // Lifetime is extended by creating unmanaged copies here and by pushing the
@@ -4614,6 +4588,14 @@ RValue SILGenFunction::emitApply(
         unmanagedCopies.push_back(arg.unmanagedCopy(*this, loc));
       }
     }
+    // similarly, we defer the emission of the breadcrumb until the result
+    // plan's finish method is called, because it must happen in the
+    // successors of the `await_async_continuation` terminator.
+    resultPlan->deferExecutorBreadcrumb(std::move(breadcrumb));
+
+  } else {
+    // In the ordinary case, we hop back to the current executor
+    breadcrumb.emit(*this, loc);
   }
 
   // Pop the argument scope.
@@ -4700,8 +4682,6 @@ RValue SILGenFunction::emitApply(
     for (auto unmanagedCopy : unmanagedCopies) {
       Cleanups.pushCleanup<FixLifetimeDestroyCleanup>(unmanagedCopy);
     }
-    // save breadcrumb as a clean-up so it is emitted in result / throw cases.
-    Cleanups.pushCleanup<EmitBreadcrumbCleanup>(std::move(breadcrumb));
   } else {
     assert(unmanagedCopies.empty());
   }

diff --git a/test/Inputs/clang-importer-sdk/usr/include/ObjCConcurrency.h b/test/Inputs/clang-importer-sdk/usr/include/ObjCConcurrency.h
@@ -310,4 +310,13 @@ struct StructWithSendableContents {
 SENDABLE id StructWithSendableContentsGetSendableComputed(struct StructWithSendableContents contents)
   __attribute__((swift_name("getter:StructWithSendableContents.sendableComputed(self:)")));
 
+@interface CostcoManager : NSObject
++ (instancetype)sharedManager;
+- (void)isCustomerEnrolledInExecutiveProgram:(NSObject *)customer completion:(void(^)(BOOL enrolled))completion;
+@end
+
+@interface Person : NSObject
++ (void)getAsCustomer:(void(^_Nonnull)(NSObject *device))completion;
+@end
+
 #pragma clang assume_nonnull end
diff --git a/test/SILGen/objc_async.swift b/test/SILGen/objc_async.swift
@@ -177,7 +177,7 @@ func testGeneric2<T: AnyObject, U>(x: GenericObject<T>, y: U) async throws {
 // CHECK-LABEL: sil {{.*}}@${{.*}}22testSlowServerFromMain
 @MainActor
 func testSlowServerFromMain(slowServer: SlowServer) async throws {
-  // CHECK: hop_to_executor %6 : $MainActor
+  // CHECK: hop_to_executor {{%.*}} : $MainActor
   // CHECK: [[RESUME_BUF:%.*]] = alloc_stack $Int
   // CHECK: [[STRINGINIT:%.*]] = function_ref @$sSS10FoundationE19_bridgeToObjectiveCSo8NSStringCyF :
   // CHECK: [[ARG:%.*]] = apply [[STRINGINIT]]
@@ -194,8 +194,8 @@ func testSlowServerFromMain(slowServer: SlowServer) async throws {
   // CHECK: destroy_value [[ARG]]
   // CHECK: await_async_continuation [[CONT]] {{.*}}, resume [[RESUME:bb[0-9]+]]
   // CHECK: [[RESUME]]:
+  // CHECK: hop_to_executor {{%.*}} : $MainActor
   // CHECK: [[RESULT:%.*]] = load [trivial] [[RESUME_BUF]]
-  // CHECK: hop_to_executor %6 : $MainActor
   // CHECK: fix_lifetime [[COPY]]
   // CHECK: destroy_value [[COPY]]
   // CHECK: dealloc_stack [[RESUME_BUF]]
@@ -223,8 +223,8 @@ func testThrowingMethodFromMain(slowServer: SlowServer) async -> String {
 // CHECK:  await_async_continuation [[RAW_CONT]] : $Builtin.RawUnsafeContinuation, resume [[RESUME:bb[0-9]+]], error [[ERROR:bb[0-9]+]]
 
 // CHECK: [[RESUME]]
-// CHECK:   {{.*}} = load [take] [[RESULT_BUF]] : $*String
 // CHECK:   hop_to_executor {{%.*}} : $MainActor
+// CHECK:   {{.*}} = load [take] [[RESULT_BUF]] : $*String
 // CHECK:   fix_lifetime [[STRING_ARG_COPY]] : $NSString
 // CHECK:   destroy_value [[STRING_ARG_COPY]] : $NSString
 // CHECK:   dealloc_stack [[RESULT_BUF]] : $*String
@@ -241,3 +241,38 @@ func testThrowingMethodFromMain(slowServer: SlowServer) async -> String {
     return "none"
   }
 }
+
+// rdar://91502776
+// CHECK-LABEL: sil hidden [ossa] @$s{{.*}}21checkCostcoMembershipSbyYaF : $@convention(thin) @async () -> Bool {
+// CHECK:    bb0:
+// CHECK:        hop_to_executor {{%.*}} : $Optional<Builtin.Executor>
+// CHECK:        [[FINAL_BUF:%.*]] = alloc_stack $Bool
+// CHECK:        [[RESULT_BUF:%.*]] = alloc_stack $NSObject
+// CHECK:        [[METH:%.*]] = objc_method {{%.*}} : $@objc_metatype Person.Type, #Person.asCustomer!foreign
+// CHECK:        get_async_continuation_addr NSObject, [[RESULT_BUF]] : $*NSObject
+// CHECK:        = apply [[METH]]
+// CHECK:        dealloc_stack {{%.*}} : $*@block_storage
+// CHECK:        await_async_continuation {{%.*}} : $Builtin.RawUnsafeContinuation, resume bb1
+// CHECK:    bb1:
+// CHECK:        hop_to_executor {{%.*}} : $Optional<Builtin.Executor>
+// CHECK:        [[RESULT:%.*]] = load [take] [[RESULT_BUF]] : $*NSObject
+// CHECK:        objc_method {{%.*}} : $CostcoManager, #CostcoManager.isCustomerEnrolled!foreign
+// CHECK:        get_async_continuation_addr Bool, [[FINAL_BUF]] : $*Bool
+// CHECK:        [[BLOCK_ARG:%.*]] = init_block_storage_header [[BLOCK_STORAGE:%.*]] : $*@block_storage
+// CHECK:        = apply {{%.*}}([[RESULT]], [[BLOCK_ARG]], [[MANAGER:%.*]]) : $@convention(objc_method)
+// CHECK:        [[EXTEND1:%.*]] = copy_value [[RESULT]] : $NSObject
+// CHECK:        [[EXTEND2:%.*]] = copy_value [[MANAGER]] : $CostcoManager
+// CHECK:        dealloc_stack [[BLOCK_STORAGE]] : $*@block_storage
+// CHECK:        await_async_continuation {{%.*}} : $Builtin.RawUnsafeContinuation, resume bb2
+// CHECK:    bb2:
+// CHECK:        hop_to_executor {{%.*}} : $Optional<Builtin.Executor>
+// CHECK:        [[ANSWER:%.*]] = load [trivial] [[FINAL_BUF]] : $*Bool
+// CHECK:        fix_lifetime [[EXTEND2]] : $CostcoManager
+// CHECK:        destroy_value [[EXTEND2]] : $CostcoManager
+// CHECK:        fix_lifetime [[EXTEND1]] : $NSObject
+// CHECK:        destroy_value [[EXTEND1]] : $NSObject
+// CHECK:        return [[ANSWER]] : $Bool
+// CHECK:  }
+func checkCostcoMembership() async -> Bool {
+  return await CostcoManager.shared().isCustomerEnrolled(inExecutiveProgram: Person.asCustomer())
+}
diff --git a/test/SILGen/objc_async_from_swift.swift b/test/SILGen/objc_async_from_swift.swift
@@ -31,8 +31,8 @@ func testSlowServing(p: SlowServing) async throws {
 
     // CHECK: objc_method {{.*}} $@convention(objc_method) <τ_0_0 where τ_0_0 : SlowServing> (@convention(block) (Int, Optional<NSString>, Optional<NSError>) -> (), τ_0_0) -> ()
     // CHECK: hop_to_executor [[GENERIC_EXECUTOR]] :
-    // CHECK:      builtin "willThrow"
-    // CHECK-NEXT: hop_to_executor [[GENERIC_EXECUTOR]] :
+    // CHECK: hop_to_executor [[GENERIC_EXECUTOR]] :
+    // CHECK-NEXT:      builtin "willThrow"
     let _: (Int, String) = try await p.tryRequestIntAndString()
 }
 
@@ -42,8 +42,8 @@ func testSlowServingAgain(p: SlowServing) async throws {
   // CHECK: hop_to_executor [[GENERIC_EXECUTOR]] :
   // CHECK: objc_method {{.*}} $@convention(objc_method) <τ_0_0 where τ_0_0 : SlowServing> (@convention(block) (Optional<NSString>, Optional<NSError>) -> (), τ_0_0) -> ()
   // CHECK: hop_to_executor [[GENERIC_EXECUTOR]] :
-  // CHECK:      builtin "willThrow"
-  // CHECK-NEXT: hop_to_executor [[GENERIC_EXECUTOR]] :
+  // CHECK: hop_to_executor [[GENERIC_EXECUTOR]] :
+  // CHECK-NEXT:      builtin "willThrow"
   let _: String = try await p.tryRequestString()
 }