[CSApply/Distributed] Identify implicitly throwing calls while applying solution

include/swift/AST/Decl.h

@@ -5300,6 +5300,9 @@ class VarDecl : public AbstractStorageDecl {
   /// Returns true if the name is the self identifier and is implicit.
   bool isSelfParameter() const;
 
+  /// Check whether the variable is the "self" of an actor method.
+  bool isActorSelf() const;
+
   /// Determine whether this property will be part of the implicit memberwise
   /// initializer.
   ///

include/swift/Sema/ConstraintSystem.h

@@ -1879,6 +1879,17 @@ class SolutionApplicationTarget {
         !expression.pattern->isImplicit();
   }
 
+  /// Check whether this is an initializaion for `async let` pattern.
+  bool isAsyncLetInitializer() const {
+    if (!(kind == Kind::expression &&
+          expression.contextualPurpose == CTP_Initialization))
+      return false;
+
+    if (auto *PBD = getInitializationPatternBindingDecl())
+      return PBD->isAsyncLet();
+    return false;
+  }
+
   /// Whether to bind the types of any variables within the pattern via
   /// one-way constraints.
   bool shouldBindPatternVarsOneWay() const {

lib/AST/Decl.cpp

@@ -6042,6 +6042,20 @@ bool VarDecl::isSelfParameter() const {
   return false;
 }
 
+bool VarDecl::isActorSelf() const {
+  if (!isSelfParameter() && !isSelfParamCapture())
+    return false;
+
+  auto *dc = getDeclContext();
+  while (!dc->isTypeContext() && !dc->isModuleScopeContext())
+    dc = dc->getParent();
+
+  // Check if this `self` parameter belogs to an actor declaration or
+  // extension.
+  auto nominal = dc->getSelfNominalTypeDecl();
+  return nominal && nominal->isActor();
+}
+
 /// Whether the given variable is the backing storage property for
 /// a declared property that is either `lazy` or has an attached
 /// property wrapper.

lib/Sema/CSApply.cpp

@@ -679,6 +679,10 @@ namespace {
                                   LookUpConformanceInModule(cs.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);
+
   public:
     /// Build a reference to the given declaration.
     Expr *buildDeclRef(SelectedOverload overload, DeclNameLoc loc,
@@ -7581,6 +7585,14 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
     apply->setArgs(args);
     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())
+        apply->setImplicitlyThrows(true);
+    }
+
     solution.setExprTypes(apply);
     Expr *result = TypeChecker::substituteInputSugarTypeForResult(apply);
     cs.cacheExprTypes(result);
@@ -7658,6 +7670,90 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
   return finishApply(apply, openedType, locator, ctorLocator);
 }
 
+/// Determine whether this closure should be treated as Sendable.
+static bool isSendableClosure(ConstraintSystem &cs,
+                              const AbstractClosureExpr *closure) {
+  if (auto fnType = cs.getType(const_cast<AbstractClosureExpr *>(closure))
+                        ->getAs<FunctionType>()) {
+    if (fnType->isSendable())
+      return true;
+  }
+
+  return false;
+}
+
+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;
+
+  auto *actor = getReferencedParamOrCapture(
+      cast<SelfApplyExpr>(context)->getBase(),
+      [&](OpaqueValueExpr *opaqueValue) -> Expr * {
+        for (const auto &existential : OpenedExistentials) {
+          if (existential.OpaqueValue == opaqueValue)
+            return existential.ExistentialValue;
+        }
+        return nullptr;
+      });
+
+  if (!actor)
+    return false;
+
+  // If this is a method reference on an potentially isolated
+  // actor then it cannot be a remote thunk.
+  if (isPotentiallyIsolatedActor(actor, [&](ParamDecl *P) {
+        return P->isIsolated() ||
+               llvm::is_contained(solution.isolatedParams, P);
+      }))
+    return false;
+
+  bool isInAsyncLetInitializer = target && target->isAsyncLetInitializer();
+
+  auto isActorInitOrDeInitContext = [&](const DeclContext *dc) {
+    return ::isActorInitOrDeInitContext(
+        dc, [&](const AbstractClosureExpr *closure) {
+          return isSendableClosure(cs, closure);
+        });
+  };
+
+  switch (ActorIsolationRestriction::forDeclaration(ref, dc)) {
+  case ActorIsolationRestriction::CrossActorSelf: {
+    // Not a thunk if it's used in actor init or de-init.
+    if (!isInAsyncLetInitializer && isActorInitOrDeInitContext(dc))
+      return false;
+
+    // Here we know that the method could be used across actors
+    // and the actor it's used on is non-isolated, which means
+    // that it could be a thunk, so we have to be conservative
+    // about it.
+    return true;
+  }
+
+  case ActorIsolationRestriction::ActorSelf: {
+    // An instance member of an actor can be referenced from an actor's
+    // designated initializer or deinitializer.
+    if (actor->isActorSelf() && !isInAsyncLetInitializer) {
+      if (auto *fn = isActorInitOrDeInitContext(dc)) {
+        if (!(isa<ConstructorDecl>(fn) &&
+              cast<ConstructorDecl>(fn)->isConvenienceInit()))
+          return false;
+      }
+    }
+
+    // Call on a non-isolated actor in async context requires
+    // implicit thunk.
+    return isInAsyncLetInitializer || cs.isAsynchronousContext(dc);
+  }
+
+  default:
+    return false;
+  }
+}
+
 // Return the precedence-yielding parent of 'expr', along with the index of
 // 'expr' as the child of that parent. The precedence-yielding parent is the
 // nearest ancestor of 'expr' which imposes a minimum precedence on 'expr'.
@@ -8279,12 +8375,9 @@ static Optional<SolutionApplicationTarget> applySolutionToInitialization(
 
   // For an async let, wrap the initializer appropriately to make it a child
   // task.
-  if (auto patternBinding = target.getInitializationPatternBindingDecl()) {
-    if (patternBinding->isAsyncLet()) {
-      resultTarget.setExpr(
-          wrapAsyncLetInitializer(
-            cs, resultTarget.getAsExpr(), resultTarget.getDeclContext()));
-    }
+  if (target.isAsyncLetInitializer()) {
+    resultTarget.setExpr(wrapAsyncLetInitializer(
+        cs, resultTarget.getAsExpr(), resultTarget.getDeclContext()));
   }
 
   return resultTarget;