[Macros] Implement peer declaration macros with expansion

Sources/_SwiftSyntaxMacros/CMakeLists.txt

@@ -13,6 +13,7 @@ add_swift_host_library(_SwiftSyntaxMacros
   Macro.swift
   MacroExpansionContext.swift
   MacroSystem.swift
+  PeerDeclarationMacro.swift
   Syntax+MacroEvaluation.swift
 )
 

Sources/_SwiftSyntaxMacros/ExpressionMacro.swift

@@ -12,7 +12,6 @@
 
 import SwiftSyntax
 import SwiftParser
-import SwiftDiagnostics
 
 /// Describes a macro that is explicitly expanded as an expression.
 public protocol ExpressionMacro: Macro {

Sources/_SwiftSyntaxMacros/FreestandingDeclarationMacro.swift

@@ -9,7 +9,6 @@
 
 import SwiftSyntax
 import SwiftParser
-import SwiftDiagnostics
 
 /// Describes a macro that forms declarations.
 public protocol FreestandingDeclarationMacro: DeclarationMacro {

Sources/_SwiftSyntaxMacros/MacroSystem.swift

@@ -120,6 +120,16 @@ class MacroApplication: SyntaxRewriter {
       // Recurse on the child node.
       let newItem = visit(item.item)
       newItems.append(item.withItem(newItem))
+
+      // Expand any peer declarations triggered by macros used as attributes.
+      if case let .decl(decl) = item.item {
+        let peers = expandPeers(of: decl)
+        newItems.append(
+          contentsOf: peers.map {
+            newDecl in CodeBlockItemSyntax(item: .decl(newDecl))
+          }
+        )
+      }
     }
 
     return CodeBlockItemListSyntax(newItems)
@@ -160,10 +170,97 @@ class MacroApplication: SyntaxRewriter {
       // Recurse on the child node.
       let newDecl = visit(item.decl)
       newItems.append(item.withDecl(newDecl))
+
+      // Expand any peer declarations triggered by macros used as attributes.
+      let peers = expandPeers(of: item.decl)
+      newItems.append(
+        contentsOf: peers.map {
+          newDecl in MemberDeclListItemSyntax(decl: newDecl)
+        }
+      )
     }
 
     return .init(newItems)
   }
+
+  override func visit(_ node: FunctionDeclSyntax) -> DeclSyntax {
+    let visitedNode = super.visit(node)
+
+    // FIXME: Generalize this to DeclSyntax, once we have attributes.
+    // Visit the node first.
+
+    guard let visitedFunc = visitedNode.as(FunctionDeclSyntax.self),
+      let attributes = visitedFunc.attributes
+    else {
+      return visitedNode
+    }
+
+    // Remove any attached attributes.
+    let newAttributes = attributes.filter {
+      guard case let .customAttribute(customAttr) = $0 else {
+        return true
+      }
+
+      guard let attributeName = customAttr.attributeName.as(SimpleTypeIdentifierSyntax.self)?.name.text,
+        let macro = macroSystem.macros[attributeName]
+      else {
+        return true
+      }
+
+      return !(macro is PeerDeclarationMacro.Type)
+    }
+
+    if newAttributes.isEmpty {
+      return DeclSyntax(visitedFunc.withAttributes(nil))
+    }
+
+    return DeclSyntax(visitedFunc.withAttributes(AttributeListSyntax(newAttributes)))
+  }
+}
+
+extension MacroApplication {
+  // If any of the custom attributes associated with the given declaration
+  // refer to "peer" declaration macros, expand them and return the resulting
+  // set of peer declarations.
+  private func expandPeers(of decl: DeclSyntax) -> [DeclSyntax] {
+    // Dig out the attribute list.
+    // FIXME: We should have a better way to get the attributes from any
+    // declaration.
+    guard
+      let attributes =
+        (decl.children(viewMode: .sourceAccurate).compactMap {
+          $0.as(AttributeListSyntax.self)
+        }).first
+    else {
+      return []
+    }
+
+    var peers: [DeclSyntax] = []
+    for attribute in attributes {
+      guard case let .customAttribute(customAttribute) = attribute,
+        let attributeName = customAttribute.attributeName.as(SimpleTypeIdentifierSyntax.self)?.name.text,
+        let macro = macroSystem.macros[attributeName],
+        let peerMacro = macro as? PeerDeclarationMacro.Type
+      else {
+        continue
+      }
+
+      do {
+        let newPeers = try peerMacro.expansion(of: customAttribute, attachedTo: decl, in: &context)
+        peers.append(contentsOf: newPeers)
+      } catch {
+        // Record the error
+        context.diagnose(
+          Diagnostic(
+            node: Syntax(attribute),
+            message: ThrownErrorDiagnostic(message: String(describing: error))
+          )
+        )
+      }
+    }
+
+    return peers
+  }
 }
 
 extension SyntaxProtocol {

Sources/_SwiftSyntaxMacros/PeerDeclarationMacro.swift

@@ -0,0 +1,25 @@
+//
+// 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
+import SwiftParser
+
+public protocol PeerDeclarationMacro: DeclarationMacro {
+  /// 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 "peer" declarations that sit alongside
+  /// the
+  static func expansion(
+    of node: CustomAttributeSyntax,
+    attachedTo declaration: DeclSyntax,
+    in context: inout MacroExpansionContext
+  ) throws -> [DeclSyntax]
+}

Tests/SwiftSyntaxMacrosTest/MacroSystemTests.swift

@@ -208,6 +208,117 @@ struct DefineBitwidthNumberedStructsMacro: FreestandingDeclarationMacro {
   }
 }
 
+public struct AddCompletionHandler: PeerDeclarationMacro {
+  public static func expansion(
+    of node: CustomAttributeSyntax,
+    attachedTo declaration: DeclSyntax,
+    in context: inout MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    // Only on functions at the moment. We could handle initializers as well
+    // with a bit of work.
+    guard let funcDecl = declaration.as(FunctionDeclSyntax.self) else {
+      throw CustomError.message("@addCompletionHandler only works on functions")
+    }
+
+    // This only makes sense for async functions.
+    if funcDecl.signature.asyncOrReasyncKeyword == nil {
+      throw CustomError.message(
+        "@addCompletionHandler requires an async function"
+      )
+    }
+
+    // Form the completion handler parameter.
+    let resultType: TypeSyntax? = funcDecl.signature.output?.returnType.withoutTrivia()
+
+    let completionHandlerParam =
+      FunctionParameterSyntax(
+        firstName: .identifier("completionHandler"),
+        colon: .colonToken(trailingTrivia: .space),
+        type: "(\(resultType ?? "")) -> Void" as TypeSyntax
+      )
+
+    // Add the completion handler parameter to the parameter list.
+    let parameterList = funcDecl.signature.input.parameterList
+    let newParameterList: FunctionParameterListSyntax
+    if let lastParam = parameterList.last {
+      // We need to add a trailing comma to the preceding list.
+      newParameterList = parameterList.removingLast()
+        .appending(
+          lastParam.withTrailingComma(
+            .commaToken(trailingTrivia: .space)
+          )
+        )
+        .appending(completionHandlerParam)
+    } else {
+      newParameterList = parameterList.appending(completionHandlerParam)
+    }
+
+    let callArguments: [String] = try parameterList.map { param in
+      guard let argName = param.secondName ?? param.firstName else {
+        throw CustomError.message(
+          "@addCompletionHandler argument must have a name"
+        )
+      }
+
+      if let paramName = param.firstName, paramName.text != "_" {
+        return "\(paramName.withoutTrivia()): \(argName.withoutTrivia())"
+      }
+
+      return "\(argName.withoutTrivia())"
+    }
+
+    let call: ExprSyntax =
+      "\(funcDecl.identifier)(\(raw: callArguments.joined(separator: ", ")))"
+
+    // FIXME: We should make CodeBlockSyntax ExpressibleByStringInterpolation,
+    // so that the full body could go here.
+    let newBody: ExprSyntax =
+      """
+
+        Task {
+          completionHandler(await \(call))
+        }
+
+      """
+
+    // Drop the @addCompletionHandler attribute from the new declaration.
+    let newAttributeList = AttributeListSyntax(
+      funcDecl.attributes?.filter {
+        guard case let .customAttribute(customAttr) = $0 else {
+          return true
+        }
+
+        return customAttr != node
+      } ?? []
+    )
+
+    let newFunc =
+      funcDecl
+      .withSignature(
+        funcDecl.signature
+          .withAsyncOrReasyncKeyword(nil)  // drop async
+          .withOutput(nil)  // drop result type
+          .withInput(  // add completion handler parameter
+            funcDecl.signature.input.withParameterList(newParameterList)
+              .withoutTrailingTrivia()
+          )
+      )
+      .withBody(
+        CodeBlockSyntax(
+          leftBrace: .leftBraceToken(leadingTrivia: .space),
+          statements: CodeBlockItemListSyntax(
+            [CodeBlockItemSyntax(item: .expr(newBody))]
+          ),
+          rightBrace: .rightBraceToken(leadingTrivia: .newline)
+        )
+      )
+      .withAttributes(newAttributeList)
+      .withLeadingTrivia(.newlines(2))
+
+    return [DeclSyntax(newFunc)]
+  }
+}
+
 // MARK: Assertion helper functions
 
 /// Assert that expanding the given macros in the original source produces
@@ -269,6 +380,7 @@ public let testMacros: [String: Macro.Type] = [
   "stringify": StringifyMacro.self,
   "myError": ErrorMacro.self,
   "bitwidthNumberedStructs": DefineBitwidthNumberedStructsMacro.self,
+  "addCompletionHandler": AddCompletionHandler.self,
 ]
 
 final class MacroSystemTests: XCTestCase {
@@ -386,4 +498,24 @@ final class MacroSystemTests: XCTestCase {
       """
     )
   }
+
+  func testAddCompletionHandler() {
+    AssertMacroExpansion(
+      macros: testMacros,
+      """
+      @addCompletionHandler
+      func f(a: Int, for b: String, _ value: Double) async -> String { }
+      """,
+      """
+
+      func f(a: Int, for b: String, _ value: Double) async -> String { }
+
+      func f(a: Int, for b: String, _ value: Double, completionHandler: (String) -> Void) {
+        Task {
+          completionHandler(await f(a: a, for: b, value))
+        }
+      }
+      """
+    )
+  }
 }