OCaml with Jane Street extensions is available from our public opam repo. Only a slice of the features described in this series are currently implemented.

Rust, OCaml, and Resource Management

Coming from OCaml, the Rust programming language has many appealing features. Rust’s system for tracking lifetime and ownership allows users to safely express patterns that are awkward in OCaml, such as:

• Stack-allocated values and custom allocation schemes.
• Managed resources that can’t be (easily) garbage collected, e.g. file descriptors or GPU memory.
• Mutable data structures in the presence of concurrency.

On the other hand, Rust’s approach comes with some trade-offs. Eschewing garbage collection requires careful consideration of lifetime and ownership throughout a codebase. Emphasizing lifetime-polymorphism can also make type inference untenable, a design choice that wouldn’t fit OCaml.

At Jane Street, we’ve been working on extending OCaml to better support these use cases, without giving up the principles that make OCaml a convenient and flexible language.

To do so, we’re introducing a system of modes, which track properties like the locality and uniqueness of OCaml values. Modes allow the compiler to emit better, lower-allocation code, empower users to write safer APIs, and with the advent of multicore, statically guarantee data race freedom—all in a lightweight way that only affects those in need.

Stack Allocation

The OCaml compiler does not statically track lifetimes. Instead, it relies on a garbage collector to figure out a suitable lifespan for each value at runtime. Values are collected only after they become unreferenced, so OCaml programs are memory-safe.

To a first approximation, this model requires allocating all values on the heap. Fortunately, OCaml’s generational GC can efficiently handle short-lived values—minor-heap allocation simply advances a ring buffer.

However, placing everything on the heap is still a pessimistic approach. Where possible, using a specialized allocator could improve performance. For example, the minor heap is typically larger than cache, so future allocations are likely to evict live values. Stack allocation would immediately re-use freed space, eliminating this concern.

Providing an alternative to heap allocation would also have other benefits:

• Every minor heap allocation brings us closer to the next minor collection cycle. A minor collection incurs some fixed overhead, but more importantly, frequent collection causes more values to be moved to the major heap. Promoted values become much costlier to collect later on.

• At Jane Street, we often write “zero-allocation” code, which must never trigger a GC cycle. A stack allocator would make it much easier to write programs that do not touch the heap.

When such performance concerns are relevant, one should arguably be using a language based on explicit memory management, like Rust. However, garbage collection is genuinely useful; explicit management is a burden on users. Ideally, a language could provide a spectrum of allocation strategies freely interoperable within a single application. With modes, users can write OCaml with all the usual GC guarantees—but when performance is paramount, opt into the consideration of lifetimes, ownership, and concurrency.

Local Variables

In OCaml, it turns out that many short-lived values can be stack-allocated. To safely refer to such values, we introduce local variables.

Determining whether a variable is local involves checking a certain condition on its lifetime. Consider the following function:

let is_int str =
    let opt = Int.of_string_opt str in
    match opt with
    | Some _ -> true
    | None -> false
;;

val is_int : string -> bool

Naively, this function incurs a heap allocation. The compiler does not know the lifetime of opt—our function could return it, or even store it in a global variable. Because opt could escape this function, the value referenced by opt may need to live forever. Therefore, it must be heap-allocated.

As the programmer, however, we can deduce that a shorter lifetime suffices. In fact, opt only needs to live until we match on it. When is_int returns, opt is no longer accessible, so it could have safely been allocated in stack memory local to is_int.

Specifically, opt is local because its lifetime does not exceed its enclosing stack frame, which we call its region. At runtime, entering is_int begins a region by saving the current stack pointer; exiting ends the region and reclaims stack-allocated memory. Since opt is only accessible within this region, it may safely be allocated in the corresponding stack frame.

Note that a stack-allocated value is not necessarily stored on the control flow stack, as seen in languages that support alloca(). In this example, we request space from a stack-based allocator backed by entirely unrelated memory.

The Locality Mode

So, local variables are those that do not escape their region. To formalize this constraint in a manner the compiler can check, we introduce modes.

• By default, variables have the global mode. A global variable has the capability to escape any region, so always references the heap.

• Variables with the new local mode cannot escape their enclosing region, so may refer to the stack.

A mode is attached to a variable upon declaration, either in a let binding or in a function parameter. In both cases, the compiler will check that the value does not escape its region.

let foo (local x) =
    let local y = 0 in
    x, y
;;

3 | x, y
    ^
Error: this value escapes its region.

A local parameter represents a promise by the callee: the function will not store a reference to the value anywhere that could be accessed after the function returns. Intuitively, it’s safe to pass a stack-allocated value to a function if we know the value’s lifetime will not be extended.

let is_empty (local str) =
    String.length str = 0
;;

val is_empty : string @ local -> bool

Here, the syntax string @ local denotes that is_empty takes its parameter “at” the local mode.

Even without explicit mode annotations, the compiler can statically determine which variables may escape their enclosing region. Such variables are assigned the global mode; all others are automatically inferred to be local. At this point, the compiler may construct values bound to local variables using stack allocation.

Local Returns

Returning a local value from a function should appear contradictory, since a function’s result has clearly escaped its region. On the other hand, if functions can only return globals, constructing fully stack-allocated values becomes difficult—they can only be built up from literals. The solution:

let local_list () =
    exclave [1; 2; 3]
;;

val local_list : unit -> int list @ local

The exclave keyword ends the current region and executes the given expression in the enclosing region. The caller receives a local variable prohibited from escaping the caller’s region. Therefore, it’s safe to allocate that value on the caller’s stack frame—the difference is simply which region the value lives in.

let bar () =
    let list = local_list () in
    list
;;

3 | list
    ^^^^
Error: this value escapes its region.

Local-returning functions are the primary method of creating stack-allocated values, as they can programmatically build up local data structures. This mechanism also allows functions to return their local parameters.

Lastly, recall that the ‘locals’ stack is distinct from the control flow stack, making this behavior easy to implement.

Locality in APIs

Locality doesn’t only facilitate stack allocation—it also lets us design safer APIs. The following code exhibits a common pattern for resource management:

Core_unix.with_file "file" ~mode:[ O_RDONLY ] ~f:(fun fd -> (* ... *))

Here, a file descriptor is opened, passed to a lambda function, and closed after the function returns. This API lets users eschew manually closing the file descriptor. However, there’s no guarantee that the descriptor is not used after it’s closed.

let stash = ref 0 in
Core_unix.with_file "file" ~mode:[ O_RDONLY ] ~f:(fun fd -> stash := fd);
Core_unix.close !stash

Exception: Unix.Unix_error(Unix.EBADF, "close", ...)

Of course, this design can be improved by making fd a local parameter. After changing the signature of with_file to the following…

val with_file : string -> mode:open_flag list -> f:(File_descr.t @ local -> 'a) -> 'a

…the callback must promise not to stash away the file descriptor. Therefore, we know the file won’t be used after the callback returns.

In this example, we’re using modes to require a promise from the caller. This usage might feel similar to local returns, and for good reason. Formally, when a parameter is used contravariantly, its mode represents a restriction on the callee, but when used covariantly (as seen here), it instead represents a restriction on the caller.

Modes vs. Types

Above, we declared a local integer x using the syntax let local. Notably, we didn’t simply add a type annotation—the local mode does not operate on types. In fact, the mode of x is entirely separate from the type of x.

Types describe data structures, that is, how to build up and take apart values. On the other hand, a mode encodes a property independent of data layout, so may be attached to a variable of any type. To illustrate this behavior, type annotations specify a variable at a mode using the syntax type @ mode.

let local x = 0
let local y = "string"
let local z = [0.0; 1.0]

val x : int @ local
val y : string @ local
val z : float list @ local

In the case of locality, the salient property is whether a value may escape its region. Variables with the global mode can escape any region, so global values are correspondingly heap allocated. Conversely, the local mode restricts a variable to its region; a local value may be stack allocated.

Locality vs. Lifetimes

Encoding locality with a mode has some advantages compared to Rust’s type-centric approach. In Rust, reference types are parameterized over specific regions represented by lifetime variables. This design is more expressive than locality, which only distinguishes values that may escape all regions from those that cannot escape any.

On the other hand, lifetime variables are a source of pervasive complexity. When references are inherently polymorphic, essentially all functions become lifetime-polymorphic as well. For example, whenever a reference lacks a lifetime annotation, an implicit lifetime variable appears:

fn print_string(s: &str);

// Is equivalent to...

fn print_string<'a>(s: &'a str);

Since Rust supports first-class functions, the result is that higher-order functions require higher-order polymorphism, for which type inference is undecidable in general.

OCaml’s modes do not affect type inference—they preserve the types of existing code, so users truly don’t need to consider modes they aren’t actively using. In OCaml, type inference, higher-order functions, and garbage collection are all important parts of the development workflow, so we consider the local mode to be a good fit.

Modes are Deep

Above, we noted that a mode describes a property independent of data layout. Such properties are deep, as opposed to the shallow layout encoded by a type. To understand this distinction, consider the following type:

type 'a list =
    | Empty
    | More of 'a * 'a list

Destructuring a value of type 'a list produces two possible outcomes: either the empty list, or a pair of a value and another list of arbitrary shape. Hence, the type only describes the value’s top-level structure.

let process list =
    match list with
    | Empty -> (* ... *)
    | More (head, remaining) -> (* ... *)
;;

Conversely, destructuring a global variable of type 'a list produces either an empty list, or a pair of a global value and another global list. That is, if the root node of the list may escape its region, the subsequent nodes clearly can too—so the entire list must be heap-allocated.

The same logic applies to the local case: destructuring a local list produces a local value and another local list. It is possible to create a local list consisting entirely of stack allocations, so we must ensure that the contents of a local list also do not escape.

Deepness enables the compiler to validate usage of local data structures:

let head (local list) =
    match list with
    | Empty -> None
    | More (head, _) -> Some head
;;

4 | | More (v, _) -> Some head
                          ^^^^
Error: This value escapes its region

If locality didn’t exhibit “deepness,” it wouldn’t be very useful—we could stack allocate the root node of a list, but we’d have no way to express that further nodes may also be stack-allocated.

Sub-Modes

Given deepness, locality might appear to be an “all or nothing” choice—so far, we’ve allocated our data structures either entirely on the stack or entirely on the heap. To break this dichotomy, we will explore another important property of modes: each mode axis admits a natural sub-typing relation.

In the case of locality, it’s intuitively safe to use a global variable as if it were local. For example, a function expecting a local parameter promises equivalent behavior whether or not the parameter actually lives on the stack. Therefore, we say global is a sub-mode of local and allow global values to be used at the local mode.

let localize x = exclave x

val localize : 'a -> 'a @ local

It is safe for a local value to reference a global, but not vice versa. At runtime, this means we can create pointers from the stack to the heap, but not from the heap to the stack. For example, we can create a local, fully stack-allocated list whose nodes refer to heap-allocated values.

let rec localize list = exclave
    match list with
    | Empty -> Empty
    | More (head, remaining) -> More (head, localize remaining)
;;

val localize : 'a list -> 'a list @ local

We could also create a local list where only the first node is stack-allocated—say, if we locally append to a global list.

let local_cons (local head) remaining = exclave
    More (head, remaining)
;;

val local_cons : 'a @ local -> 'a list -> 'a list @ local

What we cannot create is a global list containing a stack-allocated node. Again, modes are deep, so any global list must have only captured globals. This preserves the invariant that whenever a node is heap-allocated, all nodes reachable from it are also heap-allocated.

More rigorously, we could say that as we traverse a value, the current mode monotonically increases with depth. This restriction should also make intuitive sense, since any list without this property contains a pointer from the heap to the stack. Such a pointer is potential use-after-free bug—the heap node may still be reachable after the stack frame has been freed.

The above layout can be represented using Rust lifetimes, which support subtyping. However, safely manipulating such data structures requires significantly more reasoning on the programmer’s part. Lifetime variables refer to arbitrary regions, so subtyping relationships must be explicitly specified. Locality offers a compromise: considering just one lifetime—the current region—makes efficient stack allocation easy to use in many practical scenarios. Values with other lifetimes are still managed by the garbage collector.

Global Record Fields

Because modes are deep, a local record always contains local values. Such fields may be stack allocated, so must be prohibited from escaping the current region. However, since global is a sub-mode of local, inner values may also be heap allocated—and sometimes the programmer knows they always will be. In this case, locality is unnecessarily restrictive.

Therefore, we also support annotating record fields with an explicit global mode. The compiler forbids initializing a global field using a local variable. Global fields are hence allowed to escape their region:

type 'a box = { global x : 'a }

let unwrap (local box) =
    box.x
;;

val unwrap : 'a box @ local -> 'a

Explicitly mutable record fields are automatically considered global. If this were not the case, a function could leak a local variable by storing it within a local parameter, violating region safety. For example:

type 'a box = { mutable x : 'a option }

let clear (local box) =
    let local y = None in
    box.x <- y
;;

5 | box.x <- y
             ^
Error: this value escapes its region.

Locality in Practice

At Jane Street, we’ve been using locality in production for some time. Developers who work on performance sensitive systems use locals daily, and those who don’t are largely unfamiliar with the feature—which means we’ve successfully limited the costs to the users who care. Therefore, we consider locality’s expressivity and performance benefits well worth the additional language complexity.

Building on locality’s success, the compilers team is now implementing additional modes for describing ownership constraints. In part two, we will explore new mode axes representing uniqueness and linearity.