Handle module selectors in qualified lookup

Author: beccadax

include/swift/AST/NameLookup.h

@@ -486,6 +486,14 @@ class UsableFilteringDeclConsumer final : public VisibleDeclConsumer {
 /// \returns true if any declarations were removed, false otherwise.
 bool removeOverriddenDecls(SmallVectorImpl<ValueDecl*> &decls);
 
+/// Remove any declarations in the given set that do not match the
+/// module selector, if it is not empty.
+///
+/// \returns true if any declarations were removed, false otherwise.
+bool removeOutOfModuleDecls(SmallVectorImpl<ValueDecl*> &decls,
+                            Identifier moduleSelector,
+                            const DeclContext *dc);
+
 /// Remove any declarations in the given set that are shadowed by
 /// other declarations in that set.
 ///
@@ -561,6 +569,7 @@ void tryExtractDirectlyReferencedNominalTypes(
 /// Once name lookup has gathered a set of results, perform any necessary
 /// steps to prune the result set before returning it to the caller.
 void pruneLookupResultSet(const DeclContext *dc, NLOptions options,
+                          Identifier moduleSelector,
                           SmallVectorImpl<ValueDecl *> &decls);
 
 /// Do nothing if debugClient is null.

lib/AST/NameLookup.cpp

@@ -361,6 +361,42 @@ enum class ConstructorComparison {
   Better,
 };
 
+bool swift::removeOutOfModuleDecls(SmallVectorImpl<ValueDecl*> &decls,
+                                   Identifier moduleSelector,
+                                   const DeclContext *dc) {
+  if (moduleSelector.empty())
+    return false;
+
+  ASTContext &ctx = dc->getASTContext();
+
+  // 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) {
+    LLVM_DEBUG(llvm::dbgs() << "no module " << moduleSelector << "\n");
+    decls.clear();
+    return true;
+  }
+
+  bool initialCount = decls.size();
+  decls.erase(
+    std::remove_if(decls.begin(), decls.end(), [&](ValueDecl *decl) -> bool {
+      bool inScope = ctx.getImportCache().isImportedBy(decl->getModuleContext(),
+                                                       visibleFrom);
+
+      LLVM_DEBUG(decl->dumpRef(llvm::dbgs()));
+      LLVM_DEBUG(llvm::dbgs() << ": " << decl->getModuleContext()->getName()
+                              << (inScope ? " is " : " is NOT ")
+                              << "selected by " << visibleFrom->getName()
+                              << "\n");
+
+      return !inScope;
+    }),
+    decls.end());
+  return initialCount != decls.size();
+}
+
 /// Determines whether \p ctor1 is a "better" initializer than \p ctor2.
 static ConstructorComparison compareConstructors(ConstructorDecl *ctor1,
                                                  ConstructorDecl *ctor2,
@@ -2471,14 +2507,17 @@ bool namelookup::isInABIAttr(SourceFile *sourceFile, SourceLoc loc) {
 }
 
 void namelookup::pruneLookupResultSet(const DeclContext *dc, NLOptions options,
+                                      Identifier moduleSelector,
                                       SmallVectorImpl<ValueDecl *> &decls) {
   // If we're supposed to remove overridden declarations, do so now.
   if (options & NL_RemoveOverridden)
     removeOverriddenDecls(decls);
 
   // If we're supposed to remove shadowed/hidden declarations, do so now.
-  if (options & NL_RemoveNonVisible)
+  if (options & NL_RemoveNonVisible) {
+    removeOutOfModuleDecls(decls, moduleSelector, dc);
     removeShadowedDecls(decls, dc);
+  }
 
   ModuleDecl *M = dc->getParentModule();
   filterForDiscriminator(decls, M->getDebugClient());
@@ -2790,7 +2829,7 @@ QualifiedLookupRequest::evaluate(Evaluator &eval, const DeclContext *DC,
     }
   }
 
-  pruneLookupResultSet(DC, options, decls);
+  pruneLookupResultSet(DC, options, member.getModuleSelector(), decls);
   if (auto *debugClient = DC->getParentModule()->getDebugClient()) {
     debugClient->finishLookupInNominals(DC, typeDecls, member.getFullName(),
                                         options, decls);
@@ -2842,7 +2881,7 @@ ModuleQualifiedLookupRequest::evaluate(Evaluator &eval, const DeclContext *DC,
     }
   }
 
-  pruneLookupResultSet(DC, options, decls);
+  pruneLookupResultSet(DC, options, member.getModuleSelector(), decls);
 
   if (auto *debugClient = DC->getParentModule()->getDebugClient()) {
     debugClient->finishLookupInModule(DC, module, member.getFullName(),
@@ -2910,7 +2949,7 @@ AnyObjectLookupRequest::evaluate(Evaluator &evaluator, const DeclContext *dc,
       decls.push_back(decl);
   }
 
-  pruneLookupResultSet(dc, options, decls);
+  pruneLookupResultSet(dc, options, member.getModuleSelector(), decls);
   if (auto *debugClient = dc->getParentModule()->getDebugClient()) {
     debugClient->finishLookupInAnyObject(dc, member.getFullName(), options,
                                          decls);

test/NameLookup/module_selector.swift

@@ -41,6 +41,7 @@ extension A: @retroactive Swift::Equatable {
     let fn: (Swift::Int, Swift::Int) -> Swift::Int = (Swift::+)
 
     let magnitude: Int.Swift::Magnitude = main::magnitude
+    // expected-error@-1 {{cannot convert value of type 'Never' to specified type 'Int.Magnitude' (aka 'UInt')}}
 
     _ = (fn, magnitude)
 
@@ -76,18 +77,21 @@ extension B: @retroactive main::Equatable {
   // @_derivative(of:)
 
   @_dynamicReplacement(for: main::negate())
+  // FIXME improve: expected-error@-1 {{replaced function 'main::negate()' could not be found}}
 
   mutating func myNegate() {
     let fn: (main::Int, main::Int) -> main::Int =
       (main::+)
 
     let magnitude: Int.main::Magnitude = main::magnitude
+    // FIXME improve: expected-error@-1 {{'main::Magnitude' is not a member type of struct 'Swift.Int'}}
 
     _ = (fn, magnitude)
 
     if main::Bool.main::random() {
 
       main::negate()
+      // FIXME improve: expected-error@-1 {{cannot find 'main::negate' in scope}}
     }
     else {
       self = main::B(value: .main::min)
@@ -125,6 +129,7 @@ extension C: @retroactive ModuleSelectorTestingKit::Equatable {
       (ModuleSelectorTestingKit::+)
 
     let magnitude: Int.ModuleSelectorTestingKit::Magnitude = ModuleSelectorTestingKit::magnitude
+    // FIXME improve: expected-error@-1 {{'ModuleSelectorTestingKit::Magnitude' is not a member type of struct 'Swift.Int'}}
 
     _ = (fn, magnitude)
 
@@ -161,19 +166,22 @@ extension D: @retroactive Swift::Equatable {
   // @_derivative(of:)
 
   @_dynamicReplacement(for: Swift::negate())
+  // FIXME improve: expected-error@-1 {{replaced function 'Swift::negate()' could not be found}}
 
   mutating func myNegate() {
 
     let fn: (Swift::Int, Swift::Int) -> Swift::Int =
       (Swift::+)
 
     let magnitude: Int.Swift::Magnitude = Swift::magnitude
+    // expected-error@-1 {{cannot convert value of type 'Never' to specified type 'Int.Magnitude' (aka 'UInt')}}
 
     _ = (fn, magnitude)
 
     if Swift::Bool.Swift::random() {
 
       Swift::negate()
+      // FIXME improve: expected-error@-1 {{cannot find 'Swift::negate' in scope}}
     }
     else {
       self = Swift::D(value: .Swift::min)