---
title: 'A first look at the Swift macro system'
date: '2022-11-21'
tags: ['swift', 'rust']
draft: false
summary: 'Hopefully this is the beginning of the end of writing Python scripts in Swift code.'
---

<TOCInline toc={props.toc} exclude="Overview" toHeading={2} indentDepth={2} />

# Introduction

Suppose you have a Swift enum like this:
```swift
    enum MyEnum {
        case a
        case b
        case c
    }
```

And you want to write a function that takes a `MyEnum` and returns a `String`:
```swift
    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](https://docs.swift.org/swift-book/ReferenceManual/Expressions.html#grammar_literal-expression),
you see a small list of (mostly simple)[^1] builtin unary macros:

```swift
    #file
    #fileID
    #filePath
    #line
    #column
    #function
    #dsohandle
```
[^1]: The `#dsohandle` macro is the only really confusing one of these, but it's just an [ID to the parent executable / library](https://developer.apple.com/documentation/swift/dsohandle()).

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

```swift
  #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](https://nshipster.com/swift-gyb/) on GYB,
but writing `myFunction` with GYB would look like this:

```swift
    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.

```swift
// 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](https://forums.swift.org/t/a-possible-vision-for-macros-in-swift/60900) on Swift macros.

It has an attached ["pitch / vision / manifesto"](https://gist.github.com/DougGregor/4f3ba5f4eadac474ae62eae836328b71) and a slimmed-down expression macros [pitch](https://github.com/swiftlang/swift-evolution/blob/ed9dbdd5be454265497bb0a478886179598f9c82/proposals/nnnn-expression-macros.md) *(pinned to the draft as it stood — the branch is gone)*
(currently awaiting implementation).

Before I go into this proposal, it would be good to review its most recognized inspiration...
## The [Rust macro system](https://doc.rust-lang.org/book/ch19-06-macros.html)

_Note: A much more detailed primer exists [here](https://veykril.github.io/tlborm/introduction.html).
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:

```rust
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:

```rust
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](https://serde.rs/attributes.html) for automatic seralization code generation
and [vulkano](https://docs.rs/vulkano/0.19.0/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:

```swift
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](https://github.com/kishikawakatsumi/swift-syntax/commit/576669d01271d9289000991f10efd8ecc6f45023)
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

```swift
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](https://github.com/swiftlang/swift-syntax/blob/3b7d3ad1afccd53f1e2aaeb1b22437910f80fd2e/Sources/_SwiftSyntaxMacros/Macro.swift#L21).
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](https://github.com/swiftlang/swift-syntax/commit/ae619cacb4580d2928503f43e91d0ce4393062b5).
 - `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.

```swift
  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

```swift
  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
1.

```swift
    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.

2.
```swift
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](/blog/tca-state-sharing).