Improve module selector macro diagnostics

Author: beccadax

lib/Sema/CSGen.cpp

@@ -4148,6 +4148,16 @@ namespace {
       auto macros = lookupMacros(moduleIdent, macroIdent, functionRefInfo,
                                  expr->getMacroRoles());
       if (macros.empty()) {
+        // Try removing a module selector if present.
+        if (macroIdent.hasModuleSelector()) {
+          auto anyMacroIdent = DeclNameRef(macroIdent.getFullName());
+          ModuleSelectorCorrection correction(
+            lookupMacros(moduleIdent, anyMacroIdent, functionRefInfo,
+                         expr->getMacroRoles()));
+          if (correction.diagnose(ctx, expr->getMacroNameLoc(), macroIdent))
+            return Type();
+        }
+
         ctx.Diags.diagnose(expr->getMacroNameLoc(), diag::macro_undefined,
                            macroIdent)
             .highlight(expr->getMacroNameLoc().getSourceRange());

lib/Sema/TypeCheckNameLookup.cpp

@@ -1200,6 +1200,15 @@ ModuleSelectorCorrection(const SmallVectorImpl<ValueDecl *> &candidates) {
   }
 }
 
+ModuleSelectorCorrection::
+ModuleSelectorCorrection(const SmallVectorImpl<constraints::OverloadChoice> &candidates) {
+  for (auto result : candidates) {
+    auto owningModule = result.getDecl()->getModuleContext();
+    candidateModules.insert(
+      { owningModule->getNameForModuleSelector(), CandidateKind::ContextFree });
+  }
+}
+
 bool ModuleSelectorCorrection::diagnose(ASTContext &ctx,
                                         DeclNameLoc nameLoc,
                                         DeclNameRef originalName) const {

lib/Sema/TypeCheckType.cpp

@@ -1493,18 +1493,8 @@ static Type diagnoseUnknownType(const TypeResolution &resolution,
 
   // Unqualified lookup case.
   if (parentType.isNull()) {
-    // Tailored diagnostic for custom attributes.
-    if (resolution.getOptions().is(TypeResolverContext::CustomAttr)) {
-      SmallString<64> scratch;
-      llvm::raw_svector_ostream scratchOS(scratch);
-      repr->getNameRef().printPretty(scratchOS);
-
-      diags.diagnose(repr->getNameLoc(), diag::unknown_attr_name, scratch);
-
-      return ErrorType::get(ctx);
-    }
-
-    if (repr->isSimpleUnqualifiedIdentifier(ctx.Id_Self)) {
+    if (!resolution.getOptions().is(TypeResolverContext::CustomAttr)
+          && repr->isSimpleUnqualifiedIdentifier(ctx.Id_Self)) {
       DeclContext *nominalDC = nullptr;
       NominalTypeDecl *nominal = nullptr;
       if ((nominalDC = dc->getInnermostTypeContext()) &&
@@ -1570,6 +1560,17 @@ static Type diagnoseUnknownType(const TypeResolution &resolution,
       return ErrorType::get(ctx);
     }
 
+    // Tailored diagnostic for custom attributes.
+    if (resolution.getOptions().is(TypeResolverContext::CustomAttr)) {
+      SmallString<64> scratch;
+      llvm::raw_svector_ostream scratchOS(scratch);
+      repr->getNameRef().printPretty(scratchOS);
+
+      diags.diagnose(repr->getNameLoc(), diag::unknown_attr_name, scratch);
+
+      return ErrorType::get(ctx);
+    }
+
     // Fallback.
     auto L = repr->getNameLoc();
     SourceRange R = repr->getNameLoc().getSourceRange();

lib/Sema/TypeChecker.h

@@ -332,6 +332,7 @@ class ModuleSelectorCorrection {
   ModuleSelectorCorrection(const LookupResult &candidates);
   ModuleSelectorCorrection(const LookupTypeResult &candidates);
   ModuleSelectorCorrection(const SmallVectorImpl<ValueDecl *> &candidates);
+  ModuleSelectorCorrection(const SmallVectorImpl<constraints::OverloadChoice> &candidates);
 
   /// Emit an error and warnings if there were any candidates.
   ///

test/NameLookup/module_selector.swift

@@ -148,7 +148,8 @@ extension B: @retroactive main::Equatable {
     // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{none}}
 
     _ = #main::ExprMacro
-    // expected-error@-1 {{no macro named 'main::ExprMacro'}}
+    // expected-error@-1 {{'ExprMacro' is not imported through module 'main'}}
+    // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{10-14=ModuleSelectorTestingKit}}
   }
 
   @main::PeerMacro func thingy() {}
@@ -307,7 +308,8 @@ extension D: @retroactive Swift::Equatable {
     // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{none}}
 
     _ = #Swift::ExprMacro
-    // expected-error@-1 {{no macro named 'Swift::ExprMacro'}}
+    // expected-error@-1 {{'ExprMacro' is not imported through module 'Swift'}}
+    // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{10-15=ModuleSelectorTestingKit}}
   }
 
   @Swift::PeerMacro func thingy() {}
@@ -330,25 +332,29 @@ struct AvailableUser {
   @available(macOS 10.15, *) var use1: String { "foo" }
 
   @main::available() var use2
-  // FIXME improve: expected-error@-1 {{unknown attribute 'main::available'}}
-  // FIXME suppress: expected-error@-2 {{type annotation missing in pattern}}
+  // expected-error@-1 {{'available' is not imported through module 'main'}}
+  // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{4-8=ModuleSelectorTestingKit}}
+  // FIXME suppress: expected-error@-3 {{type annotation missing in pattern}}
 
   @ModuleSelectorTestingKit::available() var use4
   // no-error
 
   @Swift::available() var use5
-  // FIXME improve: expected-error@-1 {{unknown attribute 'Swift::available'}}
-  // FIXME suppress: expected-error@-2 {{type annotation missing in pattern}}
+  // expected-error@-1 {{'available' is not imported through module 'Swift'}}
+  // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{4-9=ModuleSelectorTestingKit}}
+  // FIXME suppress: expected-error@-3 {{type annotation missing in pattern}}
 }
 
 func builderUser2(@main::MyBuilder fn: () -> Void) {}
-// FIXME improve: expected-error@-1 {{unknown attribute 'main::MyBuilder'}}
+// expected-error@-1 {{'MyBuilder' is not imported through module 'main'}}
+// expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{20-24=ModuleSelectorTestingKit}}
 
 func builderUser3(@ModuleSelectorTestingKit::MyBuilder fn: () -> Void) {}
 // no-error
 
 func builderUser4(@Swift::MyBuilder fn: () -> Void) {}
-// FIXME improve: expected-error@-1 {{unknown attribute 'Swift::MyBuilder'}}
+// expected-error@-1 {{'MyBuilder' is not imported through module 'Swift'}}
+// expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{20-25=ModuleSelectorTestingKit}}
 
 // Error cases