distributed thunk flag

Author: ktoso

include/swift/AST/Expr.h

@@ -307,12 +307,13 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
     NumCaptures : 32
   );
 
-  SWIFT_INLINE_BITFIELD(ApplyExpr, Expr, 1+1+1+1+1,
+  SWIFT_INLINE_BITFIELD(ApplyExpr, Expr, 1+1+1+1+1+1,
     ThrowsIsSet : 1,
     Throws : 1,
     ImplicitlyAsync : 1,
     ImplicitlyThrows : 1,
-    NoAsync : 1
+    NoAsync : 1,
+    ShouldApplyDistributedThunk : 1
   );
 
   SWIFT_INLINE_BITFIELD_EMPTY(CallExpr, ApplyExpr);
@@ -4423,7 +4424,12 @@ class ApplyExpr : public Expr {
 
   /// Informs IRGen to that this expression should be applied as its distributed
   /// thunk, rather than invoking the function directly.
-  bool shouldApplyDistributedThunk() const;
+  bool shouldApplyDistributedThunk() const {
+    return Bits.ApplyExpr.ShouldApplyDistributedThunk;
+  }
+  void setShouldApplyDistributedThunk(bool flag) {
+    Bits.ApplyExpr.ShouldApplyDistributedThunk = flag;
+  }
 
   ValueDecl *getCalledValue() const;
 

lib/AST/Expr.cpp

@@ -1068,22 +1068,6 @@ Type OverloadSetRefExpr::getBaseType() const {
   llvm_unreachable("Unhandled overloaded set reference expression");
 }
 
-bool ApplyExpr::shouldApplyDistributedThunk() const {
-  // only a distributed decl has a chance of needing to be invoked as a thunk
-  auto func = dyn_cast<AbstractFunctionDecl>(getCalledValue());
-  if (!func || !func->isDistributed())
-    return false;
-
-  if (implicitlyThrows())
-    return true;
-
-  auto isEffectivelyAsync = func->hasAsync() || isImplicitlyAsync();
-  if (func->hasThrows() && isEffectivelyAsync)
-    return true;
-
-  return false;
-}
-
 bool OverloadSetRefExpr::hasBaseObject() const {
   if (Type BaseTy = getBaseType())
     return !BaseTy->is<AnyMetatypeType>();

lib/Sema/TypeCheckConcurrency.cpp

@@ -1003,14 +1003,16 @@ namespace {
     ///
     /// and we reach up to mark the CallExpr.
     void markNearestCallAsImplicitly(
-        Optional<ImplicitActorHopTarget> setAsync, bool setThrows = false) {
+        Optional<ImplicitActorHopTarget> setAsync,
+        bool setThrows = false, bool setDistributedThunk = false) {
       assert(applyStack.size() > 0 && "not contained within an Apply?");
 
       const auto End = applyStack.rend();
       for (auto I = applyStack.rbegin(); I != End; ++I)
         if (auto call = dyn_cast<CallExpr>(*I)) {
           if (setAsync) call->setImplicitlyAsync(*setAsync);
           if (setThrows) call->setImplicitlyThrows(true);
+          if (setDistributedThunk) call->setShouldApplyDistributedThunk(true);
           return;
         }
       llvm_unreachable("expected a CallExpr in applyStack!");
@@ -1667,7 +1669,6 @@ namespace {
     ThrowsMarkingResult tryMarkImplicitlyThrows(SourceLoc declLoc,
                                                 ConcreteDeclRef concDeclRef,
                                                 Expr* context) {
-
       ValueDecl *decl = concDeclRef.getDecl();
       ThrowsMarkingResult result = ThrowsMarkingResult::NotFound;
 
@@ -1869,8 +1870,11 @@ namespace {
       }
 
       switch (contextIsolation) {
-      case ActorIsolation::ActorInstance:
-      case ActorIsolation::DistributedActorInstance: {
+      case ActorIsolation::DistributedActorInstance:
+        markNearestCallAsImplicitly(/*setAsync*/None, /*setThrows*/false,
+                                    /*setDistributedThunk*/true);
+        LLVM_FALLTHROUGH;
+      case ActorIsolation::ActorInstance: {
         auto result = tryMarkImplicitlyAsync(
           loc, valueRef, context,
           ImplicitActorHopTarget::forGlobalActor(globalActor));
@@ -2278,6 +2282,8 @@ namespace {
               tryMarkImplicitlyAsync(memberLoc, memberRef, context,
                                      ImplicitActorHopTarget::forInstanceSelf());
               tryMarkImplicitlyThrows(memberLoc, memberRef, context);
+              markNearestCallAsImplicitly(/*setAsync*/None, /*setThrows*/false,
+                                          /*setDistributedThunk*/true);
 
             } else {
               // neither static or distributed, apply full distributed isolation
@@ -2363,8 +2369,7 @@ namespace {
           }
 
           // It wasn't a distributed func, so ban the access
-          // TODO(distributed): handle subscripts here too
-          if (auto var = dyn_cast<VarDecl>(member)) {
+          if (isPropOrSubscript(member)) {
             ctx.Diags.diagnose(
                 memberLoc, diag::distributed_actor_isolated_non_self_reference,
                 member->getDescriptiveKind(), member->getName());
@@ -2380,6 +2385,8 @@ namespace {
         if (isolation.getActorType()->isDistributedActor() &&
             !isolatedActor.isPotentiallyIsolated) {
           tryMarkImplicitlyThrows(memberLoc, memberRef, context);
+          markNearestCallAsImplicitly(/*setAsync*/None, /*setThrows*/false,
+                                      /*setDistributedThunk*/true);
         }
 
         if (implicitAsyncResult == AsyncMarkingResult::FoundAsync)

test/Distributed/Runtime/distributed_actor_remote_functions.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-distributed -parse-as-library) | %FileCheck %s --dump-input=always
+// RUN: %target-run-simple-swift(-Xfrontend -disable-availability-checking -Xfrontend -enable-experimental-distributed -parse-as-library) | %FileCheck %s --dump-input=always
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency
@@ -41,6 +41,31 @@ distributed actor SomeSpecificDistributedActor {
    "local(\(#function))"
   }
 
+  distributed func callTaskSelf_inner() async throws -> String {
+    "local(\(#function))"
+  }
+  distributed func callTaskSelf() async -> String {
+    do {
+      return try await Task {
+        let called = try await callTaskSelf_inner() // shouldn't use the distributed thunk!
+        return "local(\(#function)) -> \(called)"
+      }.value
+    } catch {
+      return "WRONG local(\(#function)) thrown(\(error))"
+    }
+  }
+
+  distributed func callDetachedSelf() async -> String {
+    do {
+      return try await Task.detached {
+        let called = try await self.callTaskSelf_inner() // shouldn't use the distributed thunk!
+        return "local(\(#function)) -> \(called)"
+      }.value
+    } catch {
+      return "WRONG local(\(#function)) thrown(\(error))"
+    }
+  }
+
   // === errors
 
   distributed func helloThrowsImplBoom() throws -> String {
@@ -74,6 +99,21 @@ extension SomeSpecificDistributedActor {
     "remote(\(#function))"
   }
 
+  @_dynamicReplacement(for:_remote_callTaskSelf())
+  nonisolated func _remote_impl_callTaskSelf() async throws -> String {
+    "remote(\(#function))"
+  }
+
+  @_dynamicReplacement(for:_remote_callDetachedSelf())
+  nonisolated func _remote_impl_callDetachedSelf() async throws -> String {
+    "remote(\(#function))"
+  }
+
+  @_dynamicReplacement(for:_remote_callTaskSelf_inner())
+  nonisolated func _remote_impl_callTaskSelf_inner() async throws -> String {
+    "remote(\(#function))"
+  }
+
   // === errors
 
   @_dynamicReplacement(for:_remote_helloThrowsImplBoom())
@@ -101,9 +141,6 @@ func __isLocalActor(_ actor: AnyObject) -> Bool {
 @available(SwiftStdlib 5.5, *)
 struct ActorAddress: ActorIdentity {
   let address: String
-  init(parse address : String) {
-    self.address = address
-  }
 }
 
 @available(SwiftStdlib 5.5, *)
@@ -120,7 +157,7 @@ struct FakeTransport: ActorTransport {
 
   func assignIdentity<Act>(_ actorType: Act.Type) -> AnyActorIdentity
       where Act: DistributedActor {
-    .init(ActorAddress(parse: ""))
+    .init(ActorAddress(address: ""))
   }
 
   func actorReady<Act>(_ actor: Act) where Act: DistributedActor {
@@ -149,6 +186,12 @@ func test_remote_invoke(address: ActorAddress, transport: ActorTransport) async
     let h4 = try! await actor.hello()
     print("\(personality) - hello: \(h4)")
 
+    let h5 = try! await actor.callTaskSelf()
+    print("\(personality) - callTaskSelf: \(h5)")
+
+    let h6 = try! await actor.callDetachedSelf()
+    print("\(personality) - callDetachedSelf: \(h6)")
+
     // error throws
     if __isRemoteActor(actor) {
       do {
@@ -180,6 +223,8 @@ func test_remote_invoke(address: ActorAddress, transport: ActorTransport) async
   // CHECK: local - helloAsync: local(helloAsync())
   // CHECK: local - helloThrows: local(helloThrows())
   // CHECK: local - hello: local(hello())
+  // CHECK: local - callTaskSelf: local(callTaskSelf()) -> local(callTaskSelf_inner())
+  // CHECK: local - callDetachedSelf: local(callDetachedSelf()) -> local(callTaskSelf_inner())
   // CHECK: local - helloThrowsImplBoom: Boom(whoFailed: "impl")
 
   print("remote isRemote: \(__isRemoteActor(remote))")
@@ -189,6 +234,8 @@ func test_remote_invoke(address: ActorAddress, transport: ActorTransport) async
   // CHECK: remote - helloAsync: remote(_remote_impl_helloAsync())
   // CHECK: remote - helloThrows: remote(_remote_impl_helloThrows())
   // CHECK: remote - hello: remote(_remote_impl_hello())
+  // CHECK: remote - callTaskSelf: remote(_remote_impl_callTaskSelf())
+  // CHECK: remote - callDetachedSelf: remote(_remote_impl_callDetachedSelf())
   // CHECK: remote - helloThrowsTransportBoom: Boom(whoFailed: "transport")
 
   print(local)
@@ -198,7 +245,7 @@ func test_remote_invoke(address: ActorAddress, transport: ActorTransport) async
 @available(SwiftStdlib 5.5, *)
 @main struct Main {
   static func main() async {
-    let address = ActorAddress(parse: "")
+    let address = ActorAddress(address: "")
     let transport = FakeTransport()
 
     await test_remote_invoke(address: address, transport: transport)