[AST/Sema] Distributed: Introduced unified way to retrieve serialization requirements for actors

xedin · xedin · commit 7b7716cda93a · 2024-03-06T13:51:14.000-08:00
diff --git a/include/swift/AST/DistributedDecl.h b/include/swift/AST/DistributedDecl.h
@@ -31,7 +31,7 @@ class DeclContext;
 class FuncDecl;
 class NominalTypeDecl;
 
-Type getAssociatedTypeOfDistributedSystemOfActor(NominalTypeDecl *actor,
+Type getAssociatedTypeOfDistributedSystemOfActor(DeclContext *actorOrExtension,
                                                  Identifier member);
 
 /// Find the concrete invocation decoder associated with the given actor.
@@ -57,12 +57,6 @@ Type getDistributedActorSystemType(NominalTypeDecl *actor);
 /// Determine the `ID` type for the given actor.
 Type getDistributedActorIDType(NominalTypeDecl *actor);
 
-/// Similar to `getDistributedSerializationRequirementType`, however, from the
-/// perspective of a concrete function. This way we're able to get the
-/// serialization requirement for specific members, also in protocols.
-Type getSerializationRequirementTypesForMember(
-    ValueDecl *member, llvm::SmallPtrSet<ProtocolDecl *, 2> &serializationRequirements);
-
 /// Get specific 'SerializationRequirement' as defined in 'nominal'
 /// type, which must conform to the passed 'protocol' which is expected
 /// to require the 'SerializationRequirement'.
@@ -76,6 +70,12 @@ AbstractFunctionDecl *
 getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
     ValueDecl *thunk);
 
+Type getDistributedActorSerializationType(DeclContext *actorOrExtension);
+
+/// Get the specific 'SerializationRequirement' type of a specific distributed
+/// actor system.
+Type getDistributedActorSystemSerializationType(NominalTypeDecl *system);
+
 /// Get the specific 'InvocationEncoder' type of a specific distributed actor
 /// system.
 Type getDistributedActorSystemInvocationEncoderType(NominalTypeDecl *system);
@@ -91,17 +91,6 @@ Type getDistributedActorSystemResultHandlerType(NominalTypeDecl *system);
 /// Get the 'ActorID' type of a specific distributed actor system.
 Type getDistributedActorSystemActorIDType(NominalTypeDecl *system);
 
-/// Get the specific protocols that the `SerializationRequirement` specifies,
-/// and all parameters / return types of distributed targets must conform to.
-///
-/// E.g. if a system declares `typealias SerializationRequirement = Codable`
-/// then this will return `{encodableProtocol, decodableProtocol}`.
-///
-/// Returns an empty set if the requirement was `Any`.
-llvm::SmallPtrSet<ProtocolDecl *, 2>
-getDistributedSerializationRequirementProtocols(
-    NominalTypeDecl *decl, ProtocolDecl* protocol);
-
 /// Check if the `allRequirements` represent *exactly* the
 /// `Encodable & Decodable` (also known as `Codable`) requirement.
 ///
diff --git a/lib/AST/DistributedDecl.cpp b/lib/AST/DistributedDecl.cpp
@@ -106,42 +106,6 @@ Type swift::getConcreteReplacementForProtocolActorSystemType(
   llvm_unreachable("Unable to fetch ActorSystem type!");
 }
 
-Type swift::getSerializationRequirementTypesForMember(
-    ValueDecl *member,
-    llvm::SmallPtrSet<ProtocolDecl *, 2> &serializationRequirements) {
-  auto &C = member->getASTContext();
-  auto *DC = member->getDeclContext();
-  auto DA = C.getDistributedActorDecl();
-
-  // === When declared inside an actor, we can get the type directly
-  if (auto classDecl = DC->getSelfClassDecl()) {
-    return getDistributedSerializationRequirementType(classDecl, C.getDistributedActorDecl());
-  }
-
-  auto SerReqAssocType = DA->getAssociatedType(C.Id_SerializationRequirement)
-      ->getDeclaredInterfaceType();
-
-  if (DC->getSelfProtocolDecl()) {
-    GenericSignature signature;
-    if (auto *genericContext = member->getAsGenericContext()) {
-      signature = genericContext->getGenericSignature();
-    } else {
-      signature = DC->getGenericSignatureOfContext();
-    }
-
-    // Also store all `SerializationRequirement : SomeProtocol` requirements
-    for (auto proto: signature->getRequiredProtocols(SerReqAssocType)) {
-      serializationRequirements.insert(proto);
-    }
-
-    // Note that this may be null, e.g. if we're a distributed func inside
-    // a protocol that did not declare a specific actor system requirement.
-    return signature->getConcreteType(SerReqAssocType);
-  }
-
-  llvm_unreachable("Unable to fetch SerializationRequirement type!");
-}
-
 Type swift::getDistributedActorSystemType(NominalTypeDecl *actor) {
   assert(!dyn_cast<ProtocolDecl>(actor) &&
          "Use getConcreteReplacementForProtocolActorSystemType instead to get"
@@ -179,6 +143,53 @@ static Type getTypeWitnessByName(NominalTypeDecl *type, ProtocolDecl *protocol,
   return conformance.getTypeWitnessByName(selfType, member);
 }
 
+Type swift::getDistributedActorSerializationType(
+    DeclContext *actorOrExtension) {
+  auto &ctx = actorOrExtension->getASTContext();
+  auto resultTy = getAssociatedTypeOfDistributedSystemOfActor(
+      actorOrExtension,
+      ctx.Id_SerializationRequirement);
+
+  // Protocols are allowed to either not provide a `SerializationRequirement`
+  // at all or provide it in a conformance requirement.
+  if ((!resultTy || resultTy->hasDependentMember()) &&
+      actorOrExtension->getSelfProtocolDecl()) {
+    auto sig = actorOrExtension->getGenericSignatureOfContext();
+
+    auto actorProtocol = ctx.getProtocol(KnownProtocolKind::DistributedActor);
+    if (!actorProtocol)
+      return Type();
+
+    auto serializationTy =
+        actorProtocol->getAssociatedType(ctx.Id_SerializationRequirement)
+            ->getDeclaredInterfaceType();
+
+    auto protocols = sig->getRequiredProtocols(serializationTy);
+    if (protocols.empty())
+      return Type();
+
+    SmallVector<Type, 2> members;
+    llvm::transform(protocols, std::back_inserter(members), [](const auto *P) {
+      return P->getDeclaredInterfaceType();
+    });
+
+    return ExistentialType::get(
+        ProtocolCompositionType::get(ctx, members,
+                                     /*inverses=*/{},
+                                     /*HasExplicitAnyObject=*/false));
+  }
+
+  return resultTy;
+}
+
+Type swift::getDistributedActorSystemSerializationType(
+    NominalTypeDecl *system) {
+  assert(!system->isDistributedActor());
+  auto &ctx = system->getASTContext();
+  return getTypeWitnessByName(system, ctx.getDistributedActorSystemDecl(),
+                              ctx.Id_SerializationRequirement);
+}
+
 Type swift::getDistributedActorSystemActorIDType(NominalTypeDecl *system) {
   assert(!system->isDistributedActor());
   auto &ctx = system->getASTContext();
@@ -248,19 +259,14 @@ swift::getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
       decoderTy->getAnyNominal());
 }
 
-Type swift::getAssociatedTypeOfDistributedSystemOfActor(NominalTypeDecl *actor,
-                                                        Identifier member) {
-  auto &ctx = actor->getASTContext();
+Type swift::getAssociatedTypeOfDistributedSystemOfActor(
+    DeclContext *actorOrExtension, Identifier member) {
+  auto &ctx = actorOrExtension->getASTContext();
 
   auto actorProtocol = ctx.getProtocol(KnownProtocolKind::DistributedActor);
   if (!actorProtocol)
     return Type();
 
-  auto actorConformance = actor->getParentModule()->lookupConformance(
-      actor->getDeclaredInterfaceType(), actorProtocol);
-  if (!actorConformance || actorConformance.isInvalid())
-    return Type();
-
   AssociatedTypeDecl *actorSystemDecl =
       actorProtocol->getAssociatedType(ctx.Id_ActorSystem);
   if (!actorSystemDecl)
@@ -275,15 +281,27 @@ Type swift::getAssociatedTypeOfDistributedSystemOfActor(NominalTypeDecl *actor,
   if (!memberTypeDecl)
     return Type();
 
-  auto depMemTy = DependentMemberType::get(
-      DependentMemberType::get(actorProtocol->getSelfInterfaceType(),
-                               actorSystemDecl),
-      memberTypeDecl);
+  Type memberTy = DependentMemberType::get(
+    DependentMemberType::get(actorProtocol->getSelfInterfaceType(),
+                             actorSystemDecl),
+    memberTypeDecl);
+
+  auto sig = actorOrExtension->getGenericSignatureOfContext();
+
+  auto *actorType = actorOrExtension->getSelfNominalTypeDecl();
+  if (isa<ProtocolDecl>(actorType))
+    return memberTy->getReducedType(sig);
+
+  auto actorConformance =
+      actorOrExtension->getParentModule()->lookupConformance(
+          actorType->getDeclaredInterfaceType(), actorProtocol);
+  if (!actorConformance || actorConformance.isInvalid())
+    return Type();
 
   auto subs = SubstitutionMap::getProtocolSubstitutions(
-      actorProtocol, actor->getDeclaredInterfaceType(), actorConformance);
+      actorProtocol, actorType->getDeclaredInterfaceType(), actorConformance);
 
-  return Type(depMemTy).subst(subs);
+  return memberTy.subst(subs)->getReducedType(sig);
 }
 
 /******************************************************************************/
diff --git a/lib/Sema/CodeSynthesisDistributedActor.cpp b/lib/Sema/CodeSynthesisDistributedActor.cpp
@@ -878,13 +878,9 @@ static bool canSynthesizeDistributedThunk(AbstractFunctionDecl *distributedTarge
     return true;
   }
 
-  SmallPtrSet<ProtocolDecl *, 2> requirementProtos;
-  if (getSerializationRequirementTypesForMember(distributedTarget,
-                                                requirementProtos)) {
-    return true;
-  }
-
-  return false;
+  auto serializationTy =
+      getDistributedActorSerializationType(distributedTarget->getDeclContext());
+  return serializationTy && !serializationTy->hasDependentMember();
 }
 
 /******************************************************************************/
diff --git a/lib/Sema/DerivedConformanceDistributedActor.cpp b/lib/Sema/DerivedConformanceDistributedActor.cpp
@@ -19,6 +19,7 @@
 #include "TypeChecker.h"
 #include "swift/Strings.h"
 #include "TypeCheckDistributed.h"
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/DistributedDecl.h"
@@ -219,7 +220,6 @@ static FuncDecl* createLocalFunc_doInvokeOnReturn(
     ParamDecl* handlerParam,
     ParamDecl* resultBufParam) {
   auto DC = parentFunc;
-  auto DAS = C.getDistributedActorSystemDecl();
   auto doInvokeLocalFuncIdent = C.getIdentifier("doInvokeOnReturn");
 
   // mock locations, we're a synthesized func and don't need real locations
@@ -243,7 +243,12 @@ static FuncDecl* createLocalFunc_doInvokeOnReturn(
       ParameterList::create(C, {resultTyParamDecl});
 
   SmallVector<Requirement, 2> requirements;
-  for (auto p : getDistributedSerializationRequirementProtocols(systemNominal, DAS)) {
+
+  auto serializationLayout =
+      getDistributedActorSystemSerializationType(systemNominal)
+          ->getExistentialLayout();
+
+  for (auto p : serializationLayout.getProtocols()) {
     auto requirement =
         Requirement(RequirementKind::Conformance,
                     resultGenericParamDecl->getDeclaredInterfaceType(),
@@ -285,7 +290,6 @@ deriveBodyDistributed_invokeHandlerOnReturn(AbstractFunctionDecl *afd,
   auto implicit = true;
   ASTContext &C = afd->getASTContext();
   auto DC = afd->getDeclContext();
-  auto DAS = C.getDistributedActorSystemDecl();
 
   // mock locations, we're a thunk and don't really need detailed locations
   const SourceLoc sloc = SourceLoc();
@@ -305,7 +309,7 @@ deriveBodyDistributed_invokeHandlerOnReturn(AbstractFunctionDecl *afd,
   auto metatypeParam = params->get(2);
 
   auto serializationRequirementTypeTy =
-      getDistributedSerializationRequirementType(nominal, DAS);
+      getDistributedActorSystemSerializationType(nominal);
 
   auto serializationRequirementMetaTypeTy =
       ExistentialMetatypeType::get(serializationRequirementTypeTy);
@@ -393,9 +397,6 @@ static FuncDecl *deriveDistributedActorSystem_invokeHandlerOnReturn(
   auto unsafeRawPointerType = C.getUnsafeRawPointerType();
   auto anyTypeType = ExistentialMetatypeType::get(C.TheAnyType); // Any.Type
 
-  //  auto serializationRequirementType =
-  //  getDistributedSerializationRequirementType(system, DAS);
-
   // params:
   // - handler: Self.ResultHandler
   // - resultBuffer:
@@ -583,7 +584,7 @@ deriveDistributedActorType_SerializationRequirement(
     return nullptr;
 
   if (auto systemNominal = systemTy->getAnyNominal())
-    return getDistributedSerializationRequirementType(systemNominal, DAS);
+    return getDistributedActorSystemSerializationType(systemNominal);
 
   return nullptr;
 }
diff --git a/lib/Sema/TypeCheckDistributed.cpp b/lib/Sema/TypeCheckDistributed.cpp
diff --git a/test/Distributed/distributed_protocols_distributed_func_serialization_requirements.swift b/test/Distributed/distributed_protocols_distributed_func_serialization_requirements.swift
