[Macros] Implement closure body macros in MacroSystem and add a test

for expanding a body macro on a closure expr.

Author: hborla

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -444,7 +444,7 @@ private func expandPreambleMacro(
 private func expandBodyMacro(
   definition: BodyMacro.Type,
   attributeNode: AttributeSyntax,
-  attachedTo decl: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+  attachedTo node: some SyntaxProtocol,
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) -> CodeBlockSyntax? {
@@ -456,7 +456,7 @@ private func expandBodyMacro(
         in: context,
         foldingWith: .standardOperators
       ),
-      declarationNode: DeclSyntax(decl.detach(in: context)),
+      declarationNode: Syntax(node.detach(in: context)),
       parentDeclNode: nil,
       extendedType: nil,
       conformanceList: nil,
@@ -473,9 +473,15 @@ private func expandBodyMacro(
   // Remove any indentation from the first line using `drop(while:)` and then
   // prepend a space when it's being introduced on a declaration that has no
   // body yet.
-  let leadingWhitespace = decl.body == nil ? " " : ""
+  let leadingWhitespace =
+    if let decl = node as? (DeclSyntaxProtocol & WithOptionalCodeBlockSyntax),
+       decl.body == nil {
+      " "
+    } else {
+      ""
+    }
   let indentedSource =
-    leadingWhitespace + expanded.indented(by: decl.indentationOfFirstLine).drop(while: { $0.isWhitespace })
+    leadingWhitespace + expanded.indented(by: node.indentationOfFirstLine).drop(while: { $0.isWhitespace })
   return "\(raw: indentedSource)"
 }
 
@@ -733,6 +739,19 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
       return AttributeRemover(removingWhere: { attributesToRemove.contains($0) }).rewrite(visitedNode)
     }
 
+    if var closureSyntax = node.as(ClosureExprSyntax.self) {
+      closureSyntax = visitClosureBodyMacros(closureSyntax)
+
+      // Visit the node, disabling the `visitAny` handling.
+      skipVisitAnyHandling.insert(Syntax(closureSyntax))
+      let visitedNode = self.visit(closureSyntax).cast(ClosureExprSyntax.self)
+      skipVisitAnyHandling.remove(Syntax(closureSyntax))
+
+      let attributesToRemove = self.macroAttributes(attachedTo: visitedNode).map(\.attributeNode)
+
+      return AttributeRemover(removingWhere: { attributesToRemove.contains($0) }).rewrite(visitedNode)
+    }
+
     return nil
   }
 
@@ -802,6 +821,41 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
     return node.with(\.body, body.with(\.statements, preamble + body.statements))
   }
 
+  func visitClosureBodyMacros(
+    _ node: ClosureExprSyntax
+  ) -> ClosureExprSyntax {
+    // Expand body macro.
+    let expandedBodies = expandMacros(
+      attachedTo: node,
+      ofType: BodyMacro.Type.self
+    ) { attributeNode, definition, _ in
+      expandBodyMacro(
+        definition: definition,
+        attributeNode: attributeNode,
+        attachedTo: node,
+        in: contextGenerator(Syntax(node)),
+        indentationWidth: indentationWidth
+      ).map { [$0] }
+    }
+
+    // Dig out the body.
+    let body: CodeBlockSyntax
+    switch expandedBodies.count {
+    case 0:
+      // Nothing changes
+      return node
+
+    case 1:
+      body = expandedBodies[0]
+
+    default:
+      contextGenerator(Syntax(node)).addDiagnostics(from: MacroExpansionError.moreThanOneBodyMacro, node: node)
+      body = expandedBodies[0]
+    }
+
+    return node.with(\.statements, body.statements)
+  }
+
   override func visit(_ node: CodeBlockItemListSyntax) -> CodeBlockItemListSyntax {
     var newItems: [CodeBlockItemSyntax] = []
     func addResult(_ node: CodeBlockItemSyntax) {
@@ -1006,9 +1060,16 @@ extension MacroApplication {
   ///
   /// The macros must be registered in `macroSystem`.
   private func macroAttributes(
-    attachedTo decl: DeclSyntax
+    attachedTo decl: some SyntaxProtocol
   ) -> [(attributeNode: AttributeSyntax, spec: MacroSpec)] {
-    guard let attributedNode = decl.asProtocol(WithAttributesSyntax.self) else {
+    let attributedNode: (any WithAttributesSyntax)?
+    if let closure = decl.as(ClosureExprSyntax.self) {
+      attributedNode = closure.signature?.asProtocol(WithAttributesSyntax.self)
+    } else {
+      attributedNode = decl.asProtocol(WithAttributesSyntax.self)
+    }
+
+    guard let attributedNode else {
       return []
     }
 
@@ -1029,7 +1090,7 @@ extension MacroApplication {
   ///
   /// The macros must be registered in `macroSystem`.
   private func macroAttributes<MacroType>(
-    attachedTo decl: DeclSyntax,
+    attachedTo decl: some SyntaxProtocol,
     ofType: MacroType.Type
   ) -> [(attributeNode: AttributeSyntax, definition: MacroType, conformanceList: InheritedTypeListSyntax)] {
     return macroAttributes(attachedTo: decl)
@@ -1049,7 +1110,7 @@ extension MacroApplication {
     ExpandedNodeCollection: Sequence<ExpandedNode>,
     MacroType
   >(
-    attachedTo decl: DeclSyntax,
+    attachedTo decl: some SyntaxProtocol,
     ofType: MacroType.Type,
     expandMacro:
       (

Tests/SwiftSyntaxMacroExpansionTest/BodyMacroTests.swift

@@ -63,6 +63,37 @@ struct RemoteBodyMacro: BodyMacro {
   }
 }
 
+struct StartTaskMacro: BodyMacro {
+  static func expansion(
+    of node: AttributeSyntax,
+    providingBodyFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax] {
+    guard let taskBody = declaration.body else {
+      return []
+    }
+    return [
+      """
+      Task \(taskBody.trimmed)
+      """
+    ]
+  }
+
+  static func expansion(
+    of node: AttributeSyntax,
+    providingBodyFor closure: ClosureExprSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax] {
+    return [
+      """
+      Task {
+      \(closure.statements.trimmed)
+      }
+      """
+    ]
+  }
+}
+
 final class BodyMacroTests: XCTestCase {
   private let indentationWidth: Trivia = .spaces(2)
 
@@ -199,4 +230,23 @@ final class BodyMacroTests: XCTestCase {
       macros: ["SourceLocationMacro": SourceLocationMacro.self]
     )
   }
+
+  func testClosureBodyExpansion() {
+    assertMacroExpansion(
+      """
+      { @StartTask in
+        a + b
+      }
+      """,
+      expandedSource: """
+        { in
+          Task {
+            a + b
+          }
+        }
+        """,
+      macros: ["StartTask": StartTaskMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
 }