[Distributed] Ensure _remote funcs synthesized before dynamic replacement (#38974)

* [Distributed] Ensure _remote funcs synthesized before dynamic replacement

* cleanup

* remove redundant synthesis cause

Author: ktoso

include/swift/AST/Decl.h

@@ -5986,6 +5986,10 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Returns 'true' if the function is distributed.
   bool isDistributed() const;
 
+  /// Get (or synthesize)  the associated remote function for this one.
+  /// For example, for `distributed func hi()` get `func _remote_hi()`.
+  AbstractFunctionDecl *getDistributedActorRemoteFuncDecl() const;
+
   PolymorphicEffectKind getPolymorphicEffectKind(EffectKind kind) const;
 
   // FIXME: Hack that provides names with keyword arguments for accessors.

include/swift/AST/TypeCheckRequests.h

@@ -933,18 +933,18 @@ class IsDistributedActorRequest :
     bool isCached() const { return true; }
 };
 
-/// Determine whether the given func is distributed.
-class IsDistributedFuncRequest :
-    public SimpleRequest<IsDistributedFuncRequest,
-        bool(FuncDecl *),
+/// Obtain the 'remote' counterpart of a 'distributed func'.
+class GetDistributedRemoteFuncRequest :
+    public SimpleRequest<GetDistributedRemoteFuncRequest,
+        AbstractFunctionDecl *(AbstractFunctionDecl *),
         RequestFlags::Cached> {
 public:
     using SimpleRequest::SimpleRequest;
 
 private:
     friend SimpleRequest;
 
-    bool evaluate(Evaluator &evaluator, FuncDecl *func) const;
+    AbstractFunctionDecl *evaluate(Evaluator &evaluator, AbstractFunctionDecl *func) const;
 
 public:
     // Caching
@@ -2046,8 +2046,6 @@ enum class ImplicitMemberAction : uint8_t {
   ResolveEncodable,
   ResolveDecodable,
   ResolveDistributedActor,
-  ResolveDistributedActorIdentity,
-  ResolveDistributedActorTransport,
 };
 
 class ResolveImplicitMemberRequest

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -100,7 +100,7 @@ SWIFT_REQUEST(TypeChecker, IsDefaultActorRequest,
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, IsDistributedActorRequest, bool(NominalTypeDecl *),
               Cached, NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, IsDistributedFuncRequest, bool(FuncDecl *),
+SWIFT_REQUEST(TypeChecker, GetDistributedRemoteFuncRequest, AbstractFunctionDecl *(AbstractFunctionDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GlobalActorInstanceRequest,
               VarDecl *(NominalTypeDecl *),

lib/AST/Decl.cpp

@@ -4222,21 +4222,12 @@ void NominalTypeDecl::synthesizeSemanticMembersIfNeeded(DeclName member) {
       if (baseName == DeclBaseName::createConstructor()) {
         if ((member.isSimpleName() || argumentNames.front() == Context.Id_from)) {
           action.emplace(ImplicitMemberAction::ResolveDecodable);
-        } else if (argumentNames.front() == Context.Id_transport) {
-          action.emplace(ImplicitMemberAction::ResolveDistributedActorTransport);
         }
       } else if (!baseName.isSpecial() &&
            baseName.getIdentifier() == Context.Id_encode &&
            (member.isSimpleName() || argumentNames.front() == Context.Id_to)) {
         action.emplace(ImplicitMemberAction::ResolveEncodable);
       }
-    } else if (member.isSimpleName() || argumentNames.size() == 2) {
-      if (baseName == DeclBaseName::createConstructor()) {
-        if (argumentNames[0] == Context.Id_resolve &&
-            argumentNames[1] == Context.Id_using) {
-          action.emplace(ImplicitMemberAction::ResolveDistributedActor);
-        }
-      }
     }
   }
 
@@ -7217,14 +7208,19 @@ bool AbstractFunctionDecl::isSendable() const {
 }
 
 bool AbstractFunctionDecl::isDistributed() const {
-  auto func = dyn_cast<FuncDecl>(this);
-  if (!func)
-    return false;
+  return this->getAttrs().hasAttribute<DistributedActorAttr>();
+}
 
-  auto mutableFunc = const_cast<FuncDecl *>(func);
-  return evaluateOrDefault(getASTContext().evaluator,
-                           IsDistributedFuncRequest{mutableFunc},
-                           false);
+AbstractFunctionDecl*
+AbstractFunctionDecl::getDistributedActorRemoteFuncDecl() const {
+  if (!this->isDistributed())
+    return nullptr;
+
+  auto mutableThis = const_cast<AbstractFunctionDecl *>(this);
+  return evaluateOrDefault(
+      getASTContext().evaluator,
+      GetDistributedRemoteFuncRequest{mutableThis},
+      nullptr);
 }
 
 BraceStmt *AbstractFunctionDecl::getBody(bool canSynthesize) const {

lib/AST/TypeCheckRequests.cpp

@@ -1035,12 +1035,6 @@ void swift::simple_display(llvm::raw_ostream &out,
   case ImplicitMemberAction::ResolveDistributedActor:
     out << "resolve DistributedActor";
     break;
-  case ImplicitMemberAction::ResolveDistributedActorIdentity:
-    out << "resolve DistributedActor.id";
-    break;
-  case ImplicitMemberAction::ResolveDistributedActorTransport:
-    out << "resolve DistributedActor.actorTransport";
-    break;
   }
 }
 

lib/SILGen/SILGenDistributed.cpp

@@ -814,7 +814,6 @@ void SILGenFunction::emitDistributedActor_resignAddress(
 void SILGenFunction::emitDistributedActorClassMemberDestruction(
     SILLocation cleanupLoc, ManagedValue selfValue, ClassDecl *cd,
     SILBasicBlock *normalMemberDestroyBB, SILBasicBlock *finishBB) {
-  ASTContext &ctx = getASTContext();
   auto selfTy = cd->getDeclaredInterfaceType();
 
   Scope scope(Cleanups, CleanupLocation(cleanupLoc));

lib/Sema/CodeSynthesis.cpp

@@ -1255,9 +1255,10 @@ ResolveImplicitMemberRequest::evaluate(Evaluator &evaluator,
 
   auto &Context = target->getASTContext();
   switch (action) {
-  case ImplicitMemberAction::ResolveImplicitInit:
+  case ImplicitMemberAction::ResolveImplicitInit: {
     TypeChecker::addImplicitConstructors(target);
     break;
+  }
   case ImplicitMemberAction::ResolveCodingKeys: {
     // CodingKeys is a special type which may be synthesized as part of
     // Encodable/Decodable conformance. If the target conforms to either
@@ -1274,17 +1275,17 @@ ResolveImplicitMemberRequest::evaluate(Evaluator &evaluator,
     if (!evaluateTargetConformanceTo(decodableProto)) {
       (void)evaluateTargetConformanceTo(encodableProto);
     }
-  }
     break;
+  }
   case ImplicitMemberAction::ResolveEncodable: {
     // encode(to:) may be synthesized as part of derived conformance to the
     // Encodable protocol.
     // If the target should conform to the Encodable protocol, check the
     // conformance here to attempt synthesis.
     auto *encodableProto = Context.getProtocol(KnownProtocolKind::Encodable);
     (void)evaluateTargetConformanceTo(encodableProto);
-  }
     break;
+  }
   case ImplicitMemberAction::ResolveDecodable: {
     // init(from:) may be synthesized as part of derived conformance to the
     // Decodable protocol.
@@ -1293,17 +1294,11 @@ ResolveImplicitMemberRequest::evaluate(Evaluator &evaluator,
     TypeChecker::addImplicitConstructors(target);
     auto *decodableProto = Context.getProtocol(KnownProtocolKind::Decodable);
     (void)evaluateTargetConformanceTo(decodableProto);
-  }
     break;
-  case ImplicitMemberAction::ResolveDistributedActor:
-  case ImplicitMemberAction::ResolveDistributedActorTransport:
-  case ImplicitMemberAction::ResolveDistributedActorIdentity: {
-    // init(transport:) and init(resolve:using:) may be synthesized as part of
-    // derived conformance to the DistributedActor protocol.
-    // If the target should conform to the DistributedActor protocol, check the
-    // conformance here to attempt synthesis.
-    // FIXME(distributed): invoke the requirement adding explicitly here
-     TypeChecker::addImplicitConstructors(target);
+  }
+case ImplicitMemberAction::ResolveDistributedActor: {
+    if (auto classDecl = dyn_cast<ClassDecl>(target))
+      swift::addImplicitDistributedActorMembers(classDecl);
     auto *distributedActorProto =
         Context.getProtocol(KnownProtocolKind::DistributedActor);
     (void)evaluateTargetConformanceTo(distributedActorProto);

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -205,7 +205,8 @@ synthesizeRemoteFuncStubBody(AbstractFunctionDecl *func, void *context) {
   auto uintInit = ctx.getIntBuiltinInitDecl(uintDecl);
 
   auto missingTransportDecl = ctx.getMissingDistributedActorTransport();
-  assert(missingTransportDecl && "Could not locate '_missingDistributedActorTransport' function");
+  assert(missingTransportDecl &&
+         "Could not locate '_missingDistributedActorTransport' function");
 
   // Create a call to _Distributed._missingDistributedActorTransport
   auto loc = func->getLoc();
@@ -243,16 +244,10 @@ synthesizeRemoteFuncStubBody(AbstractFunctionDecl *func, void *context) {
   file->setBuiltinInitializer(staticStringInit);
 
   auto startLineAndCol = SM.getPresumedLineAndColumnForLoc(distributedFunc->getStartLoc());
-//  auto *line = new (ctx) MagicIdentifierLiteralExpr(
-//      MagicIdentifierLiteralExpr::Line, loc, /*Implicit=*/true);
-//  auto *line = new (ctx) IntegerLiteralExpr(startLineAndCol.first, loc,
-//                                            /*implicit*/ true);
   auto *line = IntegerLiteralExpr::createFromUnsigned(ctx, startLineAndCol.first);
   line->setType(uintType);
   line->setBuiltinInitializer(uintInit);
 
-//  auto *column = new (ctx) MagicIdentifierLiteralExpr(
-//      MagicIdentifierLiteralExpr::Column, loc, /*Implicit=*/true);
   auto *column = IntegerLiteralExpr::createFromUnsigned(ctx, startLineAndCol.second);
   column->setType(uintType);
   column->setBuiltinInitializer(uintInit);
@@ -264,7 +259,6 @@ synthesizeRemoteFuncStubBody(AbstractFunctionDecl *func, void *context) {
 
   SmallVector<ASTNode, 2> stmts;
   stmts.push_back(call); // something() -> Never
-  // stmts.push_back(new (ctx) ReturnStmt(SourceLoc(), /*Result=*/nullptr)); // FIXME: this causes 'different types for return type: String vs. ()'
   auto body = BraceStmt::create(ctx, SourceLoc(), stmts, SourceLoc(),
                                 /*implicit=*/true);
   return { body, /*isTypeChecked=*/true };
@@ -290,7 +284,8 @@ static Identifier makeRemoteFuncIdentifier(FuncDecl* func) {
 ///
 /// and is intended to be replaced by a transport library by providing an
 /// appropriate @_dynamicReplacement function.
-static void addImplicitRemoteActorFunction(ClassDecl *decl, FuncDecl *func) {
+static FuncDecl*
+addImplicitRemoteFunction(ClassDecl *decl, FuncDecl *func) {
   auto &C = decl->getASTContext();
   auto parentDC = decl;
 
@@ -315,6 +310,10 @@ static void addImplicitRemoteActorFunction(ClassDecl *decl, FuncDecl *func) {
   remoteFuncDecl->getAttrs().add(
       new (C) DistributedActorIndependentAttr(/*IsImplicit=*/true));
 
+  // nonisolated
+  remoteFuncDecl->getAttrs().add(
+      new (C) NonisolatedAttr(/*IsImplicit=*/true));
+
   // users should never have to access this function directly;
   // it is only invoked from our distributed function thunk if the actor is remote.
   remoteFuncDecl->setUserAccessible(false);
@@ -326,16 +325,16 @@ static void addImplicitRemoteActorFunction(ClassDecl *decl, FuncDecl *func) {
   remoteFuncDecl->copyFormalAccessFrom(func, /*sourceIsParentContext=*/false);
 
   decl->addMember(remoteFuncDecl);
+  return remoteFuncDecl;
 }
 
 /// Synthesize dynamic _remote stub functions for each encountered distributed function.
-static void addImplicitRemoteActorFunctions(ClassDecl *decl) {
+static void addImplicitRemoteFunctions(ClassDecl *decl) {
   assert(decl->isDistributedActor());
-
   for (auto member : decl->getMembers()) {
-    auto func = dyn_cast<FuncDecl>(member);
+    auto func = dyn_cast<AbstractFunctionDecl>(member);
     if (func && func->isDistributed()) {
-      addImplicitRemoteActorFunction(decl, func);
+      (void) func->getDistributedActorRemoteFuncDecl();
     }
   }
 }
@@ -344,8 +343,33 @@ static void addImplicitRemoteActorFunctions(ClassDecl *decl) {
 /************************ SYNTHESIS ENTRY POINT *******************************/
 /******************************************************************************/
 
+AbstractFunctionDecl *TypeChecker::addImplicitDistributedActorRemoteFunction(
+    ClassDecl *decl, AbstractFunctionDecl *AFD) {
+  if (!decl->isDistributedActor())
+    return nullptr;
+
+  if (auto func = dyn_cast<FuncDecl>(AFD))
+    return addImplicitRemoteFunction(decl, func);
+
+  return nullptr;
+}
+
+void TypeChecker::addImplicitDistributedActorRemoteFunctions(NominalTypeDecl *decl) {
+  // Bail out if not a distributed actor definition.
+  if (!decl->isDistributedActor())
+    return;
+
+  // If the _Distributed module is missing we cannot synthesize anything.
+  if (!swift::ensureDistributedModuleLoaded(decl))
+    return;
+
+  if (auto clazz = dyn_cast<ClassDecl>(decl))
+    addImplicitRemoteFunctions(clazz);
+}
+
 /// Entry point for adding all computed members to a distributed actor decl.
-void swift::addImplicitDistributedActorMembersToClass(ClassDecl *decl) {
+// TODO(distributed): move the synthesis of protocol requirements to DerivedConformance style
+void swift::addImplicitDistributedActorMembers(ClassDecl *decl) {
   // Bail out if not a distributed actor definition.
   if (!decl->isDistributedActor())
     return;
@@ -356,5 +380,4 @@ void swift::addImplicitDistributedActorMembersToClass(ClassDecl *decl) {
 
   addFactoryResolveFunction(decl);
   addImplicitDistributedActorStoredProperties(decl);
-  addImplicitRemoteActorFunctions(decl);
 }

lib/Sema/TypeCheckAttr.cpp

@@ -2365,7 +2365,7 @@ void AttributeChecker::visitDiscardableResultAttr(DiscardableResultAttr *attr) {
   }
 }
 
-/// Lookup the replaced decl in the replacments scope.
+/// Lookup the replaced decl in the replacements scope.
 static void lookupReplacedDecl(DeclNameRef replacedDeclName,
                                const DeclAttribute  *attr,
                                const ValueDecl *replacement,
@@ -2400,9 +2400,10 @@ static void lookupReplacedDecl(DeclNameRef replacedDeclName,
   if (declCtxt->isInSpecializeExtensionContext())
     options |= NL_IncludeUsableFromInline;
 
-  if (typeCtx)
+  if (typeCtx) {
     moduleScopeCtxt->lookupQualified({typeCtx}, replacedDeclName, options,
                                      results);
+  }
 }
 
 /// Remove any argument labels from the interface type of the given value that
@@ -3498,6 +3499,19 @@ DynamicallyReplacedDeclRequest::evaluate(Evaluator &evaluator,
   if (attr->isInvalid())
     return nullptr;
 
+  auto *clazz = VD->getDeclContext()->getSelfClassDecl();
+  if (clazz && clazz->isDistributedActor()) {
+    // Since distributed actors synthesize their `_remote`
+    // counterparts to any `distributed func` declared on them, and such remote
+    // function is a popular target for dynamic function replacement - we must
+    // force the synthesis has happened here already, such that we can lookup
+    // the func for the replacement we're handling right now.
+    //
+    // This request is cached, so it won't cause wasted/duplicate work.
+    TypeChecker::addImplicitDistributedActorRemoteFunctions(clazz);
+  }
+
+
   // If we can lazily resolve the function, do so now.
   if (auto *LazyResolver = attr->Resolver) {
     auto decl = LazyResolver->loadDynamicallyReplacedFunctionDecl(

lib/Sema/TypeCheckDecl.cpp

@@ -2580,6 +2580,7 @@ static ArrayRef<Decl *> evaluateMembersRequest(
     TypeChecker::addImplicitConstructors(nominal);
   }
 
+
   // Force any conformances that may introduce more members.
   for (auto conformance : idc->getLocalConformances()) {
     auto proto = conformance->getProtocol();
@@ -2611,15 +2612,15 @@ static ArrayRef<Decl *> evaluateMembersRequest(
     TypeChecker::checkConformance(conformance->getRootNormalConformance());
   }
 
-  // If the type conforms to Encodable or Decodable, even via an extension,
-  // the CodingKeys enum is synthesized as a member of the type itself.
-  // Force it into existence.
   if (nominal) {
+    // If the type conforms to Encodable or Decodable, even via an extension,
+    // the CodingKeys enum is synthesized as a member of the type itself.
+    // Force it into existence.
     (void) evaluateOrDefault(
       ctx.evaluator,
-      ResolveImplicitMemberRequest{nominal,
-                 ImplicitMemberAction::ResolveCodingKeys},
-      {});
+        ResolveImplicitMemberRequest{nominal,
+                                     ImplicitMemberAction::ResolveCodingKeys},
+        {});
   }
 
   // If the decl has a @main attribute, we need to force synthesis of the
@@ -2638,6 +2639,10 @@ static ArrayRef<Decl *> evaluateMembersRequest(
         (void) var->getPropertyWrapperInitializerInfo();
       }
     }
+
+    if (auto *func = dyn_cast<FuncDecl>(member)) {
+      (void) func->getDistributedActorRemoteFuncDecl();
+    }
   }
 
   SortedDeclList synthesizedMembers;