[Macros] Introduce conformsTo in assertMacroExpansion

Allows testing extension macros completely

Author: soumyamahunt

Sources/SwiftSyntaxMacroExpansion/MacroSpec.swift

@@ -0,0 +1,21 @@
+import SwiftSyntax
+import SwiftSyntaxMacros
+
+/// The specification for a macro.
+public struct MacroSpec {
+  /// The type of macro.
+  let type: Macro.Type
+  /// The list of types macro needs to add conformances.
+  let conformances: [TypeSyntax]
+
+  /// 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 macro needs to add conformances.
+  public init(type: Macro.Type, conformances: [TypeSyntax] = []) {
+    self.type = type
+    self.conformances = conformances
+  }
+}

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -25,11 +25,23 @@ extension SyntaxProtocol {
     macros: [String: Macro.Type],
     in context: some MacroExpansionContext,
     indentationWidth: Trivia? = nil
+  ) -> Syntax {
+    let specs = Dictionary(uniqueKeysWithValues: macros.map { ($0.key, MacroSpec(type: $0.value)) })
+    return self.expand(macroSpecs: specs, in: context, indentationWidth: indentationWidth)
+  }
+
+  /// Expand all uses of the given set of macros with specifications within this syntax node.
+  public func expand(
+    macroSpecs: [String: MacroSpec],
+    in context: some MacroExpansionContext,
+    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.type, name: macroName)
+      let conformedTypes = InheritedTypeListSyntax(macroSpec.conformances.map { InheritedTypeSyntax(type: $0) })
+      try! system.add(conformedTypes, name: macroName)
     }
 
     let applier = MacroApplication(
@@ -316,6 +328,7 @@ private func expandExtensionMacro(
   definition: ExtensionMacro.Type,
   attributeNode: AttributeSyntax,
   attachedTo: DeclSyntax,
+  conformanceList: InheritedTypeListSyntax?,
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) throws -> CodeBlockItemListSyntax? {
@@ -336,7 +349,7 @@ private func expandExtensionMacro(
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: extendedType.detach(in: context),
-      conformanceList: [],
+      conformanceList: conformanceList ?? [],
       in: context,
       indentationWidth: indentationWidth
     )
@@ -355,11 +368,14 @@ private func expandExtensionMacro(
 enum MacroSystemError: Error {
   /// Indicates that a macro with the given name has already been defined.
   case alreadyDefined(new: Macro.Type, existing: Macro.Type)
+  /// Indicates that protocol conformances for a macro with the given name has already been defined.
+  case alreadyConforming(new: InheritedTypeListSyntax, existing: InheritedTypeListSyntax)
 }
 
 /// A system of known macros that can be expanded syntactically
 struct MacroSystem {
   var macros: [String: Macro.Type] = [:]
+  var conformanceMap: [String: InheritedTypeListSyntax] = [:]
 
   /// Create an empty macro system.
   init() {}
@@ -375,10 +391,26 @@ struct MacroSystem {
     macros[name] = macro
   }
 
+  /// Add protocol conformances for a macro to the system.
+  ///
+  /// Throws an error if there is already conformances for a macro with this name.
+  mutating func add(_ conformanceList: InheritedTypeListSyntax, name: String) throws {
+    if let knownConformanceList = conformanceMap[name] {
+      throw MacroSystemError.alreadyConforming(new: conformanceList, existing: knownConformanceList)
+    }
+
+    conformanceMap[name] = conformanceList
+  }
+
   /// Look for a macro with the given name.
   func lookup(_ macroName: String) -> Macro.Type? {
     return macros[macroName]
   }
+
+  /// Look for protocol conformances of a macro with the given name.
+  func conformaces(forMacro macroName: String) -> InheritedTypeListSyntax? {
+    return conformanceMap[macroName]
+  }
 }
 
 // MARK: - MacroApplication
@@ -729,12 +761,13 @@ 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 macro attribute, the macro definition and optional
+  /// conformance protocols list attached to `decl`.
   ///
   /// The macros must be registered in `macroSystem`.
   private func macroAttributes(
     attachedTo decl: DeclSyntax
-  ) -> [(attributeNode: AttributeSyntax, definition: Macro.Type)] {
+  ) -> [(attributeNode: AttributeSyntax, definition: Macro.Type, conformanceList: InheritedTypeListSyntax?)] {
     guard let attributedNode = decl.asProtocol(WithAttributesSyntax.self) else {
       return []
     }
@@ -747,22 +780,22 @@ extension MacroApplication {
         return nil
       }
 
-      return (attribute, macro)
+      return (attribute, macro, macroSystem.conformaces(forMacro: attributeName))
     }
   }
 
-  /// Get pairs of a macro attribute and the macro definition attached to `decl`
-  /// matching `ofType` macro type.
+  /// Get macro attribute, the macro definition and optional 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
+      .compactMap { (attributeNode: AttributeSyntax, definition: Macro.Type, conformanceList: InheritedTypeListSyntax?) in
         if let macroType = definition as? MacroType {
-          return (attributeNode, macroType)
+          return (attributeNode, macroType, conformanceList)
         } else {
           return nil
         }
@@ -778,13 +811,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 {
@@ -802,7 +835,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,
@@ -822,7 +855,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,
@@ -837,11 +870,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: context,
         indentationWidth: indentationWidth
       )
@@ -850,7 +884,7 @@ 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,
@@ -870,7 +904,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,

Sources/SwiftSyntaxMacrosTestSupport/Assertions.swift

@@ -271,6 +271,50 @@ public func assertMacroExpansion(
   indentationWidth: Trivia = .spaces(4),
   file: StaticString = #file,
   line: UInt = #line
+) {
+  let specs = Dictionary(uniqueKeysWithValues: macros.map { ($0.key, MacroSpec(type: $0.value)) })
+  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., `"stringify"`) to specification
+///     (e.g., `StringifyMacro.self`).
+///   - 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.
+///   - testFileName: The name of the test file name to use.
+///   - indentationWidth: The indentation width used in the expansion.
+public func assertMacroExpansion(
+  _ originalSource: String,
+  expandedSource expectedExpandedSource: String,
+  diagnostics: [DiagnosticSpec] = [],
+  macroSpecs: [String: MacroSpec],
+  applyFixIts: [String]? = nil,
+  fixedSource expectedFixedSource: String? = nil,
+  testModuleName: String = "TestModule",
+  testFileName: String = "test.swift",
+  indentationWidth: Trivia = .spaces(4),
+  file: StaticString = #file,
+  line: UInt = #line
 ) {
   // Parse the original source file.
   let origSourceFile = Parser.parse(source: originalSource)
@@ -280,7 +324,7 @@ public func assertMacroExpansion(
     sourceFiles: [origSourceFile: .init(moduleName: testModuleName, fullFilePath: testFileName)]
   )
 
-  let expandedSourceFile = origSourceFile.expand(macros: macros, in: context, indentationWidth: indentationWidth)
+  let expandedSourceFile = origSourceFile.expand(macroSpecs: macroSpecs, in: context, indentationWidth: indentationWidth)
   let diags = ParseDiagnosticsGenerator.diagnostics(for: expandedSourceFile)
   if !diags.isEmpty {
     XCTFail(

Tests/SwiftSyntaxMacroExpansionTest/ExtensionMacroTests.swift

@@ -128,4 +128,85 @@ final class ExtensionMacroTests: XCTestCase {
       ]
     )
   }
+
+  func testConditionalExtensionExpansion() {
+    struct CodableExtensionMacro: ExtensionMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        attachedTo: some DeclGroupSyntax,
+        providingExtensionsOf type: some TypeSyntaxProtocol,
+        conformingTo protocols: [TypeSyntax],
+        in context: some MacroExpansionContext
+      ) throws -> [ExtensionDeclSyntax] {
+        let extensions: [ExtensionDeclSyntax] = protocols.compactMap { `protocol` in
+          let decl = """
+            extension \(type.trimmed): \(`protocol`) {}
+            """ as DeclSyntax
+          return decl.as(ExtensionDeclSyntax.self)
+        }
+
+        return extensions
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      @AddCodableExtensions
+      struct MyType {
+      }
+      """,
+      expandedSource: """
+
+        struct MyType {
+        }
+
+        extension MyType: Decodable {
+        }
+
+        extension MyType: Encodable {
+        }
+        """,
+      macroSpecs: ["AddCodableExtensions": MacroSpec(type: CodableExtensionMacro.self, conformances: ["Decodable", "Encodable"])],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      struct Wrapper {
+        @AddCodableExtensions
+        struct MyType {
+        }
+      }
+      """,
+      expandedSource: """
+        struct Wrapper {
+          struct MyType {
+          }
+        }
+
+        extension MyType: Encodable {
+        }
+        """,
+      macroSpecs: ["AddCodableExtensions": MacroSpec(type: CodableExtensionMacro.self, conformances: ["Encodable"])],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      struct Wrapper {
+        @AddCodableExtensions
+        struct MyType {
+        }
+      }
+      """,
+      expandedSource: """
+        struct Wrapper {
+          struct MyType {
+          }
+        }
+        """,
+      macros: ["AddCodableExtensions": CodableExtensionMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
 }