[Distributed] Complete handling of protocol calls and witnesses using adjusted mangling scheme (#72416)

Author: ktoso

include/swift/AST/AccessorKinds.def

@@ -121,6 +121,18 @@
 /// move-only types are introduced.
 OBJC_ACCESSOR(Get, get)
 
+/// This is a "distributed thunk" getter, that is used for a "distributed var"
+/// 'get { ... }' accessor.
+///
+/// A distributed thunk is mangled as it's Storage's name.
+/// This is never invoked directly by user code, but whenever making calls on
+/// a distributed actor and the target is not known to be local, we invoke the
+/// thunk. This is equivalent to a `distributed_thunk' FuncDecl of a
+/// 'distributed func'.
+///
+/// For the most part, its behavior is exactly the same as a normal getter.
+OBJC_ACCESSOR(DistributedGet, _distributed_get)
+
 /// This is a setter: a function that is called when a value is assigned
 /// to the storage.  It takes a value of the storage type as its
 /// primary parameter, in addition to any index values that may be

include/swift/AST/Decl.h

@@ -2788,6 +2788,12 @@ class ValueDecl : public Decl {
   /// can be distributed.
   bool isDistributed() const;
 
+  /// Is this a '_distributed_get' accessor?
+  ///
+  /// These are special accessors used by distributed thunks, implementing
+  /// `distributed var get { }` accessors.
+  bool isDistributedGetAccessor() const;
+
   bool hasName() const { return bool(Name); }
   bool isOperator() const { return Name.isOperator(); }
 
@@ -8050,7 +8056,7 @@ class AccessorDecl final : public FuncDecl {
                  throws, throwsLoc, thrownTy, hasImplicitSelfDecl,
                  /*genericParams*/ nullptr, parent),
         AccessorKeywordLoc(accessorKeywordLoc), Storage(storage) {
-    assert(!async || accessorKind == AccessorKind::Get
+    assert(!async || (accessorKind == AccessorKind::Get || accessorKind == AccessorKind::DistributedGet)
            && "only get accessors can be async");
     Bits.AccessorDecl.AccessorKind = unsigned(accessorKind);
   }
@@ -8082,6 +8088,14 @@ class AccessorDecl final : public FuncDecl {
          TypeLoc thrownType, ParameterList *parameterList, Type fnRetType,
          DeclContext *parent, ClangNode clangNode = ClangNode());
 
+  static AccessorDecl *createImplicit(ASTContext &Context,
+                                  AccessorKind accessorKind,
+                                  AbstractStorageDecl *storage,
+                                  bool async,
+                                  bool throws, TypeLoc thrownType,
+                                  Type fnRetType,
+                                  DeclContext *parent);
+
   /// Create a parsed accessor.
   ///
   /// \param paramList A parameter list for e.g \c set(newValue), or \c nullptr
@@ -8110,6 +8124,7 @@ class AccessorDecl final : public FuncDecl {
   }
 
   bool isGetter() const { return getAccessorKind() == AccessorKind::Get; }
+  bool isDistributedGetter() const { return getAccessorKind() == AccessorKind::DistributedGet; }
   bool isSetter() const { return getAccessorKind() == AccessorKind::Set; }
   bool isAnyAddressor() const {
     auto kind = getAccessorKind();

include/swift/SIL/SILWitnessVisitor.h

@@ -129,6 +129,8 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
         asDerived().addMethod(SILDeclRef(accessor, SILDeclRef::Kind::Func));
         addAutoDiffDerivativeMethodsIfRequired(accessor,
                                                SILDeclRef::Kind::Func);
+        addDistributedWitnessMethodsIfRequired(accessor,
+                                               SILDeclRef::Kind::Func);
       }
     });
   }
@@ -201,10 +203,13 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
 
   void addDistributedWitnessMethodsIfRequired(AbstractFunctionDecl *AFD,
                                               SILDeclRef::Kind kind) {
-    if (!AFD || !AFD->isDistributed())
+    if (!AFD)
       return;
 
     auto thunk = AFD->getDistributedThunk();
+    if (!thunk)
+      return;
+
     SILDeclRef declRef(thunk, kind);
     asDerived().addMethod(declRef.asDistributed());
   }

lib/APIDigester/ModuleAnalyzerNodes.cpp

@@ -1745,6 +1745,7 @@ SDKContext::shouldIgnore(Decl *D, const Decl* Parent) const {
     if (!checkingABI()) {
       switch (ACC->getAccessorKind()) {
       case AccessorKind::Get:
+      case AccessorKind::DistributedGet:
       case AccessorKind::Set:
         break;
       default:

lib/AST/ASTMangler.cpp

@@ -82,6 +82,10 @@ static StringRef getCodeForAccessorKind(AccessorKind kind) {
   switch (kind) {
   case AccessorKind::Get:
     return "g";
+  case AccessorKind::DistributedGet:
+    // TODO(distributed): probably does not matter since we mangle distributed
+    //  thunk getters as the name of the Storage?
+    return "g";
   case AccessorKind::Set:
     return "s";
   case AccessorKind::WillSet:
@@ -4283,7 +4287,17 @@ void ASTMangler::appendDistributedThunk(
     appendContextOf(thunk, base);
   }
 
-  appendIdentifier(thunk->getBaseName().getIdentifier().str());
+  if (auto accessor = dyn_cast<AccessorDecl>(thunk)) {
+    assert(accessor->getAccessorKind() == AccessorKind::DistributedGet &&
+           "Only accessors marked as _distributed_get are expected to be "
+           "mangled as thunks");
+    // A distributed getter is mangled as the name of its storage (i.e. "the
+    // var")
+    appendIdentifier(accessor->getStorage()->getBaseIdentifier().str());
+  } else {
+    appendIdentifier(thunk->getBaseIdentifier().str());
+  }
+
   appendDeclType(thunk, base, FunctionMangling);
   appendOperator("F");
   appendSymbolKind(SymbolKind::DistributedThunk);

lib/AST/ASTPrinter.cpp

@@ -4011,6 +4011,7 @@ void PrintAST::visitAccessorDecl(AccessorDecl *decl) {
 
   switch (auto kind = decl->getAccessorKind()) {
   case AccessorKind::Get:
+  case AccessorKind::DistributedGet:
   case AccessorKind::Address:
   case AccessorKind::Read:
   case AccessorKind::Modify:

lib/AST/AccessRequests.cpp

@@ -52,6 +52,7 @@ AccessLevelRequest::evaluate(Evaluator &evaluator, ValueDecl *D) const {
     AbstractStorageDecl *storage = accessor->getStorage();
     switch (accessor->getAccessorKind()) {
     case AccessorKind::Get:
+    case AccessorKind::DistributedGet:
     case AccessorKind::Address:
     case AccessorKind::Read:
       return storage->getFormalAccess();

lib/AST/Decl.cpp

@@ -232,6 +232,7 @@ DescriptiveDeclKind Decl::getDescriptiveKind() const {
 
      switch (accessor->getAccessorKind()) {
      case AccessorKind::Get:
+     case AccessorKind::DistributedGet:
        return DescriptiveDeclKind::Getter;
 
      case AccessorKind::Set:
@@ -993,8 +994,10 @@ std::optional<Type> AbstractFunctionDecl::getEffectiveThrownErrorType() const {
   // has a cyclic reference if we try to get its interface type here. Find a
   // better way to express this.
   if (auto accessor = dyn_cast<AccessorDecl>(this)) {
-    if (accessor->getAccessorKind() != AccessorKind::Get)
+    if (accessor->getAccessorKind() != AccessorKind::Get &&
+        accessor->getAccessorKind() != AccessorKind::DistributedGet) {
       return std::nullopt;
+    }
   }
 
   Type interfaceType = getInterfaceType();
@@ -2787,6 +2790,7 @@ bool AbstractStorageDecl::requiresOpaqueAccessors() const {
 bool AbstractStorageDecl::requiresOpaqueAccessor(AccessorKind kind) const {
   switch (kind) {
   case AccessorKind::Get:
+  case AccessorKind::DistributedGet:
     return requiresOpaqueGetter();
   case AccessorKind::Set:
     return requiresOpaqueSetter();
@@ -6805,6 +6809,8 @@ StringRef swift::getAccessorNameForDiagnostic(AccessorKind accessorKind,
   switch (accessorKind) {
   case AccessorKind::Get:
     return article ? "a getter" : "getter";
+  case AccessorKind::DistributedGet:
+    return article ? "a distributed getter" : "distributed getter";
   case AccessorKind::Set:
     return article ? "a setter" : "setter";
   case AccessorKind::Address:
@@ -8958,6 +8964,7 @@ DeclName AbstractFunctionDecl::getEffectiveFullName() const {
     case AccessorKind::Address:
     case AccessorKind::MutableAddress:
     case AccessorKind::Get:
+    case AccessorKind::DistributedGet:
     case AccessorKind::Read:
     case AccessorKind::Modify:
       return subscript ? subscript->getName()
@@ -10148,6 +10155,25 @@ AccessorDecl *AccessorDecl::create(ASTContext &ctx, SourceLoc declLoc,
   return D;
 }
 
+AccessorDecl *AccessorDecl::createImplicit(ASTContext &ctx,
+                                           AccessorKind accessorKind,
+                                           AbstractStorageDecl *storage,
+                                           bool async, bool throws,
+                                           TypeLoc thrownType,
+                                           Type fnRetType,
+                                           DeclContext *parent) {
+  AccessorDecl *D = AccessorDecl::createImpl(
+      ctx, /*declLoc=*/SourceLoc(),
+      /*accessorKeywordLoc=*/SourceLoc(), accessorKind,
+      storage, async, /*asyncLoc=*/SourceLoc(),
+      /*throws=*/true, /*throwsLoc=*/SourceLoc(),
+      thrownType, parent,
+      /*clangNode=*/ClangNode());
+  D->setImplicit();
+  D->setResultInterfaceType(fnRetType);
+  return D;
+}
+
 AccessorDecl *AccessorDecl::createParsed(
     ASTContext &ctx, AccessorKind accessorKind, AbstractStorageDecl *storage,
     SourceLoc declLoc, SourceLoc accessorKeywordLoc, ParameterList *paramList,
@@ -10226,6 +10252,7 @@ StringRef AccessorDecl::implicitParameterNameFor(AccessorKind kind) {
   case AccessorKind::DidSet:
     return "oldValue";
   case AccessorKind::Get:
+  case AccessorKind::DistributedGet:
   case AccessorKind::Read:
   case AccessorKind::Modify:
   case AccessorKind::Address:
@@ -10237,6 +10264,7 @@ StringRef AccessorDecl::implicitParameterNameFor(AccessorKind kind) {
 bool AccessorDecl::isAssumedNonMutating() const {
   switch (getAccessorKind()) {
   case AccessorKind::Get:
+  case AccessorKind::DistributedGet:
   case AccessorKind::Address:
   case AccessorKind::Read:
     return true;
@@ -10270,6 +10298,9 @@ void AccessorDecl::printUserFacingName(raw_ostream &out) const {
   case AccessorKind::Get:
     out << "getter:";
     break;
+  case AccessorKind::DistributedGet:
+    out << "_distributed_getter:";
+    break;
   case AccessorKind::Set:
     out << "setter:";
     break;

lib/AST/DistributedDecl.cpp

@@ -1289,6 +1289,16 @@ bool ValueDecl::isDistributed() const {
   return getAttrs().hasAttribute<DistributedActorAttr>();
 }
 
+bool ValueDecl::isDistributedGetAccessor() const {
+  if (auto accessor = dyn_cast<AccessorDecl>(this)) {
+    if (accessor->getAccessorKind() == AccessorKind::DistributedGet) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 ConstructorDecl *
 NominalTypeDecl::getDistributedRemoteCallTargetInitFunction() const {
   auto mutableThis = const_cast<NominalTypeDecl *>(this);
@@ -1349,10 +1359,22 @@ AbstractFunctionDecl::getDistributedThunk() const {
   if (isDistributedThunk())
     return const_cast<FuncDecl *>(dyn_cast<FuncDecl>(this));
 
+  auto mutableThis = const_cast<AbstractFunctionDecl *>(this);
+
+  // For an accessor, get the Storage (VarDecl) and get the thunk off it.
+  //
+  // Since only 'get' computed distributed properties are allowed, we know
+  // this will be the equivalent 'get' thunk for this AccessorDecl.
+  //
+  // The AccessorDecl is not marked distributed, but the VarDecl will be.
+  if (auto accessor = dyn_cast<AccessorDecl>(mutableThis)) {
+    auto Storage = accessor->getStorage();
+    return Storage->getDistributedThunk();
+  }
+
   if (!isDistributed())
     return nullptr;
 
-  auto mutableThis = const_cast<AbstractFunctionDecl *>(this);
   return evaluateOrDefault(
       getASTContext().evaluator,
       GetDistributedThunkRequest{mutableThis},

lib/AST/PrettyStackTrace.cpp

@@ -67,6 +67,9 @@ void swift::printDeclDescription(llvm::raw_ostream &out, const Decl *D,
         case AccessorKind::Get:
           out << "getter";
           break;
+        case AccessorKind::DistributedGet:
+          out << "_distributed_getter";
+          break;
         case AccessorKind::Set:
           out << "setter";
           break;