Implementing Type Systems as Macros

14 Aug 2023

There’s a neat paper Type Systems as Macros by Chang, Knauth, and Greenman [1] that describes how to implement a typed language using an untyped host language and macro expansion. The paper is neat, but I found the code hard to follow—the paper uses a compact notation that’s convenient for print, but not so much for reproducing on one’s own. This post is my attempt to implement and explain in more accessible terms what’s presented in the paper.

The working source for the paper is available on the first author’s BitBucket, but that code uses a lot of auxiliary helper files and it can still be pretty hard to follow.

My implementation consists of a single file implemented in plain Racket. I try to keep as many of the function names the same, so hopefully it’s not too hard to follow if you’re coming from reading the paper. If you’re impatient to try running the code, the full example can be found at the bottom in § Full solution or on my SourceHut account.

A few terms #

A few terms that will crop up in this post:

Source language

Generally, this refers to the input to a compiler or interpreter—it’s the language of the source code that we want to eventually be able to write. In our specific case, this will be the simply typed lambda calculus.

Target language

Generally, the output of a compiler. For example, the target language of a compiler might be x86 assembly or WASM. In our case, we want to transform the typed source language (the simply typed lambda calculus) into the untyped target language of Racket.

Procedural macro

Procedural macros are ones that let you use the entire power of the language to operate on the AST you are transforming: you can write conditionals, loops, and even call functions on the AST under expansion. Contrast with pattern macros, like Rust’s macros made with macro_rules!, which match a pattern of syntax and transform that into new syntax just by filling holes in a pattern. Pattern macros are immensely useful, but their power is a strict subset of procedural macros.

Essentials of building a typed language with macros #

Let’s take a quick look at the language we want to write:

;; Identify function
;; λx :: a → a; in Racket notation, (→ a a)
(λ ([x : a]) x)

;; Function call
;; λx :: (b → b) → (b → b) takes an argument of type b → b (which is what λy is)
((λ ([x : (→ b b)]) x) (λ ([y : b]) y))

;; More complicated function
;; λfx :: (a → b) → a → b
(λ ([f : (→ a b)] [x : a]) (f x))

Our macros should expand, check, and rewrite this into something that looks like this:

;; Identify function again
(λ (x) x)

;; Function call
((λ (x) x) (λ (y) y))

;; More complicated function
(λ (f x) (f x))

All the types have been erased. The trick is: we want to do this so that we catch any type errors. That way, if we mistakenly write:

(λ ([f : (→ a b)] [x : c]) (f x))  ; x is the wrong type!

We should get a compile-time type error, and our macros will refuse to expand this into plain Racket.

; play_stlc.rkt:16:27: #%app: play_stlc.rkt: Function expected args with type (a), got type (c) instead

How are we going to do this? Basically, we want to interpose on every function creation and function call to check that the types at these locations line up. We’ll be working with a very simple language, but this kind of work generalizes to the kinds of languages you’re used to working with.

The overall architecture looks like this:

The important thing to remember is this: expanding erases types and attaches type information to the resulting syntax while checking that the types are consistent throughout expansion.

Let’s take a look at how we would implement that middle block.

Interposing on every function definition and call site #

The first thing we need to do is to be able to inspect every function definition and call site. In other languages, this might only be possible by writing some kind of wrapping macro that walks down the AST of your entire module and replaces ordinary function calls into a call to a macro you control.

However, we’re using Racket, and Racket is designed to make language construction easy! Racket’s reader will expand function calls like (foo bar baz) into calls to a special function #%app: (#%app foo bar baz). We can override what this symbol refers to when our code gets imported as a #lang with rename-out:

#lang racket

(require syntax/parse/define)

(provide #%module-begin #%top-interaction) ; needed for #lang

(provide (rename-out [checked-app #%app]
                     [checked-λ λ]))

The highlighted portion means to rename checked-app to #%app and checked-λ to λ in the requiring module’s scope. See the docs on #lang creation for more details.

With that in place, we can write some macros checked-app and checked-λ which will then be used when we import stlc.rkt as a #lang:

(define-syntax (checked-app stx) …)
(define-syntax (checked-λ stx) …)

#lang s-exp "stlc.rkt"

;; This λ referrs to checked-λ
(λ ([x : a]) x)

Implementing the checks #

We’ve seen a few clever tricks that we can exploit in Racket to automatically interpose on every function definition and call. Now we begin the type checking.

Core macros #

Let’s first take a look at the stars of the show: checked-λ and checked-app.

Remember: these are macros that our source language will use under the names λ and whenever there’s a normal function call. These macros need to expand to untyped Racket, albeit with the type information added as a syntax property.

First macro: checked-λ #

48
49
50
51
52
53
54
55
56

(define-syntax (checked-λ stx)
  (syntax-parse stx
    [(_ ([x (~literal :) τ_in] ...) e)                                  (pat)
     #:with [-xs -e τ_out] (comp+erase-τ/ctx #'e #'([x τ_in] ...))      (comp-erase)
     #:do [(printf "derived ~a :: ~a → ~a\n"
                   (syntax->datum stx)
                   (syntax->datum #'(τ_in ...))
                   (syntax->datum) #'τ_out)]
     (add-τ #'(λ -xs -e) #'(→ τ_in ... τ_out))]))                       (lam-return)

We’re defining a macro checked-λ, which is invoked any time a λ is used in the module requiring this one, thanks to the rename-out bit.

We’re using the excellent syntax-parse macro, introduced in [2] I believe. syntax-parse gives us fancy keywords like #:with, #:when, and #:fail-unless¹ which gives us a rich language to build macros with.²

Let’s step through this line-by-line.

The line pat describes the pattern we’re matching in stx: _ is a wildcard, and in our case will just match checked-λ. Next we have a list with elements looking like [var_name : type_name]. We bind all the variable names to x and the corresponding types to τ_in. Finally we bind some expression to the variable e. Syntax like (checked-λ ([x : a] [y : (→ a b)]) (y x)) should match this.

The next line comp-erase calls a helper function comp+erase-τ/ctx (yes that’s a mouthful) to compute the type of the body whilst erasing the type annotations at the same time. Moreover, it should do this type computation in the context of the variables x ... and their types τ_in .... This just means that when computing the type of the body of (λ ([x : Int]) {lambda body here...}), it should know that variable x has type Int.

comp+erase-τ/ctx returns three things:

-xs

the variables x ... but with the type annotations erased

-e

the body e but, again, with the type annotations erased from the source

τ_out

the derived return type of the function

We’ll dig into how this function works in Support functions: compute and erase types.

The call to printf just shows us some pretty information about the work that the macro type checker is doing. It’s not necessary and you could delete lines 52–55 with no consequences.

In the final line lam-return we return some new syntax #'(λ -xs -e). Note that this bit of syntax has no type annotations! We do add the type to this bit of syntax with add-τ—we will need this later. The type of a λ is an arrow type with arguments provided by the type annotations on the arguments, and the return type of the function supplied by the derived type of the expression from the comp-erase line.

Remember: when our macros expand, they check the types and erase them. We don’t have any checks to do here in checked-λ, but the type information we’ve added will come in handy when we check the arguments to this function in the next macro.

Second macro: checked-app #

58
59
60
61
62
63
64
65
66
67
68

(define-syntax (checked-app stx)
  (syntax-parse stx
    [(_ e_fn e_arg ...)                                                          (app-pat)
     #:with [-e_fn (_ τ_in ... τ_out)] (comp+erase-τ #'e_fn)                     (fn-type)
     #:with ([-e_arg τ_arg] ...) (stx-map comp+erase-τ #'(e_arg ...))            (arg-types)
     #:fail-unless (τ= #'(τ_arg ...) #'(τ_in ...))                               (typecheck)
     (format "~a: Function expected args with type ~a, got type ~a instead"
             (syntax-source stx)
             (syntax->datum #'(τ_in ...))
             (syntax->datum #'(τ_arg ...)))
     (add-τ #'(#%app -e_fn -e_arg ...) #'τ_out)]))                               (app-return)

Like before, let’s go through this line-by-line.

The line app-pat describes the pattern for function application: the first thing you see e_fn is a function, and the rest e_arg ... are arguments.

Next, in fn-type we expand the syntax of the function, which means the syntax we get back (-e_fn) has passed type checking and now has the type annotations erased. The erased syntax does have its type stuck on it, lam-return which we retrieve in comp+erase-τ with get-τ and store in τ_in ... τ_out.

Then we compute the types of the arguments using the same function comp+erase-τ; this time we use stx-map³ to map over the syntax object #'(e_arg …).

Once we have the type of the function (→ τ_in ... τ_out) and the types of the arguments (τ_arg ...) it becomes a simple matter to check that these types match. typecheck If they don’t, the #:fail-unless keyword will break the macro expansion with a parse error and will display a nicely-formatted message, complete with source information from syntax-source as well as expected and received types.

Like with checked-λ, we return the function application with all the type annotations erased in both the function and the arguments, but we tack on the function’s return type τ_out to the syntax we build and return: #'(#%app -e_fn -e_arg ...).

#:with ([-e_arg τ_arg] ...) (map comp+erase-τ (syntax->list #'(e_arg ...)))

Support functions: compute and erase types #

We’ve covered the major macros, but now we need to inspect the helper functions comp+erase-τ and comp+erase-τ/ctx.

Computing and erasing types: comp+erase-τ #

The first support function that we need is comp+erase-τ, which we call from the checked-app macro on lines fn-type and arg-types.

Note that we do not expand to a call to comp+erase-τ with the checked-app macro—we call this function _as part of expansion_​—i.e. this function gets called at compile-time.

Normally we define functions with define. But, because we need this function to be available for macros to call, we use define-for-syntax. This makes the function available at compile time. More about this in § Putting the phases together.

29
30
31
32

(define-for-syntax (comp+erase-τ e)
  (let* ([-e (expand-stx e)]
         [τ (get-τ -e)])
    `[,-e ,τ]))

This function is pretty simple: we take in some syntax e, expand the macros in e to get -e, which is just the same as e, but with all the type annotations erased and some type information added as a syntax property. Then we use a helper function get-τ (defined in § Convenience functions) to pull the type off of -e. Finally we return -e and its associated type τ.

Again, but with context: comp+erase-τ/ctx #

Now we need to do the same thing, but this time, we’ll be adding a context to the expansion.

What is a context? Consider the following function:

(λ ([x : Int] [y : Bool])
  (… body using variables x and y …))

Inside the body of the function, we need to somehow remember that the variable x has type Int and y has type Bool. That way, if we run into something like (+ x y) we can know that y isn’t the right type for that operation.

comp+erase-τ/ctx does the same thing and comp+erase-τ, but we add a context, which is just a list of variable ↦ type mappings. Like comp+erase-τ, we define comp+erase-τ/ctx with define-for-syntax so it’s available for checked-λ to use at compile time.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

(define-for-syntax (comp+erase-τ/ctx e ctx)
  (syntax-parse ctx
    [([x τ] ...)                                                                          (ctx-destruct)
     #:with (y ...)                                                                       (temps)
     (generate-temporaries #'(x ...))
     #:with (λ xs- e-)                                                                    (erased-version)
     (expand-stx #`(λ (y ...)
                     (let-syntax ([x (make-rename-transformer (add-τ #'y #'τ))] ...)      (rename-transformer)
                       #,e)))
     #:with τ_out (get-τ #'e-)                                                            (body-type)
     #'[xs- e- τ_out]]))

We use syntax-parse here just to destruct the list of variable ↦ type pairs [x τ] from ctx. ctx-destruct Then, for each of the variables x ..., we create a brand new variables y ... with generate-temporaries. temps

Now comes the tricky bit: we’d like to expand e, but remembering that a given variable x has type τ. To accomplish this, we do something a little sneaky: we create a lexically-scoped macro with let-syntax that has the name x. We bind this to a macro that expands to the brand new variable y that has the type information for x attached to it.

This is what the let-syntax and make-rename-transformer bit does: it makes a macro that can expand in places other than application positions, (i.e. identifier position) along with a host of other convenience functions.

So, starting with a lambda like this:

(λ ([a : Int] [b : Bool])
  (and b (zero? a)))

comp+erase-τ/ctx creates some new syntax that looks like this:

(λ (y_tmp1 y_tmp2)
  (let-syntax ([a (make-rename-transformer (add-τ #'y_tmp1 #'Int))]
               [b (make-rename-transformer (add-τ #'y_tmp2 #'Bool))])
    (and b (zero? a))))

Which gets expanded into:

(λ (y_tmp1 y_tmp2)
  (and y_tmp2 (zero? y_tmp1)))

But the symbols y_tmp2 and y_tmp1 on the second line have extra type information attached to them to help with the remainder of the type checking: add-τ simply returns the syntax you gave it in the first place, but with the second thing attached as type metadata.

Now we’ve expanded the body e of the original λ that we got but with the variables and the type information added as context. Now we simply pull off the type data for the body—i.e. the function return type—with (get-τ e-) and return the erased variables, body, and function return type.

Normally we wouldn’t have to worry about details like this, but because of how we’re using local-expand via expand-stx and the let-syntax clause inside the macro, the body of the λ gets wrapped in two layers of let-values, so we end up having to destructure all of that. The real implementation of comp+erase-τ/ctx ends up looking a little messier:

34
35
36
37
38
39
40
41
42
43
44
45
46

(define-for-syntax (comp+erase-τ/ctx e ctx)
  (syntax-parse ctx
    [([x τ] ...)
     #:with (y ...) (generate-temporaries #'(x ...))

     ;; Not sure why I need to unwrap the let-values… must be something with how
     ;; Racket automatically wraps the body of λ's.
     ;; This is consistent with the paper's implementation.
     #:with ((~literal #%plain-lambda) xs- ((~literal let-values) () ((~literal let-values) () e-)))
     (expand-stx #`(λ (y ...) (let-syntax ([x (make-rename-transformer (add-τ #'y #'τ))] ...) #,e)))

     #:with τ_out (get-τ #'e-)
     #'[xs- e- τ_out]]))

The ~literal just tells the pattern matching that we’re looking for the identifier let-values in the syntax, and we’re not trying to have some variable named let-values that we’d like to bind to whatever is there.⁴

Convenience functions: working with type information #

We are going to go from explicitly typed code to untyped code, but with type annotations on the syntax objects (AST nodes). It happens to be a convenient place to put the derived type of the data, but there’s nothing intrinsically important about this, as long as you can tie a particular syntax object back to its derived type somehow.

We define two functions: add-τ and get-τ to put and fetch data into the 'type field of a syntax object.

;; Convenience functions
(define (add-τ e t) (syntax-property e 'type t))
(define (get-τ e) (syntax-property e 'type))

We also define type equality here. For this demo, type quality is syntactic equality in the type. E.g. a ≡ a and a → b ≡ a → b, but a ≢ b and a → a ≢ b → b. This is a choice we’re making to keep things simple. It’s possible (and the authors in [1] do this) to define alternate type comparison operations to support existential types as well as subtyping.

;; Dumb equality check: just check for syntactic equality
(define (τ= t1s t2s)
  (define t1 (syntax->datum t1s))
  (define t2 (syntax->datum t2s))
  (equal? t1 t2))

We will want one small piece of convenience code in the next part:

(define-syntax-rule (expand-stx stx) (local-expand stx 'expression null))

local-expand is a function that takes some syntax and expands all the macros in it. It’s a little more complicated than that: local-expand expands the syntax in the lexical context of the currently expanding expression. What that means is that you can only call local-expand during macro expansion. This makes it hard to play with at the top-level REPL as a Lisp/Scheme is wont to do. If you are wanting to play around in the REPL, expand and expand-once are your friends. Inside our use case, expanding in the lexical context of the currently expanding macro makes this macro expansion safer.

expand-stx simply makes it so we don’t have to include so many arguments to local-expand, as they’d be the same every time we call it. See the documentation for local-expand for more details on what these arguments do.

Putting the phases together #

We’ve covered all the essential bits, but we haven’t talked about what functions need to be available when: this isn’t such a big deal in Lisp or Scheme when you can call pretty much any function from any macro, but Racket takes great pains to ensure that there is a clear dependency graph between macro expansion time and runtime. See Flatt’s Composable and Compilable Macros: You want it when? paper [3] for all the details you could ever want on this. That paper does a good job of motivating why you need clear phase separation.

That said, let’s talk about the dependency graph of the functions and macros in our code.

Normally, functions are defined at phase level 0, which is the “runtime” phase of the module. Macros are defined at phase level 1, which is the “compile time” or “expand time” phase. You can only refer to things defined at the same or higher phase level. So, normally, macros can’t call functions in the same module.

Our two main macros, checked-λ and checked-app, both depend on the functions comp+erase-τ and comp+erase-τ/ctx, along with a few other auxiliary functions. These functions in turn depend on the expand-stx macro.

The purple box is at phase level 1, since checked-λ and checked-app are both macros. Normally, the functions in the orange box would be at phase level 0, and would be unavailable for our macros. However, we can raise the phase level by using define-for-syntax instead of define, or wrap the things we need at a higher phase level in begin-for-syntax.

All the functions now at phase level 1 use the expand-stx macro, so it has to be defined at phase level 2—we can simply put the define-syntax⁵ inside of define-for-syntax.

That’s it! The full solution below puts all the functions together at the right phases.

Full solution #

The code can also be cloned from my Git repository on SourceHut.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

#lang racket

;; The Simply Typed Lambda Calculus

(require syntax/parse/define)

;; make this a sensible #lang
(provide #%module-begin #%top-interaction)
(provide (rename-out [checked-app #%app]
                     [checked-λ λ]))

(begin-for-syntax
  (require syntax/stx)                  ; needed for stx-map

  ;; Shorthand to expand macros in syntax
  ;; Yes, this is a macro for use inside other macros (or functions at phase level 1)
  (define-syntax-rule (expand-stx stx) (local-expand stx 'expression null))

  ;; Convenience functions
  (define (add-τ e t) (syntax-property e 'type t))
  (define (get-τ e) (syntax-property e 'type))

  ;; Dumb equality check: just check for syntactic equality
  (define (τ= t1s t2s)
    (define t1 (syntax->datum t1s))
    (define t2 (syntax->datum t2s))
    (equal? t1 t2)))

(define-for-syntax (comp+erase-τ e)
  (let* ([-e (expand-stx e)]
         [τ (get-τ -e)])
    `[,-e ,τ]))

(define-for-syntax (comp+erase-τ/ctx e ctx)
  (syntax-parse ctx
    [([x τ] ...)
     #:with (y ...) (generate-temporaries #'(x ...))

     ;; Not sure why I need to unwrap the let-values… must be something with how
     ;; Racket automatically wraps the body of λ's.
     ;; This is consistent with the paper's implementation.
     #:with ((~literal #%plain-lambda) xs- ((~literal let-values) () ((~literal let-values) () e-)))
     (expand-stx #`(λ (y ...) (let-syntax ([x (make-rename-transformer (add-τ #'y #'τ))] ...) #,e)))

     #:with τ_out (get-τ #'e-)
     #'[xs- e- τ_out]]))

(define-syntax (checked-λ stx)
  (syntax-parse stx
    [(_ ([x (~datum :) τ_in] ...) e)
     #:with [-xs -e τ_out] (comp+erase-τ/ctx #'e #'([x τ_in] ...))
     #:do [(printf "derived ~a :: ~a → ~a\n"
                   (syntax->datum stx)
                   (syntax->datum #'(τ_in ...))
                   (syntax->datum #'τ_out))]
     (add-τ #'(λ -xs -e) #'(→ τ_in ... τ_out))]))

(define-syntax (checked-app stx)
  (syntax-parse stx
    [(_ e_fn e_arg ...)
     #:with [-e_fn (_ τ_in ... τ_out)] (comp+erase-τ #'e_fn)
     #:with ([-e_arg τ_arg] ...) (stx-map comp+erase-τ #'(e_arg ...))
     #:fail-unless (τ= #'(τ_arg ...) #'(τ_in ...))
     (format "~a: Function expected args with type ~a, got type ~a instead"
             (syntax-source stx)
             (syntax->datum #'(τ_in ...))
             (syntax->datum #'(τ_arg ...)))
     (add-τ #'(#%app -e_fn -e_arg ...) #'τ_out)]))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

#lang s-exp "stlc.rkt"

;; Smallest example: function a → a (in Racket notation, (→ a a))
(λ ([x : a]) x)

;; a bigger example
((λ ([x : (→ b b)]) x) (λ ([y : b]) y))

;; ;; Functions!
(λ ([f : (→ a b)] [x : a]) (f x))

;; an example of an ill-typed expression
;; ((λ ([x : (→ b b)]) x) (λ ([f : (→ b c)] [y : b]) (f y)))

Acknowledgments #

Thanks to Ben Greenman for reading and providing suggestions and improvements on this post.

References #