Tail Calls and Local Regions in OCaml

This episode continues our series on local mode in the OCaml compiler branch being developed at Jane Street. I want to show how local values interact with tail calls, using a small list-processing example.

The example counts occurrences of the first element in a list. For [3; 1; 3; 8; 3; 3; 4], the result should be four. We will first implement it normally, then introduce locals and examine the resulting errors.

For an empty list, count_firsts returns zero. Otherwise, it binds the head as x and defines a recursive count_occs function over the tail. An empty tail returns one for the original head; each later match adds one, while a non-match continues without adding.

The recursive function closes over x, and the outer function calls it on the remaining list. Compiling and running the program prints four. The list in this particular example consists entirely of constants, so OCaml can avoid heap allocation already, but a general list might require it.

I annotate count_firsts to accept a local list and return an int. Compilation now reports that a local value escapes its region, with a hint that the function cannot be local because the call is a tail call. At first this is surprising: the returned integer cannot contain the list, and we have stored it in no reference.

Pattern matching a local list makes x and xs local in this example. Because count_occs closes over x, its closure must also be local: putting it on the heap could leave a heap pointer to stack data after the local region ends.

Replacing the captured x with the constant 3 makes the error disappear, confirming that the captured local variable forces the closure to be local. But we actually want to count occurrences of x, so we need to understand the tail call.

If I bind the result of count_occs to a variable and then return that variable, the call is no longer in tail position and the program compiles and runs. The distinction comes from tail-call optimization.

When the last operation in a function is another function call, tail-call optimization can transfer control directly to the callee instead of returning to the caller first. The callee eventually returns to the original caller, avoiding a call frame that would have no further work to do.

A function’s local region ends immediately before its tail call, if it has one. Thus the local count_occs closure belongs to count_firsts’ region, but using it in the tail call would let it escape that region. Binding the integer result first works because the returned integer does not retain the closure or list.

The compiler branch provides an explicit alternative: annotate the call with [@nontail]. That prevents tail-call optimization at this site, so the local region remains available for the ordinary call and the program compiles.

This annotation changes runtime behavior: the ordinary call adds a call-stack frame. For this one call the speaker says the extra allocation takes very little time, but it is real additional stack usage.

There is a subtler question: why can count_occs make a tail-recursive call with ys, even though ys comes from a local list? Its local parameter is inferred correctly, and the call works without a nontail annotation.

The explanation is a hidden mode called regional, between global and local: global is a submode of regional, and regional is a submode of local. A local parameter comes from outside the current function’s region, so it may escape that one region. The tail ys obtained by matching that parameter is regional, allowing the tail call.

To demonstrate, I define lid as an identity function taking and returning a local value. Passing ys through lid and binding the result as zs loses the information that the value was regional: the type permits the call, but its result is known only as local.

Using zs in the tail-recursive call now produces the escaping-local error. Marking that call [@nontail] makes it compile, but doing so repeatedly in recursion can use O(n) more call-stack space for a long list. Preserving regional information avoids that cost for the tail-recursive branch.

The branch that adds one after the recursive call is different: 1 + count_occs zs is not a tail call, because addition still remains to be done. Passing a local value there is therefore accepted without a nontail annotation.

The rule to remember is that the function’s region ends before a tail call. An explicit [@nontail] annotation, or binding the call result before returning it, prevents tail-call treatment at the call site.

In practice the current approach can require many nontail annotations. I think a better approach will be needed before upstreaming this work; there are ideas for heuristics, but their future remains open.