[Distributed] Correct tbd handling for distributed thunks

include/swift/IRGen/Linking.h

@@ -1608,6 +1608,7 @@ class LinkEntity {
   bool isTypeMetadataAccessFunction() const {
     return getKind() == Kind::TypeMetadataAccessFunction;
   }
+  bool isDistributedThunk() const;
   bool isDispatchThunk() const {
     return getKind() == Kind::DispatchThunk ||
            getKind() == Kind::DispatchThunkInitializer ||

include/swift/SIL/SILDeclRef.h

@@ -187,7 +187,7 @@ struct SILDeclRef {
   /// True if this references a foreign entry point for the referenced decl.
   unsigned isForeign : 1;
   /// True if this references a distributed function.
-  unsigned isDistributed : 1;
+  unsigned distributedThunk : 1;
   /// True if this references a distributed function, but it is known to be local
   unsigned isKnownToBeLocal : 1;
   /// True is this reference to function that could be looked up via a special
@@ -231,7 +231,7 @@ struct SILDeclRef {
 
   /// Produces a null SILDeclRef.
   SILDeclRef()
-      : loc(), kind(Kind::Func), isForeign(0), isDistributed(0),
+      : loc(), kind(Kind::Func), isForeign(0), distributedThunk(0),
         isKnownToBeLocal(0), isRuntimeAccessible(0),
         backDeploymentKind(BackDeploymentKind::None), defaultArgIndex(0),
         isAsyncLetClosure(0) {}
@@ -406,13 +406,13 @@ struct SILDeclRef {
   friend llvm::hash_code hash_value(const SILDeclRef &ref) {
     return llvm::hash_combine(
         ref.loc.getOpaqueValue(), static_cast<int>(ref.kind), ref.isForeign,
-        ref.isDistributed, ref.defaultArgIndex, ref.isAsyncLetClosure);
+        ref.distributedThunk, ref.defaultArgIndex, ref.isAsyncLetClosure);
   }
 
   bool operator==(SILDeclRef rhs) const {
     return loc.getOpaqueValue() == rhs.loc.getOpaqueValue() &&
            kind == rhs.kind && isForeign == rhs.isForeign &&
-           isDistributed == rhs.isDistributed &&
+           distributedThunk == rhs.distributedThunk &&
            backDeploymentKind == rhs.backDeploymentKind &&
            defaultArgIndex == rhs.defaultArgIndex && pointer == rhs.pointer &&
            isAsyncLetClosure == rhs.isAsyncLetClosure;
@@ -468,7 +468,7 @@ struct SILDeclRef {
 
   /// Returns a copy of the decl with the given back deployment kind.
   SILDeclRef asBackDeploymentKind(BackDeploymentKind backDeploymentKind) const {
-    return SILDeclRef(loc.getOpaqueValue(), kind, isForeign, isDistributed,
+    return SILDeclRef(loc.getOpaqueValue(), kind, isForeign, distributedThunk,
                       isKnownToBeLocal, isRuntimeAccessible, backDeploymentKind,
                       defaultArgIndex, isAsyncLetClosure,
                       pointer.get<AutoDiffDerivativeFunctionIdentifier *>());
@@ -511,6 +511,9 @@ struct SILDeclRef {
   /// True if the decl ref references a thunk handling potentially distributed actor functions
   bool isDistributedThunk() const;
 
+  /// True if the decl references a 'distributed' function.
+  bool isDistributed() const;
+
   /// True if the decl ref references a thunk handling a call to a function that
   /// supports back deployment.
   bool isBackDeploymentThunk() const;
@@ -605,13 +608,13 @@ struct SILDeclRef {
   friend struct llvm::DenseMapInfo<swift::SILDeclRef>;
   /// Produces a SILDeclRef from an opaque value.
   explicit SILDeclRef(void *opaqueLoc, Kind kind, bool isForeign,
-                      bool isDistributed, bool isKnownToBeLocal,
+                      bool isDistributedThunk, bool isKnownToBeLocal,
                       bool isRuntimeAccessible,
                       BackDeploymentKind backDeploymentKind,
                       unsigned defaultArgIndex, bool isAsyncLetClosure,
                       AutoDiffDerivativeFunctionIdentifier *derivativeId)
       : loc(Loc::getFromOpaqueValue(opaqueLoc)), kind(kind),
-        isForeign(isForeign), isDistributed(isDistributed),
+        isForeign(isForeign), distributedThunk(isDistributedThunk),
         isKnownToBeLocal(isKnownToBeLocal),
         isRuntimeAccessible(isRuntimeAccessible),
         backDeploymentKind(backDeploymentKind),
@@ -655,7 +658,7 @@ template<> struct DenseMapInfo<swift::SILDeclRef> {
                     : 0;
     unsigned h4 = UnsignedInfo::getHashValue(Val.isForeign);
     unsigned h5 = PointerInfo::getHashValue(Val.pointer.getOpaqueValue());
-    unsigned h6 = UnsignedInfo::getHashValue(Val.isDistributed);
+    unsigned h6 = UnsignedInfo::getHashValue(Val.distributedThunk);
     unsigned h7 = UnsignedInfo::getHashValue(unsigned(Val.backDeploymentKind));
     unsigned h8 = UnsignedInfo::getHashValue(Val.isKnownToBeLocal);
     unsigned h9 = UnsignedInfo::getHashValue(Val.isRuntimeAccessible);

lib/IRGen/GenMeta.cpp

@@ -942,8 +942,12 @@ namespace {
       SILDeclRef func(entry.getFunction());
 
       // Emit the dispatch thunk.
-      if (Resilient || IGM.getOptions().WitnessMethodElimination)
+      auto shouldEmitDispatchThunk =
+          (Resilient || IGM.getOptions().WitnessMethodElimination) &&
+          !func.isDistributed();
+      if (shouldEmitDispatchThunk) {
         IGM.emitDispatchThunk(func);
+      }
 
       {
         auto *requirement = cast<AbstractFunctionDecl>(func.getDecl());
@@ -1009,10 +1013,23 @@ namespace {
       }
 
       for (auto &entry : pi.getWitnessEntries()) {
+        if (entry.isFunction() &&
+            entry.getFunction().getDecl()->isDistributedGetAccessor()) {
+          // We avoid emitting _distributed_get accessors, as they cannot be
+          // referred to anyway
+          continue;
+        }
+
         if (Resilient) {
           if (entry.isFunction()) {
             // Define the method descriptor.
             SILDeclRef func(entry.getFunction());
+
+            /// Distributed thunks don't need resilience.
+            if (func.isDistributedThunk()) {
+              continue;
+            }
+
             auto *descriptor =
               B.getAddrOfCurrentPosition(
                 IGM.ProtocolRequirementStructTy);
@@ -1021,13 +1038,6 @@ namespace {
           }
         }
 
-        if (entry.isFunction() &&
-            entry.getFunction().getDecl()->isDistributedGetAccessor()) {
-          // We avoid emitting _distributed_get accessors, as they cannot be
-          // referred to anyway
-          continue;
-        }
-
         if (entry.isAssociatedType()) {
           auto assocType = entry.getAssociatedType();
           // Define the associated type descriptor to point to the current

lib/IRGen/GenProto.cpp

@@ -2215,6 +2215,12 @@ namespace {
               LinkEntity::forBaseConformanceDescriptor(requirement));
           B.addRelativeAddress(baseConformanceDescriptor);
         } else if (entry.getKind() == SILWitnessTable::Method) {
+          // distributed thunks don't need resilience
+          if (entry.getMethodWitness().Requirement.isDistributedThunk()) {
+            witnesses = witnesses.drop_back();
+            continue;
+          }
+
           // Method descriptor.
           auto declRef = entry.getMethodWitness().Requirement;
           auto requirement =

lib/IRGen/IRSymbolVisitor.cpp

@@ -103,6 +103,15 @@ class IRSymbolVisitorImpl : public SILSymbolVisitor {
 
   void addDispatchThunk(SILDeclRef declRef) override {
     auto entity = LinkEntity::forDispatchThunk(declRef);
+
+    // TODO: explain why
+    if (declRef.isDistributedThunk()) {
+      auto afd = declRef.getAbstractFunctionDecl();
+      if (afd && isa<ProtocolDecl>(afd->getDeclContext())) {
+        return;
+      }
+    }
+
     addLinkEntity(entity);
 
     if (declRef.getAbstractFunctionDecl()->hasAsync())
@@ -147,6 +156,14 @@ class IRSymbolVisitorImpl : public SILSymbolVisitor {
   }
 
   void addMethodDescriptor(SILDeclRef declRef) override {
+    if (declRef.isDistributedThunk()) {
+      auto afd = declRef.getAbstractFunctionDecl();
+      auto DC = afd->getDeclContext();
+      if (isa<ProtocolDecl>(DC)) {
+        return;
+      }
+    }
+
     addLinkEntity(LinkEntity::forMethodDescriptor(declRef));
   }
 

lib/IRGen/Linking.cpp

@@ -1304,6 +1304,18 @@ bool LinkEntity::isText() const {
   }
 }
 
+bool LinkEntity::isDistributedThunk() const {
+  if (!hasDecl())
+    return false;
+
+  auto value = getDecl();
+  if (auto afd = dyn_cast<AbstractFunctionDecl>(value)) {
+    return afd->isDistributedThunk();
+  }
+
+  return false;
+}
+
 bool LinkEntity::isWeakImported(ModuleDecl *module) const {
   switch (getKind()) {
   case Kind::SILGlobalVariable:

lib/SIL/IR/SILDeclRef.cpp

@@ -124,11 +124,11 @@ bool swift::requiresForeignEntryPoint(ValueDecl *vd) {
 }
 
 SILDeclRef::SILDeclRef(ValueDecl *vd, SILDeclRef::Kind kind, bool isForeign,
-                       bool isDistributed, bool isKnownToBeLocal,
+                       bool isDistributedThunk, bool isKnownToBeLocal,
                        bool isRuntimeAccessible,
                        SILDeclRef::BackDeploymentKind backDeploymentKind,
                        AutoDiffDerivativeFunctionIdentifier *derivativeId)
-    : loc(vd), kind(kind), isForeign(isForeign), isDistributed(isDistributed),
+    : loc(vd), kind(kind), isForeign(isForeign), distributedThunk(isDistributedThunk),
       isKnownToBeLocal(isKnownToBeLocal),
       isRuntimeAccessible(isRuntimeAccessible),
       backDeploymentKind(backDeploymentKind), defaultArgIndex(0),
@@ -186,7 +186,7 @@ SILDeclRef::SILDeclRef(SILDeclRef::Loc baseLoc, bool asForeign,
   }
 
   isForeign = asForeign;
-  isDistributed = asDistributed;
+  distributedThunk = asDistributed;
   isKnownToBeLocal = asDistributedKnownToBeLocal;
 }
 
@@ -1062,10 +1062,20 @@ bool SILDeclRef::isNativeToForeignThunk() const {
 }
 
 bool SILDeclRef::isDistributedThunk() const {
-  if (!isDistributed)
+  if (!distributedThunk)
     return false;
   return kind == Kind::Func;
 }
+bool SILDeclRef::isDistributed() const {
+  if (!hasFuncDecl())
+    return false;
+
+  if (auto decl = getFuncDecl()) {
+    return decl->isDistributed();
+  }
+
+  return false;
+}
 
 bool SILDeclRef::isBackDeploymentFallback() const {
   if (backDeploymentKind != BackDeploymentKind::Fallback)

lib/SIL/IR/SILSymbolVisitor.cpp

@@ -811,6 +811,7 @@ class SILSymbolVisitorImpl : public ASTVisitor<SILSymbolVisitorImpl> {
                   V.Ctx.getOpts().WitnessMethodElimination} {}
 
         void addMethod(SILDeclRef declRef) {
+          // TODO: alternatively maybe prevent adding distributed thunk here rather than inside those?
           if (Resilient || WitnessMethodElimination) {
             Visitor.addDispatchThunk(declRef);
             Visitor.addMethodDescriptor(declRef);

lib/SILGen/SILGenType.cpp

@@ -401,14 +401,6 @@ template<typename T> class SILGenWitnessTable : public SILWitnessVisitor<T> {
 
 public:
   void addMethod(SILDeclRef requirementRef) {
-    // TODO: here the requirement is thunk_decl of the protocol; it is a FUNC
-    // detect here that it is a func dec + thunk.
-    // walk up to DC, and find storage.
-    // e  requirementRef->getDecl()->dump()
-    //(func_decl implicit "distributedVariable()" interface type="<Self where Self : WorkerProtocol> (Self) -> () async throws -> String" access=internal nonisolated distributed_thunk
-    //  (parameter "self")
-    //  (parameter_list))
-
     auto reqDecl = requirementRef.getDecl();
 
     // Static functions can be witnessed by enum cases with payload
@@ -473,7 +465,6 @@ template<typename T> class SILGenWitnessTable : public SILWitnessVisitor<T> {
     // Here we notice a `distributed var` thunk requirement,
     // and witness it with the distributed thunk -- the "getter thunk".
     if (requirementRef.isDistributedThunk()) {
-
       return addMethodImplementation(
           requirementRef, getWitnessRef(requirementRef, witnessStorage->getDistributedThunk()),
           witness);

test/Distributed/Runtime/distributed_actor_protocol_call_resilient_lib.swift

@@ -0,0 +1,117 @@
+// RUN: %empty-directory(%t/src)
+// RUN: split-file %s %t/src
+
+/// Build the fake actor systems lib
+// RUN: %target-build-swift                                                    \
+// RUN:     -Xfrontend -disable-availability-checking                          \
+// RUN:     -parse-as-library -emit-library                                    \
+// RUN:     -emit-module-path %t/FakeDistributedActorSystems.swiftmodule       \
+// RUN:     -module-name FakeDistributedActorSystems                           \
+// RUN:      %S/../Inputs/FakeDistributedActorSystems.swift                    \
+// RUN:     -enable-library-evolution                                          \
+// RUN:     -o %t/%target-library-name(FakeDistributedActorSystems)
+
+/// Build the Lib
+// RUN: %target-build-swift                                                    \
+// RUN:     -Xfrontend -disable-availability-checking                          \
+// RUN:     -parse-as-library -emit-library                                    \
+// RUN:     -emit-module-path %t/ResilientLib.swiftmodule                      \
+// RUN:     -module-name ResilientLib                                          \
+// RUN:     -I %t                                                              \
+// RUN:     -L %t                                                              \
+// RUN:     %t/src/ResilientLib.swift                                          \
+// RUN:     -enable-library-evolution                                          \
+// RUN:     -o %t/%target-library-name(ResilientLib)
+
+/// Build the ActorLib
+// RUN: %target-build-swift                                                    \
+// RUN:     -Xfrontend -disable-availability-checking                          \
+// RUN:     -parse-as-library -emit-library                                    \
+// RUN:     -emit-module-path %t/ResilientActorLib.swiftmodule                 \
+// RUN:     -module-name ResilientActorLib                                     \
+// RUN:     -I %t                                                              \
+// RUN:     -L %t                                                              \
+// RUN:     %t/src/ResilientActorLib.swift                                     \
+// RUN:     -lFakeDistributedActorSystems                                      \
+// RUN:     -lResilientLib                                                     \
+// RUN:     -enable-library-evolution                                          \
+// RUN:     -o %t/%target-library-name(ResilientActorLib)
+
+/// Build the client
+// RUN: %target-build-swift                                                    \
+// RUN:     -Xfrontend -disable-availability-checking                          \
+// RUN:     -parse-as-library                                                  \
+// RUN:     -lFakeDistributedActorSystems                                      \
+// RUN:     -lResilientLib                                                     \
+// RUN:     -lResilientActorLib                                                \
+// RUN:     -module-name main                                                  \
+// RUN:     -I %t                                                              \
+// RUN:     -L %t                                                              \
+// RUN:     %s                                                                 \
+// RUN:     -enable-library-evolution                                          \
+// RUN:     -o %t/a.out
+
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// Locating the built libraries failed on Linux (construction of test case),
+// but we primarily care about macOS in this test
+// UNSUPPORTED: OS=linux-gnu
+
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+//--- ResilientLib.swift
+
+import Distributed
+
+public protocol SomeProtocol {
+  func function() async throws -> String
+}
+
+//--- ResilientActorLib.swift
+
+import ResilientLib
+
+import Distributed
+import FakeDistributedActorSystems
+
+public distributed actor Impl: SomeProtocol {
+  public typealias ActorSystem = FakeRoundtripActorSystem
+
+  public distributed func function() async throws -> String {
+    "Success!"
+  }
+}
+
+//--- Main.swift
+
+import ResilientLib
+import ResilientActorLib
+
+import Distributed
+import FakeDistributedActorSystems
+
+@main struct Main {
+  static func main() async {
+    let system = FakeRoundtripActorSystem()
+
+    print("start")
+
+    let impl = Impl(actorSystem: system)
+
+    let anyAct: any SomeProtocol = impl
+    let anyReply = try! await anyAct.function()
+    print("any reply = \(anyReply)") // CHECK: any reply = Success!
+
+    let proxy: any SomeProtocol = try! Impl.resolve(id: impl.id, using: system)
+    let proxyReply = try! await proxy.function()
+    print("proxy reply = \(proxyReply)") // CHECK: proxy reply = Success!
+
+    print("done")
+  }
+}