[Macros] Allow one attached macro declaration to fulfill multiple roles.

hborla · hborla · commit e5ba38925acf · 2023-01-23T17:56:48.000-08:00
diff --git a/include/swift/AST/SourceFile.h b/include/swift/AST/SourceFile.h
@@ -511,6 +511,13 @@ class SourceFile final : public FileUnit {
   /// macro.
   CustomAttr *getAttachedMacroAttribute() const;
 
+  /// For source files created to hold the source code created by expanding
+  /// an attached macro, this is the macro role that the expansion fulfills.
+  ///
+  /// \Returns the fulfilled macro role, or \c None if this source file is not
+  /// for a macro expansion.
+  Optional<MacroRole> getFulfilledMacroRole() 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.
diff --git a/lib/AST/Module.cpp b/lib/AST/Module.cpp
@@ -896,6 +896,34 @@ CustomAttr *SourceFile::getAttachedMacroAttribute() const {
   return genInfo.attachedMacroCustomAttr;
 }
 
+Optional<MacroRole> SourceFile::getFulfilledMacroRole() const {
+  if (Kind != SourceFileKind::MacroExpansion)
+    return None;
+
+  auto genInfo =
+      *getASTContext().SourceMgr.getGeneratedSourceInfo(*getBufferID());
+  switch (genInfo.kind) {
+  case GeneratedSourceInfo::ExpressionMacroExpansion:
+    return MacroRole::Expression;
+
+  case GeneratedSourceInfo::FreestandingDeclMacroExpansion:
+    return MacroRole::FreestandingDeclaration;
+
+  case GeneratedSourceInfo::AccessorMacroExpansion:
+    return MacroRole::Accessor;
+
+  case GeneratedSourceInfo::MemberAttributeMacroExpansion:
+    return MacroRole::MemberAttribute;
+
+  case GeneratedSourceInfo::SynthesizedMemberMacroExpansion:
+    return MacroRole::SynthesizedMembers;
+
+  case GeneratedSourceInfo::ReplacedFunctionBody:
+  case GeneratedSourceInfo::PrettyPrinted:
+    return None;
+  }
+}
+
 SourceFile *SourceFile::getEnclosingSourceFile() const {
   auto macroExpansion = getMacroExpansion();
   if (!macroExpansion)
diff --git a/lib/ASTGen/Sources/ASTGen/Macros.swift b/lib/ASTGen/Sources/ASTGen/Macros.swift
@@ -31,6 +31,14 @@ struct ExportedMacro {
   var macro: Macro.Type
 }
 
+enum MacroRole: UInt8 {
+  case Expression = 0x01
+  case FreestandingDeclaration = 0x02
+  case Accessor = 0x04
+  case MemberAttribute = 0x08
+  case SynthesizedMembers = 0x10
+}
+
 /// Resolve a reference to type metadata into a macro, if posible.
 ///
 /// Returns an unmanaged pointer to an ExportedMacro instance that describes
@@ -262,6 +270,7 @@ private func findSyntaxNodeInSourceFile<Node: SyntaxProtocol>(
 func expandAttachedMacro(
   diagEnginePtr: UnsafeMutablePointer<UInt8>,
   macroPtr: UnsafeRawPointer,
+  rawMacroRole: UInt8,
   customAttrSourceFilePtr: UnsafeRawPointer,
   customAttrSourceLocPointer: UnsafePointer<UInt8>?,
   declarationSourceFilePtr: UnsafeRawPointer,
@@ -297,6 +306,7 @@ func expandAttachedMacro(
   // Get the macro.
   let macroPtr = macroPtr.bindMemory(to: ExportedMacro.self, capacity: 1)
   let macro = macroPtr.pointee.macro
+  let macroRole = MacroRole(rawValue: rawMacroRole)
 
   // FIXME: Which source file? I don't know! This should go.
   let declarationSourceFilePtr = declarationSourceFilePtr.bindMemory(
@@ -310,8 +320,8 @@ func expandAttachedMacro(
 
   var evaluatedSyntaxStr: String
   do {
-    switch macro {
-    case let attachedMacro as AccessorDeclarationMacro.Type:
+    switch (macro, macroRole) {
+    case (let attachedMacro as AccessorDeclarationMacro.Type, .Accessor):
       let accessors = try attachedMacro.expansion(
         of: customAttrNode, attachedTo: declarationNode, in: &context
       )
@@ -321,7 +331,7 @@ func expandAttachedMacro(
         $0.withoutTrivia().description
       }.joined(separator: "\n\n")
 
-    case let attachedMacro as MemberAttributeMacro.Type:
+    case (let attachedMacro as MemberAttributeMacro.Type, .MemberAttribute):
       // Dig out the node for the parent declaration of the to-expand
       // declaration. Only member attribute macros need this.
       guard let parentDeclNode = findSyntaxNodeInSourceFile(
@@ -344,7 +354,7 @@ func expandAttachedMacro(
         $0.withoutTrivia().description
       }.joined(separator: " ")
 
-    case let attachedMacro as MemberDeclarationMacro.Type:
+    case (let attachedMacro as MemberDeclarationMacro.Type, .SynthesizedMembers):
       let members = try attachedMacro.expansion(
         of: customAttrNode,
         attachedTo: declarationNode,
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
@@ -44,6 +44,7 @@ extern "C" ptrdiff_t swift_ASTGen_evaluateMacro(
 
 extern "C" ptrdiff_t swift_ASTGen_expandAttachedMacro(
     void *diagEngine, void *macro,
+    uint32_t rawMacroRole,
     void *customAttrSourceFile,
     const void *customAttrSourceLocation,
     void *declarationSourceFile,
@@ -366,7 +367,8 @@ bool ExpandSynthesizedMemberMacroRequest::evaluate(Evaluator &evaluator,
 
 /// Determine whether the given source file is from an expansion of the given
 /// macro.
-static bool isFromExpansionOfMacro(SourceFile *sourceFile, MacroDecl *macro) {
+static bool isFromExpansionOfMacro(SourceFile *sourceFile, MacroDecl *macro,
+                                   MacroRole role) {
   while (sourceFile) {
     auto expansion = sourceFile->getMacroExpansion();
     if (!expansion)
@@ -391,7 +393,8 @@ static bool isFromExpansionOfMacro(SourceFile *sourceFile, MacroDecl *macro) {
       if (!macroDecl)
         return false;
 
-      return macroDecl == macro;
+      return macroDecl == macro &&
+             sourceFile->getFulfilledMacroRole() == role;
     } else {
       llvm_unreachable("Unknown macro expansion node kind");
     }
@@ -418,7 +421,7 @@ Expr *swift::expandMacroExpr(
 
   MacroDecl *macro = cast<MacroDecl>(macroRef.getDecl());
 
-  if (isFromExpansionOfMacro(sourceFile, macro)) {
+  if (isFromExpansionOfMacro(sourceFile, macro, MacroRole::Expression)) {
     ctx.Diags.diagnose(expr->getLoc(), diag::macro_recursive, macro->getName());
     return nullptr;
   }
@@ -599,7 +602,7 @@ bool swift::expandFreestandingDeclarationMacro(
   assert(macro->getMacroRoles()
              .contains(MacroRole::FreestandingDeclaration));
 
-  if (isFromExpansionOfMacro(sourceFile, macro)) {
+  if (isFromExpansionOfMacro(sourceFile, macro, MacroRole::FreestandingDeclaration)) {
     med->diagnose(diag::macro_recursive, macro->getName());
     return false;
   }
@@ -765,8 +768,8 @@ void swift::expandAccessors(
   // Evaluate the macro.
   NullTerminatedStringRef evaluatedSource;
 
-  if (isFromExpansionOfMacro(attrSourceFile, macro) ||
-      isFromExpansionOfMacro(declSourceFile, macro)) {
+  if (isFromExpansionOfMacro(attrSourceFile, macro, MacroRole::Accessor) ||
+      isFromExpansionOfMacro(declSourceFile, macro, MacroRole::Accessor)) {
     storage->diagnose(diag::macro_recursive, macro->getName());
     return;
   }
@@ -831,6 +834,7 @@ void swift::expandAccessors(
     swift_ASTGen_expandAttachedMacro(
         &ctx.Diags,
         externalDef.opaqueHandle,
+        static_cast<uint32_t>(MacroRole::Accessor),
         astGenAttrSourceFile, attr->AtLoc.getOpaquePointerValue(),
         astGenDeclSourceFile, searchDecl->getStartLoc().getOpaquePointerValue(),
         /*parentDeclSourceFile*/nullptr, /*parentDeclLoc*/nullptr,
@@ -943,8 +947,8 @@ bool swift::expandAttributes(CustomAttr *attr, MacroDecl *macro, Decl *member) {
   // Evaluate the macro.
   NullTerminatedStringRef evaluatedSource;
 
-  if (isFromExpansionOfMacro(attrSourceFile, macro) ||
-      isFromExpansionOfMacro(declSourceFile, macro)) {
+  if (isFromExpansionOfMacro(attrSourceFile, macro, MacroRole::MemberAttribute) ||
+      isFromExpansionOfMacro(declSourceFile, macro, MacroRole::MemberAttribute)) {
     member->diagnose(diag::macro_recursive, macro->getName());
     return false;
   }
@@ -1013,6 +1017,7 @@ bool swift::expandAttributes(CustomAttr *attr, MacroDecl *macro, Decl *member) {
     swift_ASTGen_expandAttachedMacro(
         &ctx.Diags,
         externalDef.opaqueHandle,
+        static_cast<uint32_t>(MacroRole::MemberAttribute),
         astGenAttrSourceFile, attr->AtLoc.getOpaquePointerValue(),
         astGenDeclSourceFile, searchDecl->getStartLoc().getOpaquePointerValue(),
         astGenParentDeclSourceFile, parentDecl->getStartLoc().getOpaquePointerValue(),
@@ -1127,8 +1132,8 @@ bool swift::expandSynthesizedMembers(CustomAttr *attr, MacroDecl *macro,
   // Evaluate the macro.
   NullTerminatedStringRef evaluatedSource;
 
-  if (isFromExpansionOfMacro(attrSourceFile, macro) ||
-      isFromExpansionOfMacro(declSourceFile, macro)) {
+  if (isFromExpansionOfMacro(attrSourceFile, macro, MacroRole::SynthesizedMembers) ||
+      isFromExpansionOfMacro(declSourceFile, macro, MacroRole::SynthesizedMembers)) {
     decl->diagnose(diag::macro_recursive, macro->getName());
     return false;
   }
@@ -1189,6 +1194,7 @@ bool swift::expandSynthesizedMembers(CustomAttr *attr, MacroDecl *macro,
     swift_ASTGen_expandAttachedMacro(
         &ctx.Diags,
         externalDef.opaqueHandle,
+        static_cast<uint32_t>(MacroRole::SynthesizedMembers),
         astGenAttrSourceFile, attr->AtLoc.getOpaquePointerValue(),
         astGenDeclSourceFile, decl->getStartLoc().getOpaquePointerValue(),
         /*parentDeclSourceFile*/nullptr, /*parentDeclLoc*/nullptr,
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/test/Macros/Inputs/syntax_macro_definitions.swift
@@ -274,7 +274,9 @@ public struct WrapAllProperties: MemberAttributeMacro {
   }
 }
 
-public struct TypeWrapperMacro: MemberAttributeMacro {
+public struct TypeWrapperMacro {}
+
+extension TypeWrapperMacro: MemberAttributeMacro {
   public static func expansion(
     of node: AttributeSyntax,
     attachedTo decl: DeclSyntax,
@@ -303,6 +305,23 @@ public struct TypeWrapperMacro: MemberAttributeMacro {
   }
 }
 
+extension TypeWrapperMacro: MemberDeclarationMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    attachedTo decl: DeclSyntax,
+    in context: inout MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    let storageVariable: VariableDeclSyntax =
+      """
+      private var _storage = _Storage()
+      """
+
+    return [
+      DeclSyntax(storageVariable),
+    ]
+  }
+}
+
 public struct AccessViaStorageMacro: AccessorDeclarationMacro {
   public static func expansion(
     of node: AttributeSyntax,
@@ -327,20 +346,3 @@ public struct AccessViaStorageMacro: AccessorDeclarationMacro {
     ]
   }
 }
-
-public struct TypeWrapperStorageMacro: MemberDeclarationMacro {
-  public static func expansion(
-    of node: AttributeSyntax,
-    attachedTo decl: DeclSyntax,
-    in context: inout MacroExpansionContext
-  ) throws -> [DeclSyntax] {
-    let storageVariable: VariableDeclSyntax =
-      """
-      private var _storage = _Storage()
-      """
-
-    return [
-      DeclSyntax(storageVariable),
-    ]
-  }
-}
diff --git a/test/Macros/composed_macro.swift b/test/Macros/composed_macro.swift
@@ -9,8 +9,9 @@
 // FIXME: Swift parser is not enabled on Linux CI yet.
 // REQUIRES: OS=macosx
 
-@attached(memberAttributes) macro myTypeWrapper() = #externalMacro(module: "MacroDefinition", type: "TypeWrapperMacro")
-@attached(synthesizedMembers) macro typeWrapperStorage() = #externalMacro(module: "MacroDefinition", type: "TypeWrapperStorageMacro")
+@attached(memberAttributes)
+@attached(synthesizedMembers)
+macro myTypeWrapper() = #externalMacro(module: "MacroDefinition", type: "TypeWrapperMacro")
 @attached(accessor) macro accessViaStorage() = #externalMacro(module: "MacroDefinition", type: "AccessViaStorageMacro")
 
 struct _Storage {
@@ -22,7 +23,6 @@ struct _Storage {
   }
 }
 
-@typeWrapperStorage
 @myTypeWrapper
 struct S {
   var x: Int
