[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,9 @@ 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,
+      "%kind0 requires explicit import of Distributed module",
+      (const ValueDecl *))
 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,

include/swift/AST/TypeCheckRequests.h

@@ -3025,7 +3025,7 @@ class HasCircularRawValueRequest
 
 /// Checks if the Distributed module is available.
 class DistributedModuleIsAvailableRequest
-    : public SimpleRequest<DistributedModuleIsAvailableRequest, bool(Decl *),
+    : public SimpleRequest<DistributedModuleIsAvailableRequest, bool(const ValueDecl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -3034,7 +3034,7 @@ class DistributedModuleIsAvailableRequest
   friend SimpleRequest;
 
   // Evaluation.
-  bool evaluate(Evaluator &evaluator, Decl *decl) const;
+  bool evaluate(Evaluator &evaluator, const ValueDecl *decl) const;
 
 public:
   // Cached.

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -200,7 +200,7 @@ SWIFT_REQUEST(TypeChecker, InferredGenericSignatureRequest,
                                          SourceLoc, bool, bool),
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, DistributedModuleIsAvailableRequest,
-              bool(ModuleDecl *), Cached, NoLocationInfo)
+              bool(const ValueDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, InheritedTypeRequest,
               Type(llvm::PointerUnion<const TypeDecl *, const ExtensionDecl *>,
                    unsigned, TypeResolutionStage),

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)

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"
@@ -33,7 +34,7 @@ using namespace swift;
 
 // ==== ------------------------------------------------------------------------
 
-bool swift::ensureDistributedModuleLoaded(Decl *decl) {
+bool swift::ensureDistributedModuleLoaded(const ValueDecl *decl) {
   auto &C = decl->getASTContext();
   auto moduleAvailable = evaluateOrDefault(
       C.evaluator, DistributedModuleIsAvailableRequest{decl}, false);
@@ -42,14 +43,24 @@ bool swift::ensureDistributedModuleLoaded(Decl *decl) {
 
 bool
 DistributedModuleIsAvailableRequest::evaluate(Evaluator &evaluator,
-                                              Decl *decl) const {
+                                              const ValueDecl *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);
+    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);
   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);
     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());

lib/Sema/TypeCheckDistributed.h

@@ -35,7 +35,7 @@ class NominalTypeDecl;
 /******************************************************************************/
 
 // Diagnose an error if the Distributed module is not loaded.
-bool ensureDistributedModuleLoaded(Decl *decl);
+bool ensureDistributedModuleLoaded(const ValueDecl *decl);
 
 /// Check for illegal property declarations (e.g. re-declaring transport or id)
 void checkDistributedActorProperties(const NominalTypeDecl *decl);

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 'DA' requires explicit import of Distributed module}}
 
 distributed actor class DAC {}
 // expected-error@-1{{distributed' can only be applied to 'actor' definitions, and distributed actor-isolated async functions}}
@@ -18,18 +18,24 @@ 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 'neverOk' requires explicit import of Distributed module}}
     ""
   }
 }
 
+  // expected-error@+1{{distributed actor 'DA2' requires explicit import of Distributed module}}
 distributed actor DA2 {
-  // expected-error@-1{{'Distributed' module not imported, required for 'distributed actor'}}
+
   func normal() async {}
+
+  // expected-error@+1{{distributed instance method 'dist()' requires explicit import of Distributed module}}
   distributed func dist() {}
+
+  // expected-error@+1{{distributed instance method 'distAsync()' requires explicit import of Distributed module}}
   distributed func distAsync() async {}
 
-  distributed var neverOk: String { // expected-error{{'Distributed' module not imported, required for 'distributed actor'}}
+  // expected-error@+1{{distributed property 'neverOk' requires explicit import of Distributed module}}
+  distributed var neverOk: String {
     ""
   }
 }

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 'MissingImportDistributedActor_0' requires explicit import of Distributed module}}
 
 let t: DistributedActorSystem // expected-error{{cannot find type 'DistributedActorSystem' in scope}}
 let a: ActorAddress // expected-error{{cannot find type 'ActorAddress' in scope}}