IRGen: Refer to dispatch thunk async function pointers when making calls

Part of rdar://problem/73625623.

Author: slavapestov

lib/IRGen/IRGenSIL.cpp

@@ -6133,14 +6133,22 @@ void IRGenSILFunction::visitWitnessMethodInst(swift::WitnessMethodInst *i) {
   CanType baseTy = i->getLookupType();
   ProtocolConformanceRef conformance = i->getConformance();
   SILDeclRef member = i->getMember();
+  auto fnType = IGM.getSILTypes().getConstantFunctionType(
+      IGM.getMaximalTypeExpansionContext(), member);
 
   assert(member.requiresNewWitnessTableEntry());
 
   if (IGM.isResilient(conformance.getRequirement(),
                       ResilienceExpansion::Maximal)) {
-    auto *fnPtr = IGM.getAddrOfDispatchThunk(member, NotForDefinition);
-    auto fnType = IGM.getSILTypes().getConstantFunctionType(
-        IGM.getMaximalTypeExpansionContext(), member);
+    llvm::Constant *fnPtr = IGM.getAddrOfDispatchThunk(member, NotForDefinition);
+
+    if (fnType->isAsync()) {
+      auto *fnPtrType = fnPtr->getType();
+      fnPtr = IGM.getAddrOfAsyncFunctionPointer(
+          LinkEntity::forDispatchThunk(member));
+      fnPtr = llvm::ConstantExpr::getBitCast(fnPtr, fnPtrType);
+    }
+
     auto sig = IGM.getSignature(fnType);
     auto fn = FunctionPointer::forDirect(fnType, fnPtr, sig);
 
@@ -6340,7 +6348,15 @@ void IRGenSILFunction::visitClassMethodInst(swift::ClassMethodInst *i) {
   auto *classDecl = cast<ClassDecl>(method.getDecl()->getDeclContext());
   if (IGM.hasResilientMetadata(classDecl,
                                ResilienceExpansion::Maximal)) {
-    auto *fnPtr = IGM.getAddrOfDispatchThunk(method, NotForDefinition);
+    llvm::Constant *fnPtr = IGM.getAddrOfDispatchThunk(method, NotForDefinition);
+
+    if (methodType->isAsync()) {
+      auto *fnPtrType = fnPtr->getType();
+      fnPtr = IGM.getAddrOfAsyncFunctionPointer(
+          LinkEntity::forDispatchThunk(method));
+      fnPtr = llvm::ConstantExpr::getBitCast(fnPtr, fnPtrType);
+    }
+
     auto sig = IGM.getSignature(methodType);
     FunctionPointer fn(methodType, fnPtr, sig);
 

test/IRGen/async/class_resilience.swift

@@ -27,4 +27,11 @@ open class MyBaseClass<T> {
 
 // CHECK-LABEL: @"$s16class_resilience11MyBaseClassC4waitxyYFTjTu" = {{(dllexport )?}}{{(protected )?}}global %swift.async_func_pointer
 
+// CHECK-LABEL: define {{(dllexport )?}}{{(protected )?}}swiftcc void @"$s16class_resilience14callsAwaitableyx010resilient_A09BaseClassCyxGYlF"(%swift.task* %0, %swift.executor* %1, %swift.context* swiftasync %2)
+// CHECK: %swift.async_func_pointer* @"$s15resilient_class9BaseClassC4waitxyYFTjTu"
+// CHECK: ret void
+public func callsAwaitable<T>(_ c: BaseClass<T>) async -> T {
+  return await c.wait()
+}
+
 // CHECK-LABEL: define {{(dllexport )?}}{{(protected )?}}swiftcc void @"$s16class_resilience11MyBaseClassC4waitxyYFTj"(%swift.task* %0, %swift.executor* %1, %swift.context* swiftasync %2) #0 {

test/IRGen/async/protocol_resilience.swift

@@ -18,4 +18,11 @@ public protocol MyAwaitable {
 
 // CHECK-LABEL: @"$s19protocol_resilience11MyAwaitableP4wait6ResultQzyYFTjTu" = {{(dllexport )?}}{{(protected )?}}global %swift.async_func_pointer
 
+// CHECK-LABEL: define {{(dllexport )?}}{{(protected )?}}swiftcc void @"$s19protocol_resilience14callsAwaitabley6ResultQzxY010resilient_A00D0RzlF"(%swift.task* %0, %swift.executor* %1, %swift.context* swiftasync %2)
+// CHECK: %swift.async_func_pointer* @"$s18resilient_protocol9AwaitableP4wait6ResultQzyYFTjTu"
+// CHECK: ret void
+public func callsAwaitable<T : Awaitable>(_ t: T) async -> T.Result {
+  return await t.wait()
+}
+
 // CHECK-LABEL: define {{(dllexport )?}}{{(protected )?}}swiftcc void @"$s19protocol_resilience11MyAwaitableP4wait6ResultQzyYFTj"(%swift.task* %0, %swift.executor* %1, %swift.context* swiftasync %2)