Merge pull request #63076 from hborla/macro-use-after-free

Author: hborla

include/swift/AST/Attr.h

@@ -2613,55 +2613,6 @@ class DeclAttributes {
   SourceLoc getStartLoc(bool forModifiers = false) const;
 };
 
-/// Semantic attributes that are applied to a declaration.
-///
-/// This attribute list can include attributes that are not written in
-/// source on a declaration, such as attributes that are applied during
-/// macro expansion or inferred from the declaration context.
-class SemanticDeclAttributes {
-  llvm::SmallVector<DeclAttribute *, 4> attrList;
-
-public:
-  SemanticDeclAttributes() {}
-
-  /// Add a constructed DeclAttribute to this list.
-  void add(DeclAttribute *attr) {
-    attrList.push_back(attr);
-  }
-
-  using SemanticAttrList = llvm::SmallVectorImpl<DeclAttribute *>;
-
-  template <typename AttrType, bool AllowInvalid>
-  using AttributeKindRange =
-      OptionalTransformRange<iterator_range<SemanticAttrList::const_iterator>,
-                             ToAttributeKind<AttrType, AllowInvalid>,
-                             SemanticAttrList::const_iterator>;
-
-  template <typename AttrType, bool AllowInvalid = false>
-  AttributeKindRange<AttrType, AllowInvalid> getAttributes() const {
-    return AttributeKindRange<AttrType, AllowInvalid>(
-        make_range(attrList.begin(), attrList.end()),
-                   ToAttributeKind<AttrType, AllowInvalid>());
-  }
-
-  /// Retrieve the first attribute of the given attribute class.
-  template <typename AttrType>
-  const AttrType *getAttribute(bool allowInvalid = false) const {
-    for (auto attr : attrList)
-      if (auto *specificAttr = dyn_cast<AttrType>(attr))
-        if (specificAttr->isValid() || allowInvalid)
-          return specificAttr;
-
-    return nullptr;
-  }
-
-  /// Determine whether there is an attribute with the given attribute class.
-  template <typename AttrType>
-  bool hasAttribute(bool allowInvalid = false) const {
-    return getAttribute<AttrType>(allowInvalid) != nullptr;
-  }
-};
-
 /// TypeAttributes - These are attributes that may be applied to types.
 class TypeAttributes {
   // Get a SourceLoc for every possible attribute that can be parsed in source.

include/swift/AST/Decl.h

@@ -850,11 +850,10 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     return Attrs;
   }
 
-  /// Returns the semantic attributes attached to this declaration.
-  ///
-  /// \c getSemanticAttrs is intended to be a requestified replacement
-  /// for \c getAttrs
-  SemanticDeclAttributes getSemanticAttrs() const;
+  /// Returns the semantic attributes attached to this declaration,
+  /// including attributes that are generated as the result of member
+  /// attribute macro expansion.
+  DeclAttributes getSemanticAttrs() const;
 
   /// Returns the innermost enclosing decl with an availability annotation.
   const Decl *getInnermostDeclWithAvailability() const;

include/swift/AST/SourceFile.h

@@ -500,6 +500,17 @@ class SourceFile final : public FileUnit {
   /// the \c SourceFileKind is \c MacroExpansion.
   ASTNode getMacroExpansion() const;
 
+  /// For source files created to hold the source code created by expanding
+  /// an attached macro, this is the custom attribute that describes the macro
+  /// expansion.
+  ///
+  /// The source location of this attribute is the place in the source that
+  /// triggered the creation of the macro expansion whose resulting source
+  /// code is in this source file. This will only produce a non-null value when
+  /// the \c SourceFileKind is \c MacroExpansion , and the macro is an attached
+  /// macro.
+  CustomAttr *getAttachedMacroAttribute() const;
+
   /// When this source file is enclosed within another source file, for example
   /// because it describes a macro expansion, return the source file it was
   /// enclosed in.

include/swift/AST/TypeCheckRequests.h

@@ -742,26 +742,6 @@ class AttachedPropertyWrappersRequest :
   bool isCached() const { return true; }
 };
 
-/// Request the semantic attributes attached to the given declaration.
-class AttachedSemanticAttrsRequest :
-    public SimpleRequest<AttachedSemanticAttrsRequest,
-                         SemanticDeclAttributes(Decl *),
-                         RequestFlags::Cached> {
-public:
-  using SimpleRequest::SimpleRequest;
-
-private:
-  friend SimpleRequest;
-
-  // Evaluation.
-  SemanticDeclAttributes
-  evaluate(Evaluator &evaluator, Decl *decl) const;
-
-public:
-  // Caching
-  bool isCached() const { return true; }
-};
-
 /// Request the raw (possibly unbound generic) type of the property wrapper
 /// that is attached to the given variable.
 class AttachedPropertyWrapperTypeRequest :
@@ -3840,6 +3820,24 @@ class ExpandMacroExpansionDeclRequest
   bool isCached() const { return true; }
 };
 
+/// Expand all member attribute macros attached to the given
+/// declaration.
+class ExpandMemberAttributeMacros
+    : public SimpleRequest<ExpandMemberAttributeMacros,
+                           bool(Decl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  bool evaluate(Evaluator &evaluator, Decl *decl) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 /// Resolve an external macro given its module and type name.
 class ExternalMacroDefinitionRequest
     : public SimpleRequest<ExternalMacroDefinitionRequest,

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -29,9 +29,6 @@ SWIFT_REQUEST(TypeChecker, AttachedPropertyWrapperTypeRequest,
 SWIFT_REQUEST(TypeChecker, AttachedPropertyWrappersRequest,
               llvm::TinyPtrVector<CustomAttr *>(VarDecl *), Cached,
               NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, AttachedSemanticAttrsRequest,
-              SemanticDeclAttributes(Decl *), Cached,
-              NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, CallerSideDefaultArgExprRequest,
               Expr *(DefaultArgumentExpr *), SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, CheckInconsistentImplementationOnlyImportsRequest,
@@ -455,6 +452,9 @@ SWIFT_REQUEST(TypeChecker, ExternalMacroDefinitionRequest,
 SWIFT_REQUEST(TypeChecker, ExpandMacroExpansionDeclRequest,
               ArrayRef<Decl *>(MacroExpansionDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ExpandMemberAttributeMacros,
+              bool(Decl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SynthesizeRuntimeMetadataAttrGenerator,
               Expr *(CustomAttr *, ValueDecl *),
               Cached, NoLocationInfo)

lib/AST/Decl.cpp

@@ -365,11 +365,13 @@ StringRef Decl::getDescriptiveKindName(DescriptiveDeclKind K) {
   llvm_unreachable("bad DescriptiveDeclKind");
 }
 
-SemanticDeclAttributes Decl::getSemanticAttrs() const {
+DeclAttributes Decl::getSemanticAttrs() const {
   auto mutableThis = const_cast<Decl *>(this);
-  return evaluateOrDefault(getASTContext().evaluator,
-                           AttachedSemanticAttrsRequest{mutableThis},
-                           SemanticDeclAttributes());
+  evaluateOrDefault(getASTContext().evaluator,
+                    ExpandMemberAttributeMacros{mutableThis},
+                    false);
+
+  return getAttrs();
 }
 
 const Decl *Decl::getInnermostDeclWithAvailability() const {

lib/AST/Module.cpp

@@ -887,6 +887,15 @@ ASTNode SourceFile::getMacroExpansion() const {
   return ASTNode::getFromOpaqueValue(genInfo.astNode);
 }
 
+CustomAttr *SourceFile::getAttachedMacroAttribute() const {
+  if (Kind != SourceFileKind::MacroExpansion)
+    return nullptr;
+
+  auto genInfo =
+      *getASTContext().SourceMgr.getGeneratedSourceInfo(*getBufferID());
+  return genInfo.attachedMacroCustomAttr;
+}
+
 SourceFile *SourceFile::getEnclosingSourceFile() const {
   auto macroExpansion = getMacroExpansion();
   if (!macroExpansion)

lib/Sema/TypeCheckAttr.cpp

@@ -7385,45 +7385,6 @@ ValueDecl *RenamedDeclRequest::evaluate(Evaluator &evaluator,
   return renamedDecl;
 }
 
-SemanticDeclAttributes
-AttachedSemanticAttrsRequest::evaluate(Evaluator &evaluator, Decl *decl) const {
-  // For now, this returns the same thing as 'getAttrs' using
-  // the SemanticDeclAttribtues representation.
-  SemanticDeclAttributes semanticAttrs;
-  for (auto attr : decl->getAttrs()) {
-    semanticAttrs.add(attr);
-  }
-
-  auto *parentDecl = decl->getDeclContext()->getAsDecl();
-  if (!parentDecl)
-    return semanticAttrs;
-
-  auto parentAttrs = parentDecl->getSemanticAttrs();
-  for (auto customAttrConst: parentAttrs.getAttributes<CustomAttr>()) {
-    auto customAttr = const_cast<CustomAttr *>(customAttrConst);
-    auto customAttrDecl = evaluateOrDefault(
-        evaluator,
-        CustomAttrDeclRequest{
-          customAttr,
-          parentDecl->getInnermostDeclContext()
-        },
-        nullptr);
-    if (!customAttrDecl)
-      continue;
-
-    auto macroDecl = customAttrDecl.dyn_cast<MacroDecl *>();
-    if (!macroDecl)
-      continue;
-
-    if (!macroDecl->getMacroRoles().contains(MacroRole::MemberAttribute))
-      continue;
-
-    expandAttributes(customAttr, macroDecl, decl, semanticAttrs);
-  }
-
-  return semanticAttrs;
-}
-
 template <typename ATTR>
 static void forEachCustomAttribute(
     ValueDecl *decl,