[Macros] Support module-qualified attached macro lookup

include/swift/AST/Attr.h

@@ -64,6 +64,7 @@ class ModuleDecl;
 class PatternBindingInitializer;
 class TrailingWhereClause;
 class TypeExpr;
+class IdentTypeRepr;
 
 class alignas(1 << AttrAlignInBits) AttributeBase
     : public ASTAllocated<AttributeBase> {
@@ -1769,6 +1770,15 @@ class CustomAttr final : public DeclAttribute {
   TypeRepr *getTypeRepr() const;
   Type getType() const;
 
+  /// Destructure an attribute's type repr for a macro reference.
+  ///
+  /// For a 1-level member type repr whose base and member are both identifier
+  /// types, e.g. `Foo.Bar`, return a pair of the base and the member.
+  ///
+  /// For an identifier type repr, return a pair of `nullptr` and the
+  /// identifier.
+  std::pair<IdentTypeRepr *, IdentTypeRepr *> destructureMacroRef();
+
   /// Whether the attribute has any arguments.
   bool hasArgs() const { return argList != nullptr; }
 

include/swift/AST/FreestandingMacroExpansion.h

@@ -34,6 +34,8 @@ class ArgumentList;
 /// declaration/expression nodes.
 struct MacroExpansionInfo : ASTAllocated<MacroExpansionInfo> {
   SourceLoc SigilLoc;
+  DeclNameRef ModuleName;
+  DeclNameLoc ModuleNameLoc;
   DeclNameRef MacroName;
   DeclNameLoc MacroNameLoc;
   SourceLoc LeftAngleLoc, RightAngleLoc;
@@ -43,13 +45,16 @@ struct MacroExpansionInfo : ASTAllocated<MacroExpansionInfo> {
   /// The referenced macro.
   ConcreteDeclRef macroRef;
 
-  MacroExpansionInfo(SourceLoc sigilLoc, DeclNameRef macroName,
+  MacroExpansionInfo(SourceLoc sigilLoc, DeclNameRef moduleName,
+                     DeclNameLoc moduleNameLoc, DeclNameRef macroName,
                      DeclNameLoc macroNameLoc, SourceLoc leftAngleLoc,
                      SourceLoc rightAngleLoc, ArrayRef<TypeRepr *> genericArgs,
                      ArgumentList *argList)
-      : SigilLoc(sigilLoc), MacroName(macroName), MacroNameLoc(macroNameLoc),
-        LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
-        GenericArgs(genericArgs), ArgList(argList) {}
+      : SigilLoc(sigilLoc), ModuleName(moduleName),
+        ModuleNameLoc(moduleNameLoc), MacroName(macroName),
+        MacroNameLoc(macroNameLoc), LeftAngleLoc(leftAngleLoc),
+        RightAngleLoc(rightAngleLoc), GenericArgs(genericArgs),
+        ArgList(argList) {}
 
   SourceLoc getLoc() const { return SigilLoc; }
   SourceRange getGenericArgsRange() const {
@@ -85,6 +90,10 @@ class FreestandingMacroExpansion {
 
   SourceLoc getPoundLoc() const { return getExpansionInfo()->SigilLoc; }
 
+  DeclNameLoc getModuleNameLoc() const {
+    return getExpansionInfo()->ModuleNameLoc;
+  }
+  DeclNameRef getModuleName() const { return getExpansionInfo()->ModuleName; }
   DeclNameLoc getMacroNameLoc() const {
     return getExpansionInfo()->MacroNameLoc;
   }

include/swift/AST/NameLookup.h

@@ -245,6 +245,8 @@ enum class UnqualifiedLookupFlags {
   ExcludeMacroExpansions = 1 << 6,
   /// This lookup should only return macros.
   MacroLookup            = 1 << 7,
+  /// This lookup should only return modules
+  ModuleLookup           = 1 << 8,
 };
 
 using UnqualifiedLookupOptions = OptionSet<UnqualifiedLookupFlags>;
@@ -550,8 +552,10 @@ void filterForDiscriminator(SmallVectorImpl<Result> &results,
 
 /// \returns The set of macro declarations with the given name that
 /// fulfill any of the given macro roles.
-SmallVector<MacroDecl *, 1>
-lookupMacros(DeclContext *dc, DeclNameRef macroName, MacroRoles roles);
+SmallVector<MacroDecl *, 1> lookupMacros(DeclContext *dc,
+                                         DeclNameRef moduleName,
+                                         DeclNameRef macroName,
+                                         MacroRoles roles);
 
 /// \returns Whether the given source location is inside an attached
 /// or freestanding macro argument.

include/swift/AST/TypeCheckRequests.h

@@ -3436,6 +3436,8 @@ class UnresolvedMacroReference {
   }
 
   SourceLoc getSigilLoc() const;
+  DeclNameRef getModuleName() const;
+  DeclNameLoc getModuleNameLoc() const;
   DeclNameRef getMacroName() const;
   DeclNameLoc getMacroNameLoc() const;
   SourceRange getGenericArgsRange() const;

lib/AST/Attr.cpp

@@ -2587,6 +2587,21 @@ CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
       CustomAttr(atLoc, range, type, initContext, argList, implicit);
 }
 
+std::pair<IdentTypeRepr *, IdentTypeRepr *> CustomAttr::destructureMacroRef() {
+  TypeRepr *typeRepr = getTypeRepr();
+  if (!typeRepr)
+    return {nullptr, nullptr};
+  if (auto *identType = dyn_cast<IdentTypeRepr>(typeRepr))
+    return {nullptr, identType};
+  if (auto *memType = dyn_cast<MemberTypeRepr>(typeRepr))
+    if (auto *base = dyn_cast<IdentTypeRepr>(memType->getBaseComponent()))
+      if (memType->getMemberComponents().size() == 1)
+        if (auto first =
+                dyn_cast<IdentTypeRepr>(memType->getMemberComponents().front()))
+          return {base, first};
+  return {nullptr, nullptr};
+}
+
 TypeRepr *CustomAttr::getTypeRepr() const { return typeExpr->getTypeRepr(); }
 Type CustomAttr::getType() const { return typeExpr->getInstanceType(); }
 

lib/AST/Decl.cpp

@@ -11479,11 +11479,16 @@ MacroExpansionDecl::create(
     ArgumentList *args
 ) {
   ASTContext &ctx = dc->getASTContext();
-  MacroExpansionInfo *info = new (ctx) MacroExpansionInfo{
-      poundLoc, macro, macroLoc,
-      leftAngleLoc, rightAngleLoc, genericArgs,
-      args ? args : ArgumentList::createImplicit(ctx, {})
-  };
+  MacroExpansionInfo *info = new (ctx)
+      MacroExpansionInfo{poundLoc,
+                         /*moduleName*/ DeclNameRef(),
+                         /*moduleNameLoc*/ DeclNameLoc(),
+                         macro,
+                         macroLoc,
+                         leftAngleLoc,
+                         rightAngleLoc,
+                         genericArgs,
+                         args ? args : ArgumentList::createImplicit(ctx, {})};
   return new (ctx) MacroExpansionDecl(dc, info);
 }
 

lib/AST/Expr.cpp

@@ -2769,10 +2769,15 @@ MacroExpansionExpr *MacroExpansionExpr::create(
 ) {
   ASTContext &ctx = dc->getASTContext();
   MacroExpansionInfo *info = new (ctx) MacroExpansionInfo{
-      sigilLoc, macroName, macroNameLoc,
-      leftAngleLoc, rightAngleLoc, genericArgs,
-      argList ? argList : ArgumentList::createImplicit(ctx, {})
-  };
+      sigilLoc,
+      /*moduleName*/ DeclNameRef(),
+      /*moduleNameLoc*/ DeclNameLoc(),
+      macroName,
+      macroNameLoc,
+      leftAngleLoc,
+      rightAngleLoc,
+      genericArgs,
+      argList ? argList : ArgumentList::createImplicit(ctx, {})};
   return new (ctx) MacroExpansionExpr(dc, info, roles, isImplicit, ty);
 }
 

lib/AST/NameLookup.cpp

@@ -113,6 +113,7 @@ void swift::simple_display(llvm::raw_ostream &out,
       {UnqualifiedLookupFlags::IncludeOuterResults, "IncludeOuterResults"},
       {UnqualifiedLookupFlags::TypeLookup, "TypeLookup"},
       {UnqualifiedLookupFlags::MacroLookup, "MacroLookup"},
+      {UnqualifiedLookupFlags::ModuleLookup, "ModuleLookup"},
   };
 
   auto flagsToPrint = llvm::make_filter_range(
@@ -1621,41 +1622,59 @@ static DeclName adjustLazyMacroExpansionNameKey(
   return name;
 }
 
-SmallVector<MacroDecl *, 1>
-namelookup::lookupMacros(DeclContext *dc, DeclNameRef macroName,
-                         MacroRoles roles) {
+SmallVector<MacroDecl *, 1> namelookup::lookupMacros(DeclContext *dc,
+                                                     DeclNameRef moduleName,
+                                                     DeclNameRef macroName,
+                                                     MacroRoles roles) {
   SmallVector<MacroDecl *, 1> choices;
   auto moduleScopeDC = dc->getModuleScopeContext();
   ASTContext &ctx = moduleScopeDC->getASTContext();
 
-  // When performing lookup for freestanding macro roles, only consider
-  // macro names, ignoring types.
-  bool onlyMacros = static_cast<bool>(roles & getFreestandingMacroRoles()) &&
-      !(roles - getFreestandingMacroRoles());
-
-  // Macro lookup should always exclude macro expansions; macro
-  // expansions cannot introduce new macro declarations. Note that
-  // the source location here doesn't matter.
-  UnqualifiedLookupDescriptor descriptor{
-    macroName, moduleScopeDC, SourceLoc(),
-    UnqualifiedLookupFlags::ExcludeMacroExpansions
-  };
-
-  if (onlyMacros)
-    descriptor.Options |= UnqualifiedLookupFlags::MacroLookup;
-
-  auto lookup = evaluateOrDefault(
-      ctx.evaluator, UnqualifiedLookupRequest{descriptor}, {});
-
-  for (const auto &found : lookup.allResults()) {
-    if (auto macro = dyn_cast<MacroDecl>(found.getValueDecl())) {
+  auto addChoiceIfApplicable = [&](ValueDecl *decl) {
+    if (auto macro = dyn_cast<MacroDecl>(decl)) {
       auto candidateRoles = macro->getMacroRoles();
       if ((candidateRoles && roles.contains(candidateRoles)) ||
           // FIXME: `externalMacro` should have all roles.
           macro->getBaseIdentifier().str() == "externalMacro") {
         choices.push_back(macro);
       }
     }
+  };
+
+  // When a module is specified, it's a module-qualified lookup.
+  if (moduleName) {
+    UnqualifiedLookupDescriptor moduleLookupDesc(
+        moduleName, moduleScopeDC, SourceLoc(),
+        UnqualifiedLookupFlags::ModuleLookup);
+    auto moduleLookup = evaluateOrDefault(
+        ctx.evaluator, UnqualifiedLookupRequest{moduleLookupDesc}, {});
+    auto foundTypeDecl = moduleLookup.getSingleTypeResult();
+    auto *moduleDecl = dyn_cast_or_null<ModuleDecl>(foundTypeDecl);
+    if (!moduleDecl)
+      return {};
+
+    ModuleQualifiedLookupRequest req{moduleScopeDC, moduleDecl, macroName,
+                                     SourceLoc(),
+                                     NL_ExcludeMacroExpansions | NL_OnlyMacros};
+    auto lookup = evaluateOrDefault(ctx.evaluator, req, {});
+    for (auto *found : lookup)
+      addChoiceIfApplicable(found);
+  }
+  // Otherwise it's an unqualified lookup.
+  else {
+    // Macro lookup should always exclude macro expansions; macro
+    // expansions cannot introduce new macro declarations. Note that
+    // the source location here doesn't matter.
+    UnqualifiedLookupDescriptor descriptor{
+        macroName, moduleScopeDC, SourceLoc(),
+        UnqualifiedLookupFlags::ExcludeMacroExpansions |
+            UnqualifiedLookupFlags::MacroLookup};
+
+    auto lookup = evaluateOrDefault(ctx.evaluator,
+                                    UnqualifiedLookupRequest{descriptor}, {});
+
+    for (const auto &found : lookup.allResults())
+      addChoiceIfApplicable(found.getValueDecl());
   }
 
   return choices;
@@ -1674,8 +1693,9 @@ namelookup::isInMacroArgument(SourceFile *sourceFile, SourceLoc loc) {
           inMacroArgument = true;
         } else if (auto *attr = macro.getAttr()) {
           auto *moduleScope = sourceFile->getModuleScopeContext();
-          auto results = lookupMacros(moduleScope, macro.getMacroName(),
-                                      getAttachedMacroRoles());
+          auto results =
+              lookupMacros(moduleScope, macro.getModuleName(),
+                           macro.getMacroName(), getAttachedMacroRoles());
           inMacroArgument = !results.empty();
         }
 
@@ -1696,9 +1716,9 @@ void namelookup::forEachPotentialResolvedMacro(
 ) {
   ASTContext &ctx = moduleScopeCtx->getASTContext();
   UnqualifiedLookupDescriptor lookupDesc{
-    macroName, moduleScopeCtx, SourceLoc(),
-    UnqualifiedLookupFlags::ExcludeMacroExpansions
-  };
+      macroName, moduleScopeCtx, SourceLoc(),
+      UnqualifiedLookupFlags::ExcludeMacroExpansions |
+          UnqualifiedLookupFlags::MacroLookup};
 
   auto lookup = evaluateOrDefault(
       ctx.evaluator, UnqualifiedLookupRequest{lookupDesc}, {});
@@ -3560,13 +3580,13 @@ CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
   // Look for names at module scope, so we don't trigger name lookup for
   // nested scopes. At this point, we're looking to see whether there are
   // any suitable macros.
-  if (auto *identTypeRepr =
-          dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr())) {
-    auto macros = namelookup::lookupMacros(
-        dc, identTypeRepr->getNameRef(), getAttachedMacroRoles());
-    if (!macros.empty())
-      return nullptr;
-  }
+  auto [module, macro] = attr->destructureMacroRef();
+  auto moduleName = (module) ? module->getNameRef() : DeclNameRef();
+  auto macroName = (macro) ? macro->getNameRef() : DeclNameRef();
+  auto macros = namelookup::lookupMacros(dc, moduleName, macroName,
+                                         getAttachedMacroRoles());
+  if (!macros.empty())
+    return nullptr;
 
   // Find the types referenced by the custom attribute.
   auto &ctx = dc->getASTContext();

lib/AST/TypeCheckRequests.cpp

@@ -1792,10 +1792,10 @@ DeclNameRef UnresolvedMacroReference::getMacroName() const {
   if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
     return expansion->getMacroName();
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto *identTypeRepr = dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr());
-    if (!identTypeRepr)
+    auto [_, macro] = attr->destructureMacroRef();
+    if (!macro)
       return DeclNameRef();
-    return identTypeRepr->getNameRef();
+    return macro->getNameRef();
   }
   llvm_unreachable("Unhandled case");
 }
@@ -1808,14 +1808,38 @@ SourceLoc UnresolvedMacroReference::getSigilLoc() const {
   llvm_unreachable("Unhandled case");
 }
 
+DeclNameRef UnresolvedMacroReference::getModuleName() const {
+  if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
+    return expansion->getModuleName();
+  if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
+    auto [module, _] = attr->destructureMacroRef();
+    if (!module)
+      return DeclNameRef();
+    return module->getNameRef();
+  }
+  llvm_unreachable("Unhandled case");
+}
+
+DeclNameLoc UnresolvedMacroReference::getModuleNameLoc() const {
+  if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
+    return expansion->getModuleNameLoc();
+  if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
+    auto [module, _] = attr->destructureMacroRef();
+    if (!module)
+      return DeclNameLoc();
+    return module->getNameLoc();
+  }
+  llvm_unreachable("Unhandled case");
+}
+
 DeclNameLoc UnresolvedMacroReference::getMacroNameLoc() const {
   if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
     return expansion->getMacroNameLoc();
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto *identTypeRepr = dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr());
-    if (!identTypeRepr)
+    auto [_, macro] = attr->destructureMacroRef();
+    if (!macro)
       return DeclNameLoc();
-    return identTypeRepr->getNameLoc();
+    return macro->getNameLoc();
   }
   llvm_unreachable("Unhandled case");
 }
@@ -1825,8 +1849,10 @@ SourceRange UnresolvedMacroReference::getGenericArgsRange() const {
     return expansion->getGenericArgsRange();
 
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto *typeRepr = attr->getTypeRepr();
-    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(typeRepr);
+    auto [_, macro] = attr->destructureMacroRef();
+    if (!macro)
+      return SourceRange();
+    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(macro);
     if (!genericTypeRepr)
       return SourceRange();
 
@@ -1841,8 +1867,10 @@ ArrayRef<TypeRepr *> UnresolvedMacroReference::getGenericArgs() const {
     return expansion->getGenericArgs();
 
   if (auto *attr = pointer.dyn_cast<CustomAttr *>()) {
-    auto *typeRepr = attr->getTypeRepr();
-    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(typeRepr);
+    auto [_, macro] = attr->destructureMacroRef();
+    if (!macro)
+      return {};
+    auto *genericTypeRepr = dyn_cast_or_null<GenericIdentTypeRepr>(macro);
     if (!genericTypeRepr)
       return {};
 

lib/AST/UnqualifiedLookup.cpp

@@ -281,6 +281,12 @@ void UnqualifiedLookupFactory::performUnqualifiedLookup() {
                                   "performUnqualifiedLookup",
                                   DC->getParentSourceFile());
 
+  if (options.contains(UnqualifiedLookupFlags::ModuleLookup)) {
+    lookForAModuleWithTheGivenName(DC->getModuleScopeContext());
+    recordCompletionOfAScope();
+    return;
+  }
+
   if (Loc.isValid() && DC->getParentSourceFile()) {
     // Operator lookup is always global, for the time being.
     if (!Name.isOperator())