[Distributed] Sema: Add a new distributed-thunk attribute

The attribute comes handy during solution application to
determine whether the call is using a distributed thunk.

Author: xedin

include/swift/AST/Attr.def

@@ -735,6 +735,12 @@ CONTEXTUAL_SIMPLE_DECL_ATTR(_local, KnownToBeLocal,
   APIBreakingToAdd | APIBreakingToRemove,
   130)
 
+SIMPLE_DECL_ATTR(_distributed_thunk, DistributedThunk,
+  OnFunc | UserInaccessible |
+  ABIBreakingToAdd | ABIBreakingToRemove |
+  APIBreakingToAdd | APIBreakingToRemove,
+  131)
+
 // If you're adding a new underscored attribute here, please document it in
 // docs/ReferenceGuides/UnderscoredAttributes.md.
 

include/swift/AST/Decl.h

@@ -6436,6 +6436,10 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Returns 'true' if the function is distributed.
   bool isDistributed() const;
 
+  /// Is this a thunk function used to access a distributed method outside
+  /// of its actor isolation context?
+  bool isDistributedThunk() const;
+
   /// For a 'distributed' target (func or computed property),
   /// get the 'thunk' responsible for performing the 'remoteCall'.
   ///

include/swift/AST/DiagnosticsSema.def

@@ -4604,6 +4604,9 @@ ERROR(distributed_actor_isolated_method,none,
 ERROR(distributed_local_cannot_be_used,none,
       "'local' cannot be used in user-defined code currently",
       ())
+ERROR(distributed_thunk_cannot_be_used,none,
+      "'distributed_thunk' cannot be used in user-defined code",
+      ())
 ERROR(distributed_actor_func_param_not_codable,none,
       "parameter '%0' of type %1 in %2 does not conform to serialization requirement '%3'",
       (StringRef, Type, DescriptiveDeclKind, StringRef))

lib/AST/ASTDumper.cpp

@@ -877,6 +877,9 @@ namespace {
       if (D->isDistributed()) {
         PrintWithColorRAII(OS, ExprModifierColor) << " distributed";
       }
+      if (D->isDistributedThunk()) {
+        PrintWithColorRAII(OS, ExprModifierColor) << " distributed-thunk";
+      }
 
       if (auto fac = D->getForeignAsyncConvention()) {
         OS << " foreign_async=";
@@ -1350,6 +1353,10 @@ void ValueDecl::dumpRef(raw_ostream &os) const {
     os << " known-to-be-local";
   }
 
+  if (getAttrs().hasAttribute<DistributedThunkAttr>()) {
+    os << " distributed-thunk";
+  }
+
   // Print location.
   auto &srcMgr = getASTContext().SourceMgr;
   if (getLoc().isValid()) {

lib/AST/DistributedDecl.cpp

@@ -1290,6 +1290,10 @@ bool AbstractFunctionDecl::isDistributed() const {
   return getAttrs().hasAttribute<DistributedActorAttr>();
 }
 
+bool AbstractFunctionDecl::isDistributedThunk() const {
+  return getAttrs().hasAttribute<DistributedThunkAttr>();
+}
+
 ConstructorDecl *
 NominalTypeDecl::getDistributedRemoteCallTargetInitFunction() const {
   auto mutableThis = const_cast<NominalTypeDecl *>(this);

lib/Sema/CSApply.cpp

@@ -516,10 +516,6 @@ namespace {
                                   LookUpConformanceInModule(dc->getParentModule()));
     }
 
-    /// Determine whether the given reference is to a method on
-    /// 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,
@@ -1649,6 +1645,10 @@ namespace {
 
         ConcreteDeclRef thunkRef{thunkDecl, memberRef.getSubstitutions()};
 
+        // Update member reference to point to the thunk.
+        CachedConcreteRefs[cs.getConstraintLocator(memberLocator)] =
+            thunkRef;
+
         auto declRefExpr = new (context) DeclRefExpr(
             thunkRef, memberLoc, Implicit, AccessSemantics::DirectToStorage);
 
@@ -7790,14 +7790,14 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
     cs.setType(apply, fnType->getResult());
 
     // If this is a call to a distributed method thunk,
-    // let's mark the call as implicitly throwing.
-    if (isDistributedThunk(callee, apply->getFn())) {
-      auto *FD = cast<AbstractFunctionDecl>(callee.getDecl());
-      if (!FD->hasThrows())
+    // let's mark the call as implicitly throwing/async.
+    if (isa<SelfApplyExpr>(apply->getFn())) {
+      auto *FD = dyn_cast<FuncDecl>(callee.getDecl());
+      if (FD && FD->isDistributedThunk()) {
         apply->setImplicitlyThrows(true);
-      if (!FD->hasAsync())
         apply->setImplicitlyAsync(ImplicitActorHopTarget::forInstanceSelf());
-      apply->setUsesDistributedThunk(true);
+        apply->setUsesDistributedThunk(true);
+      }
     }
 
     solution.setExprTypes(apply);
@@ -7916,14 +7916,6 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
   return ctorCall;
 }
 
-bool ExprRewriter::isDistributedThunk(ConcreteDeclRef ref, Expr *context) {
-  if (!isa<SelfApplyExpr>(context))
-    return false;
-
-  return requiresDistributedThunk(cast<SelfApplyExpr>(context)->getBase(),
-                                  context->getLoc(), ref);
-}
-
 bool ExprRewriter::requiresDistributedThunk(Expr *base, SourceLoc memberLoc,
                                             ConcreteDeclRef memberRef) {
 

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -736,6 +736,7 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
   if (isa<ClassDecl>(DC))
     thunk->getAttrs().add(new (C) FinalAttr(/*isImplicit=*/true));
 
+  thunk->getAttrs().add(new (C) DistributedThunkAttr(/*isImplicit=*/true));
   thunk->setGenericSignature(baseSignature);
   thunk->copyFormalAccessFrom(func, /*sourceIsParentContext=*/false);
   thunk->setBodySynthesizer(deriveBodyDistributed_thunk, func);

lib/Sema/TypeCheckAttr.cpp

@@ -322,6 +322,8 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
   void visitKnownToBeLocalAttr(KnownToBeLocalAttr *attr);
 
   void visitSendableAttr(SendableAttr *attr);
+
+  void visitDistributedThunkAttr(DistributedThunkAttr *attr);
 };
 
 } // end anonymous namespace
@@ -5934,6 +5936,11 @@ void AttributeChecker::visitSendableAttr(SendableAttr *attr) {
   }
 }
 
+void AttributeChecker::visitDistributedThunkAttr(DistributedThunkAttr *attr) {
+  if (!D->isImplicit())
+    diagnoseAndRemoveAttr(attr, diag::distributed_thunk_cannot_be_used);
+}
+
 void AttributeChecker::visitNonisolatedAttr(NonisolatedAttr *attr) {
   // 'nonisolated' can be applied to global and static/class variables
   // that do not have storage.

lib/Sema/TypeCheckDeclOverride.cpp

@@ -1613,6 +1613,7 @@ namespace  {
     UNINTERESTING_ATTR(UnsafeInheritExecutor)
     UNINTERESTING_ATTR(CompilerInitialized)
 
+    UNINTERESTING_ATTR(DistributedThunk)
 #undef UNINTERESTING_ATTR
 
     void visitAvailableAttr(AvailableAttr *attr) {

lib/Sema/TypeCheckEffects.cpp

@@ -2878,8 +2878,16 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
            // callee is isolated to an actor.
            auto callee = call->getCalledValue();
            if (callee) {
+             auto declKind = callee->getDescriptiveKind();
+             if (call->usesDistributedThunk()) {
+               // We need to determine whether this is a method or a property.
+               // Computed properties always form an implicit call to the thunk.
+               declKind = call->isImplicit()
+                              ? DescriptiveDeclKind::DistributedProperty
+                              : DescriptiveDeclKind::DistributedMethod;
+             }
              Ctx.Diags.diagnose(diag.expr.getStartLoc(), diag::actor_isolated_sync_func,
-                                callee->getDescriptiveKind(), callee->getName());
+                                declKind, callee->getName());
            } else {
              Ctx.Diags.diagnose(
                  diag.expr.getStartLoc(), diag::actor_isolated_sync_func_value,