[Distributed] Emit dist thunk symbol in TBD

Author: ktoso

include/swift/IRGen/Linking.h

@@ -153,6 +153,10 @@ class LinkEntity {
     /// a class.
     DispatchThunkAllocatorAsyncFunctionPointer,
 
+    /// An async function pointer for a distributed thunk.
+    /// The pointer is a FuncDecl* inside an actor (class).
+    DistributedThunkAsyncFunctionPointer,
+
     /// A method descriptor.  The pointer is a FuncDecl* inside a protocol
     /// or a class.
     MethodDescriptor,
@@ -1198,7 +1202,9 @@ class LinkEntity {
 
   static LinkEntity forAsyncFunctionPointer(SILDeclRef declRef) {
     LinkEntity entity;
-    entity.setForDecl(Kind::AsyncFunctionPointerAST,
+    entity.setForDecl(declRef.isDistributedThunk()
+                          ? Kind::DistributedThunkAsyncFunctionPointer
+                          : Kind::AsyncFunctionPointerAST,
                       declRef.getAbstractFunctionDecl());
     entity.SecondaryPointer =
         reinterpret_cast<void *>(static_cast<uintptr_t>(declRef.kind));
@@ -1264,6 +1270,8 @@ class LinkEntity {
   void mangle(llvm::raw_ostream &out) const;
   void mangle(SmallVectorImpl<char> &buffer) const;
   std::string mangleAsString() const;
+
+  SILDeclRef getSILDeclRef() const;
   SILLinkage getLinkage(ForDefinition_t isDefinition) const;
 
   bool hasDecl() const {

lib/FrontendTool/TBD.cpp

@@ -76,6 +76,7 @@ static bool validateSymbols(DiagnosticEngine &diags,
     // symbol table, so make sure to mangle IRGen names before comparing them
     // with what TBDGen created.
     auto unmangledName = nameValue.getKey();
+
     SmallString<128> name;
     llvm::Mangler::getNameWithPrefix(name, unmangledName,
                                      IRModule.getDataLayout());

lib/IRGen/Linking.cpp

@@ -460,18 +460,20 @@ std::string LinkEntity::mangleAsString() const {
     Result.append("Tu");
     return Result;
   }
+  case Kind::DistributedThunkAsyncFunctionPointer: {
+    std::string Result = getSILDeclRef().mangle();
+    Result.append("Td");
+    Result.append("Tu");
+    return Result;
+  }
   case Kind::KnownAsyncFunctionPointer: {
     std::string Result(static_cast<char *>(Pointer));
     Result.append("Tu");
     return Result;
   }
 
   case Kind::AsyncFunctionPointerAST: {
-    std::string Result;
-    Result = SILDeclRef(const_cast<ValueDecl *>(getDecl()),
-                        static_cast<SILDeclRef::Kind>(
-                            reinterpret_cast<uintptr_t>(SecondaryPointer)))
-                 .mangle();
+    std::string Result = getSILDeclRef().mangle();
     Result.append("Tu");
     return Result;
   }
@@ -483,6 +485,15 @@ std::string LinkEntity::mangleAsString() const {
   llvm_unreachable("bad entity kind!");
 }
 
+SILDeclRef LinkEntity::getSILDeclRef() const {
+  assert(getKind() == Kind::DistributedThunkAsyncFunctionPointer ||
+         getKind() == Kind::AsyncFunctionPointerAST);
+
+  return SILDeclRef(const_cast<ValueDecl *>(getDecl()),
+             static_cast<SILDeclRef::Kind>(
+                 reinterpret_cast<uintptr_t>(SecondaryPointer)));
+}
+
 SILLinkage LinkEntity::getLinkage(ForDefinition_t forDefinition) const {
   // For when `this` is a protocol conformance of some kind.
   auto getLinkageAsConformance = [&] {
@@ -717,6 +728,7 @@ SILLinkage LinkEntity::getLinkage(ForDefinition_t forDefinition) const {
     return getSILFunction()->getEffectiveSymbolLinkage();
 
   case Kind::AsyncFunctionPointerAST:
+  case Kind::DistributedThunkAsyncFunctionPointer:
     return getSILLinkage(getDeclLinkage(getDecl()), forDefinition);
 
   case Kind::DynamicallyReplaceableFunctionImpl:
@@ -779,6 +791,7 @@ bool LinkEntity::isContextDescriptor() const {
     return true;
   case Kind::AsyncFunctionPointer:
   case Kind::AsyncFunctionPointerAST:
+  case Kind::DistributedThunkAsyncFunctionPointer:
   case Kind::PropertyDescriptor:
   case Kind::DispatchThunk:
   case Kind::DispatchThunkDerivative:
@@ -962,6 +975,7 @@ llvm::Type *LinkEntity::getDefaultDeclarationType(IRGenModule &IGM) const {
   case Kind::DispatchThunkAsyncFunctionPointer:
   case Kind::DispatchThunkInitializerAsyncFunctionPointer:
   case Kind::DispatchThunkAllocatorAsyncFunctionPointer:
+  case Kind::DistributedThunkAsyncFunctionPointer:
   case Kind::PartialApplyForwarderAsyncFunctionPointer:
   case Kind::AsyncFunctionPointerAST:
   case Kind::KnownAsyncFunctionPointer:
@@ -1071,6 +1085,7 @@ bool LinkEntity::isWeakImported(ModuleDecl *module) const {
   }
 
   case Kind::AsyncFunctionPointerAST:
+  case Kind::DistributedThunkAsyncFunctionPointer:
   case Kind::DispatchThunk:
   case Kind::DispatchThunkDerivative:
   case Kind::DispatchThunkInitializer:
@@ -1165,6 +1180,7 @@ bool LinkEntity::isWeakImported(ModuleDecl *module) const {
 DeclContext *LinkEntity::getDeclContextForEmission() const {
   switch (getKind()) {
   case Kind::AsyncFunctionPointerAST:
+  case Kind::DistributedThunkAsyncFunctionPointer:
   case Kind::DispatchThunk:
   case Kind::DispatchThunkDerivative:
   case Kind::DispatchThunkInitializer:

lib/TBDGen/TBDGen.cpp

@@ -718,6 +718,7 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
 
   if (AFD->isDistributed()) {
     addSymbol(SILDeclRef(AFD).asDistributed());
+    addAsyncFunctionPointerSymbol(SILDeclRef(AFD).asDistributed());
   }
 
   // Add derivative function symbols.

test/TBD/distributed.swift

@@ -0,0 +1,26 @@
+// REQUIRES: VENDOR=apple
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend %s -enable-testing -enable-experimental-distributed -disable-availability-checking -emit-ir -o %t/test.ll -emit-tbd -emit-tbd-path %t/test.tbd
+// RUN cat %t/test.tbd | %FileCheck %s --dump-input=always
+
+import _Distributed
+
+// CHECK: @"$s4test1AC13_remote_helloyyYaKFTd" = hidden global %swift.async_func_pointer
+// CHECK: @"$s4test1AC13_remote_helloyyYaKFTdTu" = hidden global %swift.async_func_pointer
+distributed actor SomeDistributedActor {
+  distributed func hello(name: String) -> String {
+    "Hello, \(name)!"
+  }
+}
+
+// function:
+// IR unmangledName = $s4test20SomeDistributedActorC5hello4nameS2S_tF
+// function method descriptor
+// IR unmangledName = $s4test20SomeDistributedActorC5hello4nameS2S_tFTq
+// thunk, method reference
+// IR unmangledName = $s4test20SomeDistributedActorC5hello4nameS2S_tFTd
+// thunk, method reference + async function pointer
+// IR unmangledName = $s4test20SomeDistributedActorC5hello4nameS2S_tFTdTu