IRGen: Fix resilient witness tables and vtables to correctly reference async methods

Resilient witness tables and resilient class vtables are built from
descriptors. Make sure we reference the AsyncFunctionPointer of a
method implementation, and not the implementation itself, if the
method is async.

Part of rdar://problem/73625623.

Author: slavapestov

lib/IRGen/GenMeta.cpp

@@ -289,8 +289,14 @@ static void buildMethodDescriptorFields(IRGenModule &IGM,
 
   if (auto entry = VTable->getEntry(IGM.getSILModule(), fn)) {
     assert(entry->getKind() == SILVTable::Entry::Kind::Normal);
-    auto *implFn = IGM.getAddrOfSILFunction(entry->getImplementation(),
-                                            NotForDefinition);
+
+    auto *impl = entry->getImplementation();
+    llvm::Constant *implFn;
+    if (impl->isAsync())
+      implFn = IGM.getAddrOfAsyncFunctionPointer(impl);
+    else
+      implFn = IGM.getAddrOfSILFunction(impl, NotForDefinition);
+
     descriptor.addRelativeAddress(implFn);
   } else {
     // The method is removed by dead method elimination.
@@ -1757,8 +1763,14 @@ namespace {
       // The implementation of the override.
       if (auto entry = VTable->getEntry(IGM.getSILModule(), baseRef)) {
         assert(entry->getKind() == SILVTable::Entry::Kind::Override);
-        auto *implFn = IGM.getAddrOfSILFunction(entry->getImplementation(),
-                                                NotForDefinition);
+
+        auto *impl = entry->getImplementation();
+        llvm::Constant *implFn;
+        if (impl->isAsync())
+          implFn = IGM.getAddrOfAsyncFunctionPointer(impl);
+        else
+          implFn = IGM.getAddrOfSILFunction(impl, NotForDefinition);
+
         descriptor.addRelativeAddress(implFn);
       } else {
         // The method is removed by dead method elimination.

lib/IRGen/GenProto.cpp

@@ -1678,7 +1678,10 @@ void ResilientWitnessTableBuilder::collectResilientWitnesses(
     SILFunction *Func = entry.getMethodWitness().Witness;
     llvm::Constant *witness;
     if (Func) {
-      witness = IGM.getAddrOfSILFunction(Func, NotForDefinition);
+      if (Func->isAsync())
+        witness = IGM.getAddrOfAsyncFunctionPointer(Func);
+      else
+        witness = IGM.getAddrOfSILFunction(Func, NotForDefinition);
     } else {
       // The method is removed by dead method elimination.
       // It should be never called. We add a null pointer.

test/Concurrency/Runtime/Inputs/resilient_class.swift

@@ -0,0 +1,12 @@
+open class BaseClass<T> {
+  let value: T
+
+  public init(value: T) {
+    self.value = value
+  }
+
+  open func waitForNothing() async {}
+  open func wait() async -> T {
+    return value
+  }
+}

test/Concurrency/Runtime/Inputs/resilient_protocol.swift

@@ -0,0 +1,5 @@
+public protocol Awaitable {
+  associatedtype Result
+  func waitForNothing() async
+  func wait() async -> Result
+}

test/Concurrency/Runtime/class_resilience.swift

@@ -0,0 +1,47 @@
+// RUN: %empty-directory(%t)
+
+// RUN: %target-build-swift-dylib(%t/%target-library-name(resilient_class)) -Xfrontend -enable-experimental-concurrency -enable-library-evolution %S/Inputs/resilient_class.swift -emit-module -emit-module-path %t/resilient_class.swiftmodule -module-name resilient_class
+// RUN: %target-codesign %t/%target-library-name(resilient_class)
+
+// RUN: %target-build-swift -Xfrontend -enable-experimental-concurrency %s -lresilient_class -I %t -L %t -o %t/main %target-rpath(%t)
+// RUN: %target-codesign %t/main
+
+// RUN: %target-run %t/main %t/%target-library-name(resilient_class)
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+
+import StdlibUnittest
+import resilient_class
+
+class MyDerived : BaseClass<Int> {
+  override func waitForNothing() async {
+    await super.waitForNothing()
+  }
+
+  override func wait() async -> Int {
+    return await super.wait() * 2
+  }
+}
+
+func virtualWaitForNothing<T>(_ c: BaseClass<T>) async {
+  await c.waitForNothing()
+}
+
+func virtualWait<T>(_ t: BaseClass<T>) async -> T {
+  return await t.wait()
+}
+
+var AsyncVTableMethodSuite = TestSuite("ResilientClass")
+
+AsyncVTableMethodSuite.test("AsyncVTableMethod") {
+  runAsyncAndBlock {
+    let x = MyDerived(value: 321)
+
+    await virtualWaitForNothing(x)
+
+    expectEqual(642, await virtualWait(x))
+  }
+}
+
+runAllTests()

test/Concurrency/Runtime/protocol_resilience.swift

@@ -0,0 +1,45 @@
+// RUN: %empty-directory(%t)
+
+// RUN: %target-build-swift-dylib(%t/%target-library-name(resilient_protocol)) -Xfrontend -enable-experimental-concurrency -enable-library-evolution %S/Inputs/resilient_protocol.swift -emit-module -emit-module-path %t/resilient_protocol.swiftmodule -module-name resilient_protocol
+// RUN: %target-codesign %t/%target-library-name(resilient_protocol)
+
+// RUN: %target-build-swift -Xfrontend -enable-experimental-concurrency %s -lresilient_protocol -I %t -L %t -o %t/main %target-rpath(%t)
+// RUN: %target-codesign %t/main
+
+// RUN: %target-run %t/main %t/%target-library-name(resilient_protocol)
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+
+import StdlibUnittest
+import resilient_protocol
+
+struct IntAwaitable : Awaitable {
+  func waitForNothing() async {}
+
+  func wait() async -> Int {
+    return 123
+  }
+}
+
+func genericWaitForNothing<T : Awaitable>(_ t: T) async {
+  await t.waitForNothing()
+}
+
+func genericWait<T : Awaitable>(_ t: T) async -> T.Result {
+  return await t.wait()
+}
+
+var AsyncProtocolRequirementSuite = TestSuite("ResilientProtocol")
+
+AsyncProtocolRequirementSuite.test("AsyncProtocolRequirement") {
+  runAsyncAndBlock {
+    let x = IntAwaitable()
+
+    await genericWaitForNothing(x)
+
+    expectEqual(123, await genericWait(x))
+  }
+}
+
+runAllTests()

test/IRGen/async/class_resilience.swift

@@ -27,6 +27,12 @@ open class MyBaseClass<T> {
 
 // CHECK-LABEL: @"$s16class_resilience11MyBaseClassC4waitxyYFTjTu" = {{(dllexport )?}}{{(protected )?}}global %swift.async_func_pointer
 
+// CHECK-LABEL: @"$s16class_resilience11MyBaseClassCMn" = {{(dllexport )?}}{{(protected )?}}constant
+// CHECK-SAME: %swift.async_func_pointer* @"$s16class_resilience11MyBaseClassC4waitxyYFTu"
+
+// CHECK-LABEL: @"$s16class_resilience9MyDerivedCMn" = hidden constant
+// CHECK-SAME: %swift.async_func_pointer* @"$s16class_resilience9MyDerivedC4waitSiyYF010resilient_A09BaseClassCADxyYFTVTu"
+
 // 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
@@ -35,3 +41,9 @@ public func callsAwaitable<T>(_ c: BaseClass<T>) async -> T {
 }
 
 // CHECK-LABEL: define {{(dllexport )?}}{{(protected )?}}swiftcc void @"$s16class_resilience11MyBaseClassC4waitxyYFTj"(%swift.task* %0, %swift.executor* %1, %swift.context* swiftasync %2) #0 {
+
+class MyDerived : BaseClass<Int> {
+  override func wait() async -> Int {
+    return await super.wait()
+  }
+}

test/IRGen/async/protocol_resilience.swift

@@ -18,6 +18,9 @@ public protocol MyAwaitable {
 
 // CHECK-LABEL: @"$s19protocol_resilience11MyAwaitableP4wait6ResultQzyYFTjTu" = {{(dllexport )?}}{{(protected )?}}global %swift.async_func_pointer
 
+// CHECK-LABEL: @"$s19protocol_resilience19ConformsToAwaitableVyxG010resilient_A00E0AAMc" = hidden constant
+// CHECK-SAME: %swift.async_func_pointer* @"$s19protocol_resilience19ConformsToAwaitableVyxG010resilient_A00E0AaeFP4wait6ResultQzyYFTWTu"
+
 // 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
@@ -26,3 +29,11 @@ public func callsAwaitable<T : Awaitable>(_ t: T) async -> T.Result {
 }
 
 // CHECK-LABEL: define {{(dllexport )?}}{{(protected )?}}swiftcc void @"$s19protocol_resilience11MyAwaitableP4wait6ResultQzyYFTj"(%swift.task* %0, %swift.executor* %1, %swift.context* swiftasync %2)
+
+struct ConformsToAwaitable<T> : Awaitable {
+  var value: T
+
+  func wait() async -> T {
+    return value
+  }
+}

test/SILGen/async_vtable_thunk.swift

@@ -0,0 +1,13 @@
+// RUN: %target-swift-frontend -emit-silgen %s -enable-experimental-concurrency | %FileCheck %s
+// REQUIRES: concurrency
+
+class BaseClass<T> {
+  func wait() async -> T {}
+}
+
+class Derived : BaseClass<Int> {
+  override func wait() async -> Int {}
+}
+
+// CHECK-LABEL: sil private [thunk] [ossa] @$s18async_vtable_thunk7DerivedC4waitSiyYFAA9BaseClassCADxyYFTV : $@convention(method) @async (@guaranteed Derived) -> @out Int {
+