[AST] Use CustomAttr's DeclContext for ResolveMacroRequest

hamishknight · hamishknight · commit 364eba482db9 · 2025-10-16T11:21:54.000+01:00
Remove the DeclContext parameter from ResolveMacroRequest, we can now
retrieve the DeclContext either from the CustomAttr or macro expansion
expr/decl directly.
diff --git a/include/swift/AST/AnyFunctionRef.h b/include/swift/AST/AnyFunctionRef.h
@@ -265,34 +265,15 @@ class AnyFunctionRef {
     return DeclAttributes();
   }
 
-  MacroDecl *getResolvedMacro(CustomAttr *attr) const {
-    if (auto afd = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
-      return afd->getResolvedMacro(attr);
-    }
-
-    if (auto ace = TheFunction.dyn_cast<AbstractClosureExpr *>()) {
-      if (auto *ce = dyn_cast<ClosureExpr>(ace)) {
-        return ce->getResolvedMacro(attr);
-      }
-    }
-
-    return nullptr;
-  }
-
   using MacroCallback = llvm::function_ref<void(CustomAttr *, MacroDecl *)>;
 
   void
   forEachAttachedMacro(MacroRole role,
                        MacroCallback macroCallback) const {
     auto attrs = getDeclAttributes();
-    for (auto customAttrConst : attrs.getAttributes<CustomAttr>()) {
-      auto customAttr = const_cast<CustomAttr *>(customAttrConst);
-      auto *macroDecl = getResolvedMacro(customAttr);
-
-      if (!macroDecl)
-        continue;
-
-      if (!macroDecl->getMacroRoles().contains(role))
+    for (auto *customAttr : attrs.getAttributes<CustomAttr>()) {
+      auto *macroDecl = customAttr->getResolvedMacro();
+      if (!macroDecl || !macroDecl->getMacroRoles().contains(role))
         continue;
 
       macroCallback(customAttr, macroDecl);
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -67,6 +67,7 @@ class CustomAttributeInitializer;
 class GenericFunctionType;
 class LazyConformanceLoader;
 class LazyMemberLoader;
+class MacroDecl;
 class ModuleDecl;
 class NominalTypeDecl;
 class PatternBindingInitializer;
@@ -2339,6 +2340,10 @@ class CustomAttr final : public DeclAttribute {
   /// it doesn't refer to one (which can be the case for e.g macro attrs).
   NominalTypeDecl *getNominalDecl() const;
 
+  /// Retrieve the resolved macro for the CustomAttr, or \c nullptr if the
+  /// attribute does not refer to a macro.
+  MacroDecl *getResolvedMacro() 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
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -1105,10 +1105,6 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl>, public Swi
   /// declaration.
   void forEachAttachedMacro(MacroRole role, MacroCallback) const;
 
-  /// Returns the resolved macro for the given custom attribute
-  /// attached to this declaration.
-  MacroDecl *getResolvedMacro(CustomAttr *attr) const;
-
   /// Retrieve the discriminator for the given custom attribute that names
   /// an attached macro.
   unsigned getAttachedMacroDiscriminator(DeclBaseName macroName, MacroRole role,
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -4471,10 +4471,6 @@ class ClosureExpr : public AbstractClosureExpr {
     return ExplicitResultTypeAndBodyState.getPointer()->getTypeRepr();
   }
 
-  /// Returns the resolved macro for the given custom attribute
-  /// attached to this closure expression.
-  MacroDecl *getResolvedMacro(CustomAttr *customAttr);
-
   /// Determine whether the closure has a single expression for its
   /// body.
   ///
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -3728,6 +3728,8 @@ class UnresolvedMacroReference {
   ArrayRef<TypeRepr *> getGenericArgs() const;
   ArgumentList *getArgs() const;
 
+  DeclContext *getDeclContext() const;
+
   /// Returns the macro roles corresponding to this macro reference.
   MacroRoles getMacroRoles() const;
 
@@ -3753,8 +3755,7 @@ void simple_display(llvm::raw_ostream &out,
 /// Resolve a given custom attribute to an attached macro declaration.
 class ResolveMacroRequest
     : public SimpleRequest<ResolveMacroRequest,
-                           ConcreteDeclRef(UnresolvedMacroReference,
-                                           DeclContext *),
+                           ConcreteDeclRef(UnresolvedMacroReference),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -3763,8 +3764,7 @@ class ResolveMacroRequest
   friend SimpleRequest;
 
   ConcreteDeclRef evaluate(Evaluator &evaluator,
-                           UnresolvedMacroReference macroRef,
-                           DeclContext *decl) const;
+                           UnresolvedMacroReference macroRef) const;
 
 public:
   bool isCached() const { return true; }
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -400,7 +400,7 @@ SWIFT_REQUEST(TypeChecker, ResolveImplicitMemberRequest,
               evaluator::SideEffect(NominalTypeDecl *, ImplicitMemberAction),
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ResolveMacroRequest,
-              ConcreteDeclRef(UnresolvedMacroReference, const Decl *),
+              ConcreteDeclRef(UnresolvedMacroReference),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ResolveMacroConformances,
               ArrayRef<Type>(const MacroRoleAttr *, const Decl *),
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -5172,12 +5172,8 @@ class PrintAttribute : public AttributeVisitor<PrintAttribute, void, Label>,
     } else if (isTypeChecked()) {
       // If the type is null, it might be a macro reference. Try that if we're
       // dumping the fully type-checked AST.
-      auto macroRef =
-          evaluateOrDefault(const_cast<ASTContext *>(Ctx)->evaluator,
-                            ResolveMacroRequest{Attr, DC}, ConcreteDeclRef());
-      if (macroRef) {
+      if (auto macroRef = Attr->getResolvedMacro())
         printDeclRefField(macroRef, Label::always("macro"));
-      }
     }
     if (!Writer.isParsable()) {
       // The type has the semantic information we want for parsable outputs, so
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -5000,7 +5000,7 @@ void ASTMangler::appendMacroExpansionContext(
     outerExpansionLoc = decl->getLoc();
     outerExpansionDC = decl->getDeclContext();
 
-    if (auto *macroDecl = decl->getResolvedMacro(attr))
+    if (auto *macroDecl = attr->getResolvedMacro())
       baseName = macroDecl->getBaseName();
     else
       baseName = Context.getIdentifier("__unknown_macro__");
@@ -5264,7 +5264,6 @@ std::string ASTMangler::mangleAttachedMacroExpansion(
   // We don't mangle the declaration itself because doing so requires semantic
   // information (e.g., its interface type), which introduces cyclic
   // dependencies.
-  const Decl *attachedTo = decl;
   Identifier attachedToName;
   if (auto accessor = dyn_cast<AccessorDecl>(decl)) {
     auto storage = accessor->getStorage();
@@ -5294,12 +5293,6 @@ std::string ASTMangler::mangleAttachedMacroExpansion(
     }
 
     appendDeclWithName(storage, attachedToName);
-
-    // For member attribute macros, the attribute is attached to the enclosing
-    // declaration.
-    if (role == MacroRole::MemberAttribute) {
-      attachedTo = storage->getDeclContext()->getAsDecl();
-    }
   } else if (auto valueDecl = dyn_cast<ValueDecl>(decl)) {
     // Mangle the name, replacing special names with their user-facing names.
     auto name = valueDecl->getName().getBaseName();
@@ -5311,20 +5304,14 @@ std::string ASTMangler::mangleAttachedMacroExpansion(
     }
 
     appendDeclWithName(valueDecl, attachedToName);
-
-    // For member attribute macros, the attribute is attached to the enclosing
-    // declaration.
-    if (role == MacroRole::MemberAttribute) {
-      attachedTo = decl->getDeclContext()->getAsDecl();
-    }
   } else {
     appendContext(decl->getDeclContext(), nullBase, "");
     appendIdentifier("_");
   }
 
   // Determine the name of the macro.
   DeclBaseName macroName;
-  if (auto *macroDecl = attachedTo->getResolvedMacro(attr)) {
+  if (auto *macroDecl = attr->getResolvedMacro()) {
     macroName = macroDecl->getName().getBaseName();
   } else {
     macroName = decl->getASTContext().getIdentifier("__unknown_macro__");
diff --git a/lib/AST/Attr.cpp b/lib/AST/Attr.cpp
@@ -841,7 +841,7 @@ void DeclAttributes::print(ASTPrinter &Printer, const PrintOptions &Options,
     // a macro.
     if (Options.SuppressExpandedMacros) {
       if (auto customAttr = dyn_cast<CustomAttr>(DA)) {
-        if (D->getResolvedMacro(const_cast<CustomAttr *>(customAttr)))
+        if (customAttr->getResolvedMacro())
           continue;
       }
     }
@@ -3162,6 +3162,16 @@ NominalTypeDecl *CustomAttr::getNominalDecl() const {
   return evaluateOrDefault(eval, CustomAttrNominalRequest{mutThis}, nullptr);
 }
 
+MacroDecl *CustomAttr::getResolvedMacro() const {
+  auto &eval = getASTContext().evaluator;
+  auto *mutThis = const_cast<CustomAttr *>(this);
+
+  auto declRef =
+      evaluateOrDefault(eval, ResolveMacroRequest{mutThis}, ConcreteDeclRef());
+
+  return dyn_cast_or_null<MacroDecl>(declRef.getDecl());
+}
+
 TypeRepr *CustomAttr::getTypeRepr() const { return typeExpr->getTypeRepr(); }
 Type CustomAttr::getType() const { return typeExpr->getInstanceType(); }
 
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -512,28 +512,15 @@ void Decl::visitAuxiliaryDecls(
 
 void Decl::forEachAttachedMacro(MacroRole role,
                                 MacroCallback macroCallback) const {
-  for (auto customAttrConst : getExpandedAttrs().getAttributes<CustomAttr>()) {
-    auto customAttr = const_cast<CustomAttr *>(customAttrConst);
-    auto *macroDecl = getResolvedMacro(customAttr);
-
-    if (!macroDecl)
-      continue;
-
-    if (!macroDecl->getMacroRoles().contains(role))
+  for (auto *customAttr : getExpandedAttrs().getAttributes<CustomAttr>()) {
+    auto *macroDecl = customAttr->getResolvedMacro();
+    if (!macroDecl || !macroDecl->getMacroRoles().contains(role))
       continue;
 
     macroCallback(customAttr, macroDecl);
   }
 }
 
-MacroDecl *Decl::getResolvedMacro(CustomAttr *customAttr) const {
-  auto declRef = evaluateOrDefault(
-      getASTContext().evaluator,
-      ResolveMacroRequest{customAttr, getDeclContext()}, ConcreteDeclRef());
-
-  return dyn_cast_or_null<MacroDecl>(declRef.getDecl());
-}
-
 unsigned Decl::getAttachedMacroDiscriminator(DeclBaseName macroName,
                                              MacroRole role,
                                              const CustomAttr *attr) const {
@@ -12508,27 +12495,22 @@ void MissingDecl::forEachMacroExpandedDecl(MacroExpandedDeclCallback callback) {
   auto macroRef = unexpandedMacro.macroRef;
   auto *baseDecl = unexpandedMacro.baseDecl;
 
-  // If the macro itself is a macro expansion expression, it should come with
-  // a top-level code declaration that we can use for resolution. For such
-  // cases, resolve the macro to determine whether it is a declaration or
-  // code-item macro, meaning that it can produce declarations. In such cases,
-  // expand the macro and use its substituted declaration (a MacroExpansionDecl)
-  // instead.
+  // If the macro itself is a macro expansion expression, resolve it to
+  // determine whether it is a declaration or code-item macro, meaning that it
+  // can produce declarations. In such cases, expand the macro and use its
+  // substituted declaration (a MacroExpansionDecl) instead.
   if (auto freestanding = macroRef.dyn_cast<FreestandingMacroExpansion *>()) {
     if (auto expr = dyn_cast<MacroExpansionExpr>(freestanding)) {
       bool replacedWithDecl = false;
-      if (auto tlcd = dyn_cast_or_null<TopLevelCodeDecl>(baseDecl)) {
-        ASTContext &ctx = tlcd->getASTContext();
-        if (auto macro = evaluateOrDefault(
-                ctx.evaluator,
-                ResolveMacroRequest{macroRef, tlcd->getDeclContext()},
-                nullptr)) {
+      if (isa_and_nonnull<TopLevelCodeDecl>(baseDecl)) {
+        auto &eval = getASTContext().evaluator;
+        if (auto macro = evaluateOrDefault(eval, ResolveMacroRequest{macroRef},
+                                           nullptr)) {
           auto macroDecl = cast<MacroDecl>(macro.getDecl());
           auto roles = macroDecl->getMacroRoles();
           if (roles.contains(MacroRole::Declaration) ||
               roles.contains(MacroRole::CodeItem)) {
-            (void)evaluateOrDefault(ctx.evaluator,
-                                    ExpandMacroExpansionExprRequest{expr},
+            (void)evaluateOrDefault(eval, ExpandMacroExpansionExprRequest{expr},
                                     std::nullopt);
             if (auto substituted = expr->getSubstituteDecl()) {
               macroRef = substituted;
diff --git a/lib/AST/Evaluator.cpp b/lib/AST/Evaluator.cpp
@@ -17,7 +17,7 @@
 #include "swift/AST/Evaluator.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/DiagnosticEngine.h"
-#include "swift/AST/TypeCheckRequests.h" // for ResolveMacroRequest
+#include "swift/AST/DiagnosticsSema.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/LangOptions.h"
diff --git a/lib/AST/Expr.cpp b/lib/AST/Expr.cpp
@@ -2204,17 +2204,6 @@ void ClosureExpr::setExplicitResultType(Type ty) {
       ->setType(MetatypeType::get(ty));
 }
 
-MacroDecl *
-ClosureExpr::getResolvedMacro(CustomAttr *customAttr) {
-  auto &ctx = getASTContext();
-  auto declRef = evaluateOrDefault(
-      ctx.evaluator,
-      ResolveMacroRequest{customAttr, this},
-      ConcreteDeclRef());
-
-  return dyn_cast_or_null<MacroDecl>(declRef.getDecl());
-}
-
 FORWARD_SOURCE_LOCS_TO(AutoClosureExpr, Body)
 
 void AutoClosureExpr::setBody(Expr *E) {
diff --git a/lib/AST/TypeCheckRequests.cpp b/lib/AST/TypeCheckRequests.cpp
@@ -2333,6 +2333,14 @@ ArgumentList *UnresolvedMacroReference::getArgs() const {
   llvm_unreachable("Unhandled case");
 }
 
+DeclContext *UnresolvedMacroReference::getDeclContext() const {
+  if (auto *expansion = pointer.dyn_cast<FreestandingMacroExpansion *>())
+    return expansion->getDeclContext();
+  if (auto *attr = pointer.dyn_cast<CustomAttr *>())
+    return attr->getOwner().getDeclContext();
+  llvm_unreachable("Unhandled case");
+}
+
 MacroRoles UnresolvedMacroReference::getMacroRoles() const {
   if (isa<FreestandingMacroExpansion *>(pointer))
     return getFreestandingMacroRoles();
diff --git a/lib/IDE/CodeCompletion.cpp b/lib/IDE/CodeCompletion.cpp
@@ -1455,18 +1455,17 @@ void CodeCompletionCallbacksImpl::typeCheckWithLookup(
     /// The attribute might not be attached to the AST if there is no var
     /// decl it could be attached to. Type check it standalone.
 
+    auto &ctx = CurDeclContext->getASTContext();
+
     // First try to check it as an attached macro.
-    (void)evaluateOrDefault(
-        CurDeclContext->getASTContext().evaluator,
-        ResolveMacroRequest{AttrWithCompletion, CurDeclContext},
-        ConcreteDeclRef());
+    (void)AttrWithCompletion->getResolvedMacro();
 
     // If that fails, type check as a call to the attribute's type. This is
     // how, e.g., property wrappers are modelled.
     if (!Lookup.gotCallback()) {
-      ASTNode Call = CallExpr::create(
-          CurDeclContext->getASTContext(), AttrWithCompletion->getTypeExpr(),
-          AttrWithCompletion->getArgs(), /*implicit=*/true);
+      ASTNode Call =
+          CallExpr::create(ctx, AttrWithCompletion->getTypeExpr(),
+                           AttrWithCompletion->getArgs(), /*implicit=*/true);
       swift::typeCheckASTNodeAtLoc(
           TypeCheckASTNodeAtLocContext::node(CurDeclContext, Call),
           CompletionLoc);
diff --git a/lib/IDE/SourceEntityWalker.cpp b/lib/IDE/SourceEntityWalker.cpp
@@ -739,8 +739,7 @@ bool SemaAnnotator::handleCustomAttributes(Decl *D) {
       // It's a little weird that attached macros have a `TypeRepr` to begin
       // with, but given they aren't types they then don't get bound. So check
       // for a macro here and and pass a reference to it.
-      auto *mutableAttr = const_cast<CustomAttr *>(customAttr);
-      if (auto macroDecl = D->getResolvedMacro(mutableAttr)) {
+      if (auto *macroDecl = customAttr->getResolvedMacro()) {
         Type macroRefType = macroDecl->getDeclaredInterfaceType();
         auto customAttrRef =
             std::make_pair(customAttr, expansion ? cast<Decl *>(expansion) : D);
diff --git a/lib/SILGen/SILGenFunction.cpp b/lib/SILGen/SILGenFunction.cpp
@@ -369,9 +369,7 @@ static MacroInfo getMacroInfo(const GeneratedSourceInfo &Info,
   case GeneratedSourceInfo::ExtensionMacroExpansion:
   case GeneratedSourceInfo::PreambleMacroExpansion:
   case GeneratedSourceInfo::BodyMacroExpansion: {
-    auto decl = cast<Decl *>(ASTNode::getFromOpaqueValue(Info.astNode));
-    auto attr = Info.attachedMacroCustomAttr;
-    if (auto *macroDecl = decl->getResolvedMacro(attr)) {
+    if (auto *macroDecl = Info.attachedMacroCustomAttr->getResolvedMacro()) {
       Result.ExpansionLoc = RegularLocation(macroDecl);
       Result.Name = macroDecl->getBaseName().userFacingName();
       Result.Freestanding = true;
diff --git a/lib/Sema/NonisolatedNonsendingByDefaultMigration.cpp b/lib/Sema/NonisolatedNonsendingByDefaultMigration.cpp
@@ -66,12 +66,12 @@ static bool isSwiftTestingTestFunction(ValueDecl *decl) {
   if (!isa<FuncDecl>(decl))
     return false;
 
-  return llvm::any_of(decl->getAttrs(), [&decl](DeclAttribute *attr) {
+  return llvm::any_of(decl->getAttrs(), [](DeclAttribute *attr) {
     auto customAttr = dyn_cast<CustomAttr>(attr);
     if (!customAttr)
       return false;
 
-    auto *macro = decl->getResolvedMacro(customAttr);
+    auto *macro = customAttr->getResolvedMacro();
     return macro && macro->getBaseIdentifier().is("Test") &&
            macro->getParentModule()->getName().is("Testing");
   });
diff --git a/lib/Sema/TypeCheckAccess.cpp b/lib/Sema/TypeCheckAccess.cpp
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
diff --git a/lib/Sema/TypeCheckDeclPrimary.cpp b/lib/Sema/TypeCheckDeclPrimary.cpp
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
diff --git a/test/Macros/rdar162469605.swift b/test/Macros/rdar162469605.swift
