15 Jig: jumping to tail calls

Source code.

15.1 Tail Calls

With Iniquity, we’ve finally introduced some computational power via the mechanism of functions and function calls. Together with the notion of inductive data, which we have in the form of pairs, we can write fixed-sized programs that operate over arbitrarily large data.

The problem, however, is that there are a class of programs that should operate with a fixed amount of memory, but instead consume memory in proportion to the size of the data they operate on. This is unfortunate because a design flaw in our compiler now leads to asympototically bad run-times.

We can correct this problem by generating space-efficient code for function calls when those calls are in tail position.

Let’s call this language Jig.

There are no syntactic additions required: we simply will properly handling function calls.

15.2 What is a Tail Call?

A tail call is a function call that occurs in tail position. What is tail position and why is important to consider function calls made in this position?

Tail position captures the notion of “the last subexpression that needs to be computed.” If the whole program is some expression e, the e is in tail position. Computing e is the last thing (it’s the only thing!) the program needs to compute.

Let’s look at some examples to get a sense of the subexpressions in tail position. If e is in tail position and e is of the form:

• (let ((x e0)) e1): then e1 is in tail position, while e0 is not. The reason e0 is not in tail position is because after evaluating it, the e1 still needs to be evaluated. On the other hand, once e0 is evaluated, then whatever e1 evaluates to is what the whole let-expression evaluates to; it is all that remains to compute.

• (if e0 e1 e2): then both e1 and e2 are in tail position, while e0 is not. After the e0, then based on its result, either e1 or e2 is evaluated, but whichever it is determines the result of the if expression.

• (+ e0 e1): then neither e0 or e1 are in tail position because after both are evaluated, their results still must be added together.

• (f e0 ...), where f is a function (define (f x ...) e): then none of the arguments e0 ... are in tail position, because after evaluating them, the function still needs to be applied, but the body of the function, e is in tail position.

The significance of tail position is relevant to the compilation of calls. Consider the compilation of a call as described in Iniquity: function definitions and calls: a return address is pushed on the stack, arguments are pushed on the call stack, then the Jmp instruction is issued, which jumps to the called function. The function computes its result, pops its arguments, the pops the return address and jumps back to the caller.

But if the call is in tail position, what else is there to do? Nothing. So after the call, return transfers back to the caller, who then just pops their arguments and returns to their caller.

This leads to unconditional stack space consumption on every function call, even function calls that don’t need to consume space.

Consider this program:

; (Listof Number) -> Number

(define (sum xs) (sum/acc xs 0))

; (Listof Number) Number -> Number

(define (sum/acc xs a)

(if (empty? xs)

a

(sum/acc (cdr xs) (+ (car xs) a))))

The sum/acc function should operate as efficiently as a loop that iterates over the elements of a list accumulating their sum. But, as currently compiled, the function will push stack frames for each call.

Matters become worse if we were re-write this program in a seemingly benign way to locally bind a variable:

; (Listof Number) Number -> Number

(define (sum/acc xs a)

(if (empty? xs)

a

(let ((b (+ (car xs) a)))

(sum/acc (cdr xs) b))))

Now the function pushes a return point and a local binding for b on every recursive call.

But we know that whatever the recursive call produces is the answer to the overall call to sum. There’s no need for a new return point and there’s no need to keep the local binding of b since there’s no way this program can depend on it after the recursive call. Instead of pushing a new, useless, return point, we should make the call with whatever the current return address is, because that’s where control is going to jump to anyway. This is the idea of proper tail calls.

An axe to grind: the notion of proper tail calls is often referred to with misleading terminology such as tail call optimization or tail recursion. Optimization seems to imply it is a nice, but optional strategy for implementing function calls. Consequently, a large number of mainstream programming languages, most notably Java, do not properly implement tail calls. But a language without proper tail calls is fundamentally broken. It means that functions cannot reliably be designed to match the structure of the data they operate on. It means iteration cannot be expressed with function calls. There’s really no justification for it. It’s just broken. Similarly, it’s not about recursion alone (although it is critical for recursion), it really is about getting function calls, all calls, right. /rant

15.3 An Interpreter for Proper Calls

Before addressing the issue of compiling proper tail calls, let’s first think about the interpreter, starting from the interpreter we wrote for Iniquity:

iniquity/interp.rkt

#lang racket

(provide interp interp-env)

(require "ast.rkt"

"env.rkt"

"interp-prims.rkt")

;; type Answer = Value | 'err

;; type Value =

;; | Integer

;; | Boolean

;; | Character

;; | Eof

;; | Void

;; | '()

;; | (cons Value Value)

;; | (box Value)

;; | (vector Value ...)

;; | (string Char ...)

;; type REnv = (Listof (List Id Value))

;; type Defns = (Listof Defn)

;; Prog -> Answer

(define (interp p)

(match p

[(Prog ds e)

(interp-env e '() ds)]))

;; Expr Env Defns -> Answer

(define (interp-env e r ds)

(match e

[(Int i)  i]

[(Bool b) b]

[(Char c) c]

[(Eof)    eof]

[(Empty)  '()]

[(Var x)  (lookup r x)]

[(Str s)  s]

[(Prim0 'void) (void)]

[(Prim0 'read-byte) (read-byte)]

[(Prim0 'peek-byte) (peek-byte)]

[(Prim1 p e)

(match (interp-env e r ds)

['err 'err]

[v (interp-prim1 p v)])]

[(Prim2 p e1 e2)

(match (interp-env e1 r ds)

['err 'err]

[v1 (match (interp-env e2 r ds)

['err 'err]

[v2 (interp-prim2 p v1 v2)])])]

[(Prim3 p e1 e2 e3)

(match (interp-env e1 r ds)

['err 'err]

[v1 (match (interp-env e2 r ds)

['err 'err]

[v2 (match (interp-env e3 r ds)

['err 'err]

[v3 (interp-prim3 p v1 v2 v3)])])])]

[(If p e1 e2)

(match (interp-env p r ds)

['err 'err]

[v

(if v

(interp-env e1 r ds)

(interp-env e2 r ds))])]

[(Begin e1 e2)

(match (interp-env e1 r ds)

['err 'err]

[_    (interp-env e2 r ds)])]

[(Let x e1 e2)

(match (interp-env e1 r ds)

['err 'err]

[v (interp-env e2 (ext r x v) ds)])]

[(App f es)

(match (interp-env* es r ds)

['err 'err]

[vs

(match (defns-lookup ds f)

[(Defn f xs e)

; check arity matches

(if (= (length xs) (length vs))

(interp-env e (zip xs vs) ds)

'err)])])]))

;; (Listof Expr) REnv Defns -> (Listof Value) | 'err

(define (interp-env* es r ds)

(match es

['() '()]

[(cons e es)

(match (interp-env e r ds)

['err 'err]

[v (match (interp-env* es r ds)

['err 'err]

[vs (cons v vs)])])]))

;; Defns Symbol -> Defn

(define (defns-lookup ds f)

(findf (match-lambda [(Defn g _ _) (eq? f g)])

ds))

(define (zip xs ys)

(match* (xs ys)

[('() '()) '()]

[((cons x xs) (cons y ys))

(cons (list x y)

(zip xs ys))]))

What needs to be done to make it implement proper tail calls?

Well... not much. Notice how every Iniquity subexpression that is in tail position is interpreted by a call to interp-env that is itself in tail position in the Racket program!

So long as Racket implements tail calls properly, which is does, then this interpreter implements tail calls properly. The interpreter inherits the property of proper tail calls from the meta-language. This is but one reason to do tail calls correctly. Had we transliterated this program to Java, we’d be in trouble as the interpeter would inherit the lack of tail calls and we would have to re-write the interpreter, but as it is, we’re already done.

15.4 A Compiler with Proper Tail Calls

Consider the following program:

(define (f x)

(if (zero? x)

42

(f (sub1 x))))

(f 100)

It’s a silly program, but it will help illuminate what tail calls are all about and how we can make them work.

Here’s what this code will compile to, roughly:

Examples

> (asm-interp    (seq (Global 'entry)         (Label 'entry)            ; calling (f 100), so set up return address,         ; push argument, then jump         (Lea 'rax 'r1)         (Push 'rax)         (Mov 'rax 100)         (Push 'rax)         (Jmp 'f)         (Label 'r1)            ; done with (f 100), return         (Ret)            ; (define (f x) ...)         (Label 'f)         (Mov 'rax (Offset 'rsp 0))         (Cmp 'rax 0)         (Jne 'if_false)            ; if-then branch         (Mov 'rax 42)         (Jmp 'done)            ; if-else branch         (Label 'if_false)         ; calling (f (sub1 x)), so set up return address,         ; push argument, then jump         (Lea 'rax 'r2)         (Push 'rax)         (Mov 'rax (Offset 'rsp 8))         (Sub 'rax 1)         (Push 'rax)         (Jmp 'f)         (Label 'r2)            (Label 'done)         (Add 'rsp 8)  ; pop x         (Ret)))

42

Now let’s think about how this computes, paying attention to the stack.

First, the run-time system would call 'entry, so there’s going to be an address on the stack telling us where to return to when the program is done:

Next, the call to (f 100) is set up, pushing the address of 'r1 for where the call should return to, and then pushing the argument, 100. So before the Jmp to 'f, the stack looks like:

Control jumps to 'f, which asks if x is 0 by referencing the argument on the top of the stack. It is not, so control jumps to the 'if_false label, which now sets up the call (f (sub1 x)) by computing a return address for 'r2, pushing it, subtracting 1 from x, and pushing that, then jumping to 'f.

Now the stack looks like:

This asks if x is 0 by referencing the argument on the top of the stack (now: 99). It is not, so control jumps to the 'if_false label, which now sets up the call (f (sub1 x)) by computing a return address for 'r2, pushing it, subtracting 1 from x, and pushing that, then jumping to 'f.

You can see where this is going.

At this point, we make a final jump to 'f. Since x is 0, 42 is moved into 'rax, control jumps tp 'done, at which point we pop the current x off the stack, then return, which pops off the next frame of the stack. Since that frame says to jump to 'r2, that’s where control jumps to.

But 'r2 is the same as 'done! So we pop off the current x (now: 1) and return, which pops of the next frame saying jump to 'r2.

This process continues, popping two frames and jumping back to 'r2 until the stack looks like:

And we’re back at 'r2. Next we pop the current x (now: 100) and return, which pops the top frame off and jumps to 'r1. But the code following 'r1 is simply a Ret intruction, so we pop another frame (the stack is now empty) and jump back to the runtime system (with 'rax holding 42).

So to summarize, each call to f pushes two words on the stack: one for where to return and one for the argument to the function. Then, when the base case is finally reached, we spin through a loop popping all this information off.

Let’s take another look at the function:

(define (f x)

(if (zero? x)

42

(f (sub1 x))))

In the call to (f (sub1 x)), that expression is in a tail position of the function. Intuitively this means once you’ve computed the result of (f (sub1 x)) there’s nothing further to compute, you now have the answer for (f x). This suggests that you don’t need to keep the current binding for x on the stack; if there’s no further work to do, you can’t possibly need x. It also suggests there’s no need to return to the point after (f (sub1 x)); you could instead just return to the caller of (f x)!

We can modify the code to embody these ideas:

Examples

> (asm-interp    (seq (Global 'entry)         (Label 'entry)            ; calling (f 100), so set up return address,         ; push argument, then jump         (Lea 'rax 'r1)         (Push 'rax)         (Mov 'rax 100)         (Push 'rax)         (Jmp 'f)         (Label 'r1)            ; done with (f 100), return         (Ret)            ; (define (f x) ...)         (Label 'f)         (Mov 'rax (Offset 'rsp 0))         (Cmp 'rax 0)         (Jne 'if_false)            ; if-then branch         (Mov 'rax 42)         (Jmp 'done)            ; if-else branch         (Label 'if_false)         ; TAIL calling (f (sub1 x)),         ; so pop off the argument (don't need it anymore)         ; and don't push a new return address, just leave         ; our caller's return address on stack         (Mov 'rax (Offset 'rsp 0))         (Sub 'rax 1)         (Add 'rsp 8) ; pop x         (Push 'rax)  ; push arg         (Jmp 'f)            (Label 'done)         (Add 'rsp 8)  ; pop x         (Ret)))

42

Let’s step through the computation again. It starts off the same: the runtime calls 'entry, which sets up the call to (f 100), so when control jumps to 'f, the stack again looks like:

The checks if x is 0, which it is not, so jumps to 'if_false. Now the code computes x-1 and then pops x, and pushes x-1, so when we jump to 'f, the stack looks like:

Again we go through the same instructions, popping x and pushing x-1, then jumping to 'f with stack:

This continues, but the stack never grows further, until finally jumping to 'f with the stack:

At which point, 42 is moved in to 'rax and control jumps to 'done, where x is popped and then Ret pops the next frame and jumps to 'r1, which issues another Ret, popping the stack (now empty) and jumping back to the runtime system.

So this program makes 100 calls to f, but it uses a constant amount of stack space. Indeed, we could’ve made it (f 200) and it would still use the same three stack frames. The program is really computing a loop.

Moreover, we can do one better: notice that the initial call (f 100) is itself in a tail position: whatever it’s result is is the result of the whole program. We can turn this in to a tail call:

Examples

> (asm-interp    (seq (Global 'entry)         (Label 'entry)            ; TAIL calling (f 100),         ; no args to pop         ; don't push a new return address, just leave         ; our caller's return address on stack         (Mov 'rax 100)         (Push 'rax)         (Jmp 'f)            ; No need for this since we never come back:         ; (Ret)            ; (define (f x) ...)         (Label 'f)         (Mov 'rax (Offset 'rsp 0))         (Cmp 'rax 0)         (Jne 'if_false)            ; if-then branch         (Mov 'rax 42)         (Jmp 'done)            ; if-else branch         (Label 'if_false)         ; TAIL calling (f (sub1 x)),         ; so pop off the argument (don't need it anymore)         ; and don't push a new return address, just leave         ; our caller's return address on stack         (Mov 'rax (Offset 'rsp 0))         (Sub 'rax 1)         (Add 'rsp 8) ; pop x         (Push 'rax)  ; push arg         (Jmp 'f)            (Label 'done)         (Add 'rsp 8)  ; pop x         (Ret)))

42

Now the stack looks like this:

decrementing until x reaches 0 at which point 42 is put in 'rax and control jumps back to the runtime system.

In general, when a function call (f e0 ...) is in tail position, there are going to be some number of things currently pushed on the stack, which are described by the current environment c. To carry out a tail call, we need to pop all of those things described by c, then push the values of e0 ... which are the arguments for f, then jump to f.

There is a problem here, which is that we need to evaluate the subexpressions e0 ... and doing so may depend on things in the current environment, e.g. they may reference bound variables.

So we have to wait to pop the things described by c until after evaluating e0 ..., but evaluating e0 ... will need to save the values somewhere... and that somewhere is the stack.

Let’s say we have an expression that looks like this:

(let ((x 1))

(let ((y 2))

(f (+ x y) 5)))

The call to f is in tail position and it will be compiled in a compile-time environment of '(y x). The compiler will need to compile (+ x y) in that same environment, but then emit code to save the result on the stack while the next argument is evaluated.

That means by the time the arguments are evaluated and the call is ready to be made, the stack will look like:

At which point we need to remove the x and y part, but then also have the arguments 3 and 5 sitting just below the return address, i.e. we want:

To accomplish, we rely on the following helper function for generating code that moves arguments on the stack:

; Integer Integer -> Asm

(define (move-args i off)

(cond [(zero? off) (seq)]

[(zero? i)   (seq)]

[else

(seq (Mov r8 (Offset rsp (* 8 (sub1 i))))

(Mov (Offset rsp (* 8 (+ off (sub1 i)))) r8)

(move-args (sub1 i) off))]))

It moves i elements on the stack up off-positions. So if you have (length c) items in the environment and n arguments, then (move-args n (length c)) will move the arguments the right spot, just below the return address.

Once the arguments are moved to the proper spot on the stack, we can pop off the local environment and jump. So the complete code for compiling a tail call is:

; Id [Listof Expr] CEnv -> Asm

(define (compile-app-tail f es c)

(seq (compile-es es c)

(move-args (length es) (length c))

(Add rsp (* 8 (length c)))

(Jmp (symbol->label f))))

What’s left is determining when to use this strategy for a function application and when to use the prior version which pushes a return pointer.

The way we do this is to add a parameter to the expression compiler, so the new signature is:

; Expr CEnv Bool -> Asm

(define (compile-e e c t?)

...)

Calling (compile-e e c #t) signals that the expression e should be compiled assuming it is in tail position, while (compile-e e c #f) signals it is not in tail position.

If e is an application, then the compiler selects between compile-app-nontail and compile-app-tail based on t?.

If e is any other kind of expression that has sub-expressions, then the compiler function for that form also adds a t? parameter and sets t? to #f for any that are not tail positions and passes on the t? given for those in tail position. For example, here is how begin is compiled:

; Expr Expr CEnv Bool -> Asm

(define (compile-begin e1 e2 c t?)

(seq (compile-e e1 c #f)

(compile-e e2 c t?)))

There are two important places where t? is seeded to #t:

• The top-level expression is in tail position.

• The body of every function is in tail position.

The complete compiler:

jig/compile.rkt

#lang racket

(provide (all-defined-out))

(require "ast.rkt" "types.rkt" "compile-ops.rkt" a86/ast)

;; Registers used

(define rax 'rax) ; return

(define rbx 'rbx) ; heap

(define rsp 'rsp) ; stack

(define rdi 'rdi) ; arg

;; type CEnv = [Listof Variable]

;; Prog -> Asm

(define (compile p)

(match p

[(Prog ds e)

(prog (externs)

(Global 'entry)

(Label 'entry)

(Push rbx)    ; save callee-saved register

(Mov rbx rdi) ; recv heap pointer

(compile-e e '() #f)

(Pop rbx)     ; restore callee-save register

(Ret)

(compile-defines ds)

(Label 'raise_error_align)

pad-stack

(Call 'raise_error))]))

(define (externs)

(seq (Extern 'peek_byte)

(Extern 'read_byte)

(Extern 'write_byte)

(Extern 'raise_error)))

;; [Listof Defn] -> Asm

(define (compile-defines ds)

(match ds

['() (seq)]

[(cons d ds)

(seq (compile-define d)

(compile-defines ds))]))

;; Defn -> Asm

(define (compile-define d)

(match d

[(Defn f xs e)

(seq (Label (symbol->label f))

(compile-e e (reverse xs) #t)

(Add rsp (* 8 (length xs))) ; pop args

(Ret))]))

;; Expr CEnv Bool -> Asm

(define (compile-e e c t?)

(match e

[(Int i)            (compile-value i)]

[(Bool b)           (compile-value b)]

[(Char c)           (compile-value c)]

[(Eof)              (compile-value eof)]

[(Empty)            (compile-value '())]

[(Var x)            (compile-variable x c)]

[(Str s)            (compile-string s)]

[(Prim0 p)          (compile-prim0 p c)]

[(Prim1 p e)        (compile-prim1 p e c)]

[(Prim2 p e1 e2)    (compile-prim2 p e1 e2 c)]

[(Prim3 p e1 e2 e3) (compile-prim3 p e1 e2 e3 c)]

[(If e1 e2 e3)      (compile-if e1 e2 e3 c t?)]

[(Begin e1 e2)      (compile-begin e1 e2 c t?)]

[(Let x e1 e2)      (compile-let x e1 e2 c t?)]

[(App f es)         (compile-app f es c t?)]))

;; Value -> Asm

(define (compile-value v)

(seq (Mov rax (value->bits v))))

;; Id CEnv -> Asm

(define (compile-variable x c)

(let ((i (lookup x c)))

(seq (Mov rax (Offset rsp i)))))

;; String -> Asm

(define (compile-string s)

(let ((len (string-length s)))

(if (zero? len)

(seq (Mov rax type-str))

(seq (Mov rax len)

(Mov (Offset rbx 0) rax)

(compile-string-chars (string->list s) 8)

(Mov rax rbx)

(Or rax type-str)

(Add rbx

(+ 8 (* 4 (if (odd? len) (add1 len) len))))))))

;; [Listof Char] Integer -> Asm

(define (compile-string-chars cs i)

(match cs

['() (seq)]

[(cons c cs)

(seq (Mov rax (char->integer c))

(Mov (Offset rbx i) 'eax)

(compile-string-chars cs (+ 4 i)))]))

;; Op0 CEnv -> Asm

(define (compile-prim0 p c)

(compile-op0 p))

;; Op1 Expr CEnv -> Asm

(define (compile-prim1 p e c)

(seq (compile-e e c #f)

(compile-op1 p)))

;; Op2 Expr Expr CEnv -> Asm

(define (compile-prim2 p e1 e2 c)

(seq (compile-e e1 c #f)

(Push rax)

(compile-e e2 (cons #f c) #f)

(compile-op2 p)))

;; Op3 Expr Expr Expr CEnv -> Asm

(define (compile-prim3 p e1 e2 e3 c)

(seq (compile-e e1 c #f)

(Push rax)

(compile-e e2 (cons #f c) #f)

(Push rax)

(compile-e e3 (cons #f (cons #f c)) #f)

(compile-op3 p)))

;; Expr Expr Expr CEnv Bool -> Asm

(define (compile-if e1 e2 e3 c t?)

(let ((l1 (gensym 'if))

(l2 (gensym 'if)))

(seq (compile-e e1 c #f)

(Cmp rax val-false)

(Je l1)

(compile-e e2 c t?)

(Jmp l2)

(Label l1)

(compile-e e3 c t?)

(Label l2))))

;; Expr Expr CEnv Bool -> Asm

(define (compile-begin e1 e2 c t?)

(seq (compile-e e1 c #f)

(compile-e e2 c t?)))

;; Id Expr Expr CEnv Bool -> Asm

(define (compile-let x e1 e2 c t?)

(seq (compile-e e1 c #f)

(Push rax)

(compile-e e2 (cons x c) t?)

(Add rsp 8)))

;; Id [Listof Expr] CEnv Bool -> Asm

(define (compile-app f es c t?)

(if t?

(compile-app-tail f es c)

(compile-app-nontail f es c)))

;; Id [Listof Expr] CEnv -> Asm

(define (compile-app-tail f es c)

(seq (compile-es es c)

(move-args (length es) (length c))

(Add rsp (* 8 (length c)))

(Jmp (symbol->label f))))

;; Integer Integer -> Asm

(define (move-args i off)

(cond [(zero? off) (seq)]

[(zero? i)   (seq)]

[else

(seq (Mov r8 (Offset rsp (* 8 (sub1 i))))

(Mov (Offset rsp (* 8 (+ off (sub1 i)))) r8)

(move-args (sub1 i) off))]))

;; Id [Listof Expr] CEnv -> Asm

;; The return address is placed above the arguments, so callee pops

;; arguments and return address is next frame

(define (compile-app-nontail f es c)

(let ((r (gensym 'ret)))

(seq (Lea rax r)

(Push rax)

(compile-es es (cons #f c))

(Jmp (symbol->label f))

(Label r))))

;; [Listof Expr] CEnv -> Asm

(define (compile-es es c)

(match es

['() '()]

[(cons e es)

(seq (compile-e e c #f)

(Push rax)

(compile-es es (cons #f c)))]))

;; Id CEnv -> Integer

(define (lookup x cenv)

(match cenv

['() (error "undefined variable:" x)]

[(cons y rest)

(match (eq? x y)

[#t 0]

[#f (+ 8 (lookup x rest))])]))

;; Symbol -> Label

;; Produce a symbol that is a valid Nasm label

(define (symbol->label s)

(string->symbol

(string-append

"label_"

(list->string

(map (λ (c)

(if (or (char<=? #\a c #\z)

(char<=? #\A c #\Z)

(char<=? #\0 c #\9)

(memq c '(#\_ #\$ #\# #\@ #\~ #\. #\?)))

c

#\_))

(string->list (symbol->string s))))

"_"

(number->string (eq-hash-code s) 16))))