IRGen: Add support for typed throws to thunk code generation

rdar://119725769

Author: aschwaighofer

lib/IRGen/GenCall.cpp

@@ -4841,6 +4841,10 @@ Address IRGenFunction::getCalleeTypedErrorResultSlot(SILType errorType) {
   return CalleeTypedErrorResultSlot;
 }
 
+void IRGenFunction::setCalleeTypedErrorResultSlot(Address addr) {
+  CalleeTypedErrorResultSlot = addr;
+}
+
 
 /// Fetch the error result slot received from the caller.
 Address IRGenFunction::getCallerErrorResultSlot() {

lib/IRGen/GenThunk.cpp

@@ -85,6 +85,7 @@ class IRGenThunk {
   llvm::Value *indirectReturnSlot = nullptr;
   llvm::Value *selfValue = nullptr;
   llvm::Value *errorResult = nullptr;
+  llvm::Value *typedErrorIndirectErrorSlot = nullptr;
   WitnessMetadata witnessMetadata;
   Explosion params;
 
@@ -137,16 +138,31 @@ void IRGenThunk::prepareArguments() {
   SILFunctionConventions conv(origTy, IGF.getSILModule());
 
   if (origTy->hasErrorResult()) {
+    typedErrorIndirectErrorSlot = nullptr;
+
+    // Set the typed error value result slot.
+    if (conv.isTypedError() && !conv.hasIndirectSILErrorResults()) {
+      auto directTypedErrorAddr = original.takeLast();
+      auto errorType =
+        conv.getSILErrorType(IGF.IGM.getMaximalTypeExpansionContext());
+      auto &errorTI = cast<FixedTypeInfo>(IGF.getTypeInfo(errorType));
+
+     IGF.setCalleeTypedErrorResultSlot(Address(directTypedErrorAddr,
+                                               errorTI.getStorageType(),
+                                               errorTI.getFixedAlignment()));
+    } else if (conv.isTypedError()) {
+      auto directTypedErrorAddr = original.takeLast();
+      // Store for later processing when we know the argument index.
+      // (i.e. emission->setIndirectTypedErrorResultSlotArgsIndex was called)
+      typedErrorIndirectErrorSlot = directTypedErrorAddr;
+    }
+
     if (isAsync) {
       // nothing to do.
     } else {
       errorResult = original.takeLast();
-      auto errorType =
-          conv.getSILErrorType(IGF.IGM.getMaximalTypeExpansionContext());
-      auto &errorTI = cast<FixedTypeInfo>(IGF.getTypeInfo(errorType));
-
       IGF.setCallerErrorResultSlot(Address(errorResult,
-                                           errorTI.getStorageType(),
+                                           IGF.IGM.Int8PtrTy,
                                            IGF.IGM.getPointerAlignment()));
     }
   }
@@ -276,6 +292,10 @@ void IRGenThunk::emit() {
   std::unique_ptr<CallEmission> emission =
       getCallEmission(IGF, callee.getSwiftContext(), std::move(callee));
 
+  if (typedErrorIndirectErrorSlot) {
+    emission->setIndirectTypedErrorResultSlot(typedErrorIndirectErrorSlot);
+  }
+
   emission->begin();
 
   emission->setArgs(params, /*isOutlined=*/false, &witnessMetadata);

lib/IRGen/IRGenFunction.h

@@ -136,6 +136,7 @@ class IRGenFunction {
 
   Address getCallerTypedErrorResultSlot();
   Address getCalleeTypedErrorResultSlot(SILType errorType);
+  void setCalleeTypedErrorResultSlot(Address addr);
 
   /// Are we currently emitting a coroutine?
   bool isCoroutine() {

test/IRGen/typed_throws_thunks.swift

@@ -0,0 +1,67 @@
+// RUN: %target-swift-frontend -primary-file %s -emit-irgen -enable-experimental-feature TypedThrows -disable-availability-checking -runtime-compatibility-version none -target %module-target-future -enable-library-evolution | %FileCheck %s --check-prefix=CHECK
+
+// REQUIRES: PTRSIZE=64
+// UNSUPPORTED: CPU=arm64e
+
+@frozen
+public struct FixedSize : Error {
+   var x = 0
+   var y = 0
+}
+
+public protocol P {
+  associatedtype Failure: Error = any Error
+
+  func f(body: () throws -> Void) async rethrows
+  func f2(body: () throws -> Void) rethrows
+
+  func g(body: () throws(Failure) -> Void) async throws(Failure)
+  func g4(body: () throws(FixedSize) -> Void) async throws(FixedSize)
+
+  func g2(body: () throws(Failure) -> Void) throws(Failure)
+  func g3(body: () throws(FixedSize) -> Void) throws(FixedSize)
+}
+
+extension P {
+ // CHECK-LABEL: define{{.*}} swifttailcc void @"$s19typed_throws_thunks1PP1g4bodyyyy7FailureQzYKXE_tYaAGYKFTj"(ptr swiftasync %0, ptr %1, ptr %2, ptr noalias swiftself %3, ptr %4, ptr %5, ptr %6)
+ // CHECK-NOT: ret
+ // CHECK:  call { ptr, ptr } (i32, ptr, ptr, ...) @llvm.coro.suspend.async.sl_p0p0s({{.*}}, ptr %1, ptr %2, ptr %3, ptr %4, ptr %5, ptr %6)
+  public func g(body: () throws(Failure) -> Void) async throws(Failure) {
+    do {
+      return try await f(body: body)
+    } catch {
+      throw error as! Failure
+    }
+  }
+
+  // CHECK-LABEL: define{{.*}} swifttailcc void @"$s19typed_throws_thunks1PP2g44bodyyyyAA9FixedSizeVYKXE_tYaAGYKFTj"(ptr swiftasync %0, ptr %1, ptr %2, ptr noalias swiftself %3, ptr %4, ptr %5, ptr %6)
+	// CHECK-NOT: ret
+	// CHECK: call { ptr, ptr } (i32, ptr, ptr, ...) @llvm.coro.suspend.async.sl_p0p0s({{.*}} ptr %1, ptr %2, ptr %3, ptr %4, ptr %5, ptr %6)
+  public func g4(body: () throws(FixedSize) -> Void) async throws(FixedSize) {
+    do {
+      return try await f(body: body)
+    } catch {
+      throw error as! FixedSize
+    }
+  }
+
+  // CHECK-LABEL: define{{.*}} swiftcc void @"$s19typed_throws_thunks1PP2g24bodyyyy7FailureQzYKXE_tAGYKFTj"(ptr %0, ptr %1, ptr noalias swiftself %2, ptr noalias nocapture swifterror dereferenceable(8) %3, ptr %4, ptr %5, ptr %6)
+	// CHECK-NOT: ret
+  // CHECK: call swiftcc void {{.*}}(ptr %0, ptr %1, ptr noalias swiftself %2, ptr noalias nocapture swifterror dereferenceable(8) %3, ptr %4, ptr %5, ptr %6)
+
+  public func g2(body: () throws(Failure) -> Void) throws(Failure) {
+    do {
+      return try f2(body: body)
+    } catch {
+      throw error as! Failure
+    }
+  }
+
+  // CHECK-LABEL: define{{.*}} swiftcc void @"$s19typed_throws_thunks1PP2g34bodyyyyAA9FixedSizeVYKXE_tAGYKFTj"(ptr %0, ptr %1, ptr noalias swiftself %2, ptr noalias nocapture swifterror dereferenceable(8) %3, ptr %4, ptr %5, ptr %6)
+  // CHECK-NOT: ret
+  // CHECK:  call swiftcc void {{.*}}(ptr %0, ptr %1, ptr noalias swiftself %2, ptr noalias nocapture swifterror dereferenceable(8) %3, ptr %4, ptr %5, ptr %6)
+
+
+  public func g3(body: () throws(FixedSize) -> Void) throws(FixedSize) {
+  }
+}