Merge pull request #72073 from xedin/distributed-protocol-requirement-accessors

[Distributed] Emit accessors for distributed protocol requirements

Author: xedin

include/swift/IRGen/Linking.h

@@ -359,6 +359,9 @@ class LinkEntity {
     /// the metadata cache once.
     CanonicalPrespecializedGenericTypeCachingOnceToken,
 
+    /// The function used to access distributed methods and accessors.
+    DistributedAccessor,
+
     /// The same as AsyncFunctionPointer but with a different stored value, for
     /// use by TBDGen.
     /// The pointer is an AbstractFunctionDecl*.
@@ -512,8 +515,6 @@ class LinkEntity {
     /// The pointer is a const char* of the name.
     KnownAsyncFunctionPointer,
 
-    /// The pointer is SILFunction*
-    DistributedAccessor,
     /// An async function pointer for a distributed accessor (method or
     /// property).
     /// The pointer is a SILFunction*.
@@ -1341,6 +1342,10 @@ class LinkEntity {
   }
 
   static LinkEntity forDistributedTargetAccessor(SILFunction *target) {
+    return forDistributedTargetAccessor(target->getDeclContext()->getAsDecl());
+  }
+
+  static LinkEntity forDistributedTargetAccessor(Decl *target) {
     LinkEntity entity;
     entity.Pointer = target;
     entity.SecondaryPointer = nullptr;

lib/AST/ASTMangler.cpp

@@ -4246,23 +4246,22 @@ void ASTMangler::appendDistributedThunk(
          "distributed thunk to mangle must be function decl");
   assert(thunk->getContextKind() == DeclContextKind::AbstractFunctionDecl);
 
-  auto inProtocolExtensionMangleAsReference =
+  auto referenceInProtocolContextOrRequirement =
       [&thunk, asReference]() -> ProtocolDecl * {
-    if (!asReference) {
-      return nullptr;
-    }
+    auto *DC = thunk->getDeclContext();
+    if (!asReference)
+      return dyn_cast_or_null<ProtocolDecl>(DC);
 
-    if (auto extension = dyn_cast<ExtensionDecl>(thunk->getDeclContext())) {
+    if (auto extension = dyn_cast<ExtensionDecl>(DC))
       return dyn_cast_or_null<ProtocolDecl>(extension->getExtendedNominal());
-    }
+
     return nullptr;
   };
 
-  if (auto type = inProtocolExtensionMangleAsReference()) {
-    appendContext(type->getDeclContext(), base,
+  if (auto *P = referenceInProtocolContextOrRequirement()) {
+    appendContext(P->getDeclContext(), base,
                   thunk->getAlternateModuleName());
-    auto baseName = type->getBaseName();
-    appendIdentifier(Twine("$", baseName.getIdentifier().str()).str());
+    appendIdentifier(Twine("$", P->getNameStr()).str());
     appendOperator("C"); // necessary for roundtrip, though we don't use it
   } else {
     appendContextOf(thunk, base);

lib/IRGen/GenDecl.cpp

@@ -4423,8 +4423,34 @@ void IRGenModule::addProtocolConformance(ConformanceDescription &&record) {
   }
 }
 
-void IRGenModule::addAccessibleFunction(SILFunction *func) {
-  AccessibleFunctions.push_back(func);
+AccessibleFunction AccessibleFunction::forSILFunction(IRGenModule &IGM,
+                                                      SILFunction *func) {
+  assert(!func->isDistributed() && "use forDistributed(...) instead");
+
+  llvm::Constant *funcAddr = nullptr;
+  if (func->isAsync()) {
+    funcAddr = IGM.getAddrOfAsyncFunctionPointer(func);
+  } else {
+    funcAddr = IGM.getAddrOfSILFunction(func, NotForDefinition);
+  }
+
+  return AccessibleFunction(
+      /*recordName=*/LinkEntity::forAccessibleFunctionRecord(func)
+          .mangleAsString(),
+      /*funcName=*/LinkEntity::forSILFunction(func).mangleAsString(),
+      /*isDistributed=*/false, func->getLoweredFunctionType(), funcAddr);
+}
+
+AccessibleFunction AccessibleFunction::forDistributed(std::string recordName,
+                                                      std::string accessorName,
+                                                      CanSILFunctionType type,
+                                                      llvm::Constant *address) {
+  return AccessibleFunction(recordName, accessorName,
+                            /*isDistributed=*/true, type, address);
+}
+
+void IRGenModule::addAccessibleFunction(AccessibleFunction func) {
+  AccessibleFunctions.push_back(std::move(func));
 }
 
 /// Emit the protocol conformance list and return it (if asContiguousArray is
@@ -4648,17 +4674,12 @@ llvm::Constant *IRGenModule::emitTypeMetadataRecords(bool asContiguousArray) {
   return nullptr;
 }
 
-void IRGenModule::emitAccessibleFunction(
-    StringRef sectionName, SILFunction* func) {
-  std::string mangledRecordName =
-      LinkEntity::forAccessibleFunctionRecord(func).mangleAsString();
-  std::string mangledFunctionName =
-      LinkEntity::forSILFunction(func).mangleAsString();
-
+void IRGenModule::emitAccessibleFunction(StringRef sectionName,
+                                         const AccessibleFunction &func) {
   auto var = new llvm::GlobalVariable(
       Module, AccessibleFunctionRecordTy, /*isConstant=*/true,
       llvm::GlobalValue::PrivateLinkage, /*initializer=*/nullptr,
-      mangledRecordName);
+      func.getRecordName());
 
   ConstantInitBuilder builder(*this);
 
@@ -4669,47 +4690,36 @@ void IRGenModule::emitAccessibleFunction(
   // -- Field: Name (record name)
   {
     llvm::Constant *name =
-        getAddrOfGlobalString(mangledFunctionName,
+        getAddrOfGlobalString(func.getFunctionName(),
                               /*willBeRelativelyAddressed=*/true);
     fields.addRelativeAddress(name);
   }
 
   // -- Field: GenericEnvironment
   llvm::Constant *genericEnvironment = nullptr;
 
-  GenericSignature signature;
-  if (auto *env = func->getGenericEnvironment()) {
+  GenericSignature signature = func.getType()->getInvocationGenericSignature();
+  if (signature) {
     // Drop all the marker protocols because they are effect-less
     // at runtime.
-    signature = env->getGenericSignature().withoutMarkerProtocols();
+    signature = signature.withoutMarkerProtocols();
 
     genericEnvironment =
         getAddrOfGenericEnvironment(signature.getCanonicalSignature());
   }
   fields.addRelativeAddressOrNull(genericEnvironment);
 
   // -- Field: FunctionType
-  llvm::Constant *type = getTypeRef(func->getLoweredFunctionType(), signature,
-                                    MangledTypeRefRole::Metadata)
-                             .first;
+  llvm::Constant *type =
+      getTypeRef(func.getType(), signature, MangledTypeRefRole::Metadata).first;
   fields.addRelativeAddress(type);
 
   // -- Field: Function
-  llvm::Constant *funcAddr = nullptr;
-  if (func->isDistributed()) {
-    funcAddr = getAddrOfAsyncFunctionPointer(
-        LinkEntity::forDistributedTargetAccessor(func));
-  } else if (func->isAsync()) {
-    funcAddr = getAddrOfAsyncFunctionPointer(func);
-  } else {
-    funcAddr = getAddrOfSILFunction(func, NotForDefinition);
-  }
-
-  fields.addRelativeAddress(funcAddr);
+  fields.addRelativeAddress(func.getAddress());
 
   // -- Field: Flags
   AccessibleFunctionFlags flags;
-  flags.setDistributed(func->isDistributed());
+  flags.setDistributed(func.isDistributed());
   fields.addInt32(flags.getOpaqueValue());
 
   // ---- End of 'TargetAccessibleFunctionRecord' fields
@@ -4746,7 +4756,7 @@ void IRGenModule::emitAccessibleFunctions() {
     break;
   }
 
-  for (auto *func : AccessibleFunctions) {
+  for (const auto &func : AccessibleFunctions) {
     emitAccessibleFunction(fnsSectionName, func);
   }
 }