Add support for distributed functions in extensions of distributed actors

lib/SILGen/SILGen.cpp

@@ -1144,6 +1144,7 @@ static void emitOrDelayFunction(SILGenModule &SGM,
   auto linkage = constant.getLinkage(ForDefinition);
   bool mayDelay = !forceEmission &&
              (constant.isImplicit() &&
+              !constant.isDynamicallyReplaceable() &&
               !isPossiblyUsedExternally(linkage, SGM.M.isWholeModule()));
 
   // Avoid emitting a delayable definition if it hasn't already been referenced.

lib/SILGen/SILGenDistributed.cpp

@@ -951,8 +951,20 @@ void SILGenFunction::emitDistributedThunk(SILDeclRef thunk) {
     auto nativeFnSILTy = SILType::getPrimitiveObjectType(nativeMethodTy);
     auto nativeSilFnType = nativeFnSILTy.castTo<SILFunctionType>();
 
-    SILValue nativeFn = emitClassMethodRef(
+    bool isClassMethod = false;
+    if (auto classDecl = dyn_cast<ClassDecl>(fd->getDeclContext())) {
+      if (!classDecl->isFinal() && !fd->isFinal() &&
+          !fd->hasForcedStaticDispatch())
+        isClassMethod = true;
+    }
+
+    SILValue nativeFn;
+    if (isClassMethod) {
+      nativeFn = emitClassMethodRef(
         loc, params[params.size() - 1], native, nativeMethodTy);
+    } else {
+      nativeFn = emitGlobalFunctionRef(loc, native);
+    }
     auto subs = F.getForwardingSubstitutionMap();
 
     if (nativeSilFnType->hasErrorResult()) {

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -137,24 +137,23 @@ static Identifier makeRemoteFuncIdentifier(FuncDecl* distributedFunc) {
 ///
 /// and is intended to be replaced by a transport library by providing an
 /// appropriate @_dynamicReplacement function.
-AbstractFunctionDecl *TypeChecker::addImplicitDistributedActorRemoteFunction(
-    ClassDecl *decl, AbstractFunctionDecl *AFD) {
-  if (!decl->isDistributedActor())
+static AbstractFunctionDecl *addImplicitDistributedActorRemoteFunction(
+    DeclContext *parentDC, AbstractFunctionDecl *AFD) {
+  auto nominal = parentDC->getSelfNominalTypeDecl();
+  if (!nominal || !nominal->isDistributedActor())
     return nullptr;
 
   auto func = dyn_cast<FuncDecl>(AFD);
   if (!func || !func->isDistributed())
     return nullptr;
 
   // ==== if the remote func already exists, return it
-  if (auto existing = decl->lookupDirectRemoteFunc(func))
+  if (auto existing = nominal->lookupDirectRemoteFunc(func))
     return existing;
 
   // ==== Synthesize and add 'remote' func to the actor decl
 
-  auto &C = decl->getASTContext();
-  auto parentDC = decl;
-
+  auto &C = func->getASTContext();
   auto remoteFuncIdent = makeRemoteFuncIdentifier(func);
 
   auto params = ParameterList::clone(C, func->getParameters());
@@ -189,7 +188,24 @@ AbstractFunctionDecl *TypeChecker::addImplicitDistributedActorRemoteFunction(
   // same access control as the original function is fine
   remoteFuncDecl->copyFormalAccessFrom(func, /*sourceIsParentContext=*/false);
 
-  decl->addMember(remoteFuncDecl);
+  cast<IterableDeclContext>(parentDC->getAsDecl())->addMember(remoteFuncDecl);
 
   return remoteFuncDecl;
 }
+
+AbstractFunctionDecl *GetDistributedRemoteFuncRequest::evaluate(
+    Evaluator &evaluator, AbstractFunctionDecl *func) const {
+
+  if (!func->isDistributed())
+    return nullptr;
+
+  auto &C = func->getASTContext();
+  DeclContext *DC = func->getDeclContext();
+
+  // not via `ensureDistributedModuleLoaded` to avoid generating a warning,
+  // we won't be emitting the offending decl after all.
+  if (!C.getLoadedModule(C.Id_Distributed))
+    return nullptr;
+
+  return addImplicitDistributedActorRemoteFunction(DC, func);
+}

lib/Sema/TypeCheckDecl.cpp

@@ -2820,10 +2820,11 @@ static ArrayRef<Decl *> evaluateMembersRequest(
         (void) var->getPropertyWrapperAuxiliaryVariables();
         (void) var->getPropertyWrapperInitializerInfo();
       }
+    }
 
-      if (auto *func = dyn_cast<FuncDecl>(member)) {
-        (void) func->getDistributedActorRemoteFuncDecl();
-      }
+    // For a distributed function, add the remote function.
+    if (auto *func = dyn_cast<FuncDecl>(member)) {
+      (void) func->getDistributedActorRemoteFuncDecl();
     }
   }
 

lib/Sema/TypeCheckDistributed.cpp

@@ -71,36 +71,6 @@ bool IsDistributedActorRequest::evaluate(
   return classDecl->isExplicitDistributedActor();
 }
 
-AbstractFunctionDecl *GetDistributedRemoteFuncRequest::evaluate(
-    Evaluator &evaluator, AbstractFunctionDecl *func) const {
-
-  if (!func->isDistributed())
-    return nullptr;
-
-  auto &C = func->getASTContext();
-  DeclContext *DC = func->getDeclContext();
-
-  // not via `ensureDistributedModuleLoaded` to avoid generating a warning,
-  // we won't be emitting the offending decl after all.
-  if (!C.getLoadedModule(C.Id_Distributed))
-    return nullptr;
-
-  // Locate the actor decl that the member must be synthesized to.
-  // TODO(distributed): should this just be added to the extension instead when we're in one?
-  ClassDecl *decl = dyn_cast<ClassDecl>(DC);
-  if (!decl) {
-    if (auto ED = dyn_cast<ExtensionDecl>(DC)) {
-      decl = dyn_cast<ClassDecl>(ED->getExtendedNominal());
-    }
-  }
-
-  /// A distributed func cannot be added to a non-distributed actor;
-  /// If the 'decl' was not a distributed actor we must have declared and
-  /// requested it from a illegal context, let's just ignore the synthesis.
-  assert(decl && "Can't find actor detect to add implicit _remote function to");
-  return TypeChecker::addImplicitDistributedActorRemoteFunction(decl, func);
-}
-
 // ==== ------------------------------------------------------------------------
 
 /// Check whether the function is a proper distributed function
@@ -150,7 +120,7 @@ bool swift::checkDistributedFunction(FuncDecl *func, bool diagnose) {
   }
 
   // === Check _remote functions
-  ClassDecl *actorDecl = dyn_cast<ClassDecl>(func->getParent());
+  auto actorDecl = func->getParent()->getSelfNominalTypeDecl();
   assert(actorDecl && actorDecl->isDistributedActor());
 
   // _remote function for a distributed function must not be implemented by end-users,
@@ -252,10 +222,6 @@ void TypeChecker::checkDistributedActor(ClassDecl *decl) {
     // --- Check all constructors
     if (auto ctor = dyn_cast<ConstructorDecl>(member))
       checkDistributedActorConstructor(decl, ctor);
-
-    // --- synthesize _remote functions for distributed functions
-    if (auto func = dyn_cast<FuncDecl>(member))
-      (void)addImplicitDistributedActorRemoteFunction(decl, func);
   }
 
   // ==== Properties

lib/Sema/TypeChecker.h

@@ -499,14 +499,6 @@ bool checkContextualRequirements(GenericTypeDecl *decl,
 /// struct, class or actor.
 void addImplicitConstructors(NominalTypeDecl *typeDecl);
 
-/// Synthesize and add a '_remote' counterpart of the passed in `func` to `decl`.
-///
-/// \param decl the actor type to add the '_remote' definition to
-/// \param func the 'distributed func' that the '_remote' func should mirror
-/// \return the synthesized function
-AbstractFunctionDecl *addImplicitDistributedActorRemoteFunction(
-    ClassDecl* decl, AbstractFunctionDecl *func);
-
 /// Fold the given sequence expression into an (unchecked) expression
 /// tree.
 Expr *foldSequence(SequenceExpr *expr, DeclContext *dc);

test/SILGen/distributed_thunk.swift

@@ -0,0 +1,19 @@
+// RUN:  %target-swift-emit-silgen %s -enable-experimental-distributed -disable-availability-checking | %FileCheck %s --dump-input=fail
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+import _Distributed
+
+distributed actor DA { }
+
+extension DA {
+  // CHECK-LABEL: sil hidden [thunk] [ossa] @$s17distributed_thunk2DAC1fyyFTd : $@convention(method) @async (@guaranteed DA) -> @error Error
+  // CHECK: function_ref @swift_distributed_actor_is_remote
+
+  // Call the actor function
+  // CHECK: function_ref @$s17distributed_thunk2DAC1fyyF : $@convention(method) (@guaranteed DA) -> ()
+
+  // ... or the remote thunk
+  // CHECK: dynamic_function_ref @$s17distributed_thunk2DAC9_remote_fyyYaKF : $@convention(method) @async (@guaranteed DA) -> @error Error
+  distributed func f() { }
+}