[Macros] Introduce conformsTo in assertMacroExpansion

Allows testing extension macros completely

Author: soumyamahunt

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -23,6 +23,7 @@ extension SyntaxProtocol {
   /// Expand all uses of the given set of macros within this syntax node.
   public func expand(
     macros: [String: Macro.Type],
+    conformsTo conformanceMap: [String: InheritedTypeListSyntax],
     in context: some MacroExpansionContext,
     indentationWidth: Trivia? = nil
   ) -> Syntax {
@@ -31,6 +32,9 @@ extension SyntaxProtocol {
     for (macroName, macroType) in macros {
       try! system.add(macroType, name: macroName)
     }
+    for (macroName, conformanceList) in conformanceMap {
+      try! system.add(conformanceList, name: macroName)
+    }
 
     let applier = MacroApplication(
       macroSystem: system,
@@ -316,6 +320,7 @@ private func expandExtensionMacro(
   definition: ExtensionMacro.Type,
   attributeNode: AttributeSyntax,
   attachedTo: DeclSyntax,
+  conformanceList: InheritedTypeListSyntax?,
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) throws -> CodeBlockItemListSyntax? {
@@ -336,7 +341,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 +360,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 +383,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
@@ -716,12 +740,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 []
     }
@@ -734,22 +759,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
         }
@@ -765,13 +790,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 {
@@ -789,7 +814,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,
@@ -809,7 +834,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,
@@ -824,11 +849,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
       )
@@ -837,7 +863,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,
@@ -857,7 +883,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

@@ -254,13 +254,17 @@ 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`).
+///   - conformanceMap: The protocols that maco should conform to, provided as
+///     a dictionary mapping macro names (e.g., `"AddCodable"`) to protocols names
+///     (e.g., `"Decodable"`, `"Encodable"` etc.).
 ///   - testModuleName: The name of the test module to use.
 ///   - testFileName: The name of the test file name to use.
 public func assertMacroExpansion(
   _ originalSource: String,
   expandedSource expectedExpandedSource: String,
   diagnostics: [DiagnosticSpec] = [],
   macros: [String: Macro.Type],
+  conformsTo conformanceMap: [String: [TypeSyntax]] = [:],
   testModuleName: String = "TestModule",
   testFileName: String = "test.swift",
   indentationWidth: Trivia = .spaces(4),
@@ -275,7 +279,10 @@ public func assertMacroExpansion(
     sourceFiles: [origSourceFile: .init(moduleName: testModuleName, fullFilePath: testFileName)]
   )
 
-  let expandedSourceFile = origSourceFile.expand(macros: macros, in: context, indentationWidth: indentationWidth)
+  let conformanceMap: [String: InheritedTypeListSyntax] = Dictionary(
+    uniqueKeysWithValues: conformanceMap.map { ($0.key, InheritedTypeListSyntax($0.value.map { InheritedTypeSyntax(type: $0) })) }
+  )
+  let expandedSourceFile = origSourceFile.expand(macros: macros, conformsTo: conformanceMap, in: context, indentationWidth: indentationWidth)
   let diags = ParseDiagnosticsGenerator.diagnostics(for: expandedSourceFile)
   if !diags.isEmpty {
     XCTFail(

Tests/SwiftSyntaxMacroExpansionTest/ExtensionMacroTests.swift

@@ -128,4 +128,87 @@ 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 {
+        }
+        """,
+      macros: ["AddCodableExtensions": CodableExtensionMacro.self],
+      conformsTo: ["AddCodableExtensions": ["Decodable", "Encodable"]],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      struct Wrapper {
+        @AddCodableExtensions
+        struct MyType {
+        }
+      }
+      """,
+      expandedSource: """
+        struct Wrapper {
+          struct MyType {
+          }
+        }
+
+        extension MyType: Encodable {
+        }
+        """,
+      macros: ["AddCodableExtensions": CodableExtensionMacro.self],
+      conformsTo: ["AddCodableExtensions": ["Encodable"]],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      struct Wrapper {
+        @AddCodableExtensions
+        struct MyType {
+        }
+      }
+      """,
+      expandedSource: """
+        struct Wrapper {
+          struct MyType {
+          }
+        }
+        """,
+      macros: ["AddCodableExtensions": CodableExtensionMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
 }