[Macros] Support code item macro expansion (swiftlang#1442)

Author: rxwei

Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -58,6 +58,13 @@ extension CompilerPluginMessageHandler {
  let rewritten = try _openExistential(macroSyntax, do: _expand)
  expandedSource = CodeBlockItemListSyntax(rewritten.map { CodeBlockItemSyntax(item: .decl($0)) }).description
 
+ case let codeItemMacroDef as CodeItemMacro.Type:
+ func _expand<Node: FreestandingMacroExpansionSyntax>(node: Node) throws -> [CodeBlockItemSyntax] {
+ try codeItemMacroDef.expansion(of: node, in: context)
+ }
+ let rewritten = try _openExistential(macroSyntax, do: _expand)
+ expandedSource = CodeBlockItemListSyntax(rewritten).description
+
  default:
  throw MacroExpansionError.unmathedMacroRole
  }

Sources/SwiftCompilerPluginMessageHandling/PluginMessages.swift

@@ -85,12 +85,13 @@ internal enum PluginToHostMessage: Codable {
 
  enum MacroRole: String, Codable {
  case expression
- case freeStandingDeclaration
+ case declaration
  case accessor
  case memberAttribute
  case member
  case peer
  case conformance
+ case codeItem
  }
 
  struct SourceLocation: Codable {

Sources/SwiftSyntaxMacros/MacroSystem.swift

@@ -132,7 +132,13 @@ class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
  let macro = macroSystem.macros[exprExpansion.macro.text]
  {
  do {
- if let macro = macro as? DeclarationMacro.Type {
+ if let macro = macro as? CodeItemMacro.Type {
+ let expandedItemList = try macro.expansion(
+ of: exprExpansion,
+ in: context
+ )
+ newItems.append(contentsOf: expandedItemList)
+ } else if let macro = macro as? DeclarationMacro.Type {
  let expandedItemList = try macro.expansion(
  of: exprExpansion,
  in: context

Tests/SwiftSyntaxMacrosTest/MacroSystemTests.swift

@@ -618,6 +618,59 @@ extension CustomTypeWrapperMacro: AccessorMacro {
  }
 }
 
+public struct UnwrapMacro: CodeItemMacro {
+ public static func expansion(
+ of node: some FreestandingMacroExpansionSyntax,
+ in context: some MacroExpansionContext
+ ) throws -> [CodeBlockItemSyntax] {
+ guard !node.argumentList.isEmpty else {
+ throw CustomError.message("'#unwrap' requires arguments")
+ }
+ let errorThrower = node.trailingClosure
+ let identifiers = try node.argumentList.map { argument in
+ guard let tupleElement = argument.as(TupleExprElementSyntax.self),
+ let identifierExpr = tupleElement.expression.as(IdentifierExprSyntax.self)
+ else {
+ throw CustomError.message("Arguments must be identifiers")
+ }
+ return identifierExpr.identifier
+ }
+
+ func elseBlock(_ token: TokenSyntax) -> CodeBlockSyntax {
+ let expr: ExprSyntax
+ if let errorThrower {
+ expr = """
+ \(errorThrower)("\(raw: token.text)")
+ """
+ } else {
+ expr = """
+ fatalError("'\(raw: token.text)' is nil")
+ """
+ }
+ return .init(
+ statements: .init([
+ .init(
+ leadingTrivia: " ",
+ item: .expr(expr),
+ trailingTrivia: " "
+ )
+ ])
+ )
+ }
+
+ return identifiers.map { identifier in
+ CodeBlockItemSyntax(
+ item: CodeBlockItemSyntax.Item.stmt(
+ """
+
+ guard let \(raw: identifier.text) else \(elseBlock(identifier))
+ """
+ )
+ )
+ }
+ }
+}
+
 // MARK: Assertion helper functions
 
 /// Assert that expanding the given macros in the original source produces
@@ -685,6 +738,7 @@ public let testMacros: [String: Macro.Type] = [
  "wrapAllProperties": WrapAllProperties.self,
  "wrapStoredProperties": WrapStoredProperties.self,
  "customTypeWrapper": CustomTypeWrapperMacro.self,
+ "unwrap": UnwrapMacro.self,
 ]
 
 final class MacroSystemTests: XCTestCase {
@@ -976,4 +1030,36 @@ final class MacroSystemTests: XCTestCase {
  )
 
  }
+
+ func testUnwrap() {
+ assertMacroExpansion(
+ macros: testMacros,
+ #"""
+ let x: Int? = 1
+ let y: Int? = nil
+ let z: Int? = 3
+ #unwrap(x, y, z)
+ #unwrap(x, y, z) {
+ fatalError("nil is \\($0)")
+ }
+ """#,
+ #"""
+ let x: Int? = 1
+ let y: Int? = nil
+ let z: Int? = 3
+ guard let x else { fatalError("'x' is nil") }
+ guard let y else { fatalError("'y' is nil") }
+ guard let z else { fatalError("'z' is nil") }
+ guard let x else { {
+ fatalError("nil is \\($0)")
+ }("x") }
+ guard let y else { {
+ fatalError("nil is \\($0)")
+ }("y") }
+ guard let z else { {
+ fatalError("nil is \\($0)")
+ }("z") }
+ """#
+ )
+ }
 }