[Macros] Allow custom attributes to resolve to attached macro declarations

include/swift/AST/Attr.h

@@ -1712,7 +1712,7 @@ class CustomAttr final : public DeclAttribute {
   }
 
 private:
-  friend class CustomAttrNominalRequest;
+  friend class CustomAttrDeclRequest;
   void resetTypeInformation(TypeExpr *repr);
 
 private:

include/swift/AST/DiagnosticsCommon.def

@@ -221,6 +221,8 @@ NOTE(in_macro_expansion,none,
      "in expansion of macro %0 here", (DeclName))
 ERROR(macro_experimental,none,
       "macros are an experimental feature that is not enabled", ())
+ERROR(ambiguous_macro_reference,none,
+      "ambiguous reference to macro %0", (DeclName))
 
 //------------------------------------------------------------------------------
 // MARK: bridged diagnostics

include/swift/AST/MacroDeclaration.h

@@ -30,7 +30,7 @@ enum class MacroContext: uint8_t {
   /// A freestanding declaration macro.
   FreestandingDeclaration = 0x02,
   /// An attached declaration macro.
-  AttachedDeclaration = 0x03,
+  AttachedDeclaration = 0x04,
 };
 
 /// The contexts in which a particular macro declaration can be used.

include/swift/AST/NameLookupRequests.h

@@ -323,11 +323,14 @@ class TypeDeclsFromWhereClauseRequest :
                                        ExtensionDecl *ext) const;
 };
 
-/// Request the nominal type declaration to which the given custom attribute
-/// refers.
-class CustomAttrNominalRequest :
-    public SimpleRequest<CustomAttrNominalRequest,
-                         NominalTypeDecl *(CustomAttr *, DeclContext *),
+using MacroOrNominalTypeDecl =
+    llvm::PointerUnion<MacroDecl *, NominalTypeDecl *>;
+
+/// Request the macro or nominal type declaration to which the given custom
+/// attribute refers.
+class CustomAttrDeclRequest :
+    public SimpleRequest<CustomAttrDeclRequest,
+                         MacroOrNominalTypeDecl(CustomAttr *, DeclContext *),
                          RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -336,7 +339,7 @@ class CustomAttrNominalRequest :
   friend SimpleRequest;
 
   // Evaluation.
-  NominalTypeDecl *
+  MacroOrNominalTypeDecl
   evaluate(Evaluator &evaluator, CustomAttr *attr, DeclContext *dc) const;
 
 public:

include/swift/AST/NameLookupTypeIDZone.def

@@ -18,8 +18,8 @@
 SWIFT_REQUEST(NameLookup, AnyObjectLookupRequest,
               QualifiedLookupResult(const DeclContext *, DeclName, NLOptions),
               Uncached, NoLocationInfo)
-SWIFT_REQUEST(NameLookup, CustomAttrNominalRequest,
-              NominalTypeDecl *(CustomAttr *, DeclContext *), Cached,
+SWIFT_REQUEST(NameLookup, CustomAttrDeclRequest,
+              MacroOrNominalTypeDecl(CustomAttr *, DeclContext *), Cached,
               NoLocationInfo)
 SWIFT_REQUEST(NameLookup, DirectLookupRequest,
               TinyPtrVector<ValueDecl *>(DirectLookupDescriptor), Uncached,

lib/AST/Decl.cpp

@@ -6807,8 +6807,11 @@ VarDecl::getAttachedPropertyWrapperTypeInfo(unsigned i) const {
     auto attr = attrs[i];
     auto dc = getDeclContext();
     ASTContext &ctx = getASTContext();
-    nominal = evaluateOrDefault(
-        ctx.evaluator, CustomAttrNominalRequest{attr, dc}, nullptr);
+    if (auto found = evaluateOrDefault(
+           ctx.evaluator, CustomAttrDeclRequest{attr, dc}, nullptr))
+      nominal = found.dyn_cast<NominalTypeDecl *>();
+    else
+      nominal = nullptr;
   }
 
   if (!nominal)
@@ -9748,7 +9751,8 @@ NominalTypeDecl *
 ValueDecl::getRuntimeDiscoverableAttrTypeDecl(CustomAttr *attr) const {
   auto &ctx = getASTContext();
   auto *nominal = evaluateOrDefault(
-      ctx.evaluator, CustomAttrNominalRequest{attr, getDeclContext()}, nullptr);
+      ctx.evaluator, CustomAttrDeclRequest{attr, getDeclContext()}, nullptr)
+    .get<NominalTypeDecl *>();
   assert(nominal->getAttrs().hasAttribute<RuntimeMetadataAttr>());
   return nominal;
 }

lib/AST/NameLookup.cpp

@@ -3040,9 +3040,57 @@ GenericParamListRequest::evaluate(Evaluator &evaluator, GenericContext *value) c
       parsedGenericParams->getRAngleLoc());
 }
 
-NominalTypeDecl *
-CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
-                                   CustomAttr *attr, DeclContext *dc) const {
+/// Perform lookup to determine whether the given custom attribute refers to
+/// a macro declaration, and return that macro declaration.
+static MacroDecl *findMacroForCustomAttr(CustomAttr *attr, DeclContext *dc) {
+  auto *identTypeRepr = dyn_cast_or_null<IdentTypeRepr>(attr->getTypeRepr());
+  if (!identTypeRepr)
+    return nullptr;
+
+  // Look for macros at module scope. They can only occur at module scope, and
+  // we need to be sure not to trigger name lookup into type contexts along
+  // the way.
+  llvm::TinyPtrVector<MacroDecl *> macros;
+  auto moduleScopeDC = dc->getModuleScopeContext();
+  ASTContext &ctx = moduleScopeDC->getASTContext();
+  UnqualifiedLookupDescriptor descriptor(
+      identTypeRepr->getNameRef(), moduleScopeDC
+  );
+  auto lookup = evaluateOrDefault(
+      ctx.evaluator, UnqualifiedLookupRequest{descriptor}, {});
+  for (const auto &result : lookup.allResults()) {
+    // Only keep attached macros, which can be spelled as custom attributes.
+    if (auto macro = dyn_cast<MacroDecl>(result.getValueDecl()))
+      if (macro->getMacroContexts().contains(MacroContext::AttachedDeclaration))
+        macros.push_back(macro);
+  }
+
+  if (macros.empty())
+    return nullptr;
+
+  if (macros.size() > 1) {
+    ctx.Diags.diagnose(attr->getLocation(), diag::ambiguous_macro_reference,
+                       identTypeRepr->getNameRef().getFullName());
+
+    for (auto macro : macros) {
+      macro->diagnose(
+          diag::kind_declname_declared_here, macro->getDescriptiveKind(),
+          macro->getName());
+    }
+  }
+
+  return macros.front();
+}
+
+MacroOrNominalTypeDecl
+CustomAttrDeclRequest::evaluate(Evaluator &evaluator,
+                                CustomAttr *attr, DeclContext *dc) const {
+  // 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 macro = findMacroForCustomAttr(attr, dc))
+    return macro;
+
   // Find the types referenced by the custom attribute.
   auto &ctx = dc->getASTContext();
   DirectlyReferencedTypeDecls decls;

lib/Sema/TypeCheckAttr.cpp

@@ -3553,11 +3553,22 @@ void AttributeChecker::visitCustomAttr(CustomAttr *attr) {
   auto dc = D->getDeclContext();
 
   // Figure out which nominal declaration this custom attribute refers to.
-  auto nominal = evaluateOrDefault(
-    Ctx.evaluator, CustomAttrNominalRequest{attr, dc}, nullptr);
+  auto found = evaluateOrDefault(
+    Ctx.evaluator, CustomAttrDeclRequest{attr, dc}, nullptr);
+
+  // FIXME: deal with macros.
+  NominalTypeDecl *nominal = nullptr;
+  if (found) {
+    // FIXME: Do full checking of the macro arguments here by turning it into
+    // a macro expansion expression (?).
+    if (found.is<MacroDecl *>())
+      return;
+
+    nominal = found.dyn_cast<NominalTypeDecl *>();
+  }
 
   // Diagnose errors.
-  if (!nominal) {
+  if (!found) {
     auto typeRepr = attr->getTypeRepr();
 
     auto type = TypeResolution::forInterface(dc, TypeResolverContext::CustomAttr,
@@ -7377,12 +7388,15 @@ static void forEachCustomAttribute(
   for (auto *attr : decl->getAttrs().getAttributes<CustomAttr>()) {
     auto *mutableAttr = const_cast<CustomAttr *>(attr);
 
-    auto *nominal = evaluateOrDefault(
+    auto found = evaluateOrDefault(
         ctx.evaluator,
-        CustomAttrNominalRequest{mutableAttr, decl->getDeclContext()}, nullptr);
+        CustomAttrDeclRequest{mutableAttr, decl->getDeclContext()}, nullptr);
+    if (!found)
+      continue;
 
+    auto nominal = found.dyn_cast<NominalTypeDecl *>();
     if (!nominal)
-      continue;
+      continue; // FIXME: add another entry point for macros we've found
 
     if (nominal->getAttrs().hasAttribute<ATTR>())
       fn(mutableAttr, nominal);

lib/Sema/TypeCheckConcurrency.cpp

@@ -257,11 +257,15 @@ swift::checkGlobalActorAttributes(
   NominalTypeDecl *globalActorNominal = nullptr;
   for (auto attr : attrs) {
     // Figure out which nominal declaration this custom attribute refers to.
-    auto nominal = evaluateOrDefault(ctx.evaluator,
-                                     CustomAttrNominalRequest{attr, dc},
-                                     nullptr);
+    auto found = evaluateOrDefault(ctx.evaluator,
+                                   CustomAttrDeclRequest{attr, dc},
+                                    nullptr);
 
     // Ignore unresolvable custom attributes.
+    if (!found)
+      continue;
+
+    auto nominal = found.dyn_cast<NominalTypeDecl *>();
     if (!nominal)
       continue;
 

lib/Sema/TypeCheckPropertyWrapper.cpp

@@ -439,8 +439,12 @@ AttachedPropertyWrappersRequest::evaluate(Evaluator &evaluator,
   for (auto attr : var->getAttrs().getAttributes<CustomAttr>()) {
     auto mutableAttr = const_cast<CustomAttr *>(attr);
     // Figure out which nominal declaration this custom attribute refers to.
-    auto nominal = evaluateOrDefault(
-      ctx.evaluator, CustomAttrNominalRequest{mutableAttr, dc}, nullptr);
+    auto found = evaluateOrDefault(
+      ctx.evaluator, CustomAttrDeclRequest{mutableAttr, dc}, nullptr);
+
+    NominalTypeDecl *nominal = nullptr;
+    if (found)
+      nominal = found.dyn_cast<NominalTypeDecl *>();
 
     // If we didn't find a nominal type with a @propertyWrapper attribute,
     // skip this custom attribute.

lib/Sema/TypeCheckRequestFunctions.cpp

@@ -174,11 +174,15 @@ AttachedResultBuilderRequest::evaluate(Evaluator &evaluator,
   for (auto attr : decl->getAttrs().getAttributes<CustomAttr>()) {
     auto mutableAttr = const_cast<CustomAttr *>(attr);
     // Figure out which nominal declaration this custom attribute refers to.
-    auto nominal = evaluateOrDefault(ctx.evaluator,
-                                     CustomAttrNominalRequest{mutableAttr, dc},
-                                     nullptr);
+    auto found = evaluateOrDefault(ctx.evaluator,
+                                   CustomAttrDeclRequest{mutableAttr, dc},
+                                  nullptr);
 
     // Ignore unresolvable custom attributes.
+    if (!found)
+      continue;
+
+    auto nominal = found.dyn_cast<NominalTypeDecl *>();
     if (!nominal)
       continue;
 

lib/Sema/TypeCheckType.cpp

@@ -5022,7 +5022,7 @@ Type CustomAttrTypeRequest::evaluate(Evaluator &eval, CustomAttr *attr,
 
   // We always require the type to resolve to a nominal type. If the type was
   // not a nominal type, we should have already diagnosed an error via
-  // CustomAttrNominalRequest.
+  // CustomAttrDeclRequest.
   auto checkType = [](Type type) -> bool {
     while (auto *genericDecl = type->getAnyGeneric()) {
       if (isa<NominalTypeDecl>(genericDecl))

lib/Sema/TypeCheckTypeWrapper.cpp

@@ -123,9 +123,13 @@ static void getTypeWrappers(NominalTypeDecl *decl,
   // Attributes applied directly to the type.
   for (auto *attr : decl->getAttrs().getAttributes<CustomAttr>()) {
     auto *mutableAttr = const_cast<CustomAttr *>(attr);
-    auto *nominal = evaluateOrDefault(
-        ctx.evaluator, CustomAttrNominalRequest{mutableAttr, decl}, nullptr);
+    auto found = evaluateOrDefault(
+        ctx.evaluator, CustomAttrDeclRequest{mutableAttr, decl}, nullptr);
 
+    if (!found)
+      continue;
+
+    auto nominal = found.dyn_cast<NominalTypeDecl *>();
     if (!nominal)
       continue;
 

test/Macros/attached_macros_diags.swift

@@ -0,0 +1,30 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-feature Macros -module-name MacrosTest
+
+@declaration(attached) macro m1: Void = #externalMacro(module: "MyMacros", type: "Macro1")
+// expected-warning@-1{{external macro implementation type 'MyMacros.Macro1' could not be found for macro 'm1'}}
+
+@declaration(attached) macro m2(_: Int) -> Void = #externalMacro(module: "MyMacros", type: "Macro2")
+// expected-warning@-1{{external macro implementation type 'MyMacros.Macro2' could not be found for macro 'm2'}}
+// expected-note@-2 2{{macro 'm2' declared here}}
+
+@declaration(attached) macro m2(_: Double) -> Void = #externalMacro(module: "MyMacros", type: "Macro2")
+// expected-warning@-1{{external macro implementation type 'MyMacros.Macro2' could not be found for macro 'm2'}}
+// expected-note@-2 2{{macro 'm2' declared here}}
+
+@m1 struct X1 { }
+
+// FIXME: Redundant diagnostic
+@m2 struct X2 { } // expected-error 2{{ambiguous reference to macro 'm2'}}
+
+// Check for nesting rules.
+struct SkipNestedType {
+  @propertyWrapper
+  struct m1<T> {
+    init() { }
+
+    var wrappedValue: T
+  }
+
+  // We select the macro, not the property wrapper.
+  @m1 var x: Int = 0
+}