Merge pull request #3298 from calda/cal--computed-var-body-macro

Author: rintaro

Sources/SwiftCompilerPluginMessageHandling/PluginMacroExpansionContext.swift

@@ -65,6 +65,8 @@ class ParsedSyntaxRegistry {
       return Syntax(PatternSyntax.parse(from: &parser))
     case .attribute:
       return Syntax(AttributeSyntax.parse(from: &parser))
+    case .accessor:
+      return Syntax(AccessorDeclSyntax.parse(from: &parser))
     }
   }
 

Sources/SwiftCompilerPluginMessageHandling/PluginMessages.swift

@@ -256,6 +256,7 @@ public enum PluginMessage {
       case type
       case pattern
       case attribute
+      case accessor
     }
     public var kind: Kind
     public var source: String

Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift

@@ -12,10 +12,14 @@
 
 #if compiler(>=6)
 import SwiftBasicFormat
+import SwiftDiagnostics
+import SwiftIfConfig
 public import SwiftSyntax
 @_spi(MacroExpansion) @_spi(ExperimentalLanguageFeatures) public import SwiftSyntaxMacros
 #else
 import SwiftBasicFormat
+import SwiftDiagnostics
+import SwiftIfConfig
 import SwiftSyntax
 @_spi(MacroExpansion) @_spi(ExperimentalLanguageFeatures) import SwiftSyntaxMacros
 #endif
@@ -394,6 +398,38 @@ public func expandAttachedMacroWithoutCollapsing<Context: MacroExpansionContext>
           providingBodyFor: declToPass,
           in: context
         )
+      } else if let varDecl = node.as(VariableDeclSyntax.self),
+        varDecl.bindings.count == 1,
+        let binding = varDecl.bindings.first,
+        let accessorBlock = binding.accessorBlock,
+        case .getter(let stmts) = accessorBlock.accessors
+      {
+        // Create an implicit `AccessorDeclSyntax` to use for the macro expansion
+        // of this computed var decl
+        let getterDecl = AccessorDeclSyntax(
+          accessorSpecifier: .keyword(.get, presence: .missing),
+          body: CodeBlockSyntax(
+            leftBrace: accessorBlock.leftBrace,
+            statements: stmts,
+            rightBrace: accessorBlock.rightBrace
+          )
+        )
+
+        // Insert the enclosing `PatternBindingSyntax`'s var decl context
+        // into the lexical context
+        var context: MacroExpansionContext = context
+        if let varDeclLexicalContext = binding.asMacroLexicalContext() {
+          context = PrependLexicalContextWrapperContext(
+            prependLexicalContext: [varDeclLexicalContext],
+            wrapping: context
+          )
+        }
+
+        body = try attachedMacro.expansion(
+          of: attributeNode,
+          providingBodyFor: getterDecl,
+          in: context
+        )
       } else {
         // Compiler error: declaration must have a body.
         throw MacroExpansionError.declarationHasNoBody
@@ -635,3 +671,39 @@ public func collapse<Node: SyntaxProtocol>(
 
   return collapsed
 }
+
+/// Wrapper context that prepends additional `lexicalContext`
+/// to an existing `MacroExpansionContext`
+final class PrependLexicalContextWrapperContext: MacroExpansionContext {
+  init(prependLexicalContext: [Syntax], wrapping wrappedContext: MacroExpansionContext) {
+    self.prependLexicalContext = prependLexicalContext
+    self.wrappedContext = wrappedContext
+  }
+
+  let prependLexicalContext: [Syntax]
+  let wrappedContext: MacroExpansionContext
+
+  var lexicalContext: [Syntax] {
+    prependLexicalContext + wrappedContext.lexicalContext
+  }
+
+  func makeUniqueName(_ name: String) -> TokenSyntax {
+    wrappedContext.makeUniqueName(name)
+  }
+
+  func diagnose(_ diagnostic: Diagnostic) {
+    wrappedContext.diagnose(diagnostic)
+  }
+
+  func location(
+    of node: some SyntaxProtocol,
+    at position: PositionInSyntaxNode,
+    filePathMode: SourceLocationFilePathMode
+  ) -> AbstractSourceLocation? {
+    wrappedContext.location(of: node, at: position, filePathMode: filePathMode)
+  }
+
+  var buildConfiguration: (any BuildConfiguration)? {
+    wrappedContext.buildConfiguration
+  }
+}

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -707,6 +707,32 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
         let declNodeWithBody = nodeWithBody as? any DeclSyntaxProtocol & WithOptionalCodeBlockSyntax
       {
         declSyntax = DeclSyntax(visitBodyAndPreambleMacros(declNodeWithBody))
+      } else if let varDecl = node.as(VariableDeclSyntax.self),
+        varDecl.bindings.count == 1,
+        let bindingIndex = varDecl.bindings.indices.first,
+        let accessorBlock = varDecl.bindings[bindingIndex].accessorBlock,
+        case .getter = accessorBlock.accessors
+      {
+        let expandedBodies = expandMacros(
+          attachedTo: DeclSyntax(varDecl),
+          ofType: BodyMacro.Type.self
+        ) { attributeNode, definition, _ in
+          expandBodyMacro(
+            definition: definition,
+            attributeNode: attributeNode,
+            attachedTo: node,
+            in: contextGenerator(Syntax(node)),
+            indentationWidth: indentationWidth
+          ).map { [$0] }
+        }
+
+        if let newBody = expandedBodies.first {
+          let newAccessorBlock = accessorBlock.with(\.accessors, .getter(newBody.statements))
+          let newBinding = varDecl.bindings[bindingIndex].with(\.accessorBlock, newAccessorBlock)
+          var newBindings = varDecl.bindings
+          newBindings[bindingIndex] = newBinding
+          declSyntax = DeclSyntax(varDecl.with(\.bindings, newBindings))
+        }
       }
 
       // Visit the node, disabling the `visitAny` handling.

Sources/SwiftSyntaxMacros/Syntax+LexicalContext.swift

@@ -55,12 +55,24 @@ extension SyntaxProtocol {
     case is EnumCaseElementSyntax:
       return Syntax(self.detached)
 
-    // Pattern bindings have their accessors and initializer removed.
+    // Pattern bindings have their accessors and initializer removed, and wrapped in the parent variable decl.
     case var patternBinding as PatternBindingSyntax:
+      let varDecl = patternBinding.parent?.as(PatternBindingListSyntax.self)?.parent?.as(VariableDeclSyntax.self)?
+        .detached
       patternBinding = patternBinding.detached
       patternBinding.accessorBlock = nil
       patternBinding.initializer = nil
-      return Syntax(patternBinding)
+      if var varDecl {
+        varDecl.bindings = [patternBinding]
+        return Syntax(varDecl)
+      } else {
+        return Syntax(
+          VariableDeclSyntax(
+            bindingSpecifier: .keyword(.var),
+            bindings: [patternBinding]
+          )
+        )
+      }
 
     // Freestanding macros are fine as-is because if any arguments change
     // the whole macro would have to be re-evaluated.

Tests/SwiftSyntaxMacroExpansionTest/BodyMacroTests.swift

@@ -72,11 +72,32 @@ struct StartTaskMacro: BodyMacro {
     guard let taskBody = declaration.body else {
       return []
     }
-    return [
-      """
-      Task \(taskBody.trimmed)
-      """
-    ]
+
+    let taskName: String?
+    if let funcDecl = declaration.as(FunctionDeclSyntax.self) {
+      taskName = funcDecl.name.text
+    } else if declaration.is(AccessorDeclSyntax.self) {
+      taskName = context.lexicalContext
+        .compactMap { $0.as(VariableDeclSyntax.self) }
+        .first
+        .flatMap { $0.bindings.first?.pattern.as(IdentifierPatternSyntax.self)?.identifier.text }
+    } else {
+      taskName = nil
+    }
+
+    if let taskName {
+      return [
+        """
+        Task(name: \(literal: taskName)) \(taskBody.trimmed)
+        """
+      ]
+    } else {
+      return [
+        """
+        Task \(taskBody.trimmed)
+        """
+      ]
+    }
   }
 
   static func expansion(
@@ -231,6 +252,79 @@ final class BodyMacroTests: XCTestCase {
     )
   }
 
+  func testBodyExpansionOnFunc() {
+    assertMacroExpansion(
+      """
+      @StartTask
+      func x() -> Int {
+        a + b
+      }
+      """,
+      expandedSource: """
+
+        func x() -> Int {
+          Task(name: "x") {
+            a + b
+          }
+        }
+        """,
+      macros: ["StartTask": StartTaskMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testBodyExpansionOnComputedVar() {
+    assertMacroExpansion(
+      """
+      @StartTask
+      var x: Int {
+        a + b
+      }
+      """,
+      expandedSource: """
+
+        var x: Int {
+          Task(name: "x") {
+            a + b
+          }
+        }
+        """,
+      macros: ["StartTask": StartTaskMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testBodyExpansionOnAccessors() {
+    assertMacroExpansion(
+      """
+      var value: Double {
+        @StartTask get {
+          return length * width
+        }
+        @StartTask set {
+          self.value = newValue * 2
+        }
+      }
+      """,
+      expandedSource: """
+        var value: Double {
+          get {
+            Task(name: "value") {
+                return length * width
+              }
+          }
+          set {
+            Task(name: "value") {
+                self.value = newValue * 2
+              }
+          }
+        }
+        """,
+      macros: ["StartTask": StartTaskMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
+
   func testClosureBodyExpansion() {
     assertMacroExpansion(
       """

Tests/SwiftSyntaxMacroExpansionTest/LexicalContextTests.swift

@@ -151,6 +151,14 @@ extension SyntaxProtocol {
       return singleVarName
     }
 
+    // Variable declarations with a single binding.
+    if let varDecl = self.as(VariableDeclSyntax.self),
+      varDecl.bindings.count == 1,
+      let singleVarName = varDecl.bindings.first?.singleBindingName
+    {
+      return singleVarName
+    }
+
     return nil
   }
 }
@@ -303,6 +311,21 @@ final class LexicalContextTests: XCTestCase {
       indentationWidth: indentationWidth
     )
 
+    assertMacroExpansion(
+      """
+      extension A {
+        static var staticProp = #function, secondProp = #function
+      }
+      """,
+      expandedSource: """
+        extension A {
+          static var staticProp = "staticProp", secondProp = "secondProp"
+        }
+        """,
+      macros: ["function": FunctionMacro.self],
+      indentationWidth: indentationWidth
+    )
+
     assertMacroExpansion(
       """
       func f(a: Int, _: Double, c: Int) {
@@ -546,7 +569,7 @@ final class LexicalContextTests: XCTestCase {
               await _
               try _
               unsafe _
-              contextDescription: String
+              var contextDescription: String
               struct S {}
               { c in
               }