Workaround for macros not allowed to add extensions?

Background: I have several cases in my code where there is a protocol - for example, SomeAPI - that represents a category of server queries. Each of these protocols consists of several throwing functions:

protocol SomeAPI {
  func someQuery() throws -> String
  func otherQuery() throws -> String
}

For testing purposes, I make another protocol ThrowingSomeAPI that is empty, but has an extension implementing all the functions where they simply throw an error. Then for specific test cases, I make a struct that implements the ThrowingSomeAPI protocol, and has test-specific mock implementations for the API functions used in that test. That way I don't have to stub out all the functions every time.

protocol ThrowingSomeAPI: SomeAPI {}
extension ThrowingSomeAPI {
  func someQuery() throws -> String { throw TestError.unimplemented }
  func otherQuery() throws -> String { throw TestError.unimplemented }
}

struct TestSpecificMock: ThrowingSomeAPI {
  func someQuery() -> String { "test-specific data" }
}

I have several of these API protocols, so I'd like to use macros to automate all that stuff. I almost had one working, implemented as an attached macro that would go on the SomeAPI protocol, but ran into the snag that macros aren't allowed to introduce extensions.

I could require the user to manually create the extension and attach the macro to that, but at that point the macro doesn't have access to the content of the original protocol, only the name.

Has anyone found a workaround for the "no new extensions" rule that would help here? Or does this idea just have to wait until that rule is relaxed someday?

1 Like

Hey there :wave:,

I believe it's possible to achieve your goal with the current state of Swift Macros. First, let's see how to extend the SomeAPI protocol with implementations of its methods. This can be accomplished through the ExtensionMacro feature.

public enum MyExtensionMacro: ExtensionMacro {
  public static func expansion(
    of node: AttributeSyntax,
    attachedTo declaration: some DeclGroupSyntax,
    providingExtensionsOf type: some TypeSyntaxProtocol,
    conformingTo protocols: [TypeSyntax],
    in context: some MacroExpansionContext
  ) throws -> [ExtensionDeclSyntax] {
    
    // Ensure that the declaration is a protocol declaration
    guard let protocolDecl = declaration.as(ProtocolDeclSyntax.self) else { 
      // You can produce a diagnostic here if you wish.
      return []
    }

    // Retrieve all methods from the protocol
    var methods = protocolDecl.memberBlock.members
      .map(\.decl)
      .compactMap { $0.as(FunctionDeclSyntax.self) }

    // Add implementation to all methods from the protocol
    methods.indices.forEach { index in
      methods[index].body = CodeBlockSyntax {
        ExprSyntax(#"fatalError("whoops 😅")"#)
      }
    }

    let extensionDecl = ExtensionDeclSyntax(extendedType: type) {
      for method in methods {
        MemberBlockItemSyntax(decl: method)
      }
    }

    return [extensionDecl]
  }
}

Next, let's create the macro interface:

@attached(extension, names: arbitrary)
public macro MyExtensionMacro() = #externalMacro(
  module: "{your module name goes here}",
  type: "MyExtensionMacro"
)

Note that we're not using the conformances argument in the attached attribute since the macro isn't introducing any new conformances. Instead, we use names with the value arbitrary because method names can vary.

Don't forget to register this macro in the providingMacros field of your CompilerPlugin implementation, like so:

#if canImport(SwiftCompilerPlugin)
import SwiftCompilerPlugin
import SwiftSyntaxMacros

@main
struct MyPlugin: CompilerPlugin {
  let providingMacros: [Macro.Type] = [
    MyExtensionMacro.self,
  ]
}
#endif

Now we can use the macro! :rocket:

@MyExtensionMacro
protocol SomeAPI {
  func someQuery() throws -> String
  func otherQuery() throws -> String
}

which expands to:

Screenshot 2023-10-18 at 10.43.53 AM

As you can see this macro will automatically add the fatalError implementation to each of the protocol methods. You can also tailor this macro to add other default implementations based on your needs.

Now, let's discuss creating a new protocol. You're correct; you can't introduce an extension in a peer macro, but you can attach another macro to the new declaration created by a peer macro. The downside is that you'll need to copy the methods to the new protocol.

Here's how to implement it:

public enum MyPeerMacro: PeerMacro {
  public static func expansion(
    of node: AttributeSyntax,
    providingPeersOf declaration: some DeclSyntaxProtocol,
    in context: some MacroExpansionContext
  ) throws -> [DeclSyntax] {
    
    // Ensure that the declaration is a protocol declaration
    guard let protocolDecl = declaration.as(ProtocolDeclSyntax.self) else {
      // Produce a diagnostic here if necessary.
      return []
    }
    
    // Create a mutable copy of the original protocol declaration
    var newProtocolDecl = protocolDecl

    // Find and remove this macro attribute from the protocol declaration
    let index = newProtocolDecl.attributes.firstIndex(where: { $0.trimmedDescription == "@MyPeerMacro" })
    newProtocolDecl.attributes.remove(at: index!)

    // Prefix the original protocol name to create a new name for the new protocol
    newProtocolDecl.name = .identifier("MyPrefix" + protocolDecl.name.text)

    // Update the inheritance clause to inherit from the original protocol
    newProtocolDecl.inheritanceClause = InheritanceClauseSyntax(
      inheritedTypes: InheritedTypeListSyntax(
        arrayLiteral: InheritedTypeListSyntax.Element(
          type: IdentifierTypeSyntax(name: protocolDecl.name)
        )
      )
    )

    // Attach the `MyExtensionMacro` attribute to the new protocol
    newProtocolDecl.attributes.append(
      .attribute(
        AttributeSyntax(attributeName: IdentifierTypeSyntax(name: "MyExtensionMacro"))
      )
    )

    return [DeclSyntax(newProtocolDecl)]
  }
}

Since we'll always be adding a prefix to the name, the interface can be created as follows:

@attached(peer, names: prefixed(MyPrefix))
public macro MyPeerMacro() = #externalMacro(
  module: "MacroExamplesImplementation",
  type: "MyPeerMacro"
)

And, as before, don't forget to add the new macro to the providingMacros in your implementation of CompilerPlugin .

You can now happily use this setup in your unit tests! :rocket:

Hope this helps you out, and feel free to reach out if you have more questions or need further clarification. :raised_hands:

2 Likes

Thank you so much for that detailed explanation! Why does yours get around the extension limitation while mine doesn't?

Here is my implementation:

    public static func expansion(of node: SwiftSyntax.AttributeSyntax,
                                 providingPeersOf declaration: some SwiftSyntax.DeclSyntaxProtocol,
                                 in context: some SwiftSyntaxMacros.MacroExpansionContext) throws -> [SwiftSyntax.DeclSyntax] {
        guard let protocolDec = declaration.as(ProtocolDeclSyntax.self)
        else { throw Error.notAProtocol }
        guard let errorType = node.arguments?.description
        else { throw Error.missingError }
        let throwingProtoName = "Throwing\(protocolDec.name)"
        let throwingProtocol = try SwiftSyntax.ProtocolDeclSyntax(
            """
            protocol \(raw: throwingProtoName): \(protocolDec.name){}
            """)
        
        let throwingFuncs = protocolDec.functions.filter({ $0.isThrowing })
        
        let ext = try SwiftSyntax.ExtensionDeclSyntax("extension \(raw: throwingProtoName)",
                                                      membersBuilder: {
            for function in throwingFuncs {
                """
                \(raw: function.description) { throw \(raw: errorType) }
                """
            }
        })
        
        return [
            .init(fromProtocol: throwingProtocol),
            .init(ext),
        ]
    }

The key difference between our implementations lies in how we approach extensions. Your implementation directly tries to create an ExtensionDeclSyntax within the macro expansion function, which runs into Swift's restriction on adding new extensions via macros.

On the other hand, my implementation leverages peer macros (MyPeerMacro) to create a new protocol, and then attaches an additional macro (MyExtensionMacro) to this newly created protocol. This additional macro is responsible for adding the default implementations. By doing this, we sort of "de-couple" the process, making it more modular and, importantly, compliant with Swift's limitations on extensions in macros.

// Here, we're not directly adding an extension. We're marking
// the new protocol with another macro that will handle the extension part.
newProtocolDecl.attributes.append(
  .attribute(
    AttributeSyntax(attributeName: IdentifierTypeSyntax(name: "MyExtensionMacro"))
  )
)

I hope this clarifies things a bit!

3 Likes

And the reason for the restriction is that it avoids expanding every peer macro just to see if you’ve extended, say, Int. The explicit extension macro can only generate an extension for the type it’s attached to; that type just happens to be generated by another macro.

2 Likes

Thanks for the clarification. I adapted your code to my use case, and I wanted to add a couple of notes from that in case anyone else gets tripped up by these things:

  • When looking for the macro attribute to remove it from the duplicated protocol, I changed it to .trimmedDescription.hasPrefix("@MyMacro") because my macro has a parameter.
  • I'm throwing an error instead of calling fatalError(), so that meant changing ExprSyntax to StmtSyntax.
  • I created the macro package under my app project, but had to explicitly add it as a build dependency for the app target. I was confused at first because Xcode was auto-completing the macro but compiling still failed.

One other thing - how can I strip the comments in the duplicated protocol? It's obviously not essential, but it would be nice to have the expanded code a bit cleaner.

I'm running into a problem: since the ThrowingSomeAPI protocol duplicates the functions from the parent protocol, the extension counts as fulfilling only the Throwing protocol and not the parent. Thus a type that inherits from ThrowingSomeAPI isn't considered to have implemented all the SomeAPI functions, which is the whole point.

I tested this by pasting the macro expansions directly in my code, and then removing the duplicated functions from the Throwing protocol. My inheriting type was then recognized as properly implementing SomeAPI. This does surprise me since the function signatures are identical so I'd expect the implementations to satisfy both protocols.

Of course, the functions are intentionally duplicated so the expansion macro can know about them. Maybe the thing to do is include them as comments, which the expansion macro can re-parse. It's not ideal, but I'm not sure how else the expansion macro can know about them.

And to answer my above question, I managed to strip comments, as well as any non-throwing functions, with this:

throwingProtoDecl.memberBlock.members = .init {
    for member in throwingProtoDecl.memberBlock.members.filter({
        guard let function = $0.decl.as(FunctionDeclSyntax.self)
        else { return false }
        return function.isThrowing
    }) {
        {
            let commentsStripped = Trivia(pieces: member.decl.leadingTrivia.pieces.filter {
                switch $0 {
                    case .lineComment, .docLineComment, .blockComment, .docBlockComment:
                        return false
                    default: return true
                }
            })
            let linesCompacted = Trivia(pieces: commentsStripped.pieces.reduce([], {
                (partial: [TriviaPiece], next: TriviaPiece) in
                if (partial.last?.isNewline ?? false) && next.isWhitespace {
                    return partial
                }
                else {
                    var result = partial
                    result.append(next)
                    return result
                }
            }))
            var stripped = member

            stripped.leadingTrivia = linesCompacted
            return stripped
        }()
    }
}

I'm removing blank lines to clean up the ones left behind by deleting comments. This also removes any other blank lines, but I'm fine with that.

Looks like I've hit a compiler bug.

If I have struct Thing: ThrowingSomeAPI {} in the same file where I use my macro, then it works. But if I put it in a different file, I get "Type 'Thing' does not conform to protocol 'SomeAPI'". But if I manually replace the invocation of my macro with its expansion (by copying and pasting the expansion that Xcode shows), the error goes away.

I'm pretty sure that manually replacing a macro invocation with its expansion shouldn't have any effect on what the compiler produces, so I'll be filing a bug.

It seems the issue has already been filed: