[Distributed] Fix generic actors and distributed funcs returning those generics

lib/AST/DistributedDecl.cpp

@@ -207,6 +207,23 @@ Type ASTContext::getAssociatedTypeOfDistributedSystemOfActor(
     return Type();
 
   auto module = actor->getParentModule();
+
+  // In case of protocol, let's find a concrete `ActorSystem`
+  if (auto *protocol = dyn_cast<ProtocolDecl>(actor)) {
+    auto signature = protocol->getGenericSignatureOfContext();
+
+    auto systemTy =
+        signature->getConcreteType(actorSystemDecl->getDeclaredInterfaceType());
+    if (!systemTy)
+      return Type();
+
+    auto conformance = module->lookupConformance(systemTy, actorSystemProtocol);
+    if (conformance.isInvalid())
+      return Type();
+
+    return conformance.getTypeWitnessByName(systemTy, member);
+  }
+
   Type selfType = actor->getSelfInterfaceType();
   auto conformance = module->lookupConformance(selfType, actorProtocol);
   Type dependentType = actorProtocol->getSelfInterfaceType();

lib/SILGen/SILGenDistributed.cpp

@@ -118,7 +118,9 @@ static void emitDistributedIfRemoteBranch(SILGenFunction &SGF, SILLocation Loc,
                        "'Distributed' module available?");
 
   ManagedValue selfAnyObject = B.createInitExistentialRef(
-      Loc, SGF.getLoweredType(ctx.getAnyObjectType()), CanType(selfTy),
+      Loc,
+      /*existentialType=*/SGF.getLoweredType(ctx.getAnyObjectType()),
+      /*formalConcreteType=*/selfValue.getType().getASTType(),
       selfValue, {});
   auto result = SGF.emitApplyOfLibraryIntrinsic(
       Loc, isRemoteFn, SubstitutionMap(), {selfAnyObject}, SGFContext());
@@ -379,6 +381,7 @@ void SILGenFunction::emitDistributedActorFactory(FuncDecl *fd) { // TODO(distrib
   ///       having at least one call to the resolve function.
 
   auto &C = getASTContext();
+  auto DC = fd->getDeclContext();
   SILLocation loc = fd;
 
   // ==== Prepare argument references
@@ -392,12 +395,12 @@ void SILGenFunction::emitDistributedActorFactory(FuncDecl *fd) { // TODO(distrib
   ManagedValue selfArg = ManagedValue::forUnmanaged(selfArgValue);
 
   // type: SpecificDistributedActor.Type
-  auto selfArgType = F.mapTypeIntoContext(selfArg.getType().getASTType());
+  auto selfArgType = selfArg.getType().getASTType();
   auto selfMetatype = getLoweredType(selfArgType);
   SILValue selfMetatypeValue = B.createMetatype(loc, selfMetatype);
 
   // type: SpecificDistributedActor
-  auto *selfTyDecl = fd->getParent()->getSelfNominalTypeDecl();
+  auto *selfTyDecl = DC->getSelfClassDecl();
   assert(selfTyDecl->isDistributedActor());
   auto selfTy = F.mapTypeIntoContext(selfTyDecl->getDeclaredInterfaceType());
   auto returnTy = getLoweredType(selfTy);
@@ -533,8 +536,8 @@ SILGenFunction::emitConditionalResignIdentityCall(SILLocation loc,
   assert(actorDecl->isDistributedActor() &&
   "only distributed actors have actorSystem lifecycle hooks in deinit");
 
-  auto selfTy = actorDecl->getDeclaredInterfaceType();
-  
+  auto selfTy = F.mapTypeIntoContext(actorDecl->getDeclaredInterfaceType());
+
   // we only system.resignID if we are a local actor,
   // and thus the address was created by system.assignID.
   auto isRemoteBB = createBasicBlock("isRemoteBB");

lib/SILOptimizer/Utils/DistributedActor.cpp

@@ -183,7 +183,7 @@ void emitActorReadyCall(SILBuilder &B, SILLocation loc, SILValue actor,
   auto &C = F.getASTContext();
   emitDistributedActorSystemWitnessCall(
       B, loc, C.Id_actorReady, actorSystem,
-      F.mapTypeIntoContext(actor->getType()), { actor });
+      actor->getType(), { actor });
 }
 
 void emitResignIdentityCall(SILBuilder &B, SILLocation loc,

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -55,7 +55,7 @@ static VarDecl *addImplicitDistributedActorIDProperty(
   // ==== Synthesize and add 'id' property to the actor decl
   Type propertyType = getDistributedActorIDType(nominal);
 
-  VarDecl *propDecl = new (C)
+  auto *propDecl = new (C)
       VarDecl(/*IsStatic*/false, VarDecl::Introducer::Let,
               SourceLoc(), C.Id_id, nominal);
   propDecl->setImplicit();
@@ -230,20 +230,21 @@ deriveBodyDistributed_thunk(AbstractFunctionDecl *thunk, void *context) {
 
   // --- Recording invocation details
   // -- recordGenericSubstitution(s)
-  if (func->isGeneric() || nominal->isGeneric()) {
+  if (thunk->isGeneric() || nominal->isGeneric()) {
     auto recordGenericSubstitutionDecl =
         C.getRecordGenericSubstitutionOnDistributedInvocationEncoder(invocationEncoderDecl);
     assert(recordGenericSubstitutionDecl);
     auto recordGenericSubstitutionDeclRef =
         UnresolvedDeclRefExpr::createImplicit(
             C, recordGenericSubstitutionDecl->getName());
 
-    auto sig = func->getGenericSignature();
-    for (auto genParamType : sig.getGenericParams()) {
+    auto signature = thunk->getGenericSignature();
+    for (auto genParamType : signature.getGenericParams()) {
 
+      auto tyExpr = TypeExpr::createImplicit(thunk->mapTypeIntoContext(genParamType), C);
       auto subTypeExpr = new (C) DotSelfExpr(
-          TypeExpr::createImplicit(thunk->mapTypeIntoContext(genParamType), C),
-          sloc, sloc, thunk->mapTypeIntoContext(genParamType));
+          tyExpr,
+          sloc, sloc, tyExpr->getType());
 
       auto recordGenericSubArgsList =
           ArgumentList::forImplicitCallTo(
@@ -378,7 +379,8 @@ deriveBodyDistributed_thunk(AbstractFunctionDecl *thunk, void *context) {
   // -- recordReturnType
   if (!isVoidReturn) {
     // Result.self
-    auto resultType = func->getResultInterfaceType();
+    // Watch out and always map into thunk context
+    auto resultType = thunk->mapTypeIntoContext(func->getResultInterfaceType());
     auto *metaTypeRef = TypeExpr::createImplicit(resultType, C);
     auto *resultTypeExpr =
         new (C) DotSelfExpr(metaTypeRef, sloc, sloc, resultType);
@@ -426,9 +428,8 @@ deriveBodyDistributed_thunk(AbstractFunctionDecl *thunk, void *context) {
   }
 
   // === Prepare the 'RemoteCallTarget'
-  VarDecl *targetVar =
-      new (C) VarDecl(/*isStatic=*/false, VarDecl::Introducer::Let, sloc,
-                      C.Id_target, thunk);
+  auto *targetVar = new (C) VarDecl(
+      /*isStatic=*/false, VarDecl::Introducer::Let, sloc, C.Id_target, thunk);
 
   {
     // --- Mangle the thunk name
@@ -510,7 +511,8 @@ deriveBodyDistributed_thunk(AbstractFunctionDecl *thunk, void *context) {
     // -- returning: Res.Type
     if (!isVoidReturn) {
       // Result.self
-      auto resultType = func->getResultInterfaceType();
+      auto resultType =
+          func->mapTypeIntoContext(func->getResultInterfaceType());
       auto *metaTypeRef = TypeExpr::createImplicit(resultType, C);
       auto *resultTypeExpr =
           new (C) DotSelfExpr(metaTypeRef, sloc, sloc, resultType);
@@ -561,7 +563,7 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
     genericParamList = genericParams->clone(DC);
   }
 
-  GenericSignature thunkGenSig =
+  GenericSignature baseSignature =
       buildGenericSignature(C, func->getGenericSignature(),
                             /*addedParameters=*/{},
                             /*addedRequirements=*/{});
@@ -581,9 +583,9 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
     }
 
     auto paramDecl = new (C)
-        ParamDecl(SourceLoc(),
-                  /*argumentNameLoc=*/SourceLoc(), funcParam->getArgumentName(),
-                  /*parameterNameLoc=*/SourceLoc(), paramName, DC);
+         ParamDecl(SourceLoc(),
+                   /*argumentNameLoc=*/SourceLoc(), funcParam->getArgumentName(),
+                   /*parameterNameLoc=*/SourceLoc(), paramName, DC);
 
     paramDecl->setImplicit(true);
     paramDecl->setSpecifier(funcParam->getSpecifier());
@@ -593,16 +595,14 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
   }
   ParameterList *params = ParameterList::create(C, paramDecls); // = funcParams->clone(C);
 
-    auto thunk = FuncDecl::createImplicit(C, swift::StaticSpellingKind::None,
-                           thunkName, SourceLoc(),
-                           /*async=*/true, /*throws=*/true,
-                           genericParamList,
-                           params,
-                           func->getResultInterfaceType(),
-                           DC);
+  auto thunk = FuncDecl::createImplicit(
+      C, swift::StaticSpellingKind::None, thunkName, SourceLoc(),
+      /*async=*/true, /*throws=*/true,
+      genericParamList, params,
+      func->getResultInterfaceType(), DC);
   thunk->setSynthesized(true);
-  thunk->getAttrs().add(new (C) NonisolatedAttr(/*implicit=*/true));
-  thunk->setGenericSignature(thunkGenSig);
+  thunk->getAttrs().add(new (C) NonisolatedAttr(/*isImplicit=*/true));
+  thunk->setGenericSignature(baseSignature);
   thunk->copyFormalAccessFrom(func, /*sourceIsParentContext=*/false);
   thunk->setBodySynthesizer(deriveBodyDistributed_thunk, func);
 

test/Distributed/Runtime/distributed_actor_func_calls_remoteCall_genericFunc.swift

@@ -32,6 +32,12 @@ distributed actor Greeter {
 
 }
 
+distributed actor GenericGreeter<Greeting: Sendable & Codable> {
+  distributed func echo(greeting: Greeting) -> Greeting {
+    greeting
+  }
+}
+
 func test() async throws {
   let system = DefaultDistributedActorSystem()
 
@@ -62,6 +68,18 @@ func test() async throws {
   // CHECK: >> remoteCall: on:main.Greeter, target:main.Greeter.generic2(strict:_:_:), invocation:FakeInvocationEncoder(genericSubs: [Swift.String, Swift.Int], arguments: [2.0, "Caplin", [1, 2, 3]], returnType: Optional(Swift.String), errorType: nil), throwing:Swift.Never, returning:Swift.String
   print("reply: \(r2)")
   // CHECK: reply: Caplin
+
+  let gen = GenericGreeter<String>(actorSystem: system)
+  let r3 = try await gen.echo(greeting: "Hello generics!")
+  print("reply: \(r3)")
+  // CHECK: reply: Hello generics!
+
+  let genRef = try GenericGreeter<String>.resolve(id: gen.id, using: system)
+  let r32 = try await genRef.echo(greeting: "Hello generics!")
+  // CHECK: > encode generic sub: Swift.String
+  // CHECK: > encode return type: Swift.String
+  print("reply: \(r32)")
+  // CHECK: reply: Hello generics!
 }
 
 @main struct Main {

test/Distributed/Runtime/distributed_actor_hop_to.swift

@@ -24,14 +24,14 @@ protocol LifecycleWatch: DistributedActor where ActorSystem == FakeRoundtripActo
 }
 
 extension LifecycleWatch {
-  func watch() async throws {
+  func watch<T: Codable>(x: Int, _ y: T) async throws {
     // nothing here
-    print("executed: \(#function)")
+    print("executed: \(#function) - x = \(x), y = \(y)")
   }
 
-  distributed func test() async throws {
+  distributed func test<T: Codable>(x: Int, _ y: T) async throws {
     print("executed: \(#function)")
-    try await self.watch()
+    try await self.watch(x: x, y)
     print("done executed: \(#function)")
   }
 }
@@ -42,11 +42,11 @@ distributed actor Worker: LifecycleWatch {
 @main struct Main {
   static func main() async {
     let worker: any LifecycleWatch = Worker(actorSystem: DefaultDistributedActorSystem())
-    try! await worker.test()
+    try! await worker.test(x: 42, "on protocol")
 
-    // CHECK: executed: test()
-    // CHECK: executed: watch()
-    // CHECK: done executed: test()
+    // CHECK: executed: test(x:_:)
+    // CHECK: executed: watch(x:_:) - x = 42, y = on protocol
+    // CHECK: done executed: test(x:_:)
 
     print("OK") // CHECK: OK
   }

test/Distributed/Runtime/distributed_actor_remote_retains_system.swift

@@ -130,6 +130,7 @@ public struct FakeResultHandler: DistributedTargetInvocationResultHandler {
     fatalError("Not implemented: \(#function)")
   }
 }
+
 typealias DefaultDistributedActorSystem = FakeActorSystem
 
 // ==== Execute ----------------------------------------------------------------

test/Distributed/distributed_actor_generic_actor.swift

@@ -0,0 +1,28 @@
+// 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-swift-frontend -typecheck -verify -enable-experimental-distributed -disable-availability-checking -I %t 2>&1 %s
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+import Distributed
+import FakeDistributedActorSystems
+
+distributed actor Worker<Work: Sendable & Codable> {
+  typealias ActorSystem = FakeActorSystem
+
+  distributed func echo(item: Work) -> Work {
+    item
+  }
+
+  distributed func echo(items: [Work]) -> [Work] {
+    items
+  }
+
+  distributed func other<Other: Codable & Sendable>(other: Other) -> Other {
+    other
+  }
+
+  distributed func others<Other: Codable & Sendable>(other: Other) -> [Other] {
+    [other]
+  }
+}