Published on

A first look at the Swift macro system

Authors

Introduction

Suppose you have a Swift enum like this:

    enum MyEnum {
        case a
        case b
        case c
    }

And you want to write a function that takes a MyEnum and returns a String:

    func myFunction(_ myEnum: MyEnum) -> String {
        switch myEnum {
        case .a:
            return "a"
        case .b:
            return "b"
        case .c:
            return "c"
        }
    }

This is a very common pattern in Swift, and it's a bit annoying to write. You have to write the same code over and over again, and you have to remember to update it if you add a new case to the enum. It would be nice if there was a way to write this code once and then use it everywhere.

Well, there is! It's called code generation, where build tools (optimally the compiler) writes code that generates other code. In this post, I'll show you the current state of codegen in Swift and a proposal for a proper macro system for the Swift compiler (with the capability to generate the myFunction function above).

The current state of codegen

Builtin macros

Apple has actually been using macros in its Swift code since the first version. Looking at the Swift reference for literal expressions, you see a small list of (mostly simple)1 builtin unary macros:

    #file
    #fileID
    #filePath
    #line
    #column
    #function
    #dsohandle

You also have a set of macros for Apple's SDKs:

  #available
  #colorLiteral
  #imageLiteral
  #selector
  #keyPath

These aren't particularly useful to us - we want a more portable solution than a compiler patch every time we want to use a macro.

Gyb

Gyb is a Swift file with inline python code, where the output of the python is inserted inline the Swift, and Sourcery is a tool that reads Swift files and generates Swift files. NSHipster has a good article on GYB, but writing myFunction with GYB would look like this:

    func myFunction(_ myEnum: MyEnum) -> String {
        switch myEnum {
        % for case in ['a', 'b', 'c']:
        case .${case}:
            return "${case}"
        % end
        }
    }

This is a bit better than writing the same code over and over again, but it's far from ideal.

  • The code isn't checked by the compiler, so you can easily make mistakes.
  • The code is hard to read, and it's hard to see what the output will be (although it has a generated file output ready for inspection).
  • You have to run gyb at least every time the file changes, complicating the build process.
  • It's not even that much more concise! It would be nice to replace ['a', 'b', 'c'] with MyEnum.allCases, but the Python environment is completely detached from the Swift source context.

Sourcery

Sourcery takes a different approach than directly embedding Python in the Swift file. It allows for conformance to 'phantom protocols' (protocols that hold no specification) and can generate code for any type that conforms to a phantom protocol. This makes it more context-aware than GYB, but requires more work: Each phantom protocol must be implemented in its own template file, each of which is limited in its context information. Usually the code is written in a DSL called Stencil, but it can also be written in Swift or JavaScript.

// Add a conformance of AutoFunctionGeneration to MyEnum in our project files
extension MyEnum: AutoFunctionGeneration {}

// AutoFunctionGeneration.stencil

{% for enum in types.implementing.AutoFunctionGeneration|enum %}
{{ enum.accessLevel }} extension {{ enum.name }} {
  func myFunction() -> String {
    switch self {
    {% for case in enum.cases %}
    case .{{ case.name }}:
      return "{{ case.name }}"
    {% endfor %}
    }
  }
}
{% endfor %}

This is a bit better than GYB, but it's still not ideal:

  • While contextual injection points are checked by the compiler, the stencil template code still isn't, so you can easily make mistakes.
  • Creating a template file is a complicated process, and it's much more difficult to learn than just writing a Python script or a Swift function.

There are still the problems of:

  • The code is hard to read
  • You have to run Sourcery

The community has settled on Sourcery as the current state of the art, with further progress blocked by lack of first-party compiler integration. Fortunately, there's finally stirrings of a proper Swift macro system to replace these tools.

The Swift macro system

Introduction

A recent proposal just landed on the Swift Forums on Swift macros.

It has an attached "pitch / vision / manifesto" and a slimmed-down expression macros pitch (currently awaiting implementation).

Before I go into this proposal, it would be good to review its most recognized inspiration...

The Rust macro system

Note: A much more detailed primer exists here. If you're very curious about how Rust's system works, skip my explanation and read theirs.

Rust has two kinds of macros: declarative macros and procedural macros.

Declarative macros

This is the "easier" macro system. It involves matching and capturing a contained AST, much like in a regex. The matched pattern is replaced with an expansion, a block of code generated as a function of the captures. An example in the Rust macro book is a macro that multiplies by five:

macro_rules! times_five {
    ($e:expr) => { 5 * $e };
}

This macro matches any expression, and replaces it with the expansion 5 * <the expression>.

Procedural macros

This is the "harder" macro system. It involves writing a function that takes a TokenStream (a Rust AST) and returns a TokenStream. The input is the AST of the macro invocation, and the output is the AST of the expansion. An example in the Rust macro book is a macro that prints the AST of an expression:

extern crate proc_macro;
use proc_macro::TokenStream;

#[proc_macro]
pub fn print_ast(input: TokenStream) -> TokenStream {
    println!("{:#?}", input);
    input
}

This macro takes an expression, prints its AST, and returns the expression unchanged.

These are used in crates such as serde for automatic seralization code generation and vulkano for inline shader code.

The Swift macro system

The evolution proposal for the Swift macro system includes both systems, but, for now, we'll focus on the more concrete expression macros pitch. The proposal is essentially the Rust declarative macro system.

Macro expansion is a syntactic operation, which takes as input a well-formed syntax tree consisting of the full macro expansion expression (e.g., #stringify(x + y)) and produces a syntax tree as output.

There are some differences, though:

...Rust's declarative macros offer more advanced rules that match the macro arguments to a pattern and then perform a a rewrite to new syntax as described in the macro. For Swift to adopt this approach, we would likely need to invent a pattern language for matching and rewriting syntax trees.

Swift has a pattern language, but it's not as powerful as Rust's. An example is switch matching:

enum Choices {
    case a
    case v(Int)
}

let x: Int?
let y: String
let z: Choices
switch (x, y, z) {
    case (1, 2, .a):
        print(x, y)
    case let (_, y, .v(let z)):
        print(y, z)
    case (_, _, .a):
        print("a")
    case (_, _, .v):
        print("v")
}

This pattern language is useful for simple, fixed patterns, but it's not powerful enough to match arbitrary expressions, such as repetitions.

This pitch has only been out for five days, but there's already a really cool example use case! PowerAssert is a library that pretty-prints assertion failures:

#powerAssert(max(a, b) == c)
             |   |  |  |  |
             7   4  7  |  12
                       false
#powerAssert(xs.contains(4))
             |  |        |
             |  false    4
             [1, 2, 3]
#powerAssert("hello".hasPrefix("h") && "goodbye".hasSuffix("y"))
             |       |         |    |  |         |         |
             "hello" true      "h"  |  "goodbye" false     "y"
                                    false

I really enjoy this example as a proof of concept, so lets go through it step by step and see how it works!

The PowerAssert macro

Metadata and definition

struct PowerAssertMacro: ExpressionMacro {
  static var name: String { "powerAssert" }

  static var documentation: String {
    "Power Assert in Swift. Provides descriptive assertion messages."
  }

  static var genericSignature: GenericParameterClauseSyntax? = "<T>"
  static var signature: TypeSyntax = "(_ expression: @autoclosure () throws -> Bool) -> Void"

  static var owningModule: String = "Swift"

  static func apply(
    _ macro: MacroExpansionExprSyntax, in context: MacroEvaluationContext
  ) -> MacroResult<ExprSyntax> {
    let generator = CodeGenerator(macro: macro, context: context)
    return generator.generate()
  }

The metadata fields are defined in the SwiftSyntaxMacros package. For the less-straightforward fields...

  • genericSignature is the generic signature of the macro. This is given as a string, but does not have a string value, it's parsed by swift-syntax.
  • signature is the signature of the macro. This is given as a string and parsed much like the generic signature. Disambiguation between value macros and function macros happen here: if the signature is a function type, it's a function macro (e.g. #imageLiteral), otherwise it's a value macro (e.g. #line).
  • owningModule is module that must be imported to use the macro, and supplementalSignatureModules are modules that must be imported by the macro to function.

The apply function is the top-level function that is called when the macro is invoked. It takes the macro invocation AST and the context, and returns the expansion AST. This work is delegated to CodeGenerator, which holds the core logic of the macro.

CodeGenerator

The actual member variables of the struct are minimal, just the context of the invocation. The interesting logic is in the methods.

  func generate() -> MacroResult<ExprSyntax> {
    guard let expression = macro.argumentList.first else {
      if let leadingTrivia = macro.leadingTrivia {
        return MacroResult(ExprSyntax("()").withLeadingTrivia(leadingTrivia))
      }
      return MacroResult("()")
    }

    let formatted = format(expression)
    let expanded = expand(expression: formatted)

    let assertSyntax = ExprSyntax(expanded)
    if let leadingTrivia = macro.leadingTrivia {
      return MacroResult(assertSyntax.withLeadingTrivia(leadingTrivia))
    }
    return MacroResult(assertSyntax)

  }

leadingTrivia is the trivia (whitespace, comments, etc.) that comes before the macro invocation. This is used to preserve the indentation of the macro invocation. The guard statement checks that the macro invocation has an argument, and returns an empty tuple if it doesn't.

The format and expand functions are the core of the macro, so we'll get into that later.

After the macro is fully expanded, the AST is finalized into expression syntax, and the trivia is added back in.

Format

  private func format(_ expression: SyntaxProtocol) -> SyntaxProtocol {
    SourceFileSyntax(
      "\(expression.withoutTrivia())"
        .split(separator: "\n")
        .joined(separator: " ")
        .split(separator: " ")
        .joined(separator: " ")
    )
  }

All this does is remove the trivia from the expression, and then standardize whitespace into single spaces. This results in the normal spacing in the example powerAssert assertion above, regardless of whether its written as

#powerAssert("hello".hasPrefix("h") && "goodbye".hasSuffix("y"))

or

#powerAssert("hello".hasPrefix("h")        &&
"goodbye".hasSuffix("y"))

Expand

    var expressions = [Syntax]()
    parseExpression(expression, storage: &expressions)
    expressions = Array(expressions.dropFirst(2))

    let startLocation = macro.startLocation(converter: context.sourceLocationConverter)
    let endLocation = macro.macro.endLocation(converter: context.sourceLocationConverter)

    let sourceLocationConverter = SourceLocationConverter(file: "", tree: expression)
    let startColumn = endLocation.column! - startLocation.column!

This preamble takes the formatted expression and finds the start and end locations of the expression, and, implicitly, its size. It's used to find the offset of the expression, so the assertion annotations on subsequent lines can be aligned to their tokens.

return """
    PowerAssert.Assertion(#"\(macro.poundToken.withoutTrivia())\(macro.macro)(\(expression))"#, line: \(startLocation.line!))
    .assert(\(expressions.first!))
    \(
      expressions
        .reduce("") { (result, syntax) in
          let startLocation = syntax.startLocation(converter: sourceLocationConverter)
          let column = startLocation.column! + startColumn
          let syntaxType = syntax.syntaxNodeType
          if syntaxType == ArrayElementListSyntax.self
            || syntaxType == ArrayElementSyntax.self
            || // a lot more node type checks
            {
            return result
            }
            // more cases
        return result + ".capture(expression: \(syntax.withoutTrivia()), column: \(column))"
.render()
"""
.split(separator: "\n")
.joined()

This is a key difference from the AST-centric approach of the other macro systems - the output of the expansion function is a string, not an AST. The returned string is the macro expansion returned in the top-level invocation, and is parsed by swift-syntax into an AST. The actual contents of the AST just creates an array of the value and position of each subexpression in the assertion, and then passes it into a render function, which performs the final formatting at runtime (inside the string).

Conclusion

The entire macro portion of the library is only about 200 LoC, and yet runs entirely in Swift, with no need for gyb or Sourcery. I look forward to the release of the macro system and I can finally start working on a TCA state sharing macro.

Footnotes

  1. The #dsohandle macro is the only really confusing one of these, but it's just an ID to the parent executable / library. ↩

Discuss on Twitter • View on GitHub