[AST] Store owning Decl/DeclContext on CustomAttr

hamishknight · hamishknight · commit 2b8a1cccfded · 2025-10-13T13:37:29.000+01:00
Introduce CustomAttrOwner that can store either a Decl for an
attached attribute, or a DeclContext for e.g a type or closure
attribute. Store this on CustomAttr such that we can query it from
the name lookup requests.
diff --git a/include/swift/AST/ASTBridging.h b/include/swift/AST/ASTBridging.h
@@ -852,10 +852,11 @@ BridgedCDeclAttr BridgedCDeclAttr_createParsed(BridgedASTContext cContext,
                                                BridgedStringRef cName,
                                                bool underscored);
 
-SWIFT_NAME(
-    "BridgedCustomAttr.createParsed(_:atLoc:type:initContext:argumentList:)")
+SWIFT_NAME("BridgedCustomAttr.createParsed(atLoc:type:declContext:initContext:"
+           "argumentList:)")
 BridgedCustomAttr BridgedCustomAttr_createParsed(
-    BridgedASTContext cContext, swift::SourceLoc atLoc, BridgedTypeRepr cType,
+    swift::SourceLoc atLoc, BridgedTypeRepr cType,
+    BridgedDeclContext cDeclContext,
     BridgedNullableCustomAttributeInitializer cInitContext,
     BridgedNullableArgumentList cArgumentList);
 
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -68,6 +68,7 @@ class GenericFunctionType;
 class LazyConformanceLoader;
 class LazyMemberLoader;
 class ModuleDecl;
+class NominalTypeDecl;
 class PatternBindingInitializer;
 class TrailingWhereClause;
 class TypeExpr;
@@ -2272,34 +2273,63 @@ class ClangImporterSynthesizedTypeAttr : public DeclAttribute {
   }
 };
 
+/// The owning decl for a given custom attribute, or a DeclContext for a
+/// custom attribute in e.g a closure or inheritance clause.
+class CustomAttrOwner final {
+  llvm::PointerUnion<Decl *, DeclContext *> Owner;
+
+public:
+  CustomAttrOwner() : Owner(nullptr) {}
+  CustomAttrOwner(Decl *D) : Owner(D) {}
+  CustomAttrOwner(DeclContext *DC) : Owner(DC) {}
+
+  /// If the owner is a declaration, returns it, \c nullptr otherwise.
+  Decl *getAsDecl() const { return Owner.dyn_cast<Decl *>(); }
+
+  /// Retrieve the DeclContext for the CustomAttr.
+  DeclContext *getDeclContext() const;
+};
+
 /// Defines a custom attribute.
 class CustomAttr final : public DeclAttribute {
   TypeExpr *typeExpr;
   ArgumentList *argList;
+  CustomAttrOwner owner;
   CustomAttributeInitializer *initContext;
   Expr *semanticInit = nullptr;
 
   mutable unsigned isArgUnsafeBit : 1;
 
   CustomAttr(SourceLoc atLoc, SourceRange range, TypeExpr *type,
-             CustomAttributeInitializer *initContext, ArgumentList *argList,
-             bool implicit);
+             CustomAttrOwner owner, CustomAttributeInitializer *initContext,
+             ArgumentList *argList, bool implicit);
 
 public:
   static CustomAttr *create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
-                            bool implicit = false) {
-    return create(ctx, atLoc, type, /*initContext*/ nullptr,
+                            CustomAttrOwner owner, bool implicit = false) {
+    return create(ctx, atLoc, type, owner, /*initContext*/ nullptr,
                   /*argList*/ nullptr, implicit);
   }
 
   static CustomAttr *create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
+                            CustomAttrOwner owner,
                             CustomAttributeInitializer *initContext,
                             ArgumentList *argList, bool implicit = false);
 
   TypeExpr *getTypeExpr() const { return typeExpr; }
   TypeRepr *getTypeRepr() const;
   Type getType() const;
 
+  /// 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;
+
   /// Destructure an attribute's type repr for a macro reference.
   ///
   /// For a 1-level member type repr whose base and member are both identifier
@@ -2336,7 +2366,8 @@ class CustomAttr final : public DeclAttribute {
   CustomAttr *clone(ASTContext &ctx) const {
     assert(argList == nullptr &&
            "Cannot clone custom attribute with an argument list");
-    return create(ctx, AtLoc, getTypeExpr(), initContext, argList, isImplicit());
+    return create(ctx, AtLoc, getTypeExpr(), owner, initContext, argList,
+                  isImplicit());
   }
 
   bool canClone() const { return argList == nullptr; }
diff --git a/include/swift/AST/NameLookupRequests.h b/include/swift/AST/NameLookupRequests.h
@@ -350,19 +350,18 @@ class TypeDeclsFromWhereClauseRequest :
 
 /// Request the nominal type declaration to which the given custom
 /// attribute refers.
-class CustomAttrNominalRequest :
-    public SimpleRequest<CustomAttrNominalRequest,
-                         NominalTypeDecl *(CustomAttr *, DeclContext *),
-                         RequestFlags::Cached> {
+class CustomAttrNominalRequest
+    : public SimpleRequest<CustomAttrNominalRequest,
+                           NominalTypeDecl *(CustomAttr *),
+                           RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
   // Evaluation.
-  NominalTypeDecl *
-  evaluate(Evaluator &evaluator, CustomAttr *attr, DeclContext *dc) const;
+  NominalTypeDecl *evaluate(Evaluator &evaluator, CustomAttr *attr) const;
 
 public:
   // Caching
diff --git a/include/swift/AST/NameLookupTypeIDZone.def b/include/swift/AST/NameLookupTypeIDZone.def
@@ -19,8 +19,7 @@ SWIFT_REQUEST(NameLookup, AnyObjectLookupRequest,
               QualifiedLookupResult(const DeclContext *, DeclName, NLOptions),
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, CustomAttrNominalRequest,
-              NominalTypeDecl *(CustomAttr *, DeclContext *), Cached,
-              NoLocationInfo)
+              NominalTypeDecl *(CustomAttr *), Cached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, DirectLookupRequest,
               TinyPtrVector<ValueDecl *>(DirectLookupDescriptor), Uncached,
               NoLocationInfo)
diff --git a/lib/AST/Attr.cpp b/lib/AST/Attr.cpp
@@ -3053,16 +3053,26 @@ bool ImplementsAttr::isEquivalent(const ImplementsAttr *other,
           && getProtocol(DC) == other->getProtocol(DC);
 }
 
+DeclContext *CustomAttrOwner::getDeclContext() const {
+  ASSERT(!Owner.isNull());
+  if (auto *D = getAsDecl())
+    return D->getDeclContext();
+
+  return Owner.dyn_cast<DeclContext *>();
+}
+
 CustomAttr::CustomAttr(SourceLoc atLoc, SourceRange range, TypeExpr *type,
+                       CustomAttrOwner owner,
                        CustomAttributeInitializer *initContext,
                        ArgumentList *argList, bool implicit)
     : DeclAttribute(DeclAttrKind::Custom, atLoc, range, implicit),
-      typeExpr(type), argList(argList), initContext(initContext) {
+      typeExpr(type), argList(argList), owner(owner), initContext(initContext) {
   assert(type);
   isArgUnsafeBit = false;
 }
 
 CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
+                               CustomAttrOwner owner,
                                CustomAttributeInitializer *initContext,
                                ArgumentList *argList, bool implicit) {
   assert(type);
@@ -3071,7 +3081,7 @@ CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
     range.End = argList->getEndLoc();
 
   return new (ctx)
-      CustomAttr(atLoc, range, type, initContext, argList, implicit);
+      CustomAttr(atLoc, range, type, owner, initContext, argList, implicit);
 }
 
 std::pair<UnqualifiedIdentTypeRepr *, DeclRefTypeRepr *>
@@ -3091,6 +3101,16 @@ CustomAttr::destructureMacroRef() {
   return {nullptr, nullptr};
 }
 
+ASTContext &CustomAttr::getASTContext() const {
+  return getOwner().getDeclContext()->getASTContext();
+}
+
+NominalTypeDecl *CustomAttr::getNominalDecl() const {
+  auto &eval = getASTContext().evaluator;
+  auto *mutThis = const_cast<CustomAttr *>(this);
+  return evaluateOrDefault(eval, CustomAttrNominalRequest{mutThis}, nullptr);
+}
+
 TypeRepr *CustomAttr::getTypeRepr() const { return typeExpr->getTypeRepr(); }
 Type CustomAttr::getType() const { return typeExpr->getInstanceType(); }
 
diff --git a/lib/AST/Bridging/DeclAttributeBridging.cpp b/lib/AST/Bridging/DeclAttributeBridging.cpp
@@ -197,12 +197,15 @@ BridgedCDeclAttr BridgedCDeclAttr_createParsed(BridgedASTContext cContext,
 }
 
 BridgedCustomAttr BridgedCustomAttr_createParsed(
-    BridgedASTContext cContext, SourceLoc atLoc, BridgedTypeRepr cType,
+    SourceLoc atLoc, BridgedTypeRepr cType, BridgedDeclContext cDeclContext,
     BridgedNullableCustomAttributeInitializer cInitContext,
     BridgedNullableArgumentList cArgumentList) {
-  ASTContext &context = cContext.unbridged();
+  DeclContext *DC = cDeclContext.unbridged();
+  ASTContext &context = DC->getASTContext();
+  // Note we set a DeclContext as the owner, which we'll change to the attached
+  // Decl if necessary when attaching the attributes.
   return CustomAttr::create(
-      context, atLoc, new (context) TypeExpr(cType.unbridged()),
+      context, atLoc, new (context) TypeExpr(cType.unbridged()), /*owner*/ DC,
       cInitContext.unbridged(), cArgumentList.unbridged());
 }
 
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -442,6 +442,8 @@ void Decl::attachParsedAttrs(DeclAttributes attrs) {
     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.
@@ -573,8 +575,7 @@ Type Decl::getResolvedCustomAttrType(CustomAttr *attr) const {
     return ty;
 
   auto dc = getDeclContext();
-  auto *nominal = evaluateOrDefault(
-      getASTContext().evaluator, CustomAttrNominalRequest{attr, dc}, nullptr);
+  auto *nominal = attr->getNominalDecl();
   if (!nominal)
     return Type();
 
@@ -8703,11 +8704,7 @@ VarDecl::getAttachedPropertyWrapperTypeInfo(unsigned i) const {
     if (i >= attrs.size())
       return PropertyWrapperTypeInfo();
 
-    auto attr = attrs[i];
-    auto dc = getDeclContext();
-    ASTContext &ctx = getASTContext();
-    nominal = evaluateOrDefault(
-        ctx.evaluator, CustomAttrNominalRequest{attr, dc}, nullptr);
+    nominal = attrs[i]->getNominalDecl();
   }
 
   if (!nominal)
diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
@@ -3971,9 +3971,11 @@ GenericParamListRequest::evaluate(Evaluator &evaluator, GenericContext *value) c
       parsedGenericParams->getRAngleLoc());
 }
 
-NominalTypeDecl *
-CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
-                                   CustomAttr *attr, DeclContext *dc) const {
+NominalTypeDecl *CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
+                                                    CustomAttr *attr) const {
+  auto owner = attr->getOwner();
+  auto *dc = owner.getDeclContext();
+
   // 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.
diff --git a/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift b/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift
@@ -2285,9 +2285,9 @@ extension ASTGenVisitor {
     }
 
     return .createParsed(
-      self.ctx,
       atLoc: self.generateSourceLoc(node.atSign),
       type: type,
+      declContext: declContext,
       initContext: initContext.asNullable,
       argumentList: argList.asNullable
     )
diff --git a/lib/ClangImporter/ClangImporter.cpp b/lib/ClangImporter/ClangImporter.cpp
@@ -6262,8 +6262,8 @@ void cloneImportedAttributes(ValueDecl *fromDecl, ValueDecl* toDecl) {
     case DeclAttrKind::Custom: {
       CustomAttr *cAttr = cast<CustomAttr>(attr);
       attrs.add(CustomAttr::create(context, SourceLoc(), cAttr->getTypeExpr(),
-                                   cAttr->getInitContext(), cAttr->getArgs(),
-                                   true));
+                                   /*owner*/ toDecl, cAttr->getInitContext(),
+                                   cAttr->getArgs(), /*implicit*/ true));
       break;
     }
     case DeclAttrKind::DiscardableResult: {
diff --git a/lib/ClangImporter/ImportDecl.cpp b/lib/ClangImporter/ImportDecl.cpp
@@ -8885,7 +8885,10 @@ void ClangImporter::Implementation::importNontrivialAttribute(
       SmallVector<DeclAttribute *, 2> attrs(decl->getAttrs().begin(),
                                             decl->getAttrs().end());
       for (auto attr : attrs) {
-        MappedDecl->getAttrs().add(cached ? attr->clone(SwiftContext) : attr);
+        auto *newAttr = cached ? attr->clone(SwiftContext) : attr;
+        if (auto *CA = dyn_cast<CustomAttr>(newAttr))
+          CA->setOwner(MappedDecl);
+        MappedDecl->getAttrs().add(newAttr);
       }
     }
     break;
@@ -8933,7 +8936,8 @@ ClangImporter::Implementation::importSwiftAttrAttributes(Decl *MappedDecl) {
 
         if (Type mainActorType = SwiftContext.getMainActorType()) {
           auto typeExpr = TypeExpr::createImplicit(mainActorType, SwiftContext);
-          auto attr = CustomAttr::create(SwiftContext, SourceLoc(), typeExpr);
+          auto attr = CustomAttr::create(SwiftContext, SourceLoc(), typeExpr,
+                                         /*owner*/ MappedDecl);
           MappedDecl->getAttrs().add(attr);
           seenMainActorAttr = swiftAttr;
         }
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -4215,10 +4215,12 @@ ParserResult<CustomAttr> Parser::parseCustomAttribute(SourceLoc atLoc) {
            "Cannot parse a trailing closure here");
   }
 
-  // Form the attribute.
+  // Form the attribute. We set a DeclContext as the owner, which we'll change
+  // to the attached Decl if necessary when attaching the attributes.
   auto *TE = new (Context) TypeExpr(type.get());
-  auto *customAttr = CustomAttr::create(Context, atLoc, TE, initContext,
-                                        argList);
+  auto *customAttr =
+      CustomAttr::create(Context, atLoc, TE,
+                         /*owner*/ CurDeclContext, initContext, argList);
   if (status.hasCodeCompletion() && CodeCompletionCallbacks) {
     CodeCompletionCallbacks->setCompletingInAttribute(customAttr);
   }
diff --git a/lib/Sema/CSSyntacticElement.cpp b/lib/Sema/CSSyntacticElement.cpp
@@ -2595,7 +2595,8 @@ static void applySolutionToClosurePropertyWrappers(ClosureExpr *closure,
       auto &context = wrappedValueVar->getASTContext();
       auto *typeExpr = TypeExpr::createImplicit(backingType, context);
       auto *attr =
-          CustomAttr::create(context, SourceLoc(), typeExpr, /*implicit=*/true);
+          CustomAttr::create(context, SourceLoc(), typeExpr,
+                             /*owner*/ wrappedValueVar, /*implicit=*/true);
       wrappedValueVar->getAttrs().add(attr);
     }
   }
diff --git a/lib/Sema/CodeSynthesis.cpp b/lib/Sema/CodeSynthesis.cpp
@@ -268,8 +268,8 @@ static ParamDecl *createMemberwiseInitParameter(DeclContext *DC,
   // Attach a result builder attribute if needed.
   if (resultBuilderType) {
     auto typeExpr = TypeExpr::createImplicit(resultBuilderType, ctx);
-    auto attr =
-        CustomAttr::create(ctx, SourceLoc(), typeExpr, /*implicit=*/true);
+    auto attr = CustomAttr::create(ctx, SourceLoc(), typeExpr, /*owner*/ arg,
+                                   /*implicit=*/true);
     arg->getAttrs().add(attr);
   }
 
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -4619,9 +4619,7 @@ 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 *nominal = attr->getNominalDecl();
   if (!nominal) {
     if (attr->isInvalid())
       return;
@@ -8780,19 +8778,10 @@ template <typename ATTR>
 static void forEachCustomAttribute(
     Decl *decl,
     llvm::function_ref<void(CustomAttr *attr, NominalTypeDecl *)> fn) {
-  auto &ctx = decl->getASTContext();
-
   for (auto *attr : decl->getAttrs().getAttributes<CustomAttr>()) {
-    auto *mutableAttr = const_cast<CustomAttr *>(attr);
-
-    auto *nominal = evaluateOrDefault(
-        ctx.evaluator,
-        CustomAttrNominalRequest{mutableAttr, decl->getDeclContext()}, nullptr);
-    if (!nominal)
-      continue;
-
-    if (nominal->getAttrs().hasAttribute<ATTR>())
-      fn(mutableAttr, nominal);
+    auto *nominal = attr->getNominalDecl();
+    if (nominal && nominal->getAttrs().hasAttribute<ATTR>())
+      fn(attr, nominal);
   }
 }
 
diff --git a/lib/Sema/TypeCheckConcurrency.cpp b/lib/Sema/TypeCheckConcurrency.cpp
diff --git a/lib/Sema/TypeCheckConcurrency.h b/lib/Sema/TypeCheckConcurrency.h
diff --git a/lib/Sema/TypeCheckPropertyWrapper.cpp b/lib/Sema/TypeCheckPropertyWrapper.cpp
diff --git a/lib/Sema/TypeCheckRequestFunctions.cpp b/lib/Sema/TypeCheckRequestFunctions.cpp
diff --git a/lib/Sema/TypeCheckType.cpp b/lib/Sema/TypeCheckType.cpp
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp

Original file line number	Diff line number	Diff line change
`@@ -2285,9 +2285,9 @@ extension ASTGenVisitor {`
`2285`	`2285`	`}`
`2286`	`2286`
`2287`	`2287`	`return .createParsed(`
`2288`		`- self.ctx,`
`2289`	`2288`	`atLoc: self.generateSourceLoc(node.atSign),`
`2290`	`2289`	`type: type,`
	`2290`	`+ declContext: declContext,`
`2291`	`2291`	`initContext: initContext.asNullable,`
`2292`	`2292`	`argumentList: argList.asNullable`
`2293`	`2293`	`)`