[CSApply] Distributed: Inject distributed thunk when necessary

Replace distributed member references with distributed thunks
when access happens outside of distributed actor context. This
significantly simplifies distributed compute properties implementation.

Author: xedin

lib/SILGen/SILGenApply.cpp

@@ -1123,9 +1123,7 @@ class SILGenApply : public Lowering::ExprVisitor<SILGenApply> {
 
     /// Some special handling may be necessary for thunks:
     if (callSite && callSite->shouldApplyDistributedThunk()) {
-      if (auto distributedThunk = cast<AbstractFunctionDecl>(e->getDecl())->getDistributedThunk()) {
-        constant = SILDeclRef(distributedThunk).asDistributed();
-      }
+      constant = SILDeclRef(e->getDecl()).asDistributed();
     } else if (afd->isBackDeployed()) {
       // If we're calling a back deployed function then we need to call a
       // thunk instead that will handle the fallback when the original

lib/SILGen/SILGenType.cpp

@@ -1161,6 +1161,12 @@ class SILGenType : public TypeMemberVisitor<SILGenType> {
       SGM.emitPropertyWrapperBackingInitializer(vd);
     }
 
+    if (auto *thunk = vd->getDistributedThunk()) {
+      auto thunkRef = SILDeclRef(thunk).asDistributed();
+      SGM.emitFunctionDefinition(thunkRef,
+                                 SGM.getFunction(thunkRef, ForDefinition));
+    }
+
     visitAbstractStorageDecl(vd);
   }
 
@@ -1287,6 +1293,13 @@ class SILGenExtension : public TypeMemberVisitor<SILGenExtension> {
         return;
       }
     }
+
+    if (auto *thunk = vd->getDistributedThunk()) {
+      auto thunkRef = SILDeclRef(thunk).asDistributed();
+      SGM.emitFunctionDefinition(thunkRef,
+                                 SGM.getFunction(thunkRef, ForDefinition));
+    }
+
     visitAbstractStorageDecl(vd);
   }
 

lib/Sema/CSApply.cpp

@@ -520,6 +520,11 @@ namespace {
     /// a remote distributed actor in the given context.
     bool isDistributedThunk(ConcreteDeclRef ref, Expr *context);
 
+    /// Determine whether the given reference on the given
+    /// base has to be replaced with a distributed thunk instead.
+    bool requiresDistributedThunk(Expr *base, SourceLoc memberLoc,
+                                  ConcreteDeclRef memberRef);
+
   public:
     /// Build a reference to the given declaration.
     Expr *buildDeclRef(SelectedOverload overload, DeclNameLoc loc,
@@ -1615,6 +1620,64 @@ namespace {
         return forceUnwrapIfExpected(ref, memberLocator);
       }
 
+      if (requiresDistributedThunk(base, memberLoc.getStartLoc(), memberRef)) {
+        auto *decl = memberRef.getDecl();
+        FuncDecl *thunkDecl = nullptr;
+        if (auto *FD = dyn_cast<FuncDecl>(decl)) {
+          thunkDecl = FD->getDistributedThunk();
+        } else {
+          thunkDecl = cast<VarDecl>(decl)->getDistributedThunk();
+        }
+
+        if (!thunkDecl)
+          return nullptr;
+
+        auto thunkType = refTy;
+        auto thunkOpenedType = openedType;
+
+        // If this is a reference to a computed property then we need to
+        // form a function type from it with unapplied Self so
+        // (Self) -> T becomes (Self) -> () -> T
+        if (isa<VarDecl>(decl)) {
+          auto extInfo = ASTExtInfoBuilder().withAsync().withThrows().build();
+
+          thunkOpenedType =
+              FunctionType::get(/*params=*/{}, thunkOpenedType, extInfo);
+          thunkType =
+              FunctionType::get({FunctionType::Param(selfTy)}, thunkOpenedType);
+        }
+
+        ConcreteDeclRef thunkRef{thunkDecl, memberRef.getSubstitutions()};
+
+        auto declRefExpr = new (context) DeclRefExpr(
+            thunkRef, memberLoc, Implicit, AccessSemantics::DirectToStorage);
+
+        declRefExpr->setFunctionRefKind(choice.getFunctionRefKind());
+        declRefExpr->setType(thunkType);
+
+        cs.cacheType(declRefExpr);
+
+        Expr *thunkApply =
+            DotSyntaxCallExpr::create(context, declRefExpr, dotLoc, base);
+        if (Implicit)
+          thunkApply->setImplicit();
+
+        thunkApply = finishApply(cast<ApplyExpr>(thunkApply), thunkOpenedType,
+                                 locator, memberLocator);
+
+        // If this is access to a computed property, that requires
+        // implicit call.
+        if (isa<VarDecl>(decl)) {
+          auto *thunkCall = CallExpr::createImplicitEmpty(context, thunkApply);
+          thunkCall->setType(solution.simplifyType(openedType));
+          thunkCall->setShouldApplyDistributedThunk(true);
+          cs.cacheType(thunkCall);
+          return thunkCall;
+        }
+
+        return thunkApply;
+      }
+
       // For properties, build member references.
       if (auto *varDecl = dyn_cast<VarDecl>(member)) {
         if (isUnboundInstanceMember) {
@@ -1662,8 +1725,9 @@ namespace {
         refTy = refTy->replaceCovariantResultType(containerTy, 2);
 
       // Handle all other references.
-      auto declRefExpr = new (context) DeclRefExpr(memberRef, memberLoc,
-                                                   Implicit, semantics);
+      auto declRefExpr =
+          new (context) DeclRefExpr(memberRef, memberLoc, Implicit, semantics);
+
       declRefExpr->setFunctionRefKind(choice.getFunctionRefKind());
       declRefExpr->setType(refTy);
       cs.setType(declRefExpr, refTy);
@@ -7734,6 +7798,9 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
       auto *FD = cast<AbstractFunctionDecl>(callee.getDecl());
       if (!FD->hasThrows())
         apply->setImplicitlyThrows(true);
+      if (!FD->hasAsync())
+        apply->setImplicitlyAsync(ImplicitActorHopTarget::forInstanceSelf());
+      apply->setShouldApplyDistributedThunk(true);
     }
 
     solution.setExprTypes(apply);
@@ -7853,16 +7920,29 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
 }
 
 bool ExprRewriter::isDistributedThunk(ConcreteDeclRef ref, Expr *context) {
-  auto *FD = dyn_cast_or_null<AbstractFunctionDecl>(ref.getDecl());
-  if (!(FD && FD->isInstanceMember() && FD->isDistributed()))
-    return false;
-
   if (!isa<SelfApplyExpr>(context))
     return false;
 
+  return requiresDistributedThunk(cast<SelfApplyExpr>(context)->getBase(),
+                                  context->getLoc(), ref);
+}
+
+bool ExprRewriter::requiresDistributedThunk(Expr *base, SourceLoc memberLoc,
+                                            ConcreteDeclRef memberRef) {
+
+  auto *memberDecl = memberRef.getDecl();
+  assert(memberDecl);
+
+  if (auto *FD = dyn_cast<AbstractFunctionDecl>(memberDecl)) {
+    if (!(FD->isInstanceMember() && FD->isDistributed()))
+      return false;
+  } else if (auto *VD = dyn_cast<VarDecl>(memberDecl)) {
+    if (!VD->isDistributed())
+      return false;
+  }
+
   auto *actor = getReferencedParamOrCapture(
-      cast<SelfApplyExpr>(context)->getBase(),
-      [&](OpaqueValueExpr *opaqueValue) -> Expr * {
+      base, [&](OpaqueValueExpr *opaqueValue) -> Expr * {
         for (const auto &existential : OpenedExistentials) {
           if (existential.OpaqueValue == opaqueValue)
             return existential.ExistentialValue;
@@ -7921,7 +8001,7 @@ bool ExprRewriter::isDistributedThunk(ConcreteDeclRef ref, Expr *context) {
   ReferencedActor actorRef = ReferencedActor(
       actor, isPotentiallyIsolated, ReferencedActor::NonIsolatedContext);
   auto refResult = ActorReferenceResult::forReference(
-      ref, context->getLoc(), referenceDC, None, actorRef);
+      memberRef, memberLoc, referenceDC, None, actorRef);
   switch (refResult) {
   case ActorReferenceResult::ExitsActorToNonisolated:
   case ActorReferenceResult::SameConcurrencyDomain:

test/Distributed/Runtime/distributed_actor_func_calls_remoteCall_computedProperty.swift

@@ -21,7 +21,7 @@ typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
 
 distributed actor Greeter {
   distributed var theDistributedProperty: String {
-    "Patrik the Seastar"
+    "Patrick the Seastar"
   }
 }
 
@@ -32,10 +32,10 @@ func test() async throws {
   let ref = try Greeter.resolve(id: local.id, using: system)
 
   let reply = try await ref.theDistributedProperty
-  // CHECK: >> remoteCall: on:main.Greeter, target:main.Greeter.name, invocation:FakeInvocationEncoder(genericSubs: [], arguments: [], returnType: Optional(Swift.String), errorType: nil), throwing:Swift.Never, returning:Swift.String
+  // CHECK: >> remoteCall: on:main.Greeter, target:main.Greeter.theDistributedProperty(), invocation:FakeInvocationEncoder(genericSubs: [], arguments: [], returnType: Optional(Swift.String), errorType: nil), throwing:Swift.Never, returning:Swift.String
   // CHECK: << remoteCall return: Patrick the Seastar
   print("reply: \(reply)")
-  // CHECK: reply: Echo: Caplin
+  // CHECK: reply: Patrick the Seastar
 }
 
 @main struct Main {

test/Distributed/distributed_actor_var_implicitly_async_throws.swift

@@ -33,15 +33,17 @@ distributed actor D {
     "dist"
   }
 
-  // OK:
+  // FIXME: The following should be accepted.
+  /*
   distributed var distGet: String {
     get distributed {
       "okey"
     }
   }
+  */
 
-  distributed var distSetGet: String {
-    set distributed {
+  distributed var distSetGet: String { // expected-error {{'distributed' computed property 'distSetGet' cannot have setter}}
+    set distributed { // expected-error {{expected '{' to start setter definition}}
       _ = newValue
     }
     get distributed {