[Distributed] Allow requesting distributed thunks on AbstractStorageDecl

One step towards future distributed subscripts.

Author: xedin

include/swift/AST/Decl.h

@@ -5123,6 +5123,13 @@ class AbstractStorageDecl : public ValueDecl {
 
   bool hasAnyNativeDynamicAccessors() const;
 
+  /// Does this have a 'distributed' modifier?
+  bool isDistributed() const;
+
+  /// Return a distributed thunk if this computed property is marked as
+  /// 'distributed' and and nullptr otherwise.
+  FuncDecl *getDistributedThunk() const;
+
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Decl *D) {
     return D->getKind() >= DeclKind::First_AbstractStorageDecl &&
@@ -5356,13 +5363,6 @@ class VarDecl : public AbstractStorageDecl {
   /// Is this an "async let" property?
   bool isAsyncLet() const;
 
-  /// Does this have a 'distributed' modifier?
-  bool isDistributed() const;
-
-  /// Return a distributed thunk if this computed property is marked as
-  /// 'distributed' and and nullptr otherwise.
-  FuncDecl *getDistributedThunk() const;
-
   /// Is this var known to be a "local" distributed actor,
   /// if so the implicit throwing ans some isolation checks can be skipped.
   bool isKnownToBeLocal() const;

include/swift/AST/TypeCheckRequests.h

@@ -1238,11 +1238,13 @@ class GetDistributedRemoteCallArgumentInitFunctionRequest :
 /// The thunk is responsible for invoking 'remoteCall' when invoked on a remote
 /// 'distributed actor'.
 class GetDistributedThunkRequest
-    : public SimpleRequest<
-          GetDistributedThunkRequest,
-          FuncDecl *(llvm::PointerUnion<VarDecl *, AbstractFunctionDecl *>),
-          RequestFlags::Cached> {
-  using Originator = llvm::PointerUnion<VarDecl *, AbstractFunctionDecl *>;
+    : public SimpleRequest<GetDistributedThunkRequest,
+                           FuncDecl *(
+                               llvm::PointerUnion<AbstractStorageDecl *,
+                                                  AbstractFunctionDecl *>),
+                           RequestFlags::Cached> {
+  using Originator =
+      llvm::PointerUnion<AbstractStorageDecl *, AbstractFunctionDecl *>;
 
 public:
   using SimpleRequest::SimpleRequest;

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -127,7 +127,7 @@ SWIFT_REQUEST(TypeChecker, GetDistributedTargetInvocationResultHandlerOnReturnFu
               AbstractFunctionDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetDistributedThunkRequest,
-              FuncDecl *(llvm::PointerUnion<VarDecl *, AbstractFunctionDecl *>),
+              FuncDecl *(llvm::PointerUnion<AbstractStorageDecl *, AbstractFunctionDecl *>),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetDistributedActorIDPropertyRequest,
               VarDecl *(NominalTypeDecl *),

lib/AST/Decl.cpp

@@ -6455,10 +6455,6 @@ bool VarDecl::isAsyncLet() const {
   return getAttrs().hasAttribute<AsyncAttr>();
 }
 
-bool VarDecl::isDistributed() const {
-  return getAttrs().hasAttribute<DistributedActorAttr>();
-}
-
 bool VarDecl::isKnownToBeLocal() const {
   return getAttrs().hasAttribute<KnownToBeLocalAttr>();
 }

lib/AST/DistributedDecl.cpp

@@ -1290,6 +1290,10 @@ bool AbstractFunctionDecl::isDistributed() const {
   return getAttrs().hasAttribute<DistributedActorAttr>();
 }
 
+bool AbstractStorageDecl::isDistributed() const {
+  return getAttrs().hasAttribute<DistributedActorAttr>();
+}
+
 ConstructorDecl *
 NominalTypeDecl::getDistributedRemoteCallTargetInitFunction() const {
   auto mutableThis = const_cast<NominalTypeDecl *>(this);
@@ -1333,11 +1337,11 @@ AbstractFunctionDecl *ASTContext::getRemoteCallOnDistributedActorSystem(
 /********************** Distributed Actor Properties **************************/
 /******************************************************************************/
 
-FuncDecl *VarDecl::getDistributedThunk() const {
+FuncDecl *AbstractStorageDecl::getDistributedThunk() const {
   if (!isDistributed())
     return nullptr;
 
-  auto mutableThis = const_cast<VarDecl *>(this);
+  auto mutableThis = const_cast<AbstractStorageDecl *>(this);
   return evaluateOrDefault(getASTContext().evaluator,
                            GetDistributedThunkRequest{mutableThis}, nullptr);
 }

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -836,14 +836,18 @@ addDistributedActorCodableConformance(
 FuncDecl *GetDistributedThunkRequest::evaluate(Evaluator &evaluator,
                                                Originator originator) const {
   AbstractFunctionDecl *distributedTarget = nullptr;
-  if (auto *var = originator.dyn_cast<VarDecl *>()) {
-    if (!var->isDistributed())
+  if (auto *storage = originator.dyn_cast<AbstractStorageDecl *>()) {
+    if (!storage->isDistributed())
       return nullptr;
 
-    if (checkDistributedActorProperty(var, /*diagnose=*/false))
-      return nullptr;
+    if (auto *var = dyn_cast<VarDecl>(storage)) {
+      if (checkDistributedActorProperty(var, /*diagnose=*/false))
+        return nullptr;
 
-    distributedTarget = var->getAccessor(AccessorKind::Get);
+      distributedTarget = var->getAccessor(AccessorKind::Get);
+    } else {
+      llvm_unreachable("unsupported storage kind");
+    }
   } else {
     distributedTarget = originator.get<AbstractFunctionDecl *>();
     if (!distributedTarget->isDistributed())