[Distributed] AST: Make argument decoding method discoverable given distributed actor

This is going to be used by IRGen later on to make it possible to
decode arguments inside of distributed target accessor instead of
allocating additional storage.

Author: xedin

include/swift/AST/ASTContext.h

@@ -1349,6 +1349,10 @@ class ASTContext final {
   NominalTypeDecl *
   getDistributedActorInvocationDecoder(NominalTypeDecl *);
 
+  /// Find `decodeNextArgument<T>(type: T.Type) -> T` method associated with
+  /// invocation decoder of the given distributed actor.
+  FuncDecl *getDistributedActorArgumentDecodingMethod(NominalTypeDecl *);
+
 private:
   friend Decl;
 

include/swift/AST/KnownIdentifiers.def

@@ -291,6 +291,7 @@ IDENTIFIER(throwing)
 IDENTIFIER(using)
 IDENTIFIER(InvocationDecoder)
 IDENTIFIER(whenLocal)
+IDENTIFIER(decodeNextArgument)
 
 #undef IDENTIFIER
 #undef IDENTIFIER_

include/swift/AST/TypeCheckRequests.h

@@ -1109,6 +1109,25 @@ class GetDistributedActorInvocationDecoderRequest :
   bool isCached() const { return true; }
 };
 
+/// Obtain the method that could be used to decode argument values passed
+/// to a particular actor invocation type.
+class GetDistributedActorArgumentDecodingMethodRequest :
+  public SimpleRequest<GetDistributedActorArgumentDecodingMethodRequest,
+                       FuncDecl *(NominalTypeDecl *),
+                       RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  FuncDecl *evaluate(Evaluator &evaluator, NominalTypeDecl *actor) const;
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+};
+
 /// Retrieve the static "shared" property within a global actor that provides
 /// the actor instance representing the global actor.
 ///

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -115,6 +115,9 @@ SWIFT_REQUEST(TypeChecker, GetDistributedActorSystemPropertyRequest, VarDecl *(N
 SWIFT_REQUEST(TypeChecker, GetDistributedActorInvocationDecoderRequest,
               NominalTypeDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, GetDistributedActorArgumentDecodingMethodRequest,
+              FuncDecl *(NominalTypeDecl *),
+              Cache, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GlobalActorInstanceRequest,
               VarDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)

lib/Sema/TypeCheckDistributed.cpp

@@ -378,22 +378,8 @@ void TypeChecker::checkDistributedActor(ClassDecl *decl) {
   (void)decl->getDistributedActorIDProperty();
 }
 
-Type swift::getDistributedActorSystemType(NominalTypeDecl *actor) {
-  assert(actor->isDistributedActor());
-  auto &ctx = actor->getASTContext();
-
-  auto protocol = ctx.getProtocol(KnownProtocolKind::DistributedActor);
-  if (!protocol)
-    return ErrorType::get(ctx);
-
-  // Dig out the actor system type.
-  auto module = actor->getParentModule();
-  Type selfType = actor->getSelfInterfaceType();
-  auto conformance = module->lookupConformance(selfType, protocol);
-  return conformance.getTypeWitnessByName(selfType, ctx.Id_ActorSystem);
-}
-
-Type swift::getDistributedActorIDType(NominalTypeDecl *actor) {
+static Type getAssociatedTypeOfDistributedSystem(NominalTypeDecl *actor,
+                                                 Identifier member) {
   assert(actor->isDistributedActor());
   auto &ctx = actor->getASTContext();
 
@@ -410,22 +396,40 @@ Type swift::getDistributedActorIDType(NominalTypeDecl *actor) {
   if (!actorSystemProtocol)
     return ErrorType::get(ctx);
 
-  AssociatedTypeDecl *idAssocTypeDecl =
-      actorSystemProtocol->getAssociatedType(ctx.Id_ActorID);
-  if (!idAssocTypeDecl)
+  AssociatedTypeDecl *assocTypeDecl =
+      actorSystemProtocol->getAssociatedType(member);
+  if (!assocTypeDecl)
     return ErrorType::get(ctx);
 
   auto module = actor->getParentModule();
   Type selfType = actor->getSelfInterfaceType();
   auto conformance = module->lookupConformance(selfType, actorProtocol);
   Type dependentType = actorProtocol->getSelfInterfaceType();
   dependentType = DependentMemberType::get(dependentType, actorSystemDecl);
-  dependentType = DependentMemberType::get(dependentType, idAssocTypeDecl);
-  auto t = dependentType.subst(
-      SubstitutionMap::getProtocolSubstitutions(
-        actorProtocol, selfType, conformance));
+  dependentType = DependentMemberType::get(dependentType, assocTypeDecl);
 
-  return t;
+  return dependentType.subst(SubstitutionMap::getProtocolSubstitutions(
+      actorProtocol, selfType, conformance));
+}
+
+Type swift::getDistributedActorSystemType(NominalTypeDecl *actor) {
+  assert(actor->isDistributedActor());
+  auto &ctx = actor->getASTContext();
+
+  auto protocol = ctx.getProtocol(KnownProtocolKind::DistributedActor);
+  if (!protocol)
+    return ErrorType::get(ctx);
+
+  // Dig out the actor system type.
+  auto module = actor->getParentModule();
+  Type selfType = actor->getSelfInterfaceType();
+  auto conformance = module->lookupConformance(selfType, protocol);
+  return conformance.getTypeWitnessByName(selfType, ctx.Id_ActorSystem);
+}
+
+Type swift::getDistributedActorIDType(NominalTypeDecl *actor) {
+  auto &ctx = actor->getASTContext();
+  return getAssociatedTypeOfDistributedSystem(actor, ctx.Id_ActorID);
 }
 
 NominalTypeDecl *
@@ -440,36 +444,35 @@ ASTContext::getDistributedActorInvocationDecoder(NominalTypeDecl *actor) {
 NominalTypeDecl *
 GetDistributedActorInvocationDecoderRequest::evaluate(Evaluator &evaluator,
                                                       NominalTypeDecl *actor) const {
-    assert(actor->isDistributedActor());
   auto &ctx = actor->getASTContext();
+  auto decoderTy =
+      getAssociatedTypeOfDistributedSystem(actor, ctx.Id_InvocationDecoder);
+  return decoderTy->hasError() ? nullptr : decoderTy->getAnyNominal();
+}
 
-  auto actorProtocol = ctx.getProtocol(KnownProtocolKind::DistributedActor);
-  if (!actorProtocol)
+FuncDecl *ASTContext::getDistributedActorArgumentDecodingMethod(NominalTypeDecl *actor) {
+  if (!actor->isDistributedActor())
     return nullptr;
 
-  AssociatedTypeDecl *actorSystemDecl =
-      actorProtocol->getAssociatedType(ctx.Id_ActorSystem);
-  if (!actorSystemDecl)
-    return nullptr;
+  return evaluateOrDefault(
+    evaluator, GetDistributedActorArgumentDecodingMethodRequest{actor}, nullptr);
+}
 
-  auto actorSystemProtocol = ctx.getProtocol(KnownProtocolKind::DistributedActorSystem);
-  if (!actorSystemProtocol)
-    return nullptr;
 
-  AssociatedTypeDecl *decoderAssocTypeDecl =
-      actorSystemProtocol->getAssociatedType(ctx.Id_InvocationDecoder);
-  if (!decoderAssocTypeDecl)
-    return nullptr;
+FuncDecl *
+GetDistributedActorArgumentDecodingMethodRequest::evaluate(Evaluator &evaluator,
+                                                           NominalTypeDecl *actor) const {
+  auto &ctx = actor->getASTContext();
 
-  auto module = actor->getParentModule();
-  Type selfType = actor->getSelfInterfaceType();
-  auto conformance = module->lookupConformance(selfType, actorProtocol);
-  Type dependentType = actorProtocol->getSelfInterfaceType();
-  dependentType = DependentMemberType::get(dependentType, actorSystemDecl);
-  dependentType = DependentMemberType::get(dependentType, decoderAssocTypeDecl);
-  auto t = dependentType.subst(
-      SubstitutionMap::getProtocolSubstitutions(
-        actorProtocol, selfType, conformance));
+  auto *decoder = ctx.getDistributedActorInvocationDecoder(actor);
+  assert(decoder);
+
+  auto decoderTy = decoder->getInterfaceType()->getMetatypeInstanceType();
+
+  DeclName methodName(ctx, {ctx.Id_decodeNextArgument}, {ctx.Id_type});
+  auto members = TypeChecker::lookupMember(actor->getDeclContext(), decoderTy,
+                                           DeclNameRef(methodName));
 
-  return t->getAnyNominal();
+  assert(members.size() == 1);
+  return cast<FuncDecl>(members[0].getValueDecl());
 }