Merge pull request #84916 from hamishknight/mac-n-cheese

Author: hamishknight

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);

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;
@@ -619,6 +620,9 @@ class DeclAttribute : public AttributeBase {
                                      DeclAttrKind kind, SourceLoc atLoc,
                                      SourceLoc attrLoc);
 
+  /// Attaches the attribute to the given declaration.
+  void attachToDecl(Decl *D);
+
   /// Create a copy of this attribute.
   DeclAttribute *clone(ASTContext &ctx) const;
 
@@ -629,6 +633,11 @@ class DeclAttribute : public AttributeBase {
   /// would have the same effect on \p attachedTo were they attached to it. A
   /// clone should always be equivalent to the original.
   bool isEquivalent(const DeclAttribute *other, Decl *attachedTo) const;
+
+private:
+  void attachToDeclImpl(Decl *D) {
+    // Most attributes don't need any custom logic for this.
+  }
 };
 
 #define UNIMPLEMENTED_CLONE(AttrType)    \
@@ -2292,6 +2301,8 @@ class CustomAttrOwner final {
 
 /// Defines a custom attribute.
 class CustomAttr final : public DeclAttribute {
+  friend class DeclAttribute;
+
   TypeExpr *typeExpr;
   ArgumentList *argList;
   CustomAttrOwner owner;
@@ -2322,14 +2333,17 @@ class CustomAttr final : public DeclAttribute {
 
   /// Retrieve the Decl or DeclContext owner for the attribute.
   CustomAttrOwner getOwner() const { return owner; }
-  void setOwner(CustomAttrOwner newOwner) { owner = newOwner; }
 
   ASTContext &getASTContext() const;
 
   /// Retrieve the NominalTypeDecl the CustomAttr refers to, or \c nullptr if
   /// 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
@@ -2377,10 +2391,11 @@ class CustomAttr final : public DeclAttribute {
   void printCustomAttr(ASTPrinter &Printer, const PrintOptions &Options) const;
 
 private:
+  void attachToDeclImpl(Decl *D);
+
   friend class CustomAttrNominalRequest;
   void resetTypeInformation(TypeExpr *repr);
 
-private:
   friend class CustomAttrTypeRequest;
   void setType(Type ty);
 };
@@ -2517,6 +2532,7 @@ class DifferentiableAttr final
                                     ParsedAutoDiffParameter> {
   friend TrailingObjects;
   friend class DifferentiableAttributeTypeCheckRequest;
+  friend class DeclAttribute;
 
   /// The declaration on which the `@differentiable` attribute is declared.
   /// May not be a valid declaration for `@differentiable` attributes.
@@ -2576,11 +2592,9 @@ class DifferentiableAttr final
 
   Decl *getOriginalDeclaration() const { return OriginalDeclaration; }
 
-  /// Sets the original declaration on which this attribute is declared.
-  /// Should only be used by parsing and deserialization.
-  void setOriginalDeclaration(Decl *originalDeclaration);
-
 private:
+  void attachToDeclImpl(Decl *D);
+
   /// Returns true if the given `@differentiable` attribute has been
   /// type-checked.
   bool hasBeenTypeChecked() const;
@@ -2695,6 +2709,7 @@ class DerivativeAttr final
       private llvm::TrailingObjects<DerivativeAttr, ParsedAutoDiffParameter> {
   friend TrailingObjects;
   friend class DerivativeAttrOriginalDeclRequest;
+  friend class DeclAttribute;
 
   /// The declaration on which the `@derivative` attribute is declared.
   /// May not be a valid declaration for `@derivative` attributes.
@@ -2755,10 +2770,6 @@ class DerivativeAttr final
 
   Decl *getOriginalDeclaration() const { return OriginalDeclaration; }
 
-  /// Sets the original declaration on which this attribute is declared.
-  /// Should only be used by parsing and deserialization.
-  void setOriginalDeclaration(Decl *originalDeclaration);
-
   TypeRepr *getBaseTypeRepr() const { return BaseTypeRepr; }
   DeclNameRefWithLoc getOriginalFunctionName() const {
     return OriginalFunctionName;
@@ -2803,6 +2814,9 @@ class DerivativeAttr final
     // Not properly implemented (very complex and not currently needed)
     return false;
   }
+
+private:
+  void attachToDeclImpl(Decl *D);
 };
 
 /// The `@transpose(of:)` attribute registers a function as a transpose of
@@ -3538,6 +3552,8 @@ class AllowFeatureSuppressionAttr final
 
 /// Defines the @abi attribute.
 class ABIAttr : public DeclAttribute {
+  friend class DeclAttribute;
+
 public:
   ABIAttr(Decl *abiDecl, SourceLoc AtLoc, SourceRange Range, bool Implicit)
       : DeclAttribute(DeclAttrKind::ABI, AtLoc, Range, Implicit),
@@ -3564,6 +3580,9 @@ class ABIAttr : public DeclAttribute {
     // Unsupported: tricky to implement and unneeded.
     return true;
   }
+
+private:
+  void attachToDeclImpl(Decl *D);
 };
 
 /// Defines a @nonexhaustive attribute.

include/swift/AST/Decl.h

@@ -1039,6 +1039,13 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl>, public Swi
     return Attrs;
   }
 
+  /// Adds a given attribute to the Decl. This should be preferred over
+  /// \c getAttrs().add(...) since it also attaches the attribute if necessary.
+  void addAttribute(DeclAttribute *attr) {
+    attr->attachToDecl(this);
+    getAttrs().add(attr);
+  }
+
   /// Returns the attributes that were directly attached to this declaration
   /// as written in source, ie. does not include semantic attributes or ones
   /// generated by macro expansions.
@@ -1054,12 +1061,6 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl>, public Swi
   /// attribute macro expansion.
   DeclAttributes getSemanticAttrs() const;
 
-  /// Register the relationship between \c this and \p attr->abiDecl , assuming
-  /// that \p attr is attached to \c this . This is necessary for
-  /// \c ABIRoleInfo::ABIRoleInfo() to determine that \c attr->abiDecl
-  /// is ABI-only and locate its API counterpart.
-  void recordABIAttr(ABIAttr *attr);
-
   /// Set this declaration's attributes to the specified attribute list,
   /// applying any post-processing logic appropriate for attributes parsed
   /// from source code.
@@ -1104,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,

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.
   ///

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; }

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 *),

lib/AST/ASTContext.cpp

@@ -1416,7 +1416,7 @@ ASTContext::synthesizeInvertibleProtocolDecl(InvertibleProtocolKind ip) const {
   protocol->setImplicit(true);
 
   // @_marker
-  protocol->getAttrs().add(new (*this) MarkerAttr(/*implicit=*/true));
+  protocol->addAttribute(new (*this) MarkerAttr(/*implicit=*/true));
 
   // public
   protocol->setAccess(AccessLevel::Public);

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

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__");