Merge pull request swiftlang#35948 from slavapestov/async-dispatch-thunk-loadable-return

IRGen: Fix async dispatch thunk emission with loadable return value

Author: slavapestov

lib/IRGen/GenCall.cpp

@@ -4753,6 +4753,23 @@ void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
   call->setTailCall();
 }
 
+void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
+                            CanSILFunctionType fnType, Explosion &result) {
+  llvm::Value *context = IGF.getAsyncContext();
+
+  Address dataAddr = asyncLayout.emitCastTo(IGF, context);
+  for (unsigned index = 0, count = asyncLayout.getDirectReturnCount();
+       index < count; ++index) {
+    auto fieldLayout = asyncLayout.getDirectReturnLayout(index);
+    Address fieldAddr =
+        fieldLayout.project(IGF, dataAddr, /*offsets*/ llvm::None);
+    cast<LoadableTypeInfo>(fieldLayout.getType())
+        .initialize(IGF, result, fieldAddr, /*isOutlined*/ false);
+  }
+
+  emitAsyncReturn(IGF, asyncLayout, fnType);
+}
+
 FunctionPointer
 IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) {
   auto *fnTy = llvm::FunctionType::get(

lib/IRGen/GenCall.h

@@ -436,6 +436,9 @@ namespace irgen {
   void emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &layout,
                        CanSILFunctionType fnType);
 
+  void emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &layout,
+                       CanSILFunctionType fnType, Explosion &result);
+
   Address emitAutoDiffCreateLinearMapContext(
       IRGenFunction &IGF, llvm::Value *topLevelSubcontextSize);
   Address emitAutoDiffProjectTopLevelSubcontext(

lib/IRGen/GenThunk.cpp

@@ -37,6 +37,7 @@
 #include "ProtocolInfo.h"
 #include "Signature.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/PrettyStackTrace.h"
 #include "swift/IRGen/Linking.h"
 #include "swift/SIL/SILDeclRef.h"
 #include "llvm/IR/Function.h"
@@ -282,6 +283,9 @@ Callee IRGenThunk::lookupMethod() {
 }
 
 void IRGenThunk::emit() {
+  PrettyStackTraceDecl stackTraceRAII("emitting dispatch thunk for",
+                                      declRef.getDecl());
+
   GenericContextScope scope(IGF.IGM, origTy->getInvocationGenericSignature());
 
   if (isAsync) {
@@ -345,7 +349,7 @@ void IRGenThunk::emit() {
   }
 
   if (isAsync) {
-    emitAsyncReturn(IGF, *asyncLayout, origTy);
+    emitAsyncReturn(IGF, *asyncLayout, origTy, result);
     IGF.emitCoroutineOrAsyncExit();
     return;
   }

lib/IRGen/IRGenSIL.cpp

@@ -3233,19 +3233,9 @@ static void emitReturnInst(IRGenSILFunction &IGF,
     assert(!IGF.IndirectReturn.isValid() &&
            "Formally direct results should stay direct results for async "
            "functions");
-    llvm::Value *context = IGF.getAsyncContext();
-    auto layout = getAsyncContextLayout(IGF);
-
-    Address dataAddr = layout.emitCastTo(IGF, context);
-    for (unsigned index = 0, count = layout.getDirectReturnCount();
-         index < count; ++index) {
-      auto fieldLayout = layout.getDirectReturnLayout(index);
-      Address fieldAddr =
-          fieldLayout.project(IGF, dataAddr, /*offsets*/ llvm::None);
-      cast<LoadableTypeInfo>(fieldLayout.getType())
-          .initialize(IGF, result, fieldAddr, /*isOutlined*/ false);
-    }
-    emitAsyncReturn(IGF, layout, fnType);
+
+    auto asyncLayout = getAsyncContextLayout(IGF);
+    emitAsyncReturn(IGF, asyncLayout, fnType, result);
     IGF.emitCoroutineOrAsyncExit();
   } else {
     auto funcLang = IGF.CurSILFn->getLoweredFunctionType()->getLanguage();

test/Concurrency/Runtime/Inputs/resilient_class.swift

@@ -10,9 +10,15 @@ open class BaseClass<T> {
   }
 
   open func waitForNothing() async {}
+
   open func wait() async -> T {
     return value
   }
+
+  open func waitForInt() async -> Int {
+    return 123
+  }
+
   open func wait(orThrow: Bool) async throws {
     if orThrow {
       throw MyError.bad

test/Concurrency/Runtime/Inputs/resilient_protocol.swift

@@ -1,6 +1,7 @@
 public protocol Awaitable {
   associatedtype Result
   func waitForNothing() async
+  func waitForInt() async -> Int
   func wait() async -> Result
   func wait(orThrow: Bool) async throws
 }

test/Concurrency/Runtime/class_resilience.swift

@@ -26,6 +26,10 @@ class MyDerived : BaseClass<Int> {
     return await super.wait() * 2
   }
 
+  override func waitForInt() async -> Int {
+    return await super.waitForInt() * 2
+  }
+
   override func wait(orThrow: Bool) async throws {
     return try await super.wait(orThrow: orThrow)
   }
@@ -39,6 +43,10 @@ func virtualWait<T>(_ c: BaseClass<T>) async -> T {
   return await c.wait()
 }
 
+func virtualWaitForInt<T>(_ c: BaseClass<T>) async -> Int {
+  return await c.waitForInt()
+}
+
 func virtualWait<T>(orThrow: Bool, _ c: BaseClass<T>) async throws {
   return try await c.wait(orThrow: orThrow)
 }
@@ -52,6 +60,7 @@ AsyncVTableMethodSuite.test("AsyncVTableMethod") {
     await virtualWaitForNothing(x)
 
     expectEqual(642, await virtualWait(x))
+    expectEqual(246, await virtualWaitForInt(x))
 
     expectNil(try? await virtualWait(orThrow: true, x))
     try! await virtualWait(orThrow: false, x)

test/Concurrency/Runtime/protocol_resilience.swift

@@ -28,6 +28,10 @@ struct IntAwaitable : Awaitable {
     return 123
   }
 
+  func waitForInt() async -> Int {
+    return 321
+  }
+
   func wait(orThrow: Bool) async throws {
     if (orThrow) {
       throw MyError.bad
@@ -43,6 +47,10 @@ func genericWait<T : Awaitable>(_ t: T) async -> T.Result {
   return await t.wait()
 }
 
+func genericWaitForInt<T : Awaitable>(_ t: T) async -> Int {
+  return await t.waitForInt()
+}
+
 func genericWait<T : Awaitable>(orThrow: Bool, _ t: T) async throws {
   return try await t.wait(orThrow: orThrow)
 }
@@ -56,6 +64,7 @@ AsyncProtocolRequirementSuite.test("AsyncProtocolRequirement") {
     await genericWaitForNothing(x)
 
     expectEqual(123, await genericWait(x))
+    expectEqual(321, await genericWaitForInt(x))
 
     expectNil(try? await genericWait(orThrow: true, x))
     try! await genericWait(orThrow: false, x)

test/IRGen/async/class_resilience.swift

@@ -8,10 +8,18 @@ import resilient_class
 open class MyBaseClass<T> {
   var value: T
 
+  open func wait() async -> Int {
+    return 0
+  }
+
   open func wait() async -> T {
     return value
   }
 
+  open func waitThrows() async throws -> Int {
+    return 0
+  }
+
   open func waitThrows() async throws -> T {
     return value
   }

test/IRGen/async/protocol_resilience.swift

@@ -7,8 +7,13 @@ import resilient_protocol
 
 public protocol MyAwaitable {
   associatedtype Result
+
+  func wait() async -> Int
+
   func wait() async -> Result
 
+  func waitThrows() async throws -> Int
+
   func waitThrows() async throws -> Result
 
   // FIXME