Introduce a new Initializer subclass for the arguments of custom attributes

Since the introduction of custom attributes (as part of property
wrappers), we've modeled the context of expressions within these
attributes as PatternBindingInitializers. These
PatternBindingInitializers would get wired in to the variable
declarations they apply to, establishing the appropriate declaration
context hierarchy. This worked because property wrappers only every
applied to---you guessed it!---properties, so the
PatternBindingInitializer would always get filled in.

When custom attributes were extended to apply to anything for the
purposes of macros, the use of PatternBindingInitializer became less
appropriate. Specifically, the binding declaration would never get
filled in (it's always NULL), so any place in the compiler that
accesses the binding might have to deal with it being NULL, which is a
new requirement. Few did, crashes ensued.

Rather than continue to play whack-a-mole with the abused
PatternBindingInitializer, introduce a new CustomAttributeInitializer
to model the context of custom attribute arguments. When the
attributes are assigned to a declaration that has a
PatternBindingInitializer, we reparent this new initializer to the
PatternBindingInitializer. This helps separate out the logic for
custom attributes vs. actual initializers.

Fixes #76409 / rdar://136997841

Author: DougGregor

include/swift/AST/ASTBridging.h

@@ -536,6 +536,14 @@ SWIFT_NAME("getter:BridgedPatternBindingInitializer.asDeclContext(self:)")
 BridgedDeclContext BridgedPatternBindingInitializer_asDeclContext(
     BridgedPatternBindingInitializer cInit);
 
+SWIFT_NAME("BridgedCustomAttributeInitializer.create(declContext:)")
+BridgedCustomAttributeInitializer
+BridgedCustomAttributeInitializer_create(BridgedDeclContext cDeclContext);
+
+SWIFT_NAME("getter:BridgedCustomAttributeInitializer.asDeclContext(self:)")
+BridgedDeclContext BridgedCustomAttributeInitializer_asDeclContext(
+    BridgedCustomAttributeInitializer cInit);
+
 SWIFT_NAME("getter:BridgedClosureExpr.asDeclContext(self:)")
 BridgedDeclContext
 BridgedClosureExpr_asDeclContext(BridgedClosureExpr cClosure);
@@ -613,7 +621,7 @@ SWIFT_NAME(
     "BridgedCustomAttr.createParsed(_:atLoc:type:initContext:argumentList:)")
 BridgedCustomAttr BridgedCustomAttr_createParsed(
     BridgedASTContext cContext, BridgedSourceLoc cAtLoc, BridgedTypeRepr cType,
-    BridgedNullablePatternBindingInitializer cInitContext,
+    BridgedNullableCustomAttributeInitializer cInitContext,
     BridgedNullableArgumentList cArgumentList);
 
 SWIFT_NAME("BridgedDocumentationAttr.createParsed(_:atLoc:range:metadata:"

include/swift/AST/ASTBridgingWrappers.def

@@ -83,6 +83,7 @@ AST_BRIDGING_WRAPPER_NULLABLE(GenericParamList)
 AST_BRIDGING_WRAPPER_NULLABLE(TrailingWhereClause)
 AST_BRIDGING_WRAPPER_NULLABLE(ParameterList)
 AST_BRIDGING_WRAPPER_NULLABLE(PatternBindingInitializer)
+AST_BRIDGING_WRAPPER_NULLABLE(CustomAttributeInitializer)
 AST_BRIDGING_WRAPPER_NONNULL(TypeAttributes)
 AST_BRIDGING_WRAPPER_NONNULL(CustomAttribute)
 AST_BRIDGING_WRAPPER_NULLABLE(ArgumentList)

include/swift/AST/Attr.h

@@ -60,6 +60,7 @@ class AbstractFunctionDecl;
 class FuncDecl;
 class ClassDecl;
 class AccessorDecl;
+class CustomAttributeInitializer;
 class GenericFunctionType;
 class LazyConformanceLoader;
 class LazyMemberLoader;
@@ -1861,13 +1862,13 @@ class ClangImporterSynthesizedTypeAttr : public DeclAttribute {
 class CustomAttr final : public DeclAttribute {
   TypeExpr *typeExpr;
   ArgumentList *argList;
-  PatternBindingInitializer *initContext;
+  CustomAttributeInitializer *initContext;
   Expr *semanticInit = nullptr;
 
   mutable unsigned isArgUnsafeBit : 1;
 
   CustomAttr(SourceLoc atLoc, SourceRange range, TypeExpr *type,
-             PatternBindingInitializer *initContext, ArgumentList *argList,
+             CustomAttributeInitializer *initContext, ArgumentList *argList,
              bool implicit);
 
 public:
@@ -1878,7 +1879,7 @@ class CustomAttr final : public DeclAttribute {
   }
 
   static CustomAttr *create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
-                            PatternBindingInitializer *initContext,
+                            CustomAttributeInitializer *initContext,
                             ArgumentList *argList, bool implicit = false);
 
   TypeExpr *getTypeExpr() const { return typeExpr; }
@@ -1911,7 +1912,7 @@ class CustomAttr final : public DeclAttribute {
   Expr *getSemanticInit() const { return semanticInit; }
   void setSemanticInit(Expr *expr) { semanticInit = expr; }
 
-  PatternBindingInitializer *getInitContext() const { return initContext; }
+  CustomAttributeInitializer *getInitContext() const { return initContext; }
 
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DeclAttrKind::Custom;
@@ -3041,6 +3042,10 @@ class DeclAttributes {
          const_cast<const DeclAttributes *>(this)->getEffectiveSendableAttr());
   }
 
+  /// Try to find the custom attribute initializer within the list of
+  /// attributes.
+  CustomAttributeInitializer *findCustomAttributeInitializer() const;
+
 public:
   template <typename ATTR, bool AllowInvalid>
   using AttributeKindRange =

include/swift/AST/Initializer.h

@@ -22,6 +22,10 @@
 
 #include "swift/AST/DeclContext.h"
 
+namespace llvm {
+class raw_ostream;
+}
+
 namespace swift {
 class ParamDecl;
 class PatternBindingDecl;
@@ -36,6 +40,9 @@ enum class InitializerKind : uint8_t {
 
   /// A property wrapper initialization expression.
   PropertyWrapper,
+
+  /// An expression within a custom attribute.
+  CustomAttribute,
 };
 
 /// An Initializer is a kind of DeclContext used for expressions that
@@ -176,6 +183,34 @@ class PropertyWrapperInitializer : public Initializer {
   }
 };
 
+/// An expression within a custom attribute. The parent context is the
+/// context in which the attributed declaration occurs.
+class CustomAttributeInitializer : public Initializer {
+public:
+  explicit CustomAttributeInitializer(DeclContext *parent)
+      : Initializer(InitializerKind::CustomAttribute, parent) {}
+
+  static CustomAttributeInitializer *create(DeclContext *parent) {
+    return new (parent->getASTContext()) CustomAttributeInitializer(parent);
+  }
+
+  void setEnclosingInitializer(Initializer *newParent) {
+    setParent(newParent);
+  }
+
+  static bool classof(const DeclContext *DC) {
+    if (auto init = dyn_cast<Initializer>(DC))
+      return classof(init);
+    return false;
+  }
+
+  static bool classof(const Initializer *I) {
+    return I->getInitializerKind() == InitializerKind::CustomAttribute;
+  }
+};
+
+void simple_display(llvm::raw_ostream &out, Initializer *init);
+
 } // end namespace swift
 
 #endif

include/swift/Parse/Parser.h

@@ -999,7 +999,7 @@ class Parser {
   /// Parse an #if ... #endif containing only attributes.
   ParserStatus parseIfConfigAttributes(
     DeclAttributes &attributes, bool ifConfigsAreDeclAttrs,
-    PatternBindingInitializer *initContext);
+    CustomAttributeInitializer *&initContext);
 
   /// Parse a #error or #warning diagnostic.
   ParserResult<PoundDiagnosticDecl> parseDeclPoundDiagnostic();
@@ -1027,7 +1027,7 @@ class Parser {
   /// This is the inner loop, which can be called recursively.
   ParserStatus parseDeclAttributeList(DeclAttributes &Attributes,
                                       bool IfConfigsAreDeclAttrs,
-                                      PatternBindingInitializer *initContext);
+                                      CustomAttributeInitializer *&initContext);
 
   /// Parse the optional attributes before a closure declaration.
   ParserStatus parseClosureDeclAttributeList(DeclAttributes &Attributes);
@@ -1162,7 +1162,7 @@ class Parser {
   /// Parse a specific attribute.
   ParserStatus parseDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc,
                                   SourceLoc AtEndLoc,
-                                  PatternBindingInitializer *&initContext,
+                                  CustomAttributeInitializer *&initContext,
                                   bool isFromClangAttribute = false);
 
   bool isCustomAttributeArgument();
@@ -1177,7 +1177,7 @@ class Parser {
   /// will get filled in by this function. The same variable should be provided
   /// for every custom attribute within the same attribute list.
   ParserResult<CustomAttr> parseCustomAttribute(
-      SourceLoc atLoc, PatternBindingInitializer *&initContext);
+      SourceLoc atLoc, CustomAttributeInitializer *&initContext);
 
   ParserStatus parseNewDeclAttribute(DeclAttributes &Attributes,
                                      SourceLoc AtLoc, DeclAttrKind DK,
@@ -1479,7 +1479,7 @@ class Parser {
 
   ParserStatus parseTypeAttribute(TypeOrCustomAttr &result, SourceLoc AtLoc,
                                   SourceLoc AtEndLoc, ParseTypeReason reason,
-                                  PatternBindingInitializer *&initContext,
+                                  CustomAttributeInitializer *&initContext,
                                   bool justChecking = false);
 
   ParserResult<TypeRepr> parseOldStyleProtocolComposition();

lib/AST/ASTDumper.cpp

@@ -1911,17 +1911,7 @@ void swift::printContext(raw_ostream &os, DeclContext *dc) {
     break;
 
   case DeclContextKind::Initializer:
-    switch (cast<Initializer>(dc)->getInitializerKind()) {
-    case InitializerKind::PatternBinding:
-      os << "pattern binding initializer";
-      break;
-    case InitializerKind::DefaultArgument:
-      os << "default argument initializer";
-      break;
-    case InitializerKind::PropertyWrapper:
-      os << "property wrapper initializer";
-      break;
-    }
+    simple_display(os, cast<Initializer>(dc));
     break;
 
   case DeclContextKind::TopLevelCodeDecl:

lib/AST/ASTMangler.cpp

@@ -2541,6 +2541,12 @@ void ASTMangler::appendContext(const DeclContext *ctx,
       }
       return;
     }
+
+    case InitializerKind::CustomAttribute: {
+      BaseEntitySignature nullBase(nullptr);
+      appendContext(ctx->getParent(), nullBase, useModuleName);
+      return;
+    }
     }
     llvm_unreachable("bad initializer kind");
     }

lib/AST/Attr.cpp

@@ -2632,7 +2632,7 @@ ProtocolDecl *ImplementsAttr::getProtocol(DeclContext *dc) const {
 }
 
 CustomAttr::CustomAttr(SourceLoc atLoc, SourceRange range, TypeExpr *type,
-                       PatternBindingInitializer *initContext,
+                       CustomAttributeInitializer *initContext,
                        ArgumentList *argList, bool implicit)
     : DeclAttribute(DeclAttrKind::Custom, atLoc, range, implicit),
       typeExpr(type), argList(argList), initContext(initContext) {
@@ -2641,7 +2641,7 @@ CustomAttr::CustomAttr(SourceLoc atLoc, SourceRange range, TypeExpr *type,
 }
 
 CustomAttr *CustomAttr::create(ASTContext &ctx, SourceLoc atLoc, TypeExpr *type,
-                               PatternBindingInitializer *initContext,
+                               CustomAttributeInitializer *initContext,
                                ArgumentList *argList, bool implicit) {
   assert(type);
   SourceRange range(atLoc, type->getSourceRange().End);
@@ -2881,3 +2881,13 @@ bool swift::hasAttribute(const LangOptions &langOpts,
 
   return false;
 }
+
+CustomAttributeInitializer *
+DeclAttributes::findCustomAttributeInitializer() const {
+  for (auto custom : getAttributes<CustomAttr>()) {
+    if (auto initContext = custom->getInitContext())
+      return initContext;
+  }
+
+  return nullptr;
+}

lib/AST/Bridging/DeclAttributeBridging.cpp

@@ -134,7 +134,7 @@ BridgedCDeclAttr BridgedCDeclAttr_createParsed(BridgedASTContext cContext,
 
 BridgedCustomAttr BridgedCustomAttr_createParsed(
     BridgedASTContext cContext, BridgedSourceLoc cAtLoc, BridgedTypeRepr cType,
-    BridgedNullablePatternBindingInitializer cInitContext,
+    BridgedNullableCustomAttributeInitializer cInitContext,
     BridgedNullableArgumentList cArgumentList) {
   ASTContext &context = cContext.unbridged();
   return CustomAttr::create(

lib/AST/Bridging/DeclContextBridging.cpp

@@ -31,6 +31,16 @@ BridgedDeclContext BridgedPatternBindingInitializer_asDeclContext(
   return cInit.unbridged();
 }
 
+BridgedCustomAttributeInitializer
+BridgedCustomAttributeInitializer_create(BridgedDeclContext cDeclContext) {
+  return CustomAttributeInitializer::create(cDeclContext.unbridged());
+}
+
+BridgedDeclContext BridgedCustomAttributeInitializer_asDeclContext(
+    BridgedCustomAttributeInitializer cInit) {
+  return cInit.unbridged();
+}
+
 BridgedDeclContext
 BridgedClosureExpr_asDeclContext(BridgedClosureExpr cClosure) {
   return cClosure.unbridged();

lib/AST/Decl.cpp

@@ -2061,9 +2061,18 @@ PatternBindingDecl::create(ASTContext &Ctx, SourceLoc StaticLoc,
     // FIXME: We ought to reconsider this since it won't recontextualize any
     // closures/decls present in the initialization expr. This currently should
     // only affect implicit code though.
-    if (!initContext && !Parent->isLocalContext())
+    if (!initContext && !Parent->isLocalContext()) {
       initContext = PatternBindingInitializer::create(Parent);
 
+    }
+
+    if (auto firstVar = PBD->getAnchoringVarDecl(idx)) {
+      if (auto attributeInit =
+              firstVar->getAttrs().findCustomAttributeInitializer()) {
+        attributeInit->setEnclosingInitializer(initContext);
+      }
+    }
+
     // We need to call setPattern to ensure the VarDecls in the pattern have
     // the PatternBindingDecl set as their parent. We also need to call
     // setInitContext to setup the context.

lib/AST/DeclContext.cpp

@@ -869,6 +869,9 @@ unsigned DeclContext::printContext(raw_ostream &OS, const unsigned indent,
       }
       break;
     }
+    case InitializerKind::CustomAttribute:
+      OS << " CustomAttribute";
+      break;
     }
     break;
   }

lib/AST/TypeCheckRequests.cpp

@@ -1470,6 +1470,9 @@ void swift::simple_display(llvm::raw_ostream &out, Initializer *init) {
   case InitializerKind::PropertyWrapper:
     out << "property wrapper initializer";
     break;
+  case InitializerKind::CustomAttribute:
+    out << "custom attribute initializer";
+    break;
   }
 }
 

lib/ASTGen/Sources/ASTGen/Bridge.swift

@@ -34,6 +34,7 @@ extension BridgedNullableGenericParamList: /*@retroactive*/ swiftASTGen.BridgedN
 extension BridgedNullableTrailingWhereClause: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullableParameterList: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullablePatternBindingInitializer: /*@retroactive*/ swiftASTGen.BridgedNullable {}
+extension BridgedNullableCustomAttributeInitializer: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 extension BridgedNullableArgumentList: /*@retroactive*/ swiftASTGen.BridgedNullable {}
 
 extension BridgedIdentifier: /*@retroactive*/ Swift.Equatable {
@@ -80,6 +81,9 @@ extension BridgedParameterList: BridgedHasNullable {
 extension BridgedPatternBindingInitializer: BridgedHasNullable {
   typealias Nullable = BridgedNullablePatternBindingInitializer
 }
+extension BridgedCustomAttributeInitializer: BridgedHasNullable {
+  typealias Nullable = BridgedNullableCustomAttributeInitializer
+}
 extension BridgedArgumentList: BridgedHasNullable {
   typealias Nullable = BridgedNullableArgumentList
 }

lib/ASTGen/Sources/ASTGen/DeclAttrs.swift

@@ -22,7 +22,7 @@ extension ASTGenVisitor {
     var attributes: BridgedDeclAttributes
     var staticSpelling: BridgedStaticSpelling
     var staticLoc: BridgedSourceLoc
-    var initContext: BridgedPatternBindingInitializer?
+    var initContext: BridgedCustomAttributeInitializer?
   }
 
   func generateDeclAttributes(_ node: some WithAttributesSyntax & WithModifiersSyntax, allowStatic: Bool) -> DeclAttributesResult {
@@ -56,7 +56,7 @@ extension ASTGenVisitor {
 
     // '@' attributes.
     // FIXME: Factor out this and share it with 'generate(accessorDecl:)'.
-    var initContext: BridgedPatternBindingInitializer? = nil
+    var initContext: BridgedCustomAttributeInitializer? = nil
     visitIfConfigElements(node.attributes, of: AttributeSyntax.self) { element in
       switch element {
       case .ifConfigDecl(let ifConfigDecl):
@@ -100,7 +100,7 @@ extension ASTGenVisitor {
 
 // MARK: - Decl attributes
 extension ASTGenVisitor {
-  func generateDeclAttribute(attribute node: AttributeSyntax, initContext: inout BridgedPatternBindingInitializer?) -> BridgedDeclAttribute? {
+  func generateDeclAttribute(attribute node: AttributeSyntax, initContext: inout BridgedCustomAttributeInitializer?) -> BridgedDeclAttribute? {
     if let identTy = node.attributeName.as(IdentifierTypeSyntax.self) {
       let attrName = identTy.name.rawText
       let attrKind = BridgedDeclAttrKind(from: attrName.bridged)
@@ -1513,7 +1513,7 @@ extension ASTGenVisitor {
     )
   }
 
-  func generateCustomAttr(attribute node: AttributeSyntax, initContext: inout BridgedPatternBindingInitializer?) -> BridgedCustomAttr? {
+  func generateCustomAttr(attribute node: AttributeSyntax, initContext: inout BridgedCustomAttributeInitializer?) -> BridgedCustomAttr? {
     let type = self.generate(type: node.attributeName)
 
     let argList: BridgedArgumentList?
@@ -1524,7 +1524,7 @@ extension ASTGenVisitor {
       }
 
       if !self.declContext.isLocalContext && initContext == nil {
-        initContext = BridgedPatternBindingInitializer.create(declContext: self.declContext)
+        initContext = BridgedCustomAttributeInitializer.create(declContext: self.declContext)
       }
       argList = withDeclContext(initContext?.asDeclContext ?? self.declContext) {
         self.generateArgumentList(

lib/ASTGen/Sources/ASTGen/Decls.swift

@@ -487,7 +487,7 @@ extension ASTGenVisitor {
       // ensures that property initializers are correctly treated as being in a
       // local context).
       if !self.declContext.isLocalContext {
-        initContext = attrs.initContext ?? .create(declContext: self.declContext)
+        initContext = .create(declContext: self.declContext)
       }
       initExpr = withDeclContext(initContext?.asDeclContext ?? self.declContext) {
         generate(expr: initializer.value)

lib/ASTGen/Sources/ASTGen/Exprs.swift

@@ -232,7 +232,7 @@ extension ASTGenVisitor {
         return .underlying(attribute)
       }
     } body: { node in
-      var initCtx: BridgedPatternBindingInitializer?
+      var initCtx: BridgedCustomAttributeInitializer?
       if let attr = self.generateDeclAttribute(attribute: node, initContext: &initCtx) {
         result.attributes.add(attr)
       }