[Macros] Support conformsTo in assertMacroExpansion for member macros

Author: soumyamahunt

Sources/SwiftSyntaxMacroExpansion/MacroSpec.swift

@@ -1,3 +1,15 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
 
@@ -8,6 +20,15 @@ public struct MacroSpec {
   /// The list of types macro needs to add conformances.
   let conformances: [TypeSyntax]
 
+  /// The type list 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.
   ///

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -21,12 +21,14 @@ import SwiftSyntaxBuilder
 
 extension SyntaxProtocol {
   /// Expand all uses of the given set of macros within this syntax node.
+  /// - SeeAlso: ``expand(macroSpecs:in:indentationWidth:)`` to also specify
+  ///   the list of conformances passed to the macro expansion.
   public func expand(
     macros: [String: Macro.Type],
     in context: some MacroExpansionContext,
     indentationWidth: Trivia? = nil
   ) -> Syntax {
-    let specs = Dictionary(uniqueKeysWithValues: macros.map { ($0.key, MacroSpec(type: $0.value)) })
+    let specs = macros.mapValues { MacroSpec(type: $0) }
     return self.expand(macroSpecs: specs, in: context, indentationWidth: indentationWidth)
   }
 
@@ -39,9 +41,7 @@ extension SyntaxProtocol {
     // Build the macro system.
     var system = MacroSystem()
     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)
+      try! system.add(macroSpec, name: macroName)
     }
 
     let applier = MacroApplication(
@@ -140,6 +140,7 @@ private func expandMemberMacro(
   definition: MemberMacro.Type,
   attributeNode: AttributeSyntax,
   attachedTo: DeclSyntax,
+  conformanceList: InheritedTypeListSyntax,
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) throws -> MemberBlockItemListSyntax? {
@@ -151,7 +152,7 @@ private func expandMemberMacro(
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: nil,
-      conformanceList: nil,
+      conformanceList: conformanceList,
       in: context,
       indentationWidth: indentationWidth
     )
@@ -328,7 +329,7 @@ private func expandExtensionMacro(
   definition: ExtensionMacro.Type,
   attributeNode: AttributeSyntax,
   attachedTo: DeclSyntax,
-  conformanceList: InheritedTypeListSyntax?,
+  conformanceList: InheritedTypeListSyntax,
   in context: some MacroExpansionContext,
   indentationWidth: Trivia
 ) throws -> CodeBlockItemListSyntax? {
@@ -349,7 +350,7 @@ private func expandExtensionMacro(
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: extendedType.detach(in: context),
-      conformanceList: conformanceList ?? [],
+      conformanceList: conformanceList,
       in: context,
       indentationWidth: indentationWidth
     )
@@ -438,49 +439,30 @@ private func expandBodyMacro(
 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] = [:]
+  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
   }
 
-  /// 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? {
+  /// Look for a macro specification with the given name.
+  func lookup(_ macroName: String) -> MacroSpec? {
     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
@@ -898,41 +880,40 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
 // MARK:  Attached macro expansions.
 
 extension MacroApplication {
-  /// Get macro attribute, the macro definition and optional
-  /// conformance protocols list attached to `decl`.
+  /// Get macro attribute, the macro specification attached to `decl`.
   ///
   /// The macros must be registered in `macroSystem`.
   private func macroAttributes(
     attachedTo decl: DeclSyntax
-  ) -> [(attributeNode: AttributeSyntax, definition: Macro.Type, conformanceList: InheritedTypeListSyntax?)] {
+  ) -> [(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, macroSystem.conformaces(forMacro: attributeName))
+      return (attribute, macroSpec)
     }
   }
 
-  /// Get macro attribute, the macro definition and optional conformance
+  /// Get macro attribute, the macro definition and 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, conformanceList: InheritedTypeListSyntax?)] {
+  ) -> [(attributeNode: AttributeSyntax, definition: MacroType, conformanceList: InheritedTypeListSyntax)] {
     return macroAttributes(attachedTo: decl)
-      .compactMap { (attributeNode: AttributeSyntax, definition: Macro.Type, conformanceList: InheritedTypeListSyntax?) in
-        if let macroType = definition as? MacroType {
-          return (attributeNode, macroType, conformanceList)
+      .compactMap { (attributeNode: AttributeSyntax, spec: MacroSpec) in
+        if let macroType = spec.type as? MacroType {
+          return (attributeNode, macroType, spec.inheritedTypeList)
         } else {
           return nil
         }
@@ -948,7 +929,7 @@ extension MacroApplication {
   >(
     attachedTo decl: DeclSyntax,
     ofType: MacroType.Type,
-    expandMacro: (_ attributeNode: AttributeSyntax, _ definition: MacroType, _ conformanceList: InheritedTypeListSyntax?) throws -> ExpanedNodeCollection?
+    expandMacro: (_ attributeNode: AttributeSyntax, _ definition: MacroType, _ conformanceList: InheritedTypeListSyntax) throws -> ExpanedNodeCollection?
   ) -> [ExpandedNode] {
     var result: [ExpandedNode] = []
 
@@ -1026,6 +1007,7 @@ extension MacroApplication {
         definition: definition,
         attributeNode: attributeNode,
         attachedTo: decl,
+        conformanceList: conformanceList,
         in: context,
         indentationWidth: indentationWidth
       )
@@ -1147,7 +1129,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
     }

Sources/SwiftSyntaxMacrosTestSupport/Assertions.swift

@@ -300,6 +300,9 @@ func assertDiagnostic(
 ///   - 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,

Tests/SwiftSyntaxMacroExpansionTest/MemberMacroTests.swift

@@ -299,4 +299,82 @@ final class MemberMacroTests: XCTestCase {
       ]
     )
   }
+
+  func testConditionalMemberExpansion() {
+    struct CodableExtensionMacro: MemberMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        conformingTo protocols: [TypeSyntax],
+        in context: some MacroExpansionContext
+      ) throws -> [DeclSyntax] {
+        let decls: [DeclSyntax] = protocols.compactMap { `protocol` in
+          return """
+            func print(arg: some \(`protocol`)) { }
+            """ as DeclSyntax
+        }
+
+        return decls
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      @AddCodableExtensions
+      class MyType {
+      }
+      """,
+      expandedSource: """
+        class MyType {
+
+          func print(arg: some Decodable) {
+          }
+
+          func print(arg: some Encodable) {
+          }
+        }
+        """,
+      macroSpecs: ["AddCodableExtensions": MacroSpec(type: CodableExtensionMacro.self, conformances: ["Decodable", "Encodable"])],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      class Wrapper {
+        @AddCodableExtensions
+        class MyType {
+        }
+      }
+      """,
+      expandedSource: """
+        class Wrapper {
+          class MyType {
+
+            func print(arg: some Encodable) {
+            }
+          }
+        }
+        """,
+      macroSpecs: ["AddCodableExtensions": MacroSpec(type: CodableExtensionMacro.self, conformances: ["Encodable"])],
+      indentationWidth: indentationWidth
+    )
+
+    assertMacroExpansion(
+      """
+      class Wrapper {
+        @AddCodableExtensions
+        class MyType {
+        }
+      }
+      """,
+      expandedSource: """
+        class Wrapper {
+          class MyType {
+          }
+        }
+        """,
+      macros: ["AddCodableExtensions": CodableExtensionMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
 }