IRGen: Don't stored the current async frame to llvm.swift.async.context.addr twice

rdar://148826951

Author: aschwaighofer

lib/IRGen/GenCall.cpp

@@ -262,10 +262,14 @@ llvm::CallInst *IRGenFunction::emitSuspendAsyncCall(
     llvm::Function *projectFn = cast<llvm::Function>(
         (cast<llvm::Constant>(args[2])->stripPointerCasts()));
     auto *fnTy = projectFn->getFunctionType();
-
-    llvm::Value *context =
+    llvm::Value *callerContext = nullptr;
+    if (projectFn == getOrCreateResumePrjFn()) {
+      callerContext = popAysncContext(calleeContext);
+    } else {
+      callerContext =
         Builder.CreateCallWithoutDbgLoc(fnTy, projectFn, {calleeContext});
-    storeCurrentAsyncContext(context);
+    }
+    storeCurrentAsyncContext(callerContext);
   }
 
   return id;

lib/IRGen/GenFunc.cpp

@@ -2808,8 +2808,7 @@ void irgen::emitBlockHeader(IRGenFunction &IGF,
                           IGF.Builder.CreateStructGEP(headerAddr, 4, layout));
 }
 
-llvm::Value *
-IRGenFunction::emitAsyncResumeProjectContext(llvm::Value *calleeContext) {
+llvm::Value *IRGenFunction::popAysncContext(llvm::Value *calleeContext) {
   auto addr = Builder.CreateBitOrPointerCast(calleeContext, IGM.Int8PtrPtrTy);
   Address callerContextAddr(addr, IGM.Int8PtrTy, IGM.getPointerAlignment());
   llvm::Value *callerContext = Builder.CreateLoad(callerContextAddr);
@@ -2818,6 +2817,13 @@ IRGenFunction::emitAsyncResumeProjectContext(llvm::Value *calleeContext) {
         PointerAuthInfo::emit(*this, schema, addr, PointerAuthEntity());
     callerContext = emitPointerAuthAuth(*this, callerContext, authInfo);
   }
+  return callerContext;
+}
+
+llvm::Value *
+IRGenFunction::emitAsyncResumeProjectContext(llvm::Value *calleeContext) {
+  auto  callerContext = popAysncContext(calleeContext);
+
   // TODO: remove this once all platforms support lowering the intrinsic.
   // At the time of this writing only arm64 supports it.
   if (IGM.TargetInfo.canUseSwiftAsyncContextAddrIntrinsic()) {

lib/IRGen/IRGenFunction.h

@@ -196,6 +196,7 @@ class IRGenFunction {
 
   llvm::Value *emitAsyncResumeProjectContext(llvm::Value *callerContextAddr);
   llvm::Function *getOrCreateResumePrjFn();
+  llvm::Value *popAysncContext(llvm::Value *calleeContext);
   llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,
                                         ArrayRef<llvm::Value *> args);
   llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,

test/Distributed/distributed_actor_accessor_thunks_32bit.swift

@@ -145,7 +145,8 @@ public distributed actor MyOtherActor {
 // CHECK-SAME: %T27distributed_actor_accessors7MyActorC* [[ACTOR]])
 
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, %swift.error* } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT:  store ptr [[CALLER_ASYNC_CTXT]], ptr
 // CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, %swift.error* {{.*}})
 
 /// ---> Thunk and distributed method accessor for `simple2`
@@ -178,7 +179,8 @@ public distributed actor MyOtherActor {
 // CHECK-SAME: %T27distributed_actor_accessors7MyActorC* [[ACTOR]])
 
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i32, i32, i32, %swift.error* } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT:  store ptr [[CALLER_ASYNC_CTXT]], ptr
 
 /// Initialize the result buffer with values produced by the thunk
 
@@ -247,7 +249,8 @@ public distributed actor MyOtherActor {
 // CHECK-SAME: %T27distributed_actor_accessors7MyActorC* [[ACTOR]])
 
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT:  store ptr [[CALLER_ASYNC_CTXT]], ptr
 // CHECK: [[INT_RES:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 1
 // CHECK: %._value = getelementptr inbounds %TSi, %TSi* [[TYPED_RESULT_BUFF]], i32 0, i32 0
 // CHECK: store i32 [[INT_RES]], i32* %._value
@@ -292,7 +295,8 @@ public distributed actor MyOtherActor {
 
 // CHECK: [[THUNK_RESULT:%.*]] = call { i8*, i32, %swift.error* } (i32, i8*, i8*, ...) @llvm.coro.suspend.async.sl_p0i8i32p0s_swift.errorss({{.*}}, %swift.context* {{.*}}, i32 [[NATIVE_ENUM_VAL]], i32 [[NATIVE_INT_VAL]], %T27distributed_actor_accessors7MyActorC* {{.*}})
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT:  store ptr [[CALLER_ASYNC_CTXT]], ptr
 // CHECK: [[ENUM_RESULT:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 1
 // CHECK: [[NATIVE_RESULT_PTR:%.*]] = bitcast %T27distributed_actor_accessors9IndirectEO* [[TYPED_RESULT_BUFF]] to i32*
 // CHECK-NEXT: store i32 [[ENUM_RESULT]], i32* [[NATIVE_RESULT_PTR]]

test/Distributed/distributed_actor_accessor_thunks_64bit.swift

@@ -141,7 +141,8 @@ public distributed actor MyOtherActor {
 // CHECK-SAME: ptr [[ACTOR]])
 
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, ptr } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
 // CHECK: {{.*}} = call i1 (ptr, i1, ...) @llvm.coro.end.async({{.*}}, ptr {{.*}}, ptr {{.*}})
 
 /// ---> Thunk and distributed method accessor for `simple2`
@@ -172,7 +173,8 @@ public distributed actor MyOtherActor {
 // CHECK-SAME: ptr [[ACTOR]])
 
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, i64, ptr, ptr } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT:  store ptr [[CALLER_ASYNC_CTXT]], ptr
 
 /// Extract information about `String` from result and call `end`
 
@@ -222,7 +224,8 @@ public distributed actor MyOtherActor {
 // CHECK-SAME: ptr [[ACTOR]])
 
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT:  store ptr [[CALLER_ASYNC_CTXT]], ptr
 // CHECK: [[INT_RES:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 1
 // CHECK: %._value = getelementptr inbounds %TSi, ptr [[RESULT_BUFF]], i32 0, i32 0
 // CHECK: store i64 [[INT_RES]], ptr %._value
@@ -262,7 +265,8 @@ public distributed actor MyOtherActor {
 
 // CHECK: [[THUNK_RESULT:%.*]] = call { ptr, i64, ptr } (i32, ptr, ptr, ...) @llvm.coro.suspend.async.sl_p0i64p0s({{.*}}, ptr {{.*}}, i64 [[NATIVE_ENUM_VAL]], i64 [[NATIVE_INT_VAL]], ptr {{.*}})
 // CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 0
-// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
+// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
+// CHECK-NEXT:  store ptr [[CALLER_ASYNC_CTXT]], ptr
 // CHECK: [[ENUM_RESULT:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 1
 // CHECK: store i64 [[ENUM_RESULT]], ptr [[RESULT_BUFF]]