[Distributed] Handle mangling thunks in extensions with generic AS and $Stubs (#71914)

Author: ktoso

include/swift/ABI/Metadata.h

@@ -4992,24 +4992,7 @@ struct TargetAccessibleFunctionRecord {
   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

@@ -1619,8 +1619,6 @@ 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,7 +159,6 @@ class ASTMangler : public Mangler {
     DistributedThunk,
     DistributedAccessor,
     AccessibleFunctionRecord,
-    AccessibleProtocolRequirementFunctionRecord,
     BackDeploymentThunk,
     BackDeploymentFallback,
     HasSymbolQuery,
@@ -246,6 +245,12 @@ class ASTMangler : public Mangler {
                                              Type GlobalActorBound,
                                              ModuleDecl *Module);
 
+  void appendDistributedThunk(const AbstractFunctionDecl *thunk,
+                              bool asReference);
+  std::string mangleDistributedThunkRef(const AbstractFunctionDecl *thunk);
+  /// Mangling for distributed function accessible function record.
+  /// Used in Linking when emitting the record.
+  std::string mangleDistributedThunkRecord(const AbstractFunctionDecl *thunk);
   std::string mangleDistributedThunk(const AbstractFunctionDecl *thunk);
 
   /// Mangle a completion handler block implementation function, used for importing ObjC

include/swift/AST/TypeCheckRequests.h

@@ -1333,8 +1333,8 @@ class GetDistributedActorInvocationDecoderRequest :
 
 /// Obtain the method that could be used to decode argument values passed
 /// to a particular actor invocation type.
-class GetDistributedActorArgumentDecodingMethodRequest :
-  public SimpleRequest<GetDistributedActorArgumentDecodingMethodRequest,
+class GetDistributedActorConcreteArgumentDecodingMethodRequest :
+  public SimpleRequest<GetDistributedActorConcreteArgumentDecodingMethodRequest,
                        FuncDecl *(NominalTypeDecl *),
                        RequestFlags::Cached> {
 public:

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -159,7 +159,7 @@ SWIFT_REQUEST(TypeChecker, GetDistributedRemoteCallArgumentInitFunctionRequest,
 SWIFT_REQUEST(TypeChecker, GetDistributedActorInvocationDecoderRequest,
               NominalTypeDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, GetDistributedActorArgumentDecodingMethodRequest,
+SWIFT_REQUEST(TypeChecker, GetDistributedActorConcreteArgumentDecodingMethodRequest,
               FuncDecl *(NominalTypeDecl *),
               Cache, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GlobalActorInstanceRequest,

include/swift/Demangling/DemangleNodes.def

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

include/swift/IRGen/Linking.h

@@ -521,8 +521,8 @@ class LinkEntity {
 
     /// Accessible function record, which describes a function that can be
     /// looked up by name by the runtime.
+    /// The pointer is a SILFunction*.
     AccessibleFunctionRecord,
-    AccessibleProtocolRequirementFunctionRecord,
 
     /// Extended existential type shape.
     /// Pointer is the (generalized) existential type.

include/swift/Runtime/AccessibleFunction.h

@@ -24,15 +24,6 @@
 namespace swift {
 namespace runtime {
 
-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,

lib/AST/ASTMangler.cpp

@@ -922,8 +922,6 @@ 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");
@@ -3904,11 +3902,13 @@ ASTMangler::mangleOpaqueTypeDescriptorRecord(const OpaqueTypeDecl *decl) {
   return finalize();
 }
 
-std::string ASTMangler::mangleDistributedThunk(const AbstractFunctionDecl *thunk) {
+void ASTMangler::appendDistributedThunk(
+    const AbstractFunctionDecl *thunk, bool asReference) {
   // Marker protocols cannot be checked at runtime, so there is no point
   // in recording them for distributed thunks.
   llvm::SaveAndRestore<bool> savedAllowMarkerProtocols(AllowMarkerProtocols,
                                                        false);
+  // TODO: add a flag to skip class/struct information from parameter types
 
   // Since computed property SILDeclRef's refer to the "originator"
   // of the thunk, we need to mangle distributed thunks of accessors
@@ -3931,8 +3931,52 @@ std::string ASTMangler::mangleDistributedThunk(const AbstractFunctionDecl *thunk
     thunk = storage->getDistributedThunk();
     assert(thunk);
   }
+  assert(isa<AbstractFunctionDecl>(thunk) &&
+         "distributed thunk to mangle must be function decl");
+  assert(thunk->getContextKind() == DeclContextKind::AbstractFunctionDecl);
+
+  auto inProtocolExtensionMangleAsReference =
+      [&thunk, asReference]() -> ProtocolDecl * {
+    if (!asReference) {
+      return nullptr;
+    }
+
+    if (auto extension = dyn_cast<ExtensionDecl>(thunk->getDeclContext())) {
+      return dyn_cast_or_null<ProtocolDecl>(extension->getExtendedNominal());
+    }
+    return nullptr;
+  };
+
+  if (auto type = inProtocolExtensionMangleAsReference()) {
+    appendContext(type->getDeclContext(), thunk->getAlternateModuleName());
+    auto baseName = type->getBaseName();
+    appendIdentifier(Twine("$", baseName.getIdentifier().str()).str());
+    appendOperator("C"); // necessary for roundtrip, though we don't use it
+  } else {
+    appendContextOf(thunk);
+  }
 
-  return mangleEntity(thunk, SymbolKind::DistributedThunk);
+  appendIdentifier(thunk->getBaseName().getIdentifier().str());
+  appendDeclType(thunk, FunctionMangling);
+  appendOperator("F");
+  appendSymbolKind(SymbolKind::DistributedThunk);
+}
+
+std::string ASTMangler::mangleDistributedThunkRef(const AbstractFunctionDecl *thunk) {
+  beginMangling();
+  appendDistributedThunk(thunk, /*asReference=*/true);
+  return finalize();
+}
+std::string ASTMangler::mangleDistributedThunkRecord(const AbstractFunctionDecl *thunk) {
+  beginMangling();
+  appendDistributedThunk(thunk, /*asReference=*/true);
+  appendSymbolKind(SymbolKind::AccessibleFunctionRecord);
+  return finalize();
+}
+std::string ASTMangler::mangleDistributedThunk(const AbstractFunctionDecl *thunk) {
+  beginMangling();
+  appendDistributedThunk(thunk, /*asReference=*/false);
+  return finalize();
 }
 
 void ASTMangler::appendMacroExpansionContext(

lib/AST/DistributedDecl.cpp

@@ -139,7 +139,7 @@ Type swift::getSerializationRequirementTypesForMember(
     return signature->getConcreteType(SerReqAssocType);
   }
 
-  llvm_unreachable("Unable to fetch ActorSystem type!");
+  llvm_unreachable("Unable to fetch SerializationRequirement type!");
 }
 
 Type swift::getDistributedActorSystemType(NominalTypeDecl *actor) {
@@ -325,11 +325,8 @@ Type ASTContext::getAssociatedTypeOfDistributedSystemOfActor(
 
 FuncDecl*
 ASTContext::getDistributedActorArgumentDecodingMethod(NominalTypeDecl *actor) {
-  if (!actor->isDistributedActor())
-    return nullptr;
-
   return evaluateOrDefault(
-      evaluator, GetDistributedActorArgumentDecodingMethodRequest{actor}, nullptr);
+      evaluator, GetDistributedActorConcreteArgumentDecodingMethodRequest{actor}, nullptr);
 }
 
 NominalTypeDecl*