Merge pull request swiftlang#1874 from DougGregor/function-body-macros

[Macros] Implement function body macros

Author: DougGregor

Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -14,8 +14,8 @@ import SwiftBasicFormat
 import SwiftDiagnostics
 import SwiftOperators
 import SwiftSyntax
-import SwiftSyntaxMacroExpansion
-import SwiftSyntaxMacros
+@_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacroExpansion
+@_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacros
 
 extension CompilerPluginMessageHandler {
   /// Get concrete macro type from a pair of module name and type name.
@@ -166,6 +166,8 @@ private extension MacroRole {
     case .conformance: self = .extension
     case .codeItem: self = .codeItem
     case .extension: self = .extension
+    case .preamble: self = .preamble
+    case .body: self = .body
     }
   }
 }

Sources/SwiftCompilerPluginMessageHandling/PluginMessages.swift

@@ -124,6 +124,8 @@ public enum PluginMessage {
     case conformance
     case codeItem
     case `extension`
+    @_spi(ExperimentalLanguageFeature) case preamble
+    @_spi(ExperimentalLanguageFeature) case body
   }
 
   public struct SourceLocation: Codable {

Sources/SwiftSyntaxMacroExpansion/FunctionParameterUtils.swift

@@ -9,7 +9,8 @@ extension FunctionParameterSyntax {
   ///
   /// The parameter names for these three parameters are `a`, `b`, and `see`,
   /// respectively.
-  var parameterName: TokenSyntax? {
+  @_spi(Testing)
+  public var parameterName: TokenSyntax? {
     // If there were two names, the second is the parameter name.
     if let secondName {
       if secondName.text == "_" {

Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift

@@ -12,7 +12,7 @@
 
 import SwiftBasicFormat
 import SwiftSyntax
-@_spi(MacroExpansion) import SwiftSyntaxMacros
+@_spi(MacroExpansion) @_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacros
 
 public enum MacroRole {
   case expression
@@ -24,6 +24,8 @@ public enum MacroRole {
   case conformance
   case codeItem
   case `extension`
+  @_spi(ExperimentalLanguageFeature) case preamble
+  @_spi(ExperimentalLanguageFeature) case body
 }
 
 extension MacroRole {
@@ -38,18 +40,23 @@ extension MacroRole {
     case .conformance: return "ConformanceMacro"
     case .codeItem: return "CodeItemMacro"
     case .extension: return "ExtensionMacro"
+    case .preamble: return "PreambleMacro"
+    case .body: return "BodyMacro"
     }
   }
 }
 
 /// Simple diagnostic message
-private enum MacroExpansionError: Error, CustomStringConvertible {
+enum MacroExpansionError: Error, CustomStringConvertible {
   case unmatchedMacroRole(Macro.Type, MacroRole)
   case parentDeclGroupNil
   case declarationNotDeclGroup
   case declarationNotIdentified
+  case declarationHasNoBody
   case noExtendedTypeSyntax
   case noFreestandingMacroRoles(Macro.Type)
+  case moreThanOneBodyMacro
+  case preambleWithoutBody
 
   var description: String {
     switch self {
@@ -65,12 +72,20 @@ private enum MacroExpansionError: Error, CustomStringConvertible {
     case .declarationNotIdentified:
       return "declaration is not a 'Identified' syntax"
 
+    case .declarationHasNoBody:
+      return "declaration is not a type with an optional code block"
+
     case .noExtendedTypeSyntax:
       return "no extended type for extension macro"
 
     case .noFreestandingMacroRoles(let type):
       return "macro implementation type '\(type)' does not conform to any freestanding macro protocol"
 
+    case .moreThanOneBodyMacro:
+      return "function can not have more than one body macro applied to it"
+
+    case .preambleWithoutBody:
+      return "preamble macro cannot be applied to a function with no body"
     }
   }
 }
@@ -125,7 +140,7 @@ public func expandFreestandingMacro(
       expandedSyntax = Syntax(CodeBlockItemListSyntax(rewritten))
 
     case (.accessor, _), (.memberAttribute, _), (.member, _), (.peer, _), (.conformance, _), (.extension, _), (.expression, _), (.declaration, _),
-      (.codeItem, _):
+      (.codeItem, _), (.preamble, _), (.body, _):
       throw MacroExpansionError.unmatchedMacroRole(definition, macroRole)
     }
     return expandedSyntax.formattedExpansion(definition.formatMode, indentationWidth: indentationWidth)
@@ -288,6 +303,38 @@ public func expandAttachedMacroWithoutCollapsing<Context: MacroExpansionContext>
         $0.formattedExpansion(definition.formatMode, indentationWidth: indentationWidth)
       }
 
+    case (let attachedMacro as PreambleMacro.Type, .preamble):
+      guard let declToPass = Syntax(declarationNode).asProtocol(SyntaxProtocol.self) as? (DeclSyntaxProtocol & WithOptionalCodeBlockSyntax)
+      else {
+        // Compiler error: declaration must have a body.
+        throw MacroExpansionError.declarationHasNoBody
+      }
+
+      let preamble = try attachedMacro.expansion(
+        of: attributeNode,
+        providingPreambleFor: declToPass,
+        in: context
+      )
+      return preamble.map {
+        $0.formattedExpansion(definition.formatMode, indentationWidth: indentationWidth)
+      }
+
+    case (let attachedMacro as BodyMacro.Type, .body):
+      guard let declToPass = Syntax(declarationNode).asProtocol(SyntaxProtocol.self) as? (DeclSyntaxProtocol & WithOptionalCodeBlockSyntax)
+      else {
+        // Compiler error: declaration must have a body.
+        throw MacroExpansionError.declarationHasNoBody
+      }
+
+      let body = try attachedMacro.expansion(
+        of: attributeNode,
+        providingBodyFor: declToPass,
+        in: context
+      )
+      return body.map {
+        $0.formattedExpansion(definition.formatMode, indentationWidth: indentationWidth)
+      }
+
     default:
       throw MacroExpansionError.unmatchedMacroRole(definition, macroRole)
     }

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -15,7 +15,7 @@ import SwiftOperators
 @_spi(MacroExpansion) import SwiftParser
 import SwiftSyntax
 import SwiftSyntaxBuilder
-@_spi(MacroExpansion) import SwiftSyntaxMacros
+@_spi(MacroExpansion) @_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacros
 
 // MARK: - Public entry function
 
@@ -349,6 +349,72 @@ private func expandExtensionMacro(
   return "\(raw: indentedSource)"
 }
 
+/// Expand a preamble macro into a list of code items.
+private func expandPreambleMacro(
+  definition: PreambleMacro.Type,
+  attributeNode: AttributeSyntax,
+  attachedTo decl: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+  in context: some MacroExpansionContext,
+  indentationWidth: Trivia
+) -> CodeBlockItemListSyntax? {
+  guard
+    let expanded = expandAttachedMacro(
+      definition: definition,
+      macroRole: .preamble,
+      attributeNode: attributeNode.detach(
+        in: context,
+        foldingWith: .standardOperators
+      ),
+      declarationNode: DeclSyntax(decl.detach(in: context)),
+      parentDeclNode: nil,
+      extendedType: nil,
+      conformanceList: nil,
+      in: context,
+      indentationWidth: indentationWidth
+    )
+  else {
+    return []
+  }
+
+  // Match the indentation of the statements if we can, and put newlines around
+  // the preamble to separate it from the rest of the body.
+  let indentation = decl.body?.statements.indentationOfFirstLine ?? (decl.indentationOfFirstLine + indentationWidth)
+  let indentedSource = "\n" + expanded.indented(by: indentation) + "\n\n"
+  return "\(raw: indentedSource)"
+}
+
+private func expandBodyMacro(
+  definition: BodyMacro.Type,
+  attributeNode: AttributeSyntax,
+  attachedTo decl: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+  in context: some MacroExpansionContext,
+  indentationWidth: Trivia
+) -> CodeBlockSyntax? {
+  guard
+    let expanded = expandAttachedMacro(
+      definition: definition,
+      macroRole: .body,
+      attributeNode: attributeNode.detach(
+        in: context,
+        foldingWith: .standardOperators
+      ),
+      declarationNode: DeclSyntax(decl.detach(in: context)),
+      parentDeclNode: nil,
+      extendedType: nil,
+      conformanceList: nil,
+      in: context,
+      indentationWidth: indentationWidth
+    )
+  else {
+    return nil
+  }
+
+  // Wrap the body in braces.
+  let beforeBody = decl.body == nil ? " " : "";
+  let indentedSource = beforeBody + "{\n" + expanded.indented(by: decl.indentationOfFirstLine + indentationWidth) + "\n}\n"
+  return "\(raw: indentedSource)" as CodeBlockSyntax
+}
+
 // MARK: - MacroSystem
 
 /// Describes the kinds of errors that can occur within a macro system.
@@ -567,14 +633,21 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
     case .notAMacro:
       break
     }
-    if let declSyntax = node.as(DeclSyntax.self),
+    if var declSyntax = node.as(DeclSyntax.self),
       let attributedNode = node.asProtocol(WithAttributesSyntax.self),
       !attributedNode.attributes.isEmpty
     {
+      // Apply body and preamble macros.
+      if let nodeWithBody = node.asProtocol(WithOptionalCodeBlockSyntax.self),
+        let declNodeWithBody = nodeWithBody as? any DeclSyntaxProtocol & WithOptionalCodeBlockSyntax
+      {
+        declSyntax = DeclSyntax(visitBodyAndPreambleMacros(declNodeWithBody))
+      }
+
       // Visit the node, disabling the `visitAny` handling.
-      skipVisitAnyHandling.insert(node)
+      skipVisitAnyHandling.insert(Syntax(declSyntax))
       let visitedNode = self.visit(declSyntax)
-      skipVisitAnyHandling.remove(node)
+      skipVisitAnyHandling.remove(Syntax(declSyntax))
 
       let attributesToRemove = self.macroAttributes(attachedTo: visitedNode).map(\.attributeNode)
 
@@ -584,6 +657,66 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
     return nil
   }
 
+  /// Visit for both the body and preamble macros.
+  func visitBodyAndPreambleMacros<Node: DeclSyntaxProtocol & WithOptionalCodeBlockSyntax>(
+    _ node: Node
+  ) -> Node {
+    // Expand preamble macros into a set of code items.
+    let preamble = expandMacros(attachedTo: DeclSyntax(node), ofType: PreambleMacro.Type.self) { attributeNode, definition in
+      expandPreambleMacro(
+        definition: definition,
+        attributeNode: attributeNode,
+        attachedTo: node,
+        in: context,
+        indentationWidth: indentationWidth
+      )
+    }
+
+    // Expand body macro.
+    let expandedBodies = expandMacros(attachedTo: DeclSyntax(node), ofType: BodyMacro.Type.self) { attributeNode, definition in
+      expandBodyMacro(
+        definition: definition,
+        attributeNode: attributeNode,
+        attachedTo: node,
+        in: context,
+        indentationWidth: indentationWidth
+      ).map { [$0] }
+    }
+
+    // Dig out the body.
+    let body: CodeBlockSyntax
+    switch expandedBodies.count {
+    case 0 where preamble.isEmpty:
+      // Nothing changes
+      return node
+
+    case 0:
+      guard let existingBody = node.body else {
+        // Any leftover preamble statements have nowhere to go, complain and
+        // exit.
+        context.addDiagnostics(from: MacroExpansionError.preambleWithoutBody, node: node)
+
+        return node
+      }
+
+      body = existingBody
+
+    case 1:
+      body = expandedBodies[0]
+
+    default:
+      context.addDiagnostics(from: MacroExpansionError.moreThanOneBodyMacro, node: node)
+      body = expandedBodies[0]
+    }
+
+    // If there's no preamble, swap in the new body.
+    if preamble.isEmpty {
+      return node.with(\.body, body)
+    }
+
+    return node.with(\.body, body.with(\.statements, preamble + body.statements))
+  }
+
   override func visit(_ node: CodeBlockItemListSyntax) -> CodeBlockItemListSyntax {
     var newItems: [CodeBlockItemSyntax] = []
     func addResult(_ node: CodeBlockItemSyntax) {

Sources/SwiftSyntaxMacros/CMakeLists.txt

@@ -9,6 +9,7 @@
 add_swift_syntax_library(SwiftSyntaxMacros
   MacroProtocols/AccessorMacro.swift
   MacroProtocols/AttachedMacro.swift
+  MacroProtocols/BodyMacro.swift
   MacroProtocols/CodeItemMacro.swift
   MacroProtocols/DeclarationMacro.swift
   MacroProtocols/ExpressionMacro.swift
@@ -19,6 +20,7 @@ add_swift_syntax_library(SwiftSyntaxMacros
   MacroProtocols/MemberAttributeMacro.swift
   MacroProtocols/MemberMacro.swift
   MacroProtocols/PeerMacro.swift
+  MacroProtocols/PreambleMacro.swift
 
   AbstractSourceLocation.swift
   MacroExpansionContext.swift

Sources/SwiftSyntaxMacros/MacroProtocols/BodyMacro.swift

@@ -0,0 +1,26 @@
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+
+/// Describes a macro that can create the body for a function that does not
+/// have one.
+@_spi(ExperimentalLanguageFeature)
+public protocol BodyMacro: AttachedMacro {
+  /// Expand a macro described by the given custom attribute and
+  /// attached to the given declaration and evaluated within a
+  /// particular expansion context.
+  ///
+  /// The macro expansion can introduce a body for the given function.
+  static func expansion(
+    of node: AttributeSyntax,
+    providingBodyFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax]
+}

Sources/SwiftSyntaxMacros/MacroProtocols/PreambleMacro.swift

@@ -0,0 +1,28 @@
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+
+/// Describes a macro that can introduce "preamble" code into an existing
+/// function body.
+@_spi(ExperimentalLanguageFeature)
+public protocol PreambleMacro: AttachedMacro {
+  /// Expand a macro described by the given custom attribute and
+  /// attached to the given declaration and evaluated within a
+  /// particular expansion context.
+  ///
+  /// The macro expansion can introduce code items that form a preamble to
+  /// the body of the given function. The code items produced by this macro
+  /// expansion will be inserted at the beginning of the function body.
+  static func expansion(
+    of node: AttributeSyntax,
+    providingPreambleFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax]
+}