[Macros] Support code item macro expansion (#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") }
+      """#
+    )
+  }
 }