Support module selectors for cross-import overlays

include/swift/AST/Module.h

@@ -481,6 +481,11 @@ public:
   void getDeclaredCrossImportBystanders(
       SmallVectorImpl<Identifier> &bystanderNames);
 
+  /// Returns the name that should be used for this module in a module
+  /// selector. For separately-imported overlays, this will be the declaring
+  /// module's name.
+  Identifier getNameForModuleSelector();
+
   /// Retrieve the ABI name of the module, which is used for metadata and
   /// mangling.
   Identifier getABIName() const;

lib/AST/Module.cpp

@@ -2564,6 +2564,12 @@ bool ModuleDecl::getRequiredBystandersIfCrossImportOverlay(
   return false;
 }
 
+Identifier ModuleDecl::getNameForModuleSelector() {
+  if (auto declaring = getDeclaringModuleIfCrossImportOverlay())
+    return declaring->getName();
+  return this->getName();
+}
+
 bool ModuleDecl::isClangHeaderImportModule() const {
   auto importer = getASTContext().getClangModuleLoader();
   if (!importer)

lib/AST/NameLookup.cpp

@@ -372,27 +372,33 @@ bool swift::removeOutOfModuleDecls(SmallVectorImpl<ValueDecl*> &decls,
   // FIXME: Should we look this up relative to dc?
   // We'd need a new ResolutionKind.
   // FIXME: How can we diagnose this?
-  ModuleDecl *visibleFrom = ctx.getLoadedModule(moduleSelector);
-  if (!visibleFrom && moduleSelector == ctx.TheBuiltinModule->getName())
-    visibleFrom = ctx.TheBuiltinModule;
-  if (!visibleFrom) {
+  ModuleDecl *visibleFromRoot = ctx.getLoadedModule(moduleSelector);
+  if (!visibleFromRoot && moduleSelector == ctx.TheBuiltinModule->getName())
+    visibleFromRoot = ctx.TheBuiltinModule;
+  if (!visibleFromRoot) {
     LLVM_DEBUG(llvm::dbgs() << "no module " << moduleSelector << "\n");
     bool clearedAny = !decls.empty();
     decls.clear();
     return clearedAny;
   }
 
+  SmallVector<ModuleDecl *, 4> visibleFrom;
+  dc->getSeparatelyImportedOverlays(visibleFromRoot, visibleFrom);
+  if (visibleFrom.empty())
+    visibleFrom.push_back(visibleFromRoot);
+
   size_t initialCount = decls.size();
   decls.erase(
     std::remove_if(decls.begin(), decls.end(), [&](ValueDecl *decl) -> bool {
-      bool inScope = ctx.getImportCache().isImportedBy(decl->getModuleContext(),
-                                                       visibleFrom);
+      bool inScope = llvm::any_of(visibleFrom, [&](ModuleDecl *visibleFromMod) {
+        return ctx.getImportCache().isImportedBy(decl->getModuleContext(),
+                                                 visibleFromMod);
+      });
 
       LLVM_DEBUG(decl->dumpRef(llvm::dbgs()));
       LLVM_DEBUG(llvm::dbgs() << ": " << decl->getModuleContext()->getName()
                               << (inScope ? " is " : " is NOT ")
-                              << "selected by " << visibleFrom->getName()
-                              << "\n");
+                              << "selected by " << moduleSelector << "\n");
 
       return !inScope;
     }),

lib/Sema/CSDiagnostics.cpp

@@ -6378,7 +6378,8 @@ bool MemberFromWrongModuleFailure::diagnoseAsError() {
   emitDiagnosticAt(loc, diag::wrong_module_selector, Member->getName(),
                    Name.getModuleSelector());
 
-  Identifier actualModuleName = Member->getModuleContext()->getName();
+  Identifier actualModuleName =
+      Member->getModuleContext()->getNameForModuleSelector();
   ASSERT(actualModuleName != Name.getModuleSelector() &&
          "Module selector failure on member in same module?");
   emitDiagnosticAt(loc, diag::note_change_module_selector, actualModuleName)

lib/Sema/TypeCheckNameLookup.cpp

@@ -1176,7 +1176,7 @@ ModuleSelectorCorrection(const LookupResult &candidates) {
 
     auto owningModule = decl->getModuleContext();
     candidateModules.insert(
-      { owningModule->getName(), kind });
+      { owningModule->getNameForModuleSelector(), kind });
   }
 }
 
@@ -1187,7 +1187,7 @@ ModuleSelectorCorrection(const LookupTypeResult &candidates) {
   for (auto result : candidates) {
     auto owningModule = result.Member->getModuleContext();
     candidateModules.insert(
-      { owningModule->getName(), CandidateKind::ContextFree });
+      { owningModule->getNameForModuleSelector(), CandidateKind::ContextFree });
   }
 }
 

test/NameLookup/Inputs/ModuleSelectorTestingKit.swiftcrossimport/ctypes.swiftoverlay

@@ -0,0 +1,5 @@
+%YAML 1.2
+---
+version: 1
+modules:
+  - name: _ModuleSelectorTestingKit_ctypes

test/NameLookup/Inputs/_ModuleSelectorTestingKit_ctypes.swiftinterface

@@ -0,0 +1,14 @@
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name _ModuleSelectorTestingKit_ctypes -swift-version 5
+
+import Swift
+@_exported import ModuleSelectorTestingKit
+import ctypes
+
+public struct E {
+  public init(value: Int)
+  public dynamic mutating func negate()
+  public var magnitude: Swift.UInt
+}
+
+public func e() -> Void

test/NameLookup/module_selector.swift

@@ -8,7 +8,6 @@
 // * Whether X::foo finds foos in X's re-exports
 // * Whether we handle access paths correctly
 // * Interaction with ClangImporter
-// * Cross-import overlays
 // * Key path dynamic member lookup
 // * Custom type attributes (and coverage of type attrs generally is sparse)
 //
@@ -107,7 +106,6 @@ extension B: @retroactive main::Equatable {
       (main::+)
       // expected-error@-1 {{'+' is not imported through module 'main'}}
       // expected-note@-2 {{did you mean module 'Swift'?}} {{8-12=Swift}}
-      // expected-note@-3 {{did you mean module '_Concurrency'?}} {{8-12=_Concurrency}} FIXME: Accept and suggest 'Swift::' instead?
 
     let magnitude: Int.main::Magnitude = main::magnitude
     // expected-error@-1 {{'Magnitude' is not imported through module 'main'}}
@@ -195,7 +193,6 @@ extension C: @retroactive ModuleSelectorTestingKit::Equatable {
       (ModuleSelectorTestingKit::+)
       // expected-error@-1 {{'+' is not imported through module 'ModuleSelectorTestingKit'}}
       // expected-note@-2 {{did you mean module 'Swift'?}} {{8-32=Swift}}
-      // expected-note@-3 {{did you mean module '_Concurrency'?}} {{8-32=_Concurrency}} FIXME: Accept and suggest 'Swift::' instead?
 
     let magnitude: Int.ModuleSelectorTestingKit::Magnitude = ModuleSelectorTestingKit::magnitude
     // expected-error@-1 {{'Magnitude' is not imported through module 'ModuleSelectorTestingKit'}}
@@ -417,3 +414,17 @@ func badModuleNames() {
 func builtinModuleLookups(_ int: Builtin::Int64) -> Builtin::Int64 {
   return Builtin::int_bswap_Int64(int)
 }
+
+func concurrencyModuleLookups(
+  _: any Swift::Clock,
+  _: any _Concurrency::Clock,
+  _: any ModuleSelectorTestingKit::Clock
+  // expected-error@-1 {{'Clock' is not imported through module 'ModuleSelectorTestingKit'}}
+  // expected-note@-2 {{did you mean module 'Swift'?}} {{10-34=Swift}}
+) async {
+  await Swift::withTaskCancellationHandler {} onCancel: {}
+  await _Concurrency::withTaskCancellationHandler {} onCancel: {}
+  await ModuleSelectorTestingKit::withTaskCancellationHandler {} onCancel: {}
+  // expected-error@-1 {{'withTaskCancellationHandler' is not imported through module 'ModuleSelectorTestingKit'}}
+  // expected-note@-2 {{did you mean module 'Swift'?}} {{9-33=Swift}}
+}

test/NameLookup/module_selector_cross_import.swift

@@ -0,0 +1,23 @@
+// RUN: %target-typecheck-verify-swift -sdk %clang-importer-sdk -module-name main -I %S/Inputs -enable-experimental-feature ModuleSelector -enable-cross-import-overlays
+// RUN: %target-typecheck-verify-swift -sdk %clang-importer-sdk -module-name main -I %S/Inputs -enable-experimental-feature ModuleSelector -enable-experimental-feature ParserASTGen -enable-cross-import-overlays
+
+// REQUIRES: swift_feature_ModuleSelector, swift_feature_ParserASTGen
+
+import ModuleSelectorTestingKit
+import ctypes
+
+func crossImportLookups(
+  _: E,
+  _: ModuleSelectorTestingKit::E,
+  _: _ModuleSelectorTestingKit_ctypes::E,
+  _: Swift::E
+  // expected-error@-1 {{'E' is not imported through module 'Swift'}}
+  // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{6-11=ModuleSelectorTestingKit}}
+) {
+  e()
+  ModuleSelectorTestingKit::e()
+  _ModuleSelectorTestingKit_ctypes::e()
+  Swift::e()
+  // expected-error@-1 {{'e' is not imported through module 'Swift'}}
+  // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{3-8=ModuleSelectorTestingKit}}
+}