Add support for distributed functions in extensions of distributed actors

docs/ABI/Mangling.rst

@@ -217,6 +217,7 @@ types where the metadata itself has unknown layout.)
   global ::= global 'To'                 // swift-as-ObjC thunk
   global ::= global 'TD'                 // dynamic dispatch thunk
   global ::= global 'Td'                 // direct method reference thunk
+  global ::= global 'TE'                 // distributed actor thunk
   global ::= global 'TI'                 // implementation of a dynamic_replaceable function
   global ::= global 'Tu'                 // async function pointer of a function
   global ::= global 'TX'                 // function pointer of a dynamic_replaceable function

include/swift/Demangling/DemangleNodes.def

@@ -70,6 +70,7 @@ NODE(DependentProtocolConformanceAssociated)
 CONTEXT_NODE(Destructor)
 CONTEXT_NODE(DidSet)
 NODE(Directness)
+NODE(DistributedThunk)
 NODE(DynamicAttribute)
 NODE(DirectMethodReferenceAttribute)
 NODE(DynamicSelf)

lib/AST/ASTMangler.cpp

@@ -842,7 +842,7 @@ void ASTMangler::appendSymbolKind(SymbolKind SKind) {
     case SymbolKind::DynamicThunk: return appendOperator("TD");
     case SymbolKind::SwiftAsObjCThunk: return appendOperator("To");
     case SymbolKind::ObjCAsSwiftThunk: return appendOperator("TO");
-    case SymbolKind::DistributedThunk: return appendOperator("Td");
+    case SymbolKind::DistributedThunk: return appendOperator("TE");
   }
 }
 

lib/Demangling/Demangler.cpp

@@ -130,6 +130,7 @@ bool swift::Demangle::isFunctionAttr(Node::Kind kind) {
     case Node::Kind::OutlinedVariable:
     case Node::Kind::OutlinedBridgedMethod:
     case Node::Kind::MergedFunction:
+    case Node::Kind::DistributedThunk:
     case Node::Kind::DynamicallyReplaceableFunctionImpl:
     case Node::Kind::DynamicallyReplaceableFunctionKey:
     case Node::Kind::DynamicallyReplaceableFunctionVar:
@@ -2359,6 +2360,7 @@ NodePointer Demangler::demangleThunkOrSpecialization() {
     case 'O': return createNode(Node::Kind::NonObjCAttribute);
     case 'D': return createNode(Node::Kind::DynamicAttribute);
     case 'd': return createNode(Node::Kind::DirectMethodReferenceAttribute);
+    case 'E': return createNode(Node::Kind::DistributedThunk);
     case 'a': return createNode(Node::Kind::PartialApplyObjCForwarder);
     case 'A': return createNode(Node::Kind::PartialApplyForwarder);
     case 'm': return createNode(Node::Kind::MergedFunction);

lib/Demangling/NodePrinter.cpp

@@ -558,6 +558,7 @@ class NodePrinter {
     case Node::Kind::ProtocolConformanceRefInTypeModule:
     case Node::Kind::ProtocolConformanceRefInProtocolModule:
     case Node::Kind::ProtocolConformanceRefInOtherModule:
+    case Node::Kind::DistributedThunk:
     case Node::Kind::DynamicallyReplaceableFunctionKey:
     case Node::Kind::DynamicallyReplaceableFunctionImpl:
     case Node::Kind::DynamicallyReplaceableFunctionVar:
@@ -1986,6 +1987,11 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
     Printer << "opaque type symbolic reference 0x";
     Printer.writeHex(Node->getIndex());
     return nullptr;
+  case Node::Kind::DistributedThunk:
+    if (!Options.ShortenThunk) {
+      Printer << "distributed thunk for ";
+    }
+    return nullptr;
   case Node::Kind::DynamicallyReplaceableFunctionKey:
     if (!Options.ShortenThunk) {
       Printer << "dynamically replaceable key for ";

lib/Demangling/OldRemangler.cpp

@@ -729,6 +729,12 @@ ManglingError Remangler::mangleMergedFunction(Node *node, unsigned depth) {
   return ManglingError::Success;
 }
 
+ManglingError
+Remangler::mangleDistributedThunk(Node *node, unsigned depth) {
+  Buffer << "TE";
+  return ManglingError::Success;
+}
+
 ManglingError
 Remangler::mangleDynamicallyReplaceableFunctionImpl(Node *node,
                                                     unsigned depth) {

lib/Demangling/Remangler.cpp

@@ -1610,6 +1610,7 @@ ManglingError Remangler::mangleGlobal(Node *node, unsigned depth) {
       case Node::Kind::VTableAttribute:
       case Node::Kind::DirectMethodReferenceAttribute:
       case Node::Kind::MergedFunction:
+      case Node::Kind::DistributedThunk:
       case Node::Kind::DynamicallyReplaceableFunctionKey:
       case Node::Kind::DynamicallyReplaceableFunctionImpl:
       case Node::Kind::DynamicallyReplaceableFunctionVar:
@@ -2241,6 +2242,12 @@ ManglingError Remangler::mangleMergedFunction(Node *node, unsigned depth) {
   return ManglingError::Success;
 }
 
+ManglingError
+Remangler::mangleDistributedThunk(Node *node, unsigned depth) {
+  Buffer << "TE";
+  return ManglingError::Success;
+}
+
 ManglingError
 Remangler::mangleDynamicallyReplaceableFunctionImpl(Node *node,
                                                     unsigned depth) {

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/Demangle/Inputs/manglings.txt

@@ -420,3 +420,4 @@ $sSIxip6foobarP ---> foobar in Swift.DefaultIndices.subscript : A
 $s13__lldb_expr_110$10016c2d8yXZ1B10$10016c2e0LLC ---> __lldb_expr_1.(unknown context at $10016c2d8).(B in $10016c2e0)
 $s__TJO ---> $s__TJO
 $s6Foobar7Vector2VAASdRszlE10simdMatrix5scale6rotate9translateSo0C10_double3x3aACySdG_SdAJtFZ0D4TypeL_aySd__GD ---> MatrixType #1 in static (extension in Foobar):Foobar.Vector2<Swift.Double><A where A == Swift.Double>.simdMatrix(scale: Foobar.Vector2<Swift.Double>, rotate: Swift.Double, translate: Foobar.Vector2<Swift.Double>) -> __C.simd_double3x3
+$s17distributed_thunk2DAC1fyyFTE ---> distributed thunk for distributed_thunk.DA.f() -> ()

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_thunk2DAC1fyyFTE : $@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() { }
+}