[AST] Consolidate attribute attachment logic

Introduce `DeclAttribute::attachToDecl` which is the now the main
entry point for associating an attribute with a decl. Different
attributes can implement `attachToDeclImpl` to add their custom logic.
Move DifferentiableAttr, DerivativeAttr, CustomAttr, and ABIAttr over
to this new logic.

Author: hamishknight

include/swift/AST/Attr.h

@@ -619,6 +619,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 +632,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 +2300,8 @@ class CustomAttrOwner final {
 
 /// Defines a custom attribute.
 class CustomAttr final : public DeclAttribute {
+  friend class DeclAttribute;
+
   TypeExpr *typeExpr;
   ArgumentList *argList;
   CustomAttrOwner owner;
@@ -2322,7 +2332,6 @@ class CustomAttr final : public DeclAttribute {
 
   /// Retrieve the Decl or DeclContext owner for the attribute.
   CustomAttrOwner getOwner() const { return owner; }
-  void setOwner(CustomAttrOwner newOwner);
 
   ASTContext &getASTContext() const;
 
@@ -2377,10 +2386,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 +2527,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 +2587,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 +2704,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 +2765,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 +2809,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 +3547,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 +3575,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

@@ -1054,12 +1054,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.

lib/AST/Attr.cpp

@@ -349,6 +349,16 @@ bool DeclAttribute::canAttributeAppearOnDeclKind(DeclAttrKind DAK, DeclKind DK)
   llvm_unreachable("bad DeclKind");
 }
 
+void DeclAttribute::attachToDecl(Decl *D) {
+  ASSERT(D);
+  switch (getKind()) {
+#define DECL_ATTR(_, CLASS, ...)                                               \
+  case DeclAttrKind::CLASS:                                                    \
+    return static_cast<CLASS##Attr *>(this)->attachToDeclImpl(D);
+#include "swift/AST/DeclAttr.def"
+  }
+}
+
 // Ensure that every DeclAttribute subclass implements its own CloneAttr.
 static void checkDeclAttributeClones() {
 #define DECL_ATTR(_,CLASS,...) \
@@ -2460,6 +2470,35 @@ OriginallyDefinedInAttr *OriginallyDefinedInAttr::clone(ASTContext &C,
       ManglingModuleName, LinkerModuleName, Platform, MovedVersion, implicit);
 }
 
+void ABIAttr::attachToDeclImpl(Decl *D) {
+  // Register the relationship between `D` and `abiDecl`. This is necessary for
+  // `ABIRoleInfo::ABIRoleInfo()` to determine that `abiDecl` is ABI-only and
+  // locate its API counterpart.
+  Decl *owner = D;
+
+  // The ABIAttr on a VarDecl ought to point to its PBD.
+  if (auto VD = dyn_cast<VarDecl>(owner)) {
+    if (auto PBD = VD->getParentPatternBinding())
+      owner = PBD;
+  }
+
+  auto record = [&](Decl *decl) {
+    auto &evaluator = owner->getASTContext().evaluator;
+    DeclABIRoleInfoRequest(decl).recordABIOnly(evaluator, owner);
+  };
+
+  if (auto abiPBD = dyn_cast<PatternBindingDecl>(abiDecl)) {
+    // Add to *every* VarDecl in the ABI PBD, even ones that don't properly
+    // match anything in the API PBD.
+    for (auto i : range(abiPBD->getNumPatternEntries())) {
+      abiPBD->getPattern(i)->forEachVariable(record);
+    }
+    return;
+  }
+
+  record(abiDecl);
+}
+
 bool AvailableAttr::isUnconditionallyUnavailable() const {
   switch (getKind()) {
   case Kind::Default:
@@ -2820,9 +2859,10 @@ DifferentiableAttr::create(AbstractFunctionDecl *original, bool implicit,
                                       derivativeGenSig);
 }
 
-void DifferentiableAttr::setOriginalDeclaration(Decl *originalDeclaration) {
-  assert(originalDeclaration && "Original declaration must be non-null");
-  assert(!OriginalDeclaration &&
+void DifferentiableAttr::attachToDeclImpl(Decl *originalDeclaration) {
+  // FIXME: This doesn't properly handle PatternBindingDecls with multiple
+  // VarDecls.
+  assert(!OriginalDeclaration || OriginalDeclaration == originalDeclaration &&
          "Original declaration cannot have already been set");
   OriginalDeclaration = originalDeclaration;
 }
@@ -2930,9 +2970,10 @@ void DerivativeAttr::setOriginalFunctionResolver(
   ResolverContextData = resolverContextData;
 }
 
-void DerivativeAttr::setOriginalDeclaration(Decl *originalDeclaration) {
-  assert(originalDeclaration && "Original declaration must be non-null");
-  assert(!OriginalDeclaration &&
+void DerivativeAttr::attachToDeclImpl(Decl *originalDeclaration) {
+  // FIXME: This doesn't properly handle PatternBindingDecls with multiple
+  // VarDecls.
+  assert(!OriginalDeclaration || OriginalDeclaration == originalDeclaration &&
          "Original declaration cannot have already been set");
   OriginalDeclaration = originalDeclaration;
 }
@@ -3084,14 +3125,14 @@ CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
       CustomAttr(atLoc, range, type, owner, initContext, argList, implicit);
 }
 
-void CustomAttr::setOwner(CustomAttrOwner newOwner) {
+void CustomAttr::attachToDeclImpl(Decl *D) {
   // Prefer to set the parent PatternBindingDecl as the owner for a VarDecl,
   // this ensures we can handle PBDs with multiple vars.
-  if (auto *VD = dyn_cast_or_null<VarDecl>(newOwner.getAsDecl())) {
+  if (auto *VD = dyn_cast<VarDecl>(D)) {
     if (auto *PBD = VD->getParentPatternBinding())
-      newOwner = PBD;
+      D = PBD;
   }
-  owner = newOwner;
+  owner = D;
 }
 
 std::pair<UnqualifiedIdentTypeRepr *, DeclRefTypeRepr *>

lib/AST/Decl.cpp

@@ -436,15 +436,6 @@ DeclAttributes Decl::getSemanticAttrs() const {
 void Decl::attachParsedAttrs(DeclAttributes attrs) {
   ASSERT(getAttrs().isEmpty() && "attaching when there are already attrs?");
 
-  for (auto *attr : attrs.getAttributes<DifferentiableAttr>())
-    attr->setOriginalDeclaration(this);
-  for (auto *attr : attrs.getAttributes<DerivativeAttr>())
-    attr->setOriginalDeclaration(this);
-  for (auto attr : attrs.getAttributes<ABIAttr, /*AllowInvalid=*/true>())
-    recordABIAttr(attr);
-  for (auto *attr : attrs.getAttributes<CustomAttr>())
-    attr->setOwner(this);
-
   // @implementation requires an explicit @objc attribute, but
   // @_objcImplementation didn't. Insert one if necessary.
   auto implAttr = attrs.getAttribute<ObjCImplementationAttr>();
@@ -458,6 +449,9 @@ void Decl::attachParsedAttrs(DeclAttributes attrs) {
     attrs.add(objcAttr);
   }
 
+  for (auto attr : attrs)
+    attr->attachToDecl(this);
+
   getAttrs() = attrs;
 }
 
@@ -4740,33 +4734,6 @@ void ValueDecl::setRenamedDecl(const AvailableAttr *attr,
                                                 std::move(renameDecl));
 }
 
-void Decl::recordABIAttr(ABIAttr *attr) {
-  Decl *owner = this;
-
-  // The ABIAttr on a VarDecl ought to point to its PBD.
-  if (auto VD = dyn_cast<VarDecl>(owner)) {
-    if (auto PBD = VD->getParentPatternBinding()) {
-      owner = PBD;
-    }
-  }
-
-  auto record = [&](Decl *decl) {
-    auto &evaluator = owner->getASTContext().evaluator;
-    DeclABIRoleInfoRequest(decl).recordABIOnly(evaluator, owner);
-  };
-
-  if (auto abiPBD = dyn_cast<PatternBindingDecl>(attr->abiDecl)) {
-    // Add to *every* VarDecl in the ABI PBD, even ones that don't properly
-    // match anything in the API PBD.
-    for (auto i : range(abiPBD->getNumPatternEntries())) {
-      abiPBD->getPattern(i)->forEachVariable(record);
-    }
-    return;
-  }
-
-  record(attr->abiDecl);
-}
-
 void DeclABIRoleInfoRequest::recordABIOnly(Evaluator &evaluator,
                                            Decl *counterpart) {
   if (evaluator.hasCachedResult(*this))

lib/ClangImporter/ImportDecl.cpp

@@ -8926,7 +8926,7 @@ void ClangImporter::Implementation::importNontrivialAttribute(
       for (auto attr : attrs) {
         auto *newAttr = cached ? attr->clone(SwiftContext) : attr;
         if (auto *CA = dyn_cast<CustomAttr>(newAttr))
-          CA->setOwner(MappedDecl);
+          CA->attachToDecl(MappedDecl);
         MappedDecl->getAttrs().add(newAttr);
       }
     }

lib/Serialization/Deserialization.cpp

@@ -3449,7 +3449,7 @@ class DeclDeserializer {
       cast<ExtensionDecl *>(decl)->setInherited(inherited);
   }
 
-  llvm::Error finishRecursiveAttrs(Decl *decl, DeclAttribute *attrs);
+  llvm::Error finishRecursiveAttrs();
 
 public:
   DeclDeserializer(ModuleFile &MF, Serialized<Decl *> &declOrOffset)
@@ -3476,9 +3476,13 @@ class DeclDeserializer {
 
     if (DAttrs) {
       decl->getAttrs().setRawAttributeChain(DAttrs);
-      // Wire up the correct owner for any CustomAttrs we deserialized.
-      for (auto *CA : decl->getAttrs().getAttributes<CustomAttr>())
-        CA->setOwner(decl);
+      // Wire up the attached decl.
+      for (auto *attr : decl->getAttrs())
+        attr->attachToDecl(decl);
+
+      // Wire up the indices for @differentiable.
+      for (auto *attr : decl->getAttrs().getAttributes<DifferentiableAttr>())
+        attr->setParameterIndices(diffAttrParamIndicesMap[attr]);
     }
 
     if (auto value = dyn_cast<ValueDecl>(decl)) {
@@ -6774,41 +6778,24 @@ llvm::Error DeclDeserializer::deserializeDeclCommon() {
 }
 
 /// Complete attributes that contain recursive references to the decl being
-/// deserialized or to other decls. This method is called after \p decl is
+/// deserialized or to other decls. This method is called after the decl is
 /// created and stored into the \c ModuleFile::Decls table, so any cycles
 /// between mutually-referencing decls will be broken.
 ///
 /// Attributes handled here include:
 ///
-///  \li \c \@differentiable
-///  \li \c \@derivative
 ///  \li \c \@abi
-llvm::Error DeclDeserializer::finishRecursiveAttrs(Decl *decl, DeclAttribute *attrs) {
-  DeclAttributes tempAttrs;
-  tempAttrs.setRawAttributeChain(attrs);
-
-  // @differentiable and @derivative
-  for (auto *attr : tempAttrs.getAttributes<DifferentiableAttr>()) {
-    auto *diffAttr = const_cast<DifferentiableAttr *>(attr);
-    diffAttr->setOriginalDeclaration(decl);
-    diffAttr->setParameterIndices(diffAttrParamIndicesMap[diffAttr]);
-  }
-  for (auto *attr : tempAttrs.getAttributes<DerivativeAttr>()) {
-    auto *derAttr = const_cast<DerivativeAttr *>(attr);
-    derAttr->setOriginalDeclaration(decl);
-  }
-
+llvm::Error DeclDeserializer::finishRecursiveAttrs() {
   // @abi
   if (unresolvedABIAttr) {
     auto abiDeclOrError = MF.getDeclChecked(unresolvedABIAttr->second);
     if (!abiDeclOrError)
       return abiDeclOrError.takeError();
     unresolvedABIAttr->first->abiDecl = abiDeclOrError.get();
-    decl->recordABIAttr(unresolvedABIAttr->first);
   }
   if (ABIDeclCounterpartID != 0) {
     // This decl is the `abiDecl` of an `ABIAttr`. Force the decl that `ABIAttr`
-    // belongs to so that `recordABIAttr()` will be called.
+    // belongs to so that `attachToDecl()` will be called.
     auto counterpartOrError = MF.getDeclChecked(ABIDeclCounterpartID);
     if (!counterpartOrError)
       return counterpartOrError.takeError();
@@ -6865,7 +6852,7 @@ DeclDeserializer::getDeclCheckedImpl(
     if (!declOrError) \
       return declOrError; \
     declOrOffset = declOrError.get(); \
-    if (auto finishError = finishRecursiveAttrs(declOrError.get(), DAttrs)) \
+    if (auto finishError = finishRecursiveAttrs()) \
       return finishError; \
     break; \
   }