[Distributed] Target identifiers for protocol calls (#70928)

Author: ktoso

include/swift/ABI/Metadata.h

@@ -4970,7 +4970,7 @@ class DynamicReplacementScope
 /// and then called through a fully-abstracted entry point whose arguments
 /// can be constructed in code.
 template <typename Runtime>
-struct TargetAccessibleFunctionRecord final {
+struct TargetAccessibleFunctionRecord {
 public:
   /// The name of the function, which is a unique string assigned to the
   /// function so it can be looked up later.
@@ -4992,7 +4992,24 @@ struct TargetAccessibleFunctionRecord final {
   AccessibleFunctionFlags Flags;
 };
 
+/// More advanced than AccessibleFunctionRecord and contains Actor name
+template <typename Runtime>
+struct TargetAccessibleProtocolRequirementFunctionRecord
+    : public TargetAccessibleFunctionRecord<Runtime> {
+public:
+  /// The concrete Actor type for this accessor.
+  RelativeDirectPointer<const char, /*nullable*/ false> ConcreteActorName;
+
+  /// The concrete witness method mangled name.
+  /// The record name for such record is the mangled name of the protocol
+  /// method. This is the mangled name of the concrete witness method.
+  RelativeDirectPointer<const char, /*nullable*/ false>
+      ConcreteWitnessMethodName;
+};
+
 using AccessibleFunctionRecord = TargetAccessibleFunctionRecord<InProcess>;
+using AccessibleProtocolRequirementFunctionRecord =
+    TargetAccessibleProtocolRequirementFunctionRecord<InProcess>;
 
 enum class PackLifetime : uint8_t {
   OnStack = 0,

include/swift/ABI/MetadataValues.h

@@ -1574,6 +1574,8 @@ namespace SpecialPointerAuthDiscriminators {
 
   /// Functions accessible at runtime (i.e. distributed method accessors).
   const uint16_t AccessibleFunctionRecord = 0x438c; // = 17292
+  const uint16_t AccessibleProtocolRequirementFunctionRecord =
+      0xa98e; // = 43406
 
   /// C type GetExtraInhabitantTag function descriminator
   const uint16_t GetExtraInhabitantTagFunction = 0x392e; // = 14638

include/swift/AST/ASTMangler.h

@@ -159,6 +159,7 @@ class ASTMangler : public Mangler {
     DistributedThunk,
     DistributedAccessor,
     AccessibleFunctionRecord,
+    AccessibleProtocolRequirementFunctionRecord,
     BackDeploymentThunk,
     BackDeploymentFallback,
     HasSymbolQuery,

include/swift/AST/Attr.h

@@ -2603,6 +2603,28 @@ class RawLayoutAttr final : public DeclAttribute {
   }
 };
 
+/// The @_distributedThunkTarget(for:) attribute.
+class DistributedThunkTargetAttr final
+    : public DeclAttribute {
+
+  AbstractFunctionDecl *TargetFunction;
+
+public:
+  DistributedThunkTargetAttr(AbstractFunctionDecl *target)
+      : DeclAttribute(DeclAttrKind::DistributedThunkTarget, SourceLoc(),
+                      SourceRange(),
+                      /*Implicit=*/false),
+        TargetFunction(target) {}
+
+  AbstractFunctionDecl *getTargetFunction() const {
+    return TargetFunction;
+  }
+
+  static bool classof(const DeclAttribute *DA) {
+    return DA->getKind() == DeclAttrKind::DistributedThunkTarget;
+  }
+};
+
 /// Predicate used to filter MatchingAttributeRange.
 template <typename ATTR, bool AllowInvalid> struct ToAttributeKind {
   ToAttributeKind() {}

include/swift/AST/Decl.h

@@ -2773,6 +2773,12 @@ class ValueDecl : public Decl {
     Bits.ValueDecl.Synthesized = value;
   }
 
+  /// Does this have a 'distributed' modifier?
+  ///
+  /// Only member methods and computed properties of a `distributed actor`
+  /// can be distributed.
+  bool isDistributed() const;
+
   bool hasName() const { return bool(Name); }
   bool isOperator() const { return Name.isOperator(); }
 
@@ -5869,9 +5875,6 @@ class AbstractStorageDecl : public ValueDecl {
 
   bool hasAnyNativeDynamicAccessors() const;
 
-  /// Does this have a 'distributed' modifier?
-  bool isDistributed() const;
-
   /// Return a distributed thunk if this computed property is marked as
   /// 'distributed' and and nullptr otherwise.
   FuncDecl *getDistributedThunk() const;
@@ -7379,9 +7382,6 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Returns if the function is 'rethrows' or 'reasync'.
   bool hasPolymorphicEffect(EffectKind kind) const;
 
-  /// Returns 'true' if the function is distributed.
-  bool isDistributed() const;
-
   /// Is this a thunk function used to access a distributed method
   /// or computed property outside of its actor isolation context?
   bool isDistributedThunk() const {
@@ -7528,6 +7528,15 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
            getSILSynthesizeKind() == SILSynthesizeKind::DistributedActorFactory;
   }
 
+  /// Return a vector of distributed requirements that this distributed method
+  /// is implementing.
+  ///
+  /// If the method is witness to multiple requirements this is incorrect and
+  /// should be diagnosed during type-checking as it may make remoteCalls
+  /// ambiguous.
+  llvm::ArrayRef<ValueDecl *>
+  getDistributedMethodWitnessedProtocolRequirements() const;
+
   /// Determines whether this function is a 'remoteCall' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedActorSystem' protocol.

include/swift/AST/DeclAttr.def

@@ -485,7 +485,10 @@ SIMPLE_DECL_ATTR(_noExistentials, NoExistentials,
 SIMPLE_DECL_ATTR(_noObjCBridging, NoObjCBridging,
   OnAbstractFunction | OnSubscript | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   155)
-LAST_DECL_ATTR(NoObjCBridging)
+DECL_ATTR(_distributedThunkTarget, DistributedThunkTarget,
+  OnAbstractFunction | UserInaccessible | ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIStableToRemove,
+  156)
+LAST_DECL_ATTR(DistributedThunkTarget)
 
 #undef DECL_ATTR_ALIAS
 #undef CONTEXTUAL_DECL_ATTR_ALIAS

include/swift/AST/TypeCheckRequests.h

@@ -1372,6 +1372,27 @@ class GetDistributedActorArgumentDecodingMethodRequest :
   bool isCached() const { return true; }
 };
 
+/// Find out if a distributed method is implementing a distributed protocol
+/// requirement.
+class GetDistributedMethodWitnessedProtocolRequirements :
+    public SimpleRequest<GetDistributedMethodWitnessedProtocolRequirements,
+                         llvm::ArrayRef<ValueDecl *> (AbstractFunctionDecl *),
+                         RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  llvm::ArrayRef<ValueDecl *> evaluate(
+      Evaluator &evaluator,
+      AbstractFunctionDecl *nominal) 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

@@ -150,6 +150,9 @@ SWIFT_REQUEST(TypeChecker, GetDistributedActorSystemPropertyRequest,
 SWIFT_REQUEST(TypeChecker, GetDistributedRemoteCallTargetInitFunctionRequest,
               ConstructorDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, GetDistributedMethodWitnessedProtocolRequirements,
+              ArrayRef<ValueDecl *> (AbstractFunctionDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetDistributedRemoteCallArgumentInitFunctionRequest,
               ConstructorDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)

include/swift/Demangling/DemangleNodes.def

@@ -358,6 +358,9 @@ NODE(AsyncSuspendResumePartialFunction)
 NODE(AccessibleFunctionRecord)
 NODE(CompileTimeConst)
 
+// Added in Swift 5.11
+NODE(AccessibleProtocolRequirementFunctionRecord)
+
 // Added in Swift 5.7
 NODE(BackDeploymentThunk)
 NODE(BackDeploymentFallback)

include/swift/IRGen/Linking.h

@@ -382,7 +382,8 @@ class LinkEntity {
     /// A SIL global variable. The pointer is a SILGlobalVariable*.
     SILGlobalVariable,
 
-    /// An outlined read-only global object. The pointer is a SILGlobalVariable*.
+    /// An outlined read-only global object. The pointer is a
+    /// SILGlobalVariable*.
     ReadOnlyGlobalObject,
 
     // These next few are protocol-conformance kinds.
@@ -513,13 +514,15 @@ class LinkEntity {
 
     /// The pointer is SILFunction*
     DistributedAccessor,
-    /// An async function pointer for a distributed accessor (method or property).
+    /// An async function pointer for a distributed accessor (method or
+    /// property).
     /// The pointer is a SILFunction*.
     DistributedAccessorAsyncPointer,
 
     /// Accessible function record, which describes a function that can be
     /// looked up by name by the runtime.
     AccessibleFunctionRecord,
+    AccessibleProtocolRequirementFunctionRecord,
 
     /// Extended existential type shape.
     /// Pointer is the (generalized) existential type.

include/swift/RemoteInspection/RuntimeHeaders/llvm/BinaryFormat/Swift.def

@@ -30,4 +30,6 @@ HANDLE_SWIFT_SECTION(protocs, "__swift5_protos", "swift5_protocols",
                      ".sw5prt$B")
 HANDLE_SWIFT_SECTION(acfuncs, "__swift5_acfuncs", "swift5_accessible_functions",
                      ".sw5acfn$B")
+HANDLE_SWIFT_SECTION(dacfuncs, "__swift5_acpfuns", "swift5_accessible_protocol_requirement_functions",
+                     ".sw5acpfn$B")
 HANDLE_SWIFT_SECTION(mpenum, "__swift5_mpenum", "swift5_mpenum", ".sw5mpen$B")

include/swift/Runtime/AccessibleFunction.h

@@ -24,7 +24,17 @@
 namespace swift {
 namespace runtime {
 
-SWIFT_RUNTIME_STDLIB_SPI const AccessibleFunctionRecord *
+SWIFT_RUNTIME_STDLIB_SPI
+const AccessibleFunctionRecord *swift_findAccessibleFunctionForConcreteType(
+    bool findConcreteWitness,
+    // concrete target type, when performing a protocol method call
+    const char *targetActorTypeNameStart, size_t targetActorTypeNameLength,
+    // method (concrete method name, or protocol method name)
+    const char *targetNameStart, size_t targetNameLength);
+
+// DEPRECATED: Prefer 'swift_findAccessibleFunctionForConcreteType'
+SWIFT_RUNTIME_STDLIB_SPI
+const AccessibleFunctionRecord *
 swift_findAccessibleFunction(const char *targetNameStart,
                              size_t targetNameLength);
 

include/swift/Runtime/Config.h

@@ -307,6 +307,10 @@ extern uintptr_t __COMPATIBILITY_LIBRARIES_CANNOT_CHECK_THE_IS_SWIFT_BIT_DIRECTL
 #define __ptrauth_swift_accessible_function_record                             \
   __ptrauth(ptrauth_key_process_independent_data, 1,                           \
             SpecialPointerAuthDiscriminators::AccessibleFunctionRecord)
+#define __ptrauth_swift_accessible_protocol_requirement_function_record        \
+  __ptrauth(ptrauth_key_process_independent_data, 1,                           \
+            SpecialPointerAuthDiscriminators::                                 \
+                AccessibleProtocolRequirementFunctionRecord)
 #define __ptrauth_swift_objc_superclass                                        \
   __ptrauth(ptrauth_key_process_independent_data, 1,                           \
             swift::SpecialPointerAuthDiscriminators::ObjCSuperclass)
@@ -355,6 +359,7 @@ extern uintptr_t __COMPATIBILITY_LIBRARIES_CANNOT_CHECK_THE_IS_SWIFT_BIT_DIRECTL
 #define __ptrauth_swift_escalation_notification_function
 #define __ptrauth_swift_dispatch_invoke_function
 #define __ptrauth_swift_accessible_function_record
+#define __ptrauth_swift_accessible_protocol_requirement_function_record
 #define __ptrauth_swift_objc_superclass
 #define __ptrauth_swift_runtime_function_entry
 #define __ptrauth_swift_runtime_function_entry_with_key(__key)

lib/AST/ASTMangler.cpp

@@ -922,6 +922,8 @@ void ASTMangler::appendSymbolKind(SymbolKind SKind) {
     case SymbolKind::DistributedThunk: return appendOperator("TE");
     case SymbolKind::DistributedAccessor: return appendOperator("TF");
     case SymbolKind::AccessibleFunctionRecord: return appendOperator("HF");
+    case SymbolKind::AccessibleProtocolRequirementFunctionRecord:
+    return appendOperator("HpF");
     case SymbolKind::BackDeploymentThunk: return appendOperator("Twb");
     case SymbolKind::BackDeploymentFallback: return appendOperator("TwB");
     case SymbolKind::HasSymbolQuery: return appendOperator("TwS");

lib/AST/Attr.cpp

@@ -1898,6 +1898,8 @@ StringRef DeclAttribute::getAttrName() const {
     return "_section";
   case DeclAttrKind::Documentation:
     return "_documentation";
+  case DeclAttrKind::DistributedThunkTarget:
+    return "_distributedThunkTarget";
   case DeclAttrKind::Nonisolated:
     if (cast<NonisolatedAttr>(this)->isUnsafe()) {
         return "nonisolated(unsafe)";

lib/AST/Decl.cpp

@@ -11349,7 +11349,8 @@ ActorIsolation::forActorInstanceSelf(ValueDecl *decl) {
 }
 
 NominalTypeDecl *ActorIsolation::getActor() const {
-  assert(getKind() == ActorInstance);
+  assert(getKind() == ActorInstance ||
+         getKind() == GlobalActor);
 
   if (silParsed)
     return nullptr;

lib/AST/DistributedDecl.cpp

@@ -65,9 +65,20 @@ using namespace swift;
 /******************************************************************************/
 
 // TODO(distributed): make into a request
-Type swift::getConcreteReplacementForProtocolActorSystemType(ValueDecl *member) {
-  auto &C = member->getASTContext();
-  auto *DC = member->getDeclContext();
+Type swift::getConcreteReplacementForProtocolActorSystemType(
+    ValueDecl *anyValue) {
+  auto &C = anyValue->getASTContext();
+
+  // FIXME(distributed): clean this up, we want a method that gets us AS type
+  // given any value, but is this the best way?
+  DeclContext *DC;
+  if (auto nominal = dyn_cast<NominalTypeDecl>(anyValue)) {
+    DC = nominal;
+  } else if (auto extension = dyn_cast<ExtensionDecl>(anyValue)) {
+    DC = extension->getExtendedNominal();
+  } else {
+    DC = anyValue->getDeclContext();
+  }
   auto DA = C.getDistributedActorDecl();
 
   // === When declared inside an actor, we can get the type directly
@@ -78,7 +89,7 @@ Type swift::getConcreteReplacementForProtocolActorSystemType(ValueDecl *member)
   /// === Maybe the value is declared in a protocol?
   if (auto protocol = DC->getSelfProtocolDecl()) {
     GenericSignature signature;
-    if (auto *genericContext = member->getAsGenericContext()) {
+    if (auto *genericContext = anyValue->getAsGenericContext()) {
       signature = genericContext->getGenericSignature();
     } else {
       signature = DC->getGenericSignatureOfContext();
@@ -312,7 +323,6 @@ Type ASTContext::getAssociatedTypeOfDistributedSystemOfActor(
 /******** Functions on DistributedActorSystem and friends *********************/
 /******************************************************************************/
 
-
 FuncDecl*
 ASTContext::getDistributedActorArgumentDecodingMethod(NominalTypeDecl *actor) {
   if (!actor->isDistributedActor())
@@ -380,6 +390,20 @@ bool swift::checkDistributedSerializationRequirementIsExactlyCodable(
       std::count(protocols.begin(), protocols.end(), decodable) == 1;
 }
 
+llvm::ArrayRef<ValueDecl *>
+AbstractFunctionDecl::getDistributedMethodWitnessedProtocolRequirements() const {
+  auto mutableThis = const_cast<AbstractFunctionDecl *>(this);
+
+  // Only a 'distributed' decl can witness 'distributed' protocol
+  if (!isDistributed()) {
+    return llvm::ArrayRef<ValueDecl *>();
+  }
+
+  return evaluateOrDefault(
+      getASTContext().evaluator,
+      GetDistributedMethodWitnessedProtocolRequirements(mutableThis), {});
+}
+
 /******************************************************************************/
 /********************* Ad-hoc protocol requirement checks *********************/
 /******************************************************************************/
@@ -1257,11 +1281,7 @@ AbstractFunctionDecl::isDistributedTargetInvocationResultHandlerOnReturn() const
 /********************** Distributed Functions *********************************/
 /******************************************************************************/
 
-bool AbstractFunctionDecl::isDistributed() const {
-  return getAttrs().hasAttribute<DistributedActorAttr>();
-}
-
-bool AbstractStorageDecl::isDistributed() const {
+bool ValueDecl::isDistributed() const {
   return getAttrs().hasAttribute<DistributedActorAttr>();
 }
 

lib/ASTGen/Sources/ASTGen/DeclAttrs.swift

@@ -253,7 +253,8 @@ extension ASTGenVisitor {
         .usableFromInline,
         .used,
         .warnUnqualifiedAccess,
-        .weakLinked:
+        .weakLinked,
+        .distributedThunkTarget:
 
         return self.generateSimpleDeclAttr(attribute: node, kind: attrKind)
 

lib/Demangling/NodePrinter.cpp

@@ -630,6 +630,7 @@ class NodePrinter {
     case Node::Kind::AsyncAwaitResumePartialFunction:
     case Node::Kind::AsyncSuspendResumePartialFunction:
     case Node::Kind::AccessibleFunctionRecord:
+    case Node::Kind::AccessibleProtocolRequirementFunctionRecord:
     case Node::Kind::BackDeploymentThunk:
     case Node::Kind::BackDeploymentFallback:
     case Node::Kind::ExtendedExistentialTypeShape:
@@ -2294,6 +2295,11 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
       Printer << "accessible function runtime record for ";
     }
     return nullptr;
+  case Node::Kind::AccessibleProtocolRequirementFunctionRecord:
+    if (!Options.ShortenThunk) {
+      Printer << "accessible distributed function runtime record for ";
+    }
+    return nullptr;
   case Node::Kind::DynamicallyReplaceableFunctionKey:
     if (!Options.ShortenThunk) {
       Printer << "dynamically replaceable key for ";