Merge pull request swiftlang#1871 from hborla/refine-extension-macros

hborla · web-flow · commit 8828c298a3be · 2023-07-02T20:47:13.000-07:00
[Macros] Pass a list of protocols to conform to when expanding an `extension` macro.
diff --git a/Examples/Sources/ExamplePlugin/Macros.swift b/Examples/Sources/ExamplePlugin/Macros.swift
@@ -91,8 +91,13 @@ public struct SendableExtensionMacro: ExtensionMacro {
     of node: AttributeSyntax,
     attachedTo: some DeclGroupSyntax,
     providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
     in context: some MacroExpansionContext
   ) throws -> [ExtensionDeclSyntax] {
+    if protocols.isEmpty {
+      return []
+    }
+
     let sendableExtension: DeclSyntax =
       """
       extension \(type.trimmed): Sendable {}
diff --git a/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift b/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
@@ -122,15 +122,25 @@ extension CompilerPluginMessageHandler {
         expandingSyntax: expandingSyntax
       )
 
-    case .expandAttachedMacro(let macro, let macroRole, let discriminator, let attributeSyntax, let declSyntax, let parentDeclSyntax, let extendedTypeSyntax):
+    case .expandAttachedMacro(
+      let macro,
+      let macroRole,
+      let discriminator,
+      let attributeSyntax,
+      let declSyntax,
+      let parentDeclSyntax,
+      let extendedTypeSyntax,
+      let conformanceListSyntax
+    ):
       try expandAttachedMacro(
         macro: macro,
         macroRole: macroRole,
         discriminator: discriminator,
         attributeSyntax: attributeSyntax,
         declSyntax: declSyntax,
         parentDeclSyntax: parentDeclSyntax,
-        extendedTypeSyntax: extendedTypeSyntax
+        extendedTypeSyntax: extendedTypeSyntax,
+        conformanceListSyntax: conformanceListSyntax
       )
 
     case .loadPluginLibrary(let libraryPath, let moduleName):
diff --git a/Sources/SwiftCompilerPluginMessageHandling/Macros.swift b/Sources/SwiftCompilerPluginMessageHandling/Macros.swift
@@ -87,7 +87,8 @@ extension CompilerPluginMessageHandler {
     attributeSyntax: PluginMessage.Syntax,
     declSyntax: PluginMessage.Syntax,
     parentDeclSyntax: PluginMessage.Syntax?,
-    extendedTypeSyntax: PluginMessage.Syntax?
+    extendedTypeSyntax: PluginMessage.Syntax?,
+    conformanceListSyntax: PluginMessage.Syntax?
   ) throws {
     let sourceManager = SourceManager()
     let context = PluginMacroExpansionContext(
@@ -104,6 +105,10 @@ extension CompilerPluginMessageHandler {
     let extendedType = extendedTypeSyntax.map {
       sourceManager.add($0).cast(TypeSyntax.self)
     }
+    let conformanceList = conformanceListSyntax.map {
+      let placeholderStruct = sourceManager.add($0).cast(StructDeclSyntax.self)
+      return placeholderStruct.inheritanceClause!.inheritedTypeCollection
+    }
 
     // TODO: Make this a 'String?' and remove non-'hasExpandMacroResult' branches.
     let expandedSources: [String]?
@@ -120,6 +125,7 @@ extension CompilerPluginMessageHandler {
         declarationNode: declarationNode,
         parentDeclNode: parentDeclNode,
         extendedType: extendedType,
+        conformanceList: conformanceList,
         in: context
       )
       if let expansions, hostCapability.hasExpandMacroResult {
diff --git a/Sources/SwiftCompilerPluginMessageHandling/PluginMessages.swift b/Sources/SwiftCompilerPluginMessageHandling/PluginMessages.swift
@@ -34,7 +34,8 @@
     attributeSyntax: PluginMessage.Syntax,
     declSyntax: PluginMessage.Syntax,
     parentDeclSyntax: PluginMessage.Syntax?,
-    extendedTypeSyntax: PluginMessage.Syntax?
+    extendedTypeSyntax: PluginMessage.Syntax?,
+    conformanceListSyntax: PluginMessage.Syntax?
   )
 
   /// Optionally implemented message to load a dynamic link library.
@@ -76,7 +77,7 @@
 }
 
 @_spi(PluginMessage) public enum PluginMessage {
-  public static var PROTOCOL_VERSION_NUMBER: Int { 6 }  // Added 'expandMacroResult'.
+  public static var PROTOCOL_VERSION_NUMBER: Int { 7 }  // Pass extension protocol list
 
   public struct HostCapability: Codable {
     var protocolVersion: Int
diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift b/Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift
@@ -185,6 +185,7 @@ public func expandAttachedMacroWithoutCollapsing<Context: MacroExpansionContext>
   declarationNode: DeclSyntax,
   parentDeclNode: DeclSyntax?,
   extendedType: TypeSyntax?,
+  conformanceList: InheritedTypeListSyntax?,
   in context: Context
 ) -> [String]? {
   do {
@@ -266,42 +267,6 @@ public func expandAttachedMacroWithoutCollapsing<Context: MacroExpansionContext>
         $0.formattedExpansion(definition.formatMode)
       }
 
-    case (let attachedMacro as ConformanceMacro.Type, .conformance):
-      guard let declGroup = declarationNode.asProtocol(DeclGroupSyntax.self) else {
-        // Compiler error: type mismatch.
-        throw MacroExpansionError.declarationNotDeclGroup
-      }
-      guard let identified = declarationNode.asProtocol(IdentifiedDeclSyntax.self)
-      else {
-        // Compiler error: type mismatch.
-        throw MacroExpansionError.declarationNotIdentified
-      }
-
-      // Local function to expand a conformance macro once we've opened up
-      // the existential.
-      func expandConformanceMacro(
-        _ node: some DeclGroupSyntax
-      ) throws -> [(TypeSyntax, GenericWhereClauseSyntax?)] {
-        return try attachedMacro.expansion(
-          of: attributeNode,
-          providingConformancesOf: node,
-          in: context
-        )
-      }
-
-      let conformances = try _openExistential(
-        declGroup,
-        do: expandConformanceMacro
-      )
-
-      // Form a buffer of extension declarations to return to the caller.
-      return conformances.map { typeSyntax, whereClause in
-        let typeName = identified.identifier.trimmedDescription
-        let protocolName = typeSyntax.trimmedDescription
-        let whereClause = whereClause?.trimmedDescription ?? ""
-        return "extension \(typeName) : \(protocolName) \(whereClause) {}"
-      }
-
     case (let attachedMacro as ExtensionMacro.Type, .extension):
       guard let declGroup = declarationNode.asProtocol(DeclGroupSyntax.self) else {
         // Compiler error: type mismatch.
@@ -312,6 +277,8 @@ public func expandAttachedMacroWithoutCollapsing<Context: MacroExpansionContext>
         throw MacroExpansionError.noExtendedTypeSyntax
       }
 
+      let protocols = conformanceList?.map(\.typeName) ?? []
+
       // Local function to expand an extension macro once we've opened up
       // the existential.
       func expandExtensionMacro(
@@ -321,6 +288,7 @@ public func expandAttachedMacroWithoutCollapsing<Context: MacroExpansionContext>
           of: attributeNode,
           attachedTo: node,
           providingExtensionsOf: extendedType,
+          conformingTo: protocols,
           in: context
         )
       }
@@ -364,6 +332,7 @@ public func expandAttachedMacro<Context: MacroExpansionContext>(
   declarationNode: DeclSyntax,
   parentDeclNode: DeclSyntax?,
   extendedType: TypeSyntax?,
+  conformanceList: InheritedTypeListSyntax?,
   in context: Context
 ) -> String? {
   let expandedSources = expandAttachedMacroWithoutCollapsing(
@@ -373,6 +342,7 @@ public func expandAttachedMacro<Context: MacroExpansionContext>(
     declarationNode: declarationNode,
     parentDeclNode: parentDeclNode,
     extendedType: extendedType,
+    conformanceList: conformanceList,
     in: context
   )
   return expandedSources.map {
diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
@@ -136,7 +136,6 @@ class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
           || macro is MemberMacro.Type
           || macro is AccessorMacro.Type
           || macro is MemberAttributeMacro.Type
-          || macro is ConformanceMacro.Type
           || macro is ExtensionMacro.Type)
       }
 
@@ -213,7 +212,7 @@ class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
 
         if let declGroup = decl.asProtocol(DeclGroupSyntax.self) {
           newItems.append(
-            contentsOf: expandConformances(of: declGroup).map {
+            contentsOf: expandExtensions(of: declGroup).map {
               newDecl in CodeBlockItemSyntax(item: .decl(newDecl))
             }
           )
@@ -429,45 +428,29 @@ extension MacroApplication {
   // If any of the custom attributes associated with the given declaration
   // refer to conformance macros, expand them and return the resulting
   // set of extension declarations.
-  private func expandConformances(of decl: DeclGroupSyntax) -> [DeclSyntax] {
-    let extendedType: Syntax
+  private func expandExtensions(of decl: DeclGroupSyntax) -> [DeclSyntax] {
+    let extendedType: TypeSyntax
     if let identified = decl.asProtocol(IdentifiedDeclSyntax.self) {
-      extendedType = Syntax(identified.identifier.trimmed)
+      extendedType = "\(identified.identifier.trimmed)"
     } else if let ext = decl.as(ExtensionDeclSyntax.self) {
-      extendedType = Syntax(ext.extendedType.trimmed)
+      extendedType = "\(ext.extendedType.trimmed)"
     } else {
       return []
     }
 
     var extensions: [DeclSyntax] = []
-    let macroAttributes = getMacroAttributes(attachedTo: decl.as(DeclSyntax.self)!, ofType: ConformanceMacro.Type.self)
-    for (attribute, conformanceMacro) in macroAttributes {
-      do {
-        let newConformances = try conformanceMacro.expansion(of: attribute, providingConformancesOf: decl, in: context)
-
-        for (type, whereClause) in newConformances {
-          var ext: DeclSyntax = """
-            extension \(extendedType): \(type) { }
-            """
-          if let whereClause {
-            ext = DeclSyntax((ext.cast(ExtensionDeclSyntax.self)).with(\.genericWhereClause, whereClause))
-          }
-
-          extensions.append(DeclSyntax(ext))
-        }
-      } catch {
-        context.addDiagnostics(from: error, node: attribute)
-      }
-    }
 
     let extensionMacroAttrs = getMacroAttributes(attachedTo: decl.as(DeclSyntax.self)!, ofType: ExtensionMacro.Type.self)
-    let extendedTypeSyntax = TypeSyntax("\(extendedType.trimmed)")
     for (attribute, extensionMacro) in extensionMacroAttrs {
       do {
+        // FIXME: We need a way for unit tests of extension macros to
+        // specify protocols already stated in source (e.g. as arguments
+        // to `assertMacroExpansion`).
         let newExtensions = try extensionMacro.expansion(
           of: attribute,
           attachedTo: decl,
-          providingExtensionsOf: extendedTypeSyntax,
+          providingExtensionsOf: extendedType,
+          conformingTo: [],
           in: context
         )
 
diff --git a/Sources/SwiftSyntaxMacros/MacroProtocols/ConformanceMacro.swift b/Sources/SwiftSyntaxMacros/MacroProtocols/ConformanceMacro.swift
@@ -14,7 +14,7 @@ import SwiftSyntax
 
 /// Describes a macro that can add conformances to the declaration it's
 /// attached to.
-public protocol ConformanceMacro: AttachedMacro {
+public protocol ConformanceMacro: ExtensionMacro {
   /// Expand an attached conformance macro to produce a set of conformances.
   ///
   /// - Parameters:
@@ -31,3 +31,38 @@ public protocol ConformanceMacro: AttachedMacro {
     in context: some MacroExpansionContext
   ) throws -> [(TypeSyntax, GenericWhereClauseSyntax?)]
 }
+
+extension ConformanceMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    attachedTo declaration: some DeclGroupSyntax,
+    providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
+    in context: some MacroExpansionContext
+  ) throws -> [ExtensionDeclSyntax] {
+
+    let newConformances = try expansion(
+      of: node,
+      providingConformancesOf: declaration,
+      in: context
+    )
+
+    var extensions: [ExtensionDeclSyntax] = []
+    for (proto, whereClause) in newConformances {
+      let decl: DeclSyntax =
+        """
+        extension \(type.trimmed): \(proto) {}
+        """
+
+      var extensionDecl = decl.cast(ExtensionDeclSyntax.self)
+
+      if let whereClause {
+        extensionDecl = extensionDecl.with(\.genericWhereClause, whereClause)
+      }
+
+      extensions.append(extensionDecl)
+    }
+
+    return extensions
+  }
+}
diff --git a/Sources/SwiftSyntaxMacros/MacroProtocols/ExtensionMacro.swift b/Sources/SwiftSyntaxMacros/MacroProtocols/ExtensionMacro.swift
@@ -21,6 +21,9 @@ public protocol ExtensionMacro: AttachedMacro {
   ///   - node: The custom attribute describing the attached macro.
   ///   - declaration: The declaration the macro attribute is attached to.
   ///   - type: The type to provide extensions of.
+  ///   - protocols: The list of protocols to add conformances to. These will
+  ///     always be protocols that `type` does not already state a conformance
+  ///     to.
   ///   - context: The context in which to perform the macro expansion.
   ///
   /// - Returns: the set of extension declarations introduced by the macro,
@@ -30,6 +33,7 @@ public protocol ExtensionMacro: AttachedMacro {
     of node: AttributeSyntax,
     attachedTo declaration: some DeclGroupSyntax,
     providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
     in context: some MacroExpansionContext
   ) throws -> [ExtensionDeclSyntax]
 }
diff --git a/Sources/SwiftSyntaxMacros/MacroSystem.swift b/Sources/SwiftSyntaxMacros/MacroSystem.swift
@@ -115,7 +115,6 @@ class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
           || macro is MemberMacro.Type
           || macro is AccessorMacro.Type
           || macro is MemberAttributeMacro.Type
-          || macro is ConformanceMacro.Type
           || macro is ExtensionMacro.Type)
       }
 
@@ -192,7 +191,7 @@ class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
 
         if let declGroup = decl.asProtocol(DeclGroupSyntax.self) {
           newItems.append(
-            contentsOf: expandConformances(of: declGroup).map {
+            contentsOf: expandExtensions(of: declGroup).map {
               newDecl in CodeBlockItemSyntax(item: .decl(newDecl))
             }
           )
@@ -408,45 +407,29 @@ extension MacroApplication {
   // If any of the custom attributes associated with the given declaration
   // refer to conformance macros, expand them and return the resulting
   // set of extension declarations.
-  private func expandConformances(of decl: DeclGroupSyntax) -> [DeclSyntax] {
-    let extendedType: Syntax
+  private func expandExtensions(of decl: DeclGroupSyntax) -> [DeclSyntax] {
+    let extendedType: TypeSyntax
     if let identified = decl.asProtocol(IdentifiedDeclSyntax.self) {
-      extendedType = Syntax(identified.identifier.trimmed)
+      extendedType = "\(identified.identifier.trimmed)"
     } else if let ext = decl.as(ExtensionDeclSyntax.self) {
-      extendedType = Syntax(ext.extendedType.trimmed)
+      extendedType = "\(ext.extendedType.trimmed)"
     } else {
       return []
     }
 
     var extensions: [DeclSyntax] = []
-    let macroAttributes = getMacroAttributes(attachedTo: decl.as(DeclSyntax.self)!, ofType: ConformanceMacro.Type.self)
-    for (attribute, conformanceMacro) in macroAttributes {
-      do {
-        let newConformances = try conformanceMacro.expansion(of: attribute, providingConformancesOf: decl, in: context)
-
-        for (type, whereClause) in newConformances {
-          var ext: DeclSyntax = """
-            extension \(extendedType): \(type) { }
-            """
-          if let whereClause {
-            ext = DeclSyntax((ext.cast(ExtensionDeclSyntax.self)).with(\.genericWhereClause, whereClause))
-          }
-
-          extensions.append(DeclSyntax(ext))
-        }
-      } catch {
-        context.addDiagnostics(from: error, node: attribute)
-      }
-    }
 
     let extensionMacroAttrs = getMacroAttributes(attachedTo: decl.as(DeclSyntax.self)!, ofType: ExtensionMacro.Type.self)
-    let extendedTypeSyntax = TypeSyntax("\(extendedType.trimmed)")
     for (attribute, extensionMacro) in extensionMacroAttrs {
       do {
+        // FIXME: We need a way for unit tests of extension macros to
+        // specify protocols already stated in source (e.g. as arguments
+        // to `assertMacroExpansion`).
         let newExtensions = try extensionMacro.expansion(
           of: attribute,
           attachedTo: decl,
-          providingExtensionsOf: extendedTypeSyntax,
+          providingExtensionsOf: extendedType,
+          conformingTo: [],
           in: context
         )
 
diff --git a/Tests/SwiftSyntaxMacrosTest/MacroSystemTests.swift b/Tests/SwiftSyntaxMacrosTest/MacroSystemTests.swift
@@ -687,6 +687,7 @@ public struct SendableExtensionMacro: ExtensionMacro {
     of node: AttributeSyntax,
     attachedTo: some DeclGroupSyntax,
     providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
     in context: some MacroExpansionContext
   ) throws -> [ExtensionDeclSyntax] {
     let sendableExtension: DeclSyntax =
diff --git a/lit_tests/compiler_plugin_basic.swift b/lit_tests/compiler_plugin_basic.swift
