[Distributed] Diagnose missing import also for funcs in extensions

Resolves rdar://125813581

Author: ktoso

include/swift/AST/Decl.h

@@ -173,11 +173,13 @@ enum class DescriptiveDeclKind : uint8_t {
   Struct,
   Class,
   Actor,
+  DistributedActor,
   Protocol,
   GenericEnum,
   GenericStruct,
   GenericClass,
   GenericActor,
+  GenericDistributedActor,
   GenericType,
   Subscript,
   StaticSubscript,

include/swift/AST/DiagnosticsSema.def

@@ -5431,9 +5431,12 @@ ERROR(isolated_default_argument_context,none,
 ERROR(conflicting_default_argument_isolation,none,
       "default argument cannot be both %0 and %1",
       (ActorIsolation, ActorIsolation))
-ERROR(distributed_actor_needs_explicit_distributed_import,none,
-      "'Distributed' module not imported, required for 'distributed actor'",
-      ())
+ERROR(distributed_decl_needs_explicit_distributed_import,none,
+      "%0 declaration requires importing Distributed module explicitly",
+      (DescriptiveDeclKind))
+ERROR(distributed_func_needs_explicit_distributed_import,none,
+      "%0 declaration requires importing Distributed module explicitly",
+      (DescriptiveDeclKind))
 NOTE(distributed_func_cannot_overload_on_async_only,none,
      "%0 previously declared here, cannot overload distributed methods on effect only", (const ValueDecl *))
 NOTE(distributed_func_other_ambiguous_overload_here,none,

lib/AST/Decl.cpp

@@ -187,12 +187,23 @@ DescriptiveDeclKind Decl::getDescriptiveKind() const {
               : DescriptiveDeclKind::Struct;
 
    case DeclKind::Class: {
-     bool isActor = cast<ClassDecl>(this)->isActor();
-     return cast<ClassDecl>(this)->getGenericParams()
-                ? (isActor ? DescriptiveDeclKind::GenericActor
-                           : DescriptiveDeclKind::GenericClass)
-                : (isActor ? DescriptiveDeclKind::Actor
-                           : DescriptiveDeclKind::Class);
+     auto clazz = cast<ClassDecl>(this);
+     bool isAnyActor = clazz->isAnyActor();
+     bool isGeneric = clazz->getGenericParams();
+
+     auto kind = isGeneric ? DescriptiveDeclKind::GenericClass
+                           : DescriptiveDeclKind::Class;
+
+     if (isAnyActor) {
+       if (clazz->isDistributedActor()) {
+         kind = isGeneric ? DescriptiveDeclKind::GenericDistributedActor
+                          : DescriptiveDeclKind::DistributedActor;
+       } else {
+         kind = isGeneric ? DescriptiveDeclKind::GenericActor
+                          : DescriptiveDeclKind::Actor;
+       }
+     }
+     return kind;
    }
 
    case DeclKind::Var: {
@@ -332,11 +343,13 @@ StringRef Decl::getDescriptiveKindName(DescriptiveDeclKind K) {
   ENTRY(Struct, "struct");
   ENTRY(Class, "class");
   ENTRY(Actor, "actor");
+  ENTRY(DistributedActor, "distributed actor");
   ENTRY(Protocol, "protocol");
   ENTRY(GenericEnum, "generic enum");
   ENTRY(GenericStruct, "generic struct");
   ENTRY(GenericClass, "generic class");
   ENTRY(GenericActor, "generic actor");
+  ENTRY(GenericDistributedActor, "generic distributed actor");
   ENTRY(GenericType, "generic type");
   ENTRY(Subscript, "subscript");
   ENTRY(StaticSubscript, "static subscript");

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -231,6 +231,7 @@ deriveBodyDistributed_thunk(AbstractFunctionDecl *thunk, void *context) {
 
   // === Type:
   StructDecl *RCT = C.getRemoteCallTargetDecl();
+  assert(RCT && "Missing RemoteCalLTarget declaration");
   Type remoteCallTargetTy = RCT->getDeclaredInterfaceType();
 
   // === __isRemoteActor(self)
@@ -742,8 +743,13 @@ static FuncDecl *createSameSignatureDistributedThunkDecl(DeclContext *DC,
 }
 
 static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
+  auto &C = func->getASTContext();
   auto DC = func->getDeclContext();
 
+  if (C.getLoadedModule(C.Id_Distributed)) {
+    return nullptr;
+  }
+
   FuncDecl *thunk =
       createSameSignatureDistributedThunkDecl(DC, func);
   assert(thunk && "couldn't create a distributed thunk");

lib/Sema/TypeCheckDistributed.cpp

@@ -25,6 +25,7 @@
 #include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeVisitor.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/Basic/Defer.h"
 #include "swift/AST/ASTPrinter.h"
@@ -45,11 +46,21 @@ DistributedModuleIsAvailableRequest::evaluate(Evaluator &evaluator,
                                               Decl *decl) const {
   auto &C = decl->getASTContext();
 
-  if (C.getLoadedModule(C.Id_Distributed))
+  auto DistributedModule = C.getLoadedModule(C.Id_Distributed);
+  if (!DistributedModule) {
+    decl->diagnose(diag::distributed_decl_needs_explicit_distributed_import,
+                   decl->getDescriptiveKind());
+    return false;
+  }
+
+  auto &importCache = C.getImportCache();
+  if (importCache.isImportedBy(DistributedModule, decl->getDeclContext())) {
     return true;
+  }
 
   // seems we're missing the Distributed module, ask to import it explicitly
-  decl->diagnose(diag::distributed_actor_needs_explicit_distributed_import);
+  decl->diagnose(diag::distributed_decl_needs_explicit_distributed_import,
+                 decl->getDescriptiveKind());
   return false;
 }
 
@@ -502,6 +513,10 @@ bool swift::checkDistributedFunction(AbstractFunctionDecl *func) {
   if (!func->isDistributed())
     return false;
 
+  // ==== Ensure the Distributed module is available,
+  if (!swift::ensureDistributedModuleLoaded(func))
+    return true;
+
   auto &C = func->getASTContext();
   return evaluateOrDefault(C.evaluator,
                            CheckDistributedFunctionRequest{func},
@@ -521,13 +536,11 @@ bool CheckDistributedFunctionRequest::evaluate(
   auto module = func->getParentModule();
 
   /// If no distributed module is available, then no reason to even try checks.
-  if (!C.getLoadedModule(C.Id_Distributed))
+  if (!C.getLoadedModule(C.Id_Distributed)) {
+    func->diagnose(diag::distributed_decl_needs_explicit_distributed_import,
+                   func->getDescriptiveKind());
     return true;
-
-//  // No checking for protocol requirements because they are not required
-//  // to have `SerializationRequirement`.
-//  if (isa<ProtocolDecl>(func->getDeclContext()))
-//    return false;
+  }
 
   Type serializationReqType =
       getDistributedActorSerializationType(func->getDeclContext());

test/Distributed/Inputs/FakeDistributedActorSystems.swift

@@ -14,6 +14,13 @@
 
 import Distributed
 
+// ==== Example Distributed Actors ----------------------------------------------
+
+@available(SwiftStdlib 5.7, *)
+public distributed actor FakeRoundtripActorSystemDistributedActor {
+  public typealias ActorSystem = FakeRoundtripActorSystem
+}
+
 // ==== Fake Address -----------------------------------------------------------
 
 public struct ActorAddress: Hashable, Sendable, Codable {

test/Distributed/Runtime/distributed_actor_func_calls_remoteCall_extension.swift

@@ -0,0 +1,46 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/../Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-build-swift -module-name main -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s %S/../Inputs/FakeDistributedActorSystems.swift -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s --dump-input=always
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+// UNSUPPORTED: OS=windows-msvc
+
+import Distributed
+import FakeDistributedActorSystems
+
+typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
+
+extension FakeRoundtripActorSystemDistributedActor {
+  distributed func echo(name: String) -> String {
+    return "Echo: \(name)"
+  }
+}
+
+func test() async throws {
+  let system = DefaultDistributedActorSystem()
+
+  let local = FakeRoundtripActorSystemDistributedActor(actorSystem: system)
+  let ref = try FakeRoundtripActorSystemDistributedActor.resolve(id: local.id, using: system)
+
+  let reply = try await ref.echo(name: "Caplin")
+  // CHECK: >> remoteCall: on:main.FakeRoundtripActorSystemDistributedActor, target:main.FakeRoundtripActorSystemDistributedActor.echo(name:), invocation:FakeInvocationEncoder(genericSubs: [], arguments: ["Caplin"], returnType: Optional(Swift.String), errorType: nil), throwing:Swift.Never, returning:Swift.String
+
+  // CHECK: << remoteCall return: Echo: Caplin
+  print("reply: \(reply)")
+  // CHECK: reply: Echo: Caplin
+}
+
+@main struct Main {
+  static func main() async {
+    try! await test()
+  }
+}

test/Distributed/distributed_actor_derived_conformances.swift

@@ -11,4 +11,4 @@ typealias DefaultDistributedActorSystem = FakeActorSystem
 
 distributed actor DA: Comparable {}
 // expected-error@-1 {{type 'DA' does not conform to protocol 'Comparable'}}
-// expected-note@-2 {{automatic synthesis of 'Comparable' is not supported for actor declarations}}
+// expected-note@-2 {{automatic synthesis of 'Comparable' is not supported for distributed actor declarations}}

test/Distributed/distributed_actor_extension_without_import.swift

@@ -0,0 +1,15 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-swift-frontend -typecheck -verify -verify-ignore-unknown -disable-availability-checking -I %t 2>&1 %s
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// On purpose missing 'import Distributed'
+import FakeDistributedActorSystems
+
+extension FakeRoundtripActorSystemDistributedActor {
+  // expected-error@+1{{}}
+  distributed func echo(name: String) -> String {
+    return "Echo: \(name)"
+  }
+}

test/Distributed/distributed_actor_is_experimental_enabled_missing_import.swift renamed to test/Distributed/distributed_actor_missing_distributed_import.swift

@@ -5,7 +5,7 @@
 actor SomeActor {}
 
 distributed actor DA {}
-// expected-error@-1{{'Distributed' module not imported, required for 'distributed actor'}}
+// expected-error@-1{{distributed actor declaration requires importing Distributed module explicitly}}
 
 distributed actor class DAC {}
 // expected-error@-1{{distributed' can only be applied to 'actor' definitions, and distributed actor-isolated async functions}}
@@ -18,18 +18,20 @@ actor A {
 }
 
 actor B {
-  distributed var neverOk: String { // expected-error{{'Distributed' module not imported, required for 'distributed actor'}}
+  distributed var neverOk: String { // expected-error{{distributed property declaration requires importing Distributed module explicitly}}
     ""
   }
 }
 
 distributed actor DA2 {
-  // expected-error@-1{{'Distributed' module not imported, required for 'distributed actor'}}
+  // expected-error@-1{{distributed actor declaration requires importing Distributed module explicitly}}
   func normal() async {}
+  // expected-error@+1{{distributed instance method declaration requires importing Distributed module explicitly}}
   distributed func dist() {}
+  // expected-error@+1{{distributed instance method declaration requires importing Distributed module explicitly}}
   distributed func distAsync() async {}
 
-  distributed var neverOk: String { // expected-error{{'Distributed' module not imported, required for 'distributed actor'}}
+  distributed var neverOk: String { // expected-error{{distributed property declaration requires importing Distributed module explicitly}}
     ""
   }
 }

test/Distributed/distributed_missing_import.swift

@@ -5,7 +5,7 @@
 actor SomeActor { }
 
 distributed actor MissingImportDistributedActor_0 { }
-// expected-error@-1{{'Distributed' module not imported, required for 'distributed actor'}}
+// expected-error@-1{{distributed actor declaration requires importing Distributed module explicitly}}
 
 let t: DistributedActorSystem // expected-error{{cannot find type 'DistributedActorSystem' in scope}}
 let a: ActorAddress // expected-error{{cannot find type 'ActorAddress' in scope}}