[Macros] Copy attrs/modifiers on MacroExpansionDecl to expanded decls

rintaro · rintaro · commit 06e6de2bbea4 · 2023-05-15T16:12:57.000-07:00
As per the SE-0397 amendment, copy attributes and modifiers on
MacroExpansionDecl to the expanded declarations.
diff --git a/CodeGeneration/Sources/generate-swiftsyntax/templates/swiftsyntax/SyntaxCollectionsFile.swift b/CodeGeneration/Sources/generate-swiftsyntax/templates/swiftsyntax/SyntaxCollectionsFile.swift
@@ -19,6 +19,8 @@ let syntaxCollectionsFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
   DeclSyntax(
     """
     public protocol SyntaxCollection: SyntaxProtocol, Sequence where Element: SyntaxProtocol {
+      /// Creates a new collection with the elements.
+      init(_ children: [Element])
       /// The number of elements, `present` or `missing`, in this collection.
       var count: Int { get }
     }
diff --git a/Sources/SwiftParser/CMakeLists.txt b/Sources/SwiftParser/CMakeLists.txt
@@ -27,6 +27,7 @@ add_swift_host_library(SwiftParser
   Specifiers.swift
   Statements.swift
   StringLiterals.swift
+  StringLiteralRepresentedLiteralValue.swift
   SyntaxUtils.swift
   TokenConsumer.swift
   TokenPrecedence.swift
diff --git a/Sources/SwiftSyntax/generated/SyntaxCollections.swift b/Sources/SwiftSyntax/generated/SyntaxCollections.swift
@@ -13,6 +13,9 @@
 //===----------------------------------------------------------------------===//
 
 public protocol SyntaxCollection: SyntaxProtocol, Sequence where Element: SyntaxProtocol {
+  /// Creates a new collection with the elements.
+  init(_ children: [Element])
+  
   /// The number of elements, `present` or `missing`, in this collection.
   var count: Int {
     get
diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift b/Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift
@@ -31,7 +31,17 @@ public func expandFreestandingMacro(
         expandedSyntax = try Syntax(exprMacroDef.expansion(of: node, in: context))
 
       case let declMacroDef as DeclarationMacro.Type:
-        let rewritten = try declMacroDef.expansion(of: node, in: context)
+        var rewritten = try declMacroDef.expansion(of: node, in: context)
+        // Copy attributes and modifiers to the generated decls.
+        if let expansionDecl = node.as(MacroExpansionDeclSyntax.self) {
+          rewritten = rewritten.map {
+            applyingAttrsModifiers(
+              to: $0,
+              attributes: expansionDecl.attributes,
+              modifiers: expansionDecl.modifiers
+            )
+          }
+        }
         expandedSyntax = Syntax(
           CodeBlockItemListSyntax(
             rewritten.map {
@@ -201,3 +211,35 @@ fileprivate extension SyntaxProtocol {
     return formatted.trimmedDescription(matching: { $0.isWhitespace })
   }
 }
+
+func applyingAttrsModifiers(
+  to node: DeclSyntax,
+  attributes: AttributeListSyntax?,
+  modifiers: ModifierListSyntax?
+) -> DeclSyntax {
+  func _combine<C: SyntaxCollection>(_ left: C, _ right: C?) -> C? {
+    guard let right = right else { return left }
+    var elems: [C.Element] = []
+    elems.append(contentsOf: left)
+    elems.append(contentsOf: right)
+    return C(elems)
+  }
+  var node = node
+  if let attributes = attributes,
+    let withAttrs = node.asProtocol(WithAttributesSyntax.self)
+  {
+    node = withAttrs.with(
+      \.attributes,
+      _combine(attributes, withAttrs.attributes)
+    ).as(DeclSyntax.self)!
+  }
+  if let modifiers = modifiers,
+    let withModifiers = node.asProtocol(WithModifiersSyntax.self)
+  {
+    node = withModifiers.with(
+      \.modifiers,
+      _combine(modifiers, withModifiers.modifiers)
+    ).as(DeclSyntax.self)!
+  }
+  return node
+}
diff --git a/Sources/SwiftSyntaxMacros/MacroSystem.swift b/Sources/SwiftSyntaxMacros/MacroSystem.swift
@@ -125,36 +125,41 @@ class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
   override func visit(_ node: CodeBlockItemListSyntax) -> CodeBlockItemListSyntax {
     var newItems: [CodeBlockItemSyntax] = []
     for item in node {
-      // Expand declaration macros that were parsed as macro expansion
-      // expressions in this context.
-      if case let .expr(exprItem) = item.item,
-        let exprExpansion = exprItem.as(MacroExpansionExprSyntax.self),
-        let macro = macroSystem.macros[exprExpansion.macro.text]
+      if let expansion = item.item.asProtocol(FreestandingMacroExpansionSyntax.self),
+        let macro = macroSystem.macros[expansion.macro.text]
       {
-        do {
+        func _expand(expansion: FreestandingMacroExpansionSyntax) throws {
           if let macro = macro as? CodeItemMacro.Type {
             let expandedItemList = try macro.expansion(
-              of: exprExpansion,
+              of: expansion,
               in: context
             )
             newItems.append(contentsOf: expandedItemList)
           } else if let macro = macro as? DeclarationMacro.Type {
-            let expandedItemList = try macro.expansion(
-              of: exprExpansion,
+            var expandedItemList = try macro.expansion(
+              of: expansion,
               in: context
             )
+            if let declExpansion = expansion.as(MacroExpansionDeclSyntax.self) {
+              expandedItemList = expandedItemList.map {
+                applyingAttrsModifiers(to: $0, attributes: declExpansion.attributes, modifiers: declExpansion.modifiers)
+              }
+            }
             newItems.append(
               contentsOf: expandedItemList.map {
                 CodeBlockItemSyntax(item: .decl($0))
               }
             )
           } else if let macro = macro as? ExpressionMacro.Type {
             let expandedExpr = try macro.expansion(
-              of: exprExpansion,
+              of: expansion,
               in: context
             )
             newItems.append(CodeBlockItemSyntax(item: .init(expandedExpr)))
           }
+        }
+        do {
+          try _openExistential(expansion, do: _expand)
         } catch {
           context.addDiagnostics(from: error, node: node)
         }
@@ -189,10 +194,13 @@ class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
         let freestandingMacro = macro as? DeclarationMacro.Type
       {
         do {
-          let expandedList = try freestandingMacro.expansion(
+          var expandedList = try freestandingMacro.expansion(
             of: declExpansion,
             in: context
           )
+          expandedList = expandedList.map {
+            applyingAttrsModifiers(to: $0, attributes: declExpansion.attributes, modifiers: declExpansion.modifiers)
+          }
 
           newItems.append(
             contentsOf: expandedList.map { decl in
@@ -468,6 +476,40 @@ extension MacroApplication {
   }
 }
 
+// FIXME: This was copied from 'SwiftSyntaxMacroExpansion' module.
+// All expansion logics should be unified.
+private func applyingAttrsModifiers(
+  to node: DeclSyntax,
+  attributes: AttributeListSyntax?,
+  modifiers: ModifierListSyntax?
+) -> DeclSyntax {
+  func _combine<C: SyntaxCollection>(_ left: C, _ right: C?) -> C? {
+    guard let right = right else { return left }
+    var elems: [C.Element] = []
+    elems.append(contentsOf: left)
+    elems.append(contentsOf: right)
+    return C(elems)
+  }
+  var node = node
+  if let attributes = attributes,
+    let withAttrs = node.asProtocol(WithAttributesSyntax.self)
+  {
+    node = withAttrs.with(
+      \.attributes,
+      _combine(attributes, withAttrs.attributes)
+    ).as(DeclSyntax.self)!
+  }
+  if let modifiers = modifiers,
+    let withModifiers = node.asProtocol(WithModifiersSyntax.self)
+  {
+    node = withModifiers.with(
+      \.modifiers,
+      _combine(modifiers, withModifiers.modifiers)
+    ).as(DeclSyntax.self)!
+  }
+  return node
+}
+
 extension SyntaxProtocol {
   /// Expand all uses of the given set of macros within this syntax
   /// node.
diff --git a/Tests/SwiftSyntaxMacrosTest/MacroSystemTests.swift b/Tests/SwiftSyntaxMacrosTest/MacroSystemTests.swift
@@ -651,6 +651,27 @@ public struct UnwrapMacro: CodeItemMacro {
   }
 }
 
+public struct DeclsFromStringsMacro: DeclarationMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    var strings: [String] = []
+    for arg in node.argumentList {
+      guard
+        let value = arg.expression.as(StringLiteralExprSyntax.self)?.representedLiteralValue
+      else {
+        continue
+      }
+      strings.append(value)
+    }
+
+    return strings.map {
+      "\(raw: $0)"
+    }
+  }
+}
+
 // MARK: Tests
 
 /// The set of test macros we use here.
@@ -1037,4 +1058,39 @@ final class MacroSystemTests: XCTestCase {
       indentationWidth: indentationWidth
     )
   }
+
+  func testDeclsFromStringLiterals() {
+    assertMacroExpansion(
+      #"""
+      struct S {
+        public #decls(
+          """
+          static func foo() {
+          print("value") }
+          """,
+          "struct Inner {\n\n}"
+        )
+        @attr static #decls("var value1 = 1", "typealias A = B")
+      }
+      @attribute
+      @otherAttribute(x: 1) #decls("@moreAttibute var global = 42")
+      """#,
+      expandedSource: """
+        struct S {
+          public static func foo() {
+            print("value")
+          }
+          public struct Inner {
+
+          }
+          @attr static var value1 = 1
+          @attr typealias A = B
+        }
+        @attribute
+        @otherAttribute(x: 1) @moreAttibute var global = 42
+        """,
+      macros: ["decls": DeclsFromStringsMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,8 @@ let syntaxCollectionsFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {`
`19`	`19`	`DeclSyntax(`
`20`	`20`	`"""`
`21`	`21`	`public protocol SyntaxCollection: SyntaxProtocol, Sequence where Element: SyntaxProtocol {`
	`22`	`+ /// Creates a new collection with the elements.`
	`23`	`+ init(_ children: [Element])`
`22`	`24`	/// The number of elements, `present` or `missing`, in this collection.
`23`	`25`	`var count: Int { get }`
`24`	`26`	`}`