[IRGen] Distributed: Dispatch decodeNextArgument via witness thunk if concrete decoder is not available

Since distributed actors can now be generic over actor system
it's no longer guaranteed for concrete decoder to be present,
we need to handle that case specifically during accessor emission.

Author: xedin

lib/IRGen/GenDistributed.cpp

@@ -81,44 +81,48 @@ struct ArgumentDecoderInfo {
   /// The type of `decodeNextArgument` method.
   CanSILFunctionType MethodType;
 
-  /// Protocol requirements associated with the generic
-  /// parameter `Argument` of this decode method.
-  GenericSignature::RequiredProtocols ProtocolRequirements;
-
-  // Witness metadata for conformance to DistributedTargetInvocationDecoder
-  // protocol.
+  /// Witness metadata for conformance to DistributedTargetInvocationDecoder
+  /// protocol.
   WitnessMetadata Witness;
 
+  /// Indicates whether `decodeNextArgument` is referenced through
+  /// a protocol witness thunk.
+  bool UsesWitnessDispatch;
+
   ArgumentDecoderInfo(llvm::Value *decoder, llvm::Value *decoderType,
                       llvm::Value *decoderWitnessTable,
                       FunctionPointer decodeNextArgumentPtr,
-                      CanSILFunctionType decodeNextArgumentTy)
+                      CanSILFunctionType decodeNextArgumentTy,
+                      bool usesWitnessDispatch)
       : Decoder(decoder), MethodPtr(decodeNextArgumentPtr),
         MethodType(decodeNextArgumentTy),
-        ProtocolRequirements(findProtocolRequirements(decodeNextArgumentTy)) {
+        UsesWitnessDispatch(usesWitnessDispatch) {
     Witness.SelfMetadata = decoderType;
     Witness.SelfWitnessTable = decoderWitnessTable;
   }
 
   CanSILFunctionType getMethodType() const { return MethodType; }
 
-  ArrayRef<ProtocolDecl *> getProtocolRequirements() const {
-    return ProtocolRequirements;
+  WitnessMetadata *getWitnessMetadata() const {
+    return const_cast<WitnessMetadata *>(&Witness);
   }
 
-  /// Form a callee to a decode method - `decodeNextArgument`.
-  Callee getCallee() const;
+  /// Protocol requirements associated with the generic
+  /// parameter `Argument` of this decode method.
+  GenericSignature::RequiredProtocols getProtocolRequirements() const {
+    if (UsesWitnessDispatch)
+      return {};
 
-private:
-  static GenericSignature::RequiredProtocols
-  findProtocolRequirements(CanSILFunctionType decodeMethodTy) {
-    auto signature = decodeMethodTy->getInvocationGenericSignature();
+    auto signature = MethodType->getInvocationGenericSignature();
     auto genericParams = signature.getGenericParams();
 
     // func decodeNextArgument<Arg : #SerializationRequirement#>() throws -> Arg
     assert(genericParams.size() == 1);
     return signature->getRequiredProtocols(genericParams.front());
   }
+
+  /// Form a callee to a decode method - `decodeNextArgument`.
+  Callee getCallee() const;
 };
 
 class DistributedAccessor {
@@ -415,7 +419,8 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
   emission->begin();
   {
     emission->setArgs(decodeArgs, /*isOutlined=*/false,
-                      /*witnessMetadata=*/nullptr);
+                      decoder.UsesWitnessDispatch ? decoder.getWitnessMetadata()
+                                                  : nullptr);
 
     Explosion result;
     emission->emitToExplosion(result, /*isOutlined=*/false);
@@ -519,6 +524,9 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
 void DistributedAccessor::lookupWitnessTables(
     llvm::Value *value, ArrayRef<ProtocolDecl *> protocols,
     Explosion &witnessTables) {
+  if (protocols.empty())
+    return;
+
   auto conformsToProtocol = IGM.getConformsToProtocolFunctionPointer();
 
   for (auto *protocol : protocols) {
@@ -785,11 +793,32 @@ DistributedAccessor::getCalleeForDistributedTarget(llvm::Value *self) const {
 
 ArgumentDecoderInfo DistributedAccessor::findArgumentDecoder(
     llvm::Value *decoder, llvm::Value *decoderTy, llvm::Value *witnessTable) {
+  auto &C = IGM.Context;
   auto *actor = getDistributedActorOf(Target);
   auto expansionContext = IGM.getMaximalTypeExpansionContext();
 
-  auto *decodeFn = IGM.Context.getDistributedActorArgumentDecodingMethod(actor);
-  assert(decodeFn && "no suitable decoder?");
+  auto *decodeFn = C.getDistributedActorArgumentDecodingMethod(actor);
+
+  // If distributed actor is generic over actor system, we have to
+  // use witness to reference `decodeNextArgument`.
+  if (!decodeFn) {
+    auto decoderProtocol = C.getDistributedTargetInvocationDecoderDecl();
+    auto decodeNextArgRequirement =
+        decoderProtocol->getSingleRequirement(C.Id_decodeNextArgument);
+    assert(decodeNextArgRequirement);
+    SILDeclRef decodeNextArgumentRef(decodeNextArgRequirement);
+
+    llvm::Constant *fnPtr =
+        IGM.getAddrOfDispatchThunk(decodeNextArgumentRef, NotForDefinition);
+    auto fnType = IGM.getSILTypes().getConstantFunctionType(
+        IGM.getMaximalTypeExpansionContext(), decodeNextArgumentRef);
+
+    auto sig = IGM.getSignature(fnType);
+    auto fn = FunctionPointer::forDirect(fnType, fnPtr,
+                                         /*secondaryValue=*/nullptr, sig, true);
+    return {decoder, decoderTy, witnessTable,
+            fn,      fnType,    /*usesWitnessDispatch=*/true};
+  }
 
   auto methodTy = IGM.getSILTypes().getConstantFunctionType(
       expansionContext, SILDeclRef(decodeFn));
@@ -848,7 +877,8 @@ ArgumentDecoderInfo DistributedAccessor::findArgumentDecoder(
         /*secondaryValue=*/nullptr, signature);
   }
 
-  return {decoder, decoderTy, witnessTable, methodPtr, methodTy};
+  return {decoder,   decoderTy, witnessTable,
+          methodPtr, methodTy,  /*usesWitnessDispatch=*/false};
 }
 
 SILType DistributedAccessor::getResultType() const {

lib/Sema/TypeCheckDistributed.cpp

@@ -869,7 +869,10 @@ GetDistributedActorArgumentDecodingMethodRequest::evaluate(Evaluator &evaluator,
   auto &ctx = actor->getASTContext();
 
   auto *decoder = ctx.getDistributedActorInvocationDecoder(actor);
-  assert(decoder);
+  // If distributed actor is generic over actor system, there is not
+  // going to be a concrete decoder.
+  if (!decoder)
+    return nullptr;
 
   auto decoderTy = decoder->getDeclaredInterfaceType();
 
@@ -942,4 +945,4 @@ GetDistributedMethodWitnessedProtocolRequirements::evaluate(
   }
 
   return C.AllocateCopy(result);
-}
+}

test/Distributed/Runtime/distributed_actor_localSystem_generic.swift

@@ -0,0 +1,38 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -module-name main -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s -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
+
+// rdar://90373022
+// UNSUPPORTED: OS=watchos
+
+import Distributed
+
+distributed actor Worker<ActorSystem> where ActorSystem: DistributedActorSystem<any Codable>, ActorSystem.ActorID: Codable {
+  distributed func hi(name: String) {
+    print("Hi, \(name)!")
+  }
+
+  nonisolated var description: Swift.String {
+    "Worker(\(id))"
+  }
+}
+
+// ==== Execute ----------------------------------------------------------------
+@main struct Main {
+  static func main() async throws {
+    let system = LocalTestingDistributedActorSystem()
+
+    let actor = Worker(actorSystem: system)
+    try await actor.hi(name: "P") // local calls should still just work
+    // CHECK: Hi, P!
+  }
+}