[Distributed] Naive impl of whenLocal for distributed actors (#39654)

* [Distributed] Hacky impl of whenLocal for distributed actors

* cleanup & fix one test

* disable test on windows

Author: ktoso

include/swift/AST/KnownIdentifiers.def

@@ -264,6 +264,7 @@ IDENTIFIER(identifier)
 IDENTIFIER(_distributedActorRemoteInitialize)
 IDENTIFIER(_distributedActorDestroy)
 IDENTIFIER(__isRemoteActor)
+IDENTIFIER(whenLocal)
 
 #undef IDENTIFIER
 #undef IDENTIFIER_

lib/Sema/TypeCheckConcurrency.cpp

@@ -1432,6 +1432,11 @@ namespace {
       bool isPotentiallyIsolated = false;
       if (!var) {
         isPotentiallyIsolated = false;
+      } else if (var->getName().str().equals("__secretlyKnownToBeLocal")) {
+        // FIXME(distributed): we did a dynamic check and know that this actor is local,
+        //  but we can't express that to the type system; the real implementation
+        //  will have to mark 'self' as "known to be local" after an is-local check.
+        isPotentiallyIsolated = true;
       } else if (auto param = dyn_cast<ParamDecl>(var)) {
         isPotentiallyIsolated = param->isIsolated();
       } else if (var->isSelfParamCapture()) {

stdlib/public/Distributed/DistributedActor.swift

@@ -116,6 +116,28 @@ extension DistributedActor {
   }
 }
 
+// ==== Local actor special handling -------------------------------------------
+
+@available(SwiftStdlib 5.5, *)
+extension DistributedActor {
+
+  /// Executes the passed 'body' only when the distributed actor is local instance.
+  ///
+  /// The `Self` passed to the the body closure is isolated, meaning that the
+  /// closure can be used to call non-distributed functions, or even access actor
+  /// state.
+  ///
+  /// When the actor is remote, the closure won't be executed and this function will return nil.
+  public nonisolated func whenLocal<T>(_ body: @Sendable (isolated Self) async throws -> T)
+    async rethrows -> T? where T: Sendable {
+    if __isLocalActor(self) {
+       return try await body(self)
+    } else {
+      return nil
+    }
+  }
+}
+
 /******************************************************************************/
 /***************************** Actor Identity *********************************/
 /******************************************************************************/

test/Distributed/Runtime/distributed_actor_deinit.swift

@@ -7,7 +7,7 @@
 // UNSUPPORTED: use_os_stdlib
 // UNSUPPORTED: back_deployment_runtime
 
-// temporary non-support. tracked in rdar://82593574
+// FIXME(distributed): Distributed actors currently have some issues on windows, isRemote always returns false. rdar://82593574
 // UNSUPPORTED: windows
 
 import _Distributed

test/Distributed/Runtime/distributed_actor_whenLocal.swift

@@ -0,0 +1,83 @@
+// 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
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+// FIXME(distributed): Distributed actors currently have some issues on windows, isRemote always returns false. rdar://82593574
+// UNSUPPORTED: windows
+
+import _Distributed
+
+distributed actor Capybara {
+  // only the local capybara can do this!
+  func eat() -> String {
+    "watermelon"
+  }
+}
+
+
+// ==== Fake Transport ---------------------------------------------------------
+@available(SwiftStdlib 5.5, *)
+struct ActorAddress: ActorIdentity {
+  let address: String
+  init(parse address: String) {
+    self.address = address
+  }
+}
+
+@available(SwiftStdlib 5.5, *)
+struct FakeTransport: ActorTransport {
+  func decodeIdentity(from decoder: Decoder) throws -> AnyActorIdentity {
+    fatalError("not implemented:\(#function)")
+  }
+
+  func resolve<Act>(_ identity: AnyActorIdentity, as actorType: Act.Type)
+  throws -> Act? where Act: DistributedActor {
+    return nil
+  }
+
+  func assignIdentity<Act>(_ actorType: Act.Type) -> AnyActorIdentity
+          where Act: DistributedActor {
+    let id = ActorAddress(parse: "xxx")
+    return .init(id)
+  }
+
+  func actorReady<Act>(_ actor: Act) where Act: DistributedActor {
+  }
+
+  func resignIdentity(_ id: AnyActorIdentity) {
+  }
+}
+
+func test() async throws {
+  let transport = FakeTransport()
+
+  let local = Capybara(transport: transport)
+  // await local.eat() // SHOULD ERROR
+  let valueWhenLocal: String? = await local.whenLocal { __secretlyKnownToBeLocal in
+    __secretlyKnownToBeLocal.eat()
+  }
+
+  // CHECK: valueWhenLocal: watermelon
+  print("valueWhenLocal: \(valueWhenLocal ?? "nil")")
+
+  let remote = try Capybara.resolve(local.id, using: transport)
+  let valueWhenRemote: String? = await remote.whenLocal { __secretlyKnownToBeLocal in
+    __secretlyKnownToBeLocal.eat()
+  }
+  
+  // CHECK: valueWhenRemote: nil
+  print("valueWhenRemote: \(valueWhenRemote ?? "nil")")
+}
+
+@available(SwiftStdlib 5.5, *)
+@main struct Main {
+  static func main() async {
+    try! await test()
+  }
+}

test/Distributed/actor_protocols.swift

@@ -52,9 +52,8 @@ actor A2: DistributedActor {
   }
 }
 
-class C2: DistributedActor {
+final class C2: DistributedActor {
   // expected-error@-1{{non-actor type 'C2' cannot conform to the 'Actor' protocol}}
-  // expected-error@-2{{non-final class 'C2' cannot conform to 'Sendable'; use '@unchecked Sendable'}}
   nonisolated var id: AnyActorIdentity {
     fatalError()
   }