Merge pull request #78030 from DougGregor/custom-attribute-initializer-declcontext

DougGregor · web-flow · commit c1fa7db17167 · 2024-12-07T07:32:28.000-08:00
diff --git a/include/swift/AST/ASTBridging.h b/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:"
diff --git a/include/swift/AST/ASTBridgingWrappers.def b/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)
diff --git a/include/swift/AST/Attr.h b/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;
diff --git a/include/swift/AST/Initializer.h b/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
diff --git a/include/swift/Parse/Parser.h b/include/swift/Parse/Parser.h
@@ -998,8 +998,7 @@ class Parser {
 
   /// Parse an #if ... #endif containing only attributes.
   ParserStatus parseIfConfigAttributes(
-    DeclAttributes &attributes, bool ifConfigsAreDeclAttrs,
-    PatternBindingInitializer *initContext);
+    DeclAttributes &attributes, bool ifConfigsAreDeclAttrs);
 
   /// Parse a #error or #warning diagnostic.
   ParserResult<PoundDiagnosticDecl> parseDeclPoundDiagnostic();
@@ -1022,13 +1021,6 @@ class Parser {
   ParserStatus parseDeclAttributeList(DeclAttributes &Attributes,
                                       bool IfConfigsAreDeclAttrs = false);
 
-  /// Parse the optional attributes before a declaration.
-  ///
-  /// This is the inner loop, which can be called recursively.
-  ParserStatus parseDeclAttributeList(DeclAttributes &Attributes,
-                                      bool IfConfigsAreDeclAttrs,
-                                      PatternBindingInitializer *initContext);
-
   /// Parse the optional attributes before a closure declaration.
   ParserStatus parseClosureDeclAttributeList(DeclAttributes &Attributes);
 
@@ -1162,7 +1154,6 @@ class Parser {
   /// Parse a specific attribute.
   ParserStatus parseDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc,
                                   SourceLoc AtEndLoc,
-                                  PatternBindingInitializer *&initContext,
                                   bool isFromClangAttribute = false);
 
   bool isCustomAttributeArgument();
@@ -1171,13 +1162,7 @@ class Parser {
   /// Parse a custom attribute after the initial '@'.
   ///
   /// \param atLoc The location of the already-parsed '@'.
-  ///
-  /// \param initContext A reference to the initializer context used
-  /// for the set of custom attributes. This should start as nullptr, and
-  /// 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);
+  ParserResult<CustomAttr> parseCustomAttribute(SourceLoc atLoc);
 
   ParserStatus parseNewDeclAttribute(DeclAttributes &Attributes,
                                      SourceLoc AtLoc, DeclAttrKind DK,
@@ -1479,7 +1464,6 @@ class Parser {
 
   ParserStatus parseTypeAttribute(TypeOrCustomAttr &result, SourceLoc AtLoc,
                                   SourceLoc AtEndLoc, ParseTypeReason reason,
-                                  PatternBindingInitializer *&initContext,
                                   bool justChecking = false);
 
   ParserResult<TypeRepr> parseOldStyleProtocolComposition();
diff --git a/lib/AST/ASTDumper.cpp b/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:
diff --git a/lib/AST/ASTMangler.cpp b/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");
     }
diff --git a/lib/AST/Attr.cpp b/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);
diff --git a/lib/AST/Bridging/DeclAttributeBridging.cpp b/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(
diff --git a/lib/AST/Bridging/DeclContextBridging.cpp b/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();
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -2061,8 +2061,17 @@ 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);
+    }
+
+    // Set up the custom attribute contexts for each variable in this pattern.
+    PBD->getPattern(idx)->forEachVariable([&](VarDecl *var) {
+      for (auto custom : var->getAttrs().getAttributes<CustomAttr>()) {
+        if (auto attributeInit = custom->getInitContext())
+          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
diff --git a/lib/AST/DeclContext.cpp b/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;
   }
diff --git a/lib/AST/TypeCheckRequests.cpp b/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;
   }
 }
 
diff --git a/lib/ASTGen/Sources/ASTGen/Bridge.swift b/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
 }
diff --git a/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift b/lib/ASTGen/Sources/ASTGen/DeclAttrs.swift
diff --git a/lib/ASTGen/Sources/ASTGen/Decls.swift b/lib/ASTGen/Sources/ASTGen/Decls.swift
diff --git a/lib/ASTGen/Sources/ASTGen/Exprs.swift b/lib/ASTGen/Sources/ASTGen/Exprs.swift
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
diff --git a/lib/Parse/ParseIfConfig.cpp b/lib/Parse/ParseIfConfig.cpp
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
diff --git a/test/Macros/macro_expand_peers.swift b/test/Macros/macro_expand_peers.swift

Original file line number	Diff line number	Diff line change
`@@ -2541,6 +2541,12 @@ void ASTMangler::appendContext(const DeclContext *ctx,`
`2541`	`2541`	`}`
`2542`	`2542`	`return;`
`2543`	`2543`	`}`
	`2544`	`+`
	`2545`	`+ case InitializerKind::CustomAttribute: {`
	`2546`	`+ BaseEntitySignature nullBase(nullptr);`
	`2547`	`+ appendContext(ctx->getParent(), nullBase, useModuleName);`
	`2548`	`+ return;`
	`2549`	`+ }`
`2544`	`2550`	`}`
`2545`	`2551`	`llvm_unreachable("bad initializer kind");`
`2546`	`2552`	`}`