swiftlang · ahoppen · Feb 13, 2024 · Nov 3, 2023
@@ -49,6 +49,11 @@
   - Description: `TriviaPiece` now has a computed property `isComment` that returns `true` if the trivia piece is a comment.
   - Pull Request: https://github.com/apple/swift-syntax/pull/2469
 
+- New `assertMacroExpansion` API with option to specify macro specifications with `macroSpecs` argument
+  - Description: `macroSpecs` can have additional specifications like conformances provided by member or extension macro that can be used for macro expansion.
+  - Issue: https://github.com/apple/swift-syntax/issues/2031
+  - Pull Request: https://github.com/apple/swift-syntax/pull/2327
+
 ## API Behavior Changes
 
 ## Deprecations

@@ -5,6 +5,7 @@ add_swift_syntax_library(SwiftSyntaxMacroExpansion
   MacroExpansion.swift
   MacroExpansionDiagnosticMessages.swift
   MacroReplacement.swift
+  MacroSpec.swift
   MacroSystem.swift
 )
 

@@ -0,0 +1,44 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// 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 SwiftSyntaxMacros
+
+/// The information of a macro declaration, to be used with `assertMacroExpansion`.
+///
+/// In addition to specifying the macro’s type, this allows the specification of conformances that will be passed to the macro’s `expansion` function.
+public struct MacroSpec {
+  /// The type of macro.
+  let type: Macro.Type
+  /// The list of types for which the macro needs to add conformances.
+  let conformances: [TypeSyntax]
+
+  /// An `InheritedTypeListSytnax` containing all the types for which the macro needs to add conformances.
+  var inheritedTypeList: InheritedTypeListSyntax {
+    return InheritedTypeListSyntax {
+      for conformance in conformances {
+        InheritedTypeSyntax(type: conformance)
+      }
+    }
+  }
+
+  /// Creates a new specification from provided macro type
+  /// and optional list of generated conformances.
+  ///
+  /// - Parameters:
+  ///   - type: The type of macro.
+  ///   - conformances: The list of types that will be passed to the macro’s `expansion` function.
+  public init(type: Macro.Type, conformances: [TypeSyntax] = []) {
+    self.type = type
+    self.conformances = conformances
+  }
+}
@@ -35,15 +35,30 @@ extension SyntaxProtocol {
   }
 
   /// Expand all uses of the given set of macros within this syntax node.
+  /// - SeeAlso: ``expand(macroSpecs:contextGenerator:indentationWidth:)``
+  ///   to also specify the list of conformances passed to the macro expansion.
   public func expand<Context: MacroExpansionContext>(
     macros: [String: Macro.Type],
     contextGenerator: @escaping (Syntax) -> Context,
     indentationWidth: Trivia? = nil
+  ) -> Syntax {
+    return expand(
+      macroSpecs: macros.mapValues { MacroSpec(type: $0) },
+      contextGenerator: contextGenerator,
+      indentationWidth: indentationWidth
+    )
+  }
+
+  /// Expand all uses of the given set of macros with specifications within this syntax node.
+  public func expand<Context: MacroExpansionContext>(
+    macroSpecs: [String: MacroSpec],
+    contextGenerator: @escaping (Syntax) -> Context,
+    indentationWidth: Trivia? = nil
   ) -> Syntax {
     // Build the macro system.
     var system = MacroSystem()
-    for (macroName, macroType) in macros {
-      try! system.add(macroType, name: macroName)
+    for (macroName, macroSpec) in macroSpecs {
+      try! system.add(macroSpec, name: macroName)
     }
 
     let applier = MacroApplication(
@@ -142,6 +157,7 @@ private func expandMemberMacro(
   definition: MemberMacro.Type,
   attributeNode: AttributeSyntax,
   attachedTo: DeclSyntax,
+  conformanceList: InheritedTypeListSyntax,
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) throws -> MemberBlockItemListSyntax? {
@@ -153,7 +169,7 @@ private func expandMemberMacro(
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: nil,
-      conformanceList: nil,
+      conformanceList: conformanceList,
       in: context,
       indentationWidth: indentationWidth
     )
@@ -330,6 +346,7 @@ private func expandExtensionMacro(
   definition: ExtensionMacro.Type,
   attributeNode: AttributeSyntax,
   attachedTo: DeclSyntax,
+  conformanceList: InheritedTypeListSyntax,
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) throws -> CodeBlockItemListSyntax? {
@@ -349,7 +366,7 @@ private func expandExtensionMacro(
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: extendedType.detach(in: context),
-      conformanceList: [],
+      conformanceList: conformanceList,
       in: context,
       indentationWidth: indentationWidth
     )
@@ -442,24 +459,24 @@ enum MacroSystemError: Error {
 
 /// A system of known macros that can be expanded syntactically
 struct MacroSystem {
-  var macros: [String: Macro.Type] = [:]
+  var macros: [String: MacroSpec] = [:]
 
   /// Create an empty macro system.
   init() {}
 
-  /// Add a macro to the system.
+  /// Add a macro specification to the system.
   ///
   /// Throws an error if there is already a macro with this name.
-  mutating func add(_ macro: Macro.Type, name: String) throws {
-    if let knownMacro = macros[name] {
-      throw MacroSystemError.alreadyDefined(new: macro, existing: knownMacro)
+  mutating func add(_ macroSpec: MacroSpec, name: String) throws {
+    if let knownMacroSpec = macros[name] {
+      throw MacroSystemError.alreadyDefined(new: macroSpec.type, existing: knownMacroSpec.type)
     }
 
-    macros[name] = macro
+    macros[name] = macroSpec
   }
 
-  /// Look for a macro with the given name.
-  func lookup(_ macroName: String) -> Macro.Type? {
+  /// Look for a macro specification with the given name.
+  func lookup(_ macroName: String) -> MacroSpec? {
     return macros[macroName]
   }
 }
@@ -680,7 +697,7 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
     _ 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
+    let preamble = expandMacros(attachedTo: DeclSyntax(node), ofType: PreambleMacro.Type.self) { attributeNode, definition, _ in
       expandPreambleMacro(
         definition: definition,
         attributeNode: attributeNode,
@@ -691,7 +708,7 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
     }
 
     // Expand body macro.
-    let expandedBodies = expandMacros(attachedTo: DeclSyntax(node), ofType: BodyMacro.Type.self) { attributeNode, definition in
+    let expandedBodies = expandMacros(attachedTo: DeclSyntax(node), ofType: BodyMacro.Type.self) { attributeNode, definition, _ in
       expandBodyMacro(
         definition: definition,
         attributeNode: attributeNode,
@@ -900,40 +917,40 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
 // MARK:  Attached macro expansions.
 
 extension MacroApplication {
-  /// Get pairs of a macro attribute and the macro definition attached to `decl`.
+  /// Get pairs of a macro attribute and the macro specification attached to `decl`.
   ///
   /// The macros must be registered in `macroSystem`.
   private func macroAttributes(
     attachedTo decl: DeclSyntax
-  ) -> [(attributeNode: AttributeSyntax, definition: Macro.Type)] {
+  ) -> [(attributeNode: AttributeSyntax, spec: MacroSpec)] {
     guard let attributedNode = decl.asProtocol(WithAttributesSyntax.self) else {
       return []
     }
 
     return attributedNode.attributes.compactMap {
       guard case let .attribute(attribute) = $0,
         let attributeName = attribute.attributeName.as(IdentifierTypeSyntax.self)?.name.text,
-        let macro = macroSystem.lookup(attributeName)
+        let macroSpec = macroSystem.lookup(attributeName)
       else {
         return nil
       }
 
-      return (attribute, macro)
+      return (attribute, macroSpec)
     }
   }
 
-  /// Get pairs of a macro attribute and the macro definition attached to `decl`
-  /// matching `ofType` macro type.
+  /// Get a list of the macro attribute, the macro definition and the conformance
+  /// protocols list attached to `decl` matching `ofType` macro type.
   ///
   /// The macros must be registered in `macroSystem`.
   private func macroAttributes<MacroType>(
     attachedTo decl: DeclSyntax,
     ofType: MacroType.Type
-  ) -> [(attributeNode: AttributeSyntax, definition: MacroType)] {
+  ) -> [(attributeNode: AttributeSyntax, definition: MacroType, conformanceList: InheritedTypeListSyntax)] {
     return macroAttributes(attachedTo: decl)
-      .compactMap { (attributeNode: AttributeSyntax, definition: Macro.Type) in
-        if let macroType = definition as? MacroType {
-          return (attributeNode, macroType)
+      .compactMap { (attributeNode: AttributeSyntax, spec: MacroSpec) in
+        if let macroType = spec.type as? MacroType {
+          return (attributeNode, macroType, spec.inheritedTypeList)
         } else {
           return nil
         }
@@ -949,13 +966,13 @@ extension MacroApplication {
   >(
     attachedTo decl: DeclSyntax,
     ofType: MacroType.Type,
-    expandMacro: (_ attributeNode: AttributeSyntax, _ definition: MacroType) throws -> ExpanedNodeCollection?
+    expandMacro: (_ attributeNode: AttributeSyntax, _ definition: MacroType, _ conformanceList: InheritedTypeListSyntax) throws -> ExpanedNodeCollection?
   ) -> [ExpandedNode] {
     var result: [ExpandedNode] = []
 
     for macroAttribute in macroAttributes(attachedTo: decl, ofType: ofType) {
       do {
-        if let expanded = try expandMacro(macroAttribute.attributeNode, macroAttribute.definition) {
+        if let expanded = try expandMacro(macroAttribute.attributeNode, macroAttribute.definition, macroAttribute.conformanceList) {
           result += expanded
         }
       } catch {
@@ -973,7 +990,7 @@ extension MacroApplication {
   ///
   /// - Returns: The macro-synthesized peers
   private func expandMemberDeclPeers(of decl: DeclSyntax) -> [MemberBlockItemSyntax] {
-    return expandMacros(attachedTo: decl, ofType: PeerMacro.Type.self) { attributeNode, definition in
+    return expandMacros(attachedTo: decl, ofType: PeerMacro.Type.self) { attributeNode, definition, conformanceList in
       return try expandPeerMacroMember(
         definition: definition,
         attributeNode: attributeNode,
@@ -993,7 +1010,7 @@ extension MacroApplication {
   ///
   /// - Returns: The macro-synthesized peers
   private func expandCodeBlockPeers(of decl: DeclSyntax) -> [CodeBlockItemSyntax] {
-    return expandMacros(attachedTo: decl, ofType: PeerMacro.Type.self) { attributeNode, definition in
+    return expandMacros(attachedTo: decl, ofType: PeerMacro.Type.self) { attributeNode, definition, conformanceList in
       return try expandPeerMacroCodeItem(
         definition: definition,
         attributeNode: attributeNode,
@@ -1008,11 +1025,12 @@ extension MacroApplication {
   ///
   /// - Returns: The macro-synthesized extensions
   private func expandExtensions(of decl: DeclSyntax) -> [CodeBlockItemSyntax] {
-    return expandMacros(attachedTo: decl, ofType: ExtensionMacro.Type.self) { attributeNode, definition in
+    return expandMacros(attachedTo: decl, ofType: ExtensionMacro.Type.self) { attributeNode, definition, conformanceList in
       return try expandExtensionMacro(
         definition: definition,
         attributeNode: attributeNode,
         attachedTo: decl,
+        conformanceList: conformanceList,
         in: contextGenerator(Syntax(decl)),
         indentationWidth: indentationWidth
       )
@@ -1021,11 +1039,12 @@ extension MacroApplication {
 
   /// Expand all 'member' macros attached to `decl`.
   private func expandMembers(of decl: DeclSyntax) -> [MemberBlockItemSyntax] {
-    return expandMacros(attachedTo: decl, ofType: MemberMacro.Type.self) { attributeNode, definition in
+    return expandMacros(attachedTo: decl, ofType: MemberMacro.Type.self) { attributeNode, definition, conformanceList in
       return try expandMemberMacro(
         definition: definition,
         attributeNode: attributeNode,
         attachedTo: decl,
+        conformanceList: conformanceList,
         in: contextGenerator(Syntax(decl)),
         indentationWidth: indentationWidth
       )
@@ -1041,7 +1060,7 @@ extension MacroApplication {
     of decl: DeclSyntax,
     parentDecl: DeclSyntax
   ) -> [AttributeListSyntax.Element] {
-    return expandMacros(attachedTo: parentDecl, ofType: MemberAttributeMacro.Type.self) { attributeNode, definition in
+    return expandMacros(attachedTo: parentDecl, ofType: MemberAttributeMacro.Type.self) { attributeNode, definition, conformanceList in
       return try expandMemberAttributeMacro(
         definition: definition,
         attributeNode: attributeNode,
@@ -1147,7 +1166,7 @@ extension MacroApplication {
     expandMacro: (_ macro: Macro.Type, _ node: any FreestandingMacroExpansionSyntax) throws -> ExpandedMacroType?
   ) -> MacroExpansionResult<ExpandedMacroType> {
     guard let node,
-      let macro = macroSystem.lookup(node.macroName.text)
+      let macro = macroSystem.lookup(node.macroName.text)?.type
     else {
       return .notAMacro
     }

@@ -299,6 +299,52 @@ func assertDiagnostic(
 ///   - macros: The macros that should be expanded, provided as a dictionary
 ///     mapping macro names (e.g., `"stringify"`) to implementation types
 ///     (e.g., `StringifyMacro.self`).
+///   - testModuleName: The name of the test module to use.
+///   - testFileName: The name of the test file name to use.
+///   - indentationWidth: The indentation width used in the expansion.
+///
+/// - SeeAlso: ``assertMacroExpansion(_:expandedSource:diagnostics:macroSpecs:applyFixIts:fixedSource:testModuleName:testFileName:indentationWidth:file:line:)``
+///   to also specify the list of conformances passed to the macro expansion.
+public func assertMacroExpansion(
+  _ originalSource: String,
+  expandedSource expectedExpandedSource: String,
+  diagnostics: [DiagnosticSpec] = [],
+  macros: [String: Macro.Type],
+  applyFixIts: [String]? = nil,
+  fixedSource expectedFixedSource: String? = nil,
+  testModuleName: String = "TestModule",
+  testFileName: String = "test.swift",
+  indentationWidth: Trivia = .spaces(4),
+  file: StaticString = #file,
+  line: UInt = #line
+) {
+  let specs = macros.mapValues { MacroSpec(type: $0) }
+  assertMacroExpansion(
+    originalSource,
+    expandedSource: expectedExpandedSource,
+    diagnostics: diagnostics,
+    macroSpecs: specs,
+    applyFixIts: applyFixIts,
+    fixedSource: expectedFixedSource,
+    testModuleName: testModuleName,
+    testFileName: testFileName,
+    indentationWidth: indentationWidth
+  )
+}
+
+/// Assert that expanding the given macros in the original source produces
+/// the given expanded source code.
+///
+/// - Parameters:
+///   - originalSource: The original source code, which is expected to contain
+///     macros in various places (e.g., `#stringify(x + y)`).
+///   - expectedExpandedSource: The source code that we expect to see after
+///     performing macro expansion on the original source.
+///   - diagnostics: The diagnostics when expanding any macro
+///   - macroSpecs: The macros that should be expanded, provided as a dictionary
+///     mapping macro names (e.g., `"CodableMacro"`) to specification with macro type
+///     (e.g., `CodableMacro.self`) and a list of conformances macro provides
+///     (e.g., `["Decodable", "Encodable"]`).
 ///   - applyFixIts: If specified, filters the Fix-Its that are applied to generate `fixedSource` to only those whose message occurs in this array. If `nil`, all Fix-Its from the diagnostics are applied.
 ///   - fixedSource: If specified, asserts that the source code after applying Fix-Its matches this string.
 ///   - testModuleName: The name of the test module to use.
@@ -308,7 +354,7 @@ public func assertMacroExpansion(
   _ originalSource: String,
   expandedSource expectedExpandedSource: String,
   diagnostics: [DiagnosticSpec] = [],
-  macros: [String: Macro.Type],
+  macroSpecs: [String: MacroSpec],
   applyFixIts: [String]? = nil,
   fixedSource expectedFixedSource: String? = nil,
   testModuleName: String = "TestModule",
@@ -329,7 +375,7 @@ public func assertMacroExpansion(
     return BasicMacroExpansionContext(sharingWith: context, lexicalContext: syntax.allMacroLexicalContexts())
   }
 
-  let expandedSourceFile = origSourceFile.expand(macros: macros, contextGenerator: contextGenerator, indentationWidth: indentationWidth)
+  let expandedSourceFile = origSourceFile.expand(macroSpecs: macroSpecs, contextGenerator: contextGenerator, indentationWidth: indentationWidth)
   let diags = ParseDiagnosticsGenerator.diagnostics(for: expandedSourceFile)
   if !diags.isEmpty {
     XCTFail(