[Async CC] Find wtable in async context in thunk.

Previously, when looking up a protocol method, the witness table was
always exepcted to be the final argument passed to the function.

That is true for sync protocol witnesses but not for async witnesses.
In the async case, the witness table is embedded in the async context.

Here, the witness table is dug out of the async context.

rdar://problem/71491604

Author: nate-chandler

lib/IRGen/GenThunk.cpp

@@ -64,34 +64,49 @@ IRGenModule::getAddrOfDispatchThunk(SILDeclRef declRef,
 static FunctionPointer lookupMethod(IRGenFunction &IGF, SILDeclRef declRef) {
   auto expansionContext = IGF.IGM.getMaximalTypeExpansionContext();
   auto *decl = cast<AbstractFunctionDecl>(declRef.getDecl());
+  auto funcTy = IGF.IGM.getSILModule().Types.getConstantFunctionType(
+      expansionContext, declRef);
+
+  auto getAsyncContextLayout = [&]() {
+    auto originalType = funcTy;
+    auto forwardingSubstitutionMap =
+        decl->getGenericEnvironment()
+            ? decl->getGenericEnvironment()->getForwardingSubstitutionMap()
+            : SubstitutionMap();
+    auto substitutedType = originalType->substGenericArgs(
+        IGF.IGM.getSILModule(), forwardingSubstitutionMap,
+        IGF.IGM.getMaximalTypeExpansionContext());
+    auto layout = irgen::getAsyncContextLayout(
+        IGF.IGM, originalType, substitutedType, forwardingSubstitutionMap);
+    return layout;
+  };
 
   // Protocol case.
   if (isa<ProtocolDecl>(decl->getDeclContext())) {
     // Find the witness table.
-    llvm::Value *wtable = (IGF.CurFn->arg_end() - 1);
+    llvm::Value *wtable;
+    if (funcTy->isAsync()) {
+      auto layout = getAsyncContextLayout();
+      assert(layout.hasTrailingWitnesses());
+      auto context = layout.emitCastTo(IGF, IGF.getAsyncContext());
+      auto wtableAddr =
+          layout.getSelfWitnessTableLayout().project(IGF, context, llvm::None);
+      wtable = IGF.Builder.CreateLoad(wtableAddr);
+    } else {
+      wtable = (IGF.CurFn->arg_end() - 1);
+    }
 
     // Find the witness we're interested in.
     return emitWitnessMethodValue(IGF, wtable, declRef);
   }
 
   // Class case.
-  auto funcTy = IGF.IGM.getSILModule().Types.getConstantFunctionType(
-      expansionContext, declRef);
 
   // Load the metadata, or use the 'self' value if we have a static method.
   llvm::Value *self;
 
   if (funcTy->isAsync()) {
-    auto originalType = funcTy;
-    auto forwardingSubstitutionMap =
-        decl->getGenericEnvironment()
-            ? decl->getGenericEnvironment()->getForwardingSubstitutionMap()
-            : SubstitutionMap();
-    auto substitutedType = originalType->substGenericArgs(
-        IGF.IGM.getSILModule(), forwardingSubstitutionMap,
-        IGF.IGM.getMaximalTypeExpansionContext());
-    auto layout = getAsyncContextLayout(IGF.IGM, originalType, substitutedType,
-                                        forwardingSubstitutionMap);
+    auto layout = getAsyncContextLayout();
     assert(layout.hasLocalContext());
     auto context = layout.emitCastTo(IGF, IGF.getAsyncContext());
     auto localContextAddr =

test/IRGen/async/Inputs/protocol-1instance-void_to_void.swift

@@ -0,0 +1,13 @@
+import _Concurrency
+
+public protocol Protokol {
+  func protocolinstanceVoidToVoid() async
+}
+
+public struct Impl : Protokol {
+  public init() {}
+  public func protocolinstanceVoidToVoid() async {
+    print(self)
+  }
+}
+

test/IRGen/async/run-call-resilient-protocolinstance-void-to-void.swift

@@ -0,0 +1,30 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift-dylib(%t/%target-library-name(PrintShims)) %S/../../Inputs/print-shims.swift -module-name PrintShims -emit-module -emit-module-path %t/PrintShims.swiftmodule
+// RUN: %target-codesign %t/%target-library-name(PrintShims)
+// RUN: %target-build-swift-dylib(%t/%target-library-name(ResilientProtocol)) %S/Inputs/protocol-1instance-void_to_void.swift -Xfrontend -enable-experimental-concurrency -module-name ResilientProtocol -emit-module -emit-module-path %t/ResilientProtocol.swiftmodule
+// RUN: %target-codesign %t/%target-library-name(ResilientProtocol)
+// RUN: %target-build-swift -Xfrontend -enable-experimental-concurrency %s -emit-ir -I %t -L %t -lPrintShim -lResilientProtocol | %FileCheck %s --check-prefix=CHECK-LL
+// RUN: %target-build-swift -Xfrontend -enable-experimental-concurrency %s -module-name main -o %t/main -I %t -L %t -lPrintShims -lResilientProtocol %target-rpath(%t) 
+// RUN: %target-codesign %t/main
+// RUN: %target-run %t/main %t/%target-library-name(PrintShims) %t/%target-library-name(ResilientProtocol) | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: swift_test_mode_optimize_none
+// REQUIRES: concurrency
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: CPU=arm64e
+
+import _Concurrency
+import ResilientProtocol
+
+func call<T : Protokol>(_ t: T) async {
+  await t.protocolinstanceVoidToVoid()
+}
+
+// CHECK-LL: define hidden swiftcc void @"$s4main4callyyxY17ResilientProtocol8ProtokolRzlF"(%swift.task* {{%[0-9]+}}, %swift.executor* {{%[0-9]+}}, %swift.context* {{%[0-9]+}}) {{#[0-9]*}} {
+func test_case() async {
+  let impl = Impl()
+  await call(impl) // CHECK: Impl()
+}
+
+_Concurrency.runAsync(test_case)

validation-test/compiler_crashers_2_fixed/rdar71491604.swift

@@ -0,0 +1,6 @@
+// RUN: %target-swift-frontend %s -emit-ir -enable-library-evolution -enable-experimental-concurrency
+// REQUIRES: concurrency
+
+public protocol P {
+  func f() async
+}