Merge pull request #34784 from nate-chandler/concurrency/irgen/rdar71491604

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

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
+}