[Distributed] Support complex generic distributed actors in thunk gen (#70842)

Author: ktoso

lib/IRGen/GenDistributed.cpp

@@ -221,7 +221,7 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
   // A generic parameter that represents instance of invocation decoder.
   auto *decoderType =
       GenericTypeParamType::get(/*isParameterPack=*/false,
-                                /*depth=*/1, /*index=*/0, Context);
+                                /*depth=*/0, /*index=*/0, Context);
 
   // decoder
   parameters.push_back(GenericFunctionType::Param(
@@ -249,13 +249,15 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
   parameters.push_back(GenericFunctionType::Param(Context.getUIntType()));
 
   // actor
-  {
-    auto targetTy = Target->getLoweredFunctionType();
-    auto actorLoc = targetTy->getParameters().back();
 
+  auto actorTypeParam =
+      GenericTypeParamType::get(/*isParameterPack=*/false,
+                                /*depth=*/0, /*index=*/1, Context);
     parameters.push_back(
-        GenericFunctionType::Param(actorLoc.getInterfaceType()));
-  }
+        GenericFunctionType::Param(actorTypeParam));
+  auto distributedActorTy =
+        Context.getDistributedActorDecl()
+        ->getDeclaredInterfaceType();
 
   auto decoderProtocolTy =
       Context
@@ -268,11 +270,8 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
     SmallVector<GenericTypeParamType *, 4> genericParams;
     SmallVector<Requirement, 4> genericRequirements;
 
-    auto *actor = getDistributedActorOf(Target);
-    assert(actor);
-
-    for (auto *genericParam : actor->getInnermostGenericParamTypes())
-      genericParams.push_back(genericParam);
+    assert(getDistributedActorOf(Target) &&
+           "target must be declared inside distributed actor");
 
     // Add a generic parameter `D` which stands for decoder type in the
     // accessor signature - `inout D`.
@@ -281,7 +280,14 @@ static CanSILFunctionType getAccessorType(IRGenModule &IGM,
     genericRequirements.push_back(
         {RequirementKind::Conformance, decoderType, decoderProtocolTy});
 
-    signature = GenericSignature::get(genericParams, genericRequirements);
+    genericRequirements.push_back(
+        {RequirementKind::Conformance, actorTypeParam, distributedActorTy});
+
+    genericParams.push_back(actorTypeParam);
+
+    signature = buildGenericSignature(Context, GenericSignature(),
+                          std::move(genericParams),
+                          std::move(genericRequirements));
   }
 
   auto accessorTy = GenericFunctionType::get(
@@ -650,14 +656,13 @@ void DistributedAccessor::emit() {
   // Metadata that represents passed in the invocation decoder.
   auto *decoderType = params.claimNext();
 
-  // If the distributed thunk is declared in a protocol that conforms
-  // to `DistributedActor` protocol, there is an extract parameter that
-  // represents a type of protocol witness.
-  if (isa<ProtocolDecl>(actor))
-    (void)params.claimNext();
+  // Metadata that represents the actor the invocation is on.
+  auto *actorType = params.claimNext();
+  (void)actorType;
 
   // Witness table for decoder conformance to DistributedTargetInvocationDecoder
   auto *decoderProtocolWitness = params.claimNext();
+  auto *distributedActorWitness = params.claimNext();
 
   // Preliminary: Setup async context for this accessor.
   {

stdlib/public/Distributed/DistributedActor.cpp

@@ -61,7 +61,10 @@ using TargetExecutorSignature =
                         /*witnessTables=*/void **,
                         /*numWitnessTables=*/size_t,
                         /*decoderType=*/Metadata *,
-                        /*decoderWitnessTable=*/void **),
+                        /*actorType=*/Metadata *,
+                        /*decoderWitnessTable=*/void **,
+                        /*distributedActorWitnessTable=*/void **
+                        ),
                    /*throws=*/true>;
 
 SWIFT_CC(swiftasync)
@@ -88,7 +91,9 @@ using DistributedAccessorSignature =
                         /*numWitnessTables=*/size_t,
                         /*actor=*/HeapObject *,
                         /*decoderType=*/Metadata *,
-                        /*decoderWitnessTable=*/void **),
+                        /*actorType=*/Metadata *,
+                        /*decoderWitnessTable=*/void **,
+                        /*actorWitnessTable=*/void **),
                    /*throws=*/true>;
 
 SWIFT_CC(swiftasync)
@@ -124,7 +129,9 @@ void swift_distributed_execute_target(
     void **witnessTables,
     size_t numWitnessTables,
     Metadata *decoderType,
-    void **decoderWitnessTable) {
+    Metadata *actorType,
+    void **decoderWitnessTable,
+    void **actorWitnessTable) {
   auto *accessor = findDistributedAccessor(targetNameStart, targetNameLength);
   if (!accessor) {
     SwiftError *error =
@@ -159,5 +166,7 @@ void swift_distributed_execute_target(
                 numWitnessTables,
                 actor,
                 decoderType,
-                decoderWitnessTable);
+                actorType,
+                decoderWitnessTable,
+                actorWitnessTable);
 }

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -670,8 +670,8 @@ public struct RemoteCallTarget: CustomStringConvertible, Hashable {
 
 @available(SwiftStdlib 5.7, *)
 @_silgen_name("swift_distributed_execute_target")
-func _executeDistributedTarget<D: DistributedTargetInvocationDecoder>(
-  on actor: AnyObject, // DistributedActor
+func _executeDistributedTarget<D: DistributedTargetInvocationDecoder, DA: DistributedActor>(
+  on actor: DA, // DistributedActor
   _ targetName: UnsafePointer<UInt8>, _ targetNameLength: UInt,
   argumentDecoder: inout D,
   argumentTypes: Builtin.RawPointer,

test/Distributed/Runtime/distributed_actor_func_calls_decoded_args_deinit.swift

@@ -2,7 +2,7 @@
 // RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/../Inputs/FakeDistributedActorSystems.swift
 // RUN: %target-build-swift -module-name main -O -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s %S/../Inputs/FakeDistributedActorSystems.swift -o %t/a.out
 // RUN: %target-codesign %t/a.out
-// RUN: %target-run %t/a.out | %FileCheck %s --color --dump-input=always
+// RUN: %target-run %t/a.out | %FileCheck %s --color 
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency

test/Distributed/Runtime/distributed_actor_generic_constraint_issue_115497090.swift

@@ -0,0 +1,60 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/../Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-build-swift -module-name main -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s %S/../Inputs/FakeDistributedActorSystems.swift -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s --color
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+// FIXME(distributed): Distributed has some issues on windows currently, see also: rdar://82593574
+// UNSUPPORTED: OS=windows-msvc
+
+import Distributed
+import FakeDistributedActorSystems
+
+protocol Key {
+  static var isInteger: Bool { get }
+}
+
+distributed actor TestActor<Object> where Object: Codable & Identifiable, Object.ID: Key {
+  public init(actorSystem: ActorSystem) {
+    self.actorSystem = actorSystem
+  }
+
+  public distributed func handleObject(_ object: Object) async throws {
+    print("self.id = \(object.id)")
+  }
+}
+
+struct SomeKey: Codable, Key, Hashable {
+  static var isInteger: Bool { false }
+}
+
+struct Something: Codable, Identifiable {
+  typealias ID = SomeKey
+  var id: SomeKey = .init()
+}
+
+typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
+
+// ==== Execute ----------------------------------------------------------------
+@main struct Main {
+  static func main() async {
+    let system = DefaultDistributedActorSystem()
+
+    let actor: TestActor<Something> = TestActor(actorSystem: system)
+    let resolved: TestActor<Something> = try! .resolve(id: actor.id, using: system)
+    try! await resolved.handleObject(Something())
+
+    // CHECK: self.id = SomeKey()
+
+    // CHECK: OK
+    print("OK")
+  }
+}

test/Distributed/Runtime/distributed_parameter_validation.swift

@@ -2,7 +2,7 @@
 // RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/../Inputs/FakeDistributedActorSystems.swift
 // RUN: %target-build-swift -module-name main  -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s %S/../Inputs/FakeDistributedActorSystems.swift -o %t/a.out
 // RUN: %target-codesign %t/a.out
-// RUN: %target-run %t/a.out | %FileCheck %s --color --dump-input=always
+// RUN: %target-run %t/a.out | %FileCheck %s --color
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency

test/Distributed/distributed_actor_accessor_thunks_32bit.swift

@@ -94,7 +94,7 @@ public distributed actor MyOtherActor {
 
 // CHECK: define hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7simple1yySiYaKFTE"
 
-// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7simple1yySiYaKFTETF"(%swift.context* swiftasync %0, %swift.opaque* nocapture %1, i8* %2, i8* %3, {{.*}}, %T27distributed_actor_accessors7MyActorC* [[ACTOR:%.*]], %swift.type* [[DECODER_TYPE:%.*]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]])
+// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7simple1yySiYaKFTETF"(%swift.context* swiftasync %0, %swift.opaque* nocapture %1, i8* %2, i8* %3, {{.*}}, %T27distributed_actor_accessors7MyActorC* [[ACTOR:%.*]], %swift.type* [[DECODER_TYPE:%.*]], %swift.type* [[ACTOR_TYPE:%.*]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]], i8** [[ACTOR_WITNESS:%.*]])
 
 /// Read the current offset and cast an element to `Int`
 
@@ -194,7 +194,7 @@ public distributed actor MyOtherActor {
 // CHECK: define hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7simple3ySiSSYaKFTE"
 
 /// !!! in `simple3` interesting bits are: argument value extraction (because string is exploded into N arguments) and call to distributed thunk
-// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7simple3ySiSSYaKFTETF"(%swift.context* swiftasync %0, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR]], %swift.type* [[DECODER_TYPE:%.*]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]])
+// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7simple3ySiSSYaKFTETF"(%swift.context* swiftasync %0, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR]], %swift.type* [[DECODER_TYPE:%.*]], %swift.type* [[ACTOR_TYPE:%.*]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]], i8** [[ACTOR_WITNESS:%.*]])
 
 
 // CHECK: [[TYPED_RESULT_BUFF:%.*]] = bitcast i8* [[RESULT_BUFF]] to %TSi*
@@ -257,7 +257,7 @@ public distributed actor MyOtherActor {
 
 // CHECK: define hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC16single_case_enumyAA7SimpleEOAFYaKFTE"
 
-// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC16single_case_enumyAA7SimpleEOAFYaKFTETF"(%swift.context* swiftasync %0, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR]], %swift.type* [[DECODER_TYPE]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]])
+// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC16single_case_enumyAA7SimpleEOAFYaKFTETF"(%swift.context* swiftasync %0, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR]], %swift.type* [[DECODER_TYPE]], %swift.type* [[ACTOR_TYPE]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]], i8** [[ACTOR_WITNESS:%.*]])
 
 /// Let's check that the call doesn't have any arguments and returns nothing.
 
@@ -303,7 +303,7 @@ public distributed actor MyOtherActor {
 
 // CHECK: define hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7complexyAA11LargeStructVSaySiG_AA3ObjCSSSgAFtYaKFTE"
 
-// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7complexyAA11LargeStructVSaySiG_AA3ObjCSSSgAFtYaKFTETF"(%swift.context* swiftasync {{.*}}, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR]], %swift.type* [[DECODER_TYPE:%.*]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]])
+// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC7complexyAA11LargeStructVSaySiG_AA3ObjCSSSgAFtYaKFTETF"(%swift.context* swiftasync {{.*}}, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR]], %swift.type* [[DECODER_TYPE:%.*]], %swift.type* [[ACTOR_TYPE]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]], i8** [[ACTOR_WITNESS:%.*]])
 
 /// First, let's check that all of the different argument types here are loaded correctly.
 
@@ -362,7 +362,7 @@ public distributed actor MyOtherActor {
 
 /// ---> Accessor for `genericArgs`
 
-// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC11genericArgsyyx_Sayq_GtYaKSeRzSERzSeR_SER_r0_lFTETF"(%swift.context* swiftasync %0, %swift.opaque*  nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUF:%.*]], i8* [[GENERIC_SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR:%.*]], %swift.type* [[DECODER_TYPE:%.*]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]])
+// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors7MyActorC11genericArgsyyx_Sayq_GtYaKSeRzSERzSeR_SER_r0_lFTETF"(%swift.context* swiftasync %0, %swift.opaque*  nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUF:%.*]], i8* [[GENERIC_SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors7MyActorC* [[ACTOR:%.*]], %swift.type* [[DECODER_TYPE:%.*]], %swift.type* [[ACTOR_TYPE]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]], i8** [[ACTOR_WITNESS:%.*]])
 
 /// ---> Load `T`
 
@@ -418,7 +418,7 @@ public distributed actor MyOtherActor {
 
 /// Let's check that there is argument decoding since parameter list is empty
 
-// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors12MyOtherActorC5emptyyyYaKFTETF"(%swift.context* swiftasync {{.*}}, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors12MyOtherActorC* {{.*}}, %swift.type* [[DECODER_TYPE:%.*]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]])
+// CHECK: define linkonce_odr hidden swift{{(tail)?}}cc void @"$s27distributed_actor_accessors12MyOtherActorC5emptyyyYaKFTETF"(%swift.context* swiftasync {{.*}}, %swift.opaque* nocapture [[ARG_DECODER:%.*]], i8* [[ARG_TYPES:%.*]], i8* [[RESULT_BUFF:%.*]], i8* [[SUBS:%.*]], i8* [[WITNESS_TABLES:%.*]], i32 [[NUM_WITNESS_TABLES:%.*]], %T27distributed_actor_accessors12MyOtherActorC* {{.*}}, %swift.type* [[DECODER_TYPE:%.*]], %swift.type* [[ACTOR_TYPE]], i8** [[DECODER_PROTOCOL_WITNESS:%.*]], i8** [[ACTOR_WITNESS:%.*]])
 // CHECK-NEXT: entry:
 // CHECK-NEXT: {{.*}} = alloca %swift.context*
 // CHECK-NEXT: %swifterror = alloca %swift.error*