OxCaml - Stack Allocations and Locality

One of the key language extensions focuses on stack allocation and locality of data. This is critical for latency-sensitive operations, minimising GC pauses due to reduced heap allocation. The OxCaml docs on stack allocation have detailed information.

At runtime, stack allocations do not take place on the function call stack, but on a separately-allocated stack that follows the same layout as the OCaml minor heap. Every function, loop, and module-level binding (top-level binding, which has to be globally accessible) represents a new region, corresponding to a separate stack frame. Values annotated with the stack_ keyword are allocated in the particular stack frame in which they appear and are local to that region, whereas heap-allocated values are global. Stack-allocated values cannot be referenced outside of their region, as the stack frame is popped upon a return, while global values are allowed to escape as they are heap allocated and their lifetimes exceed that of the stack frame.

However, some types such as an int can mode-cross the locality axis. Since an int is register allocated and is an immediate value, its lifetime exceeds that of the current stack frame. A local int can mode-cross to a global value and be returned from the function body.

All branches of expressions such as if or match ... with must have the same locality, either local or global.

Nested regions result from the nesting of functions and other constructs that create new regions. Within a nested region, values could be referenced from the same region that are local to it, along with values from the outer region which are described as outer-local.

A region is not the same as a scope. For instance, refer to this example from the docs:

let f () =
  let counter =
    let r = stack_ (ref 42) in
    incr r;
    r
  in
  ...

While the scope of r is within counter, the stack-allocated value can be accessed and is alive throughout f.

The compiler cannot infer modes across files and relies on annotations in the interface (.mli files). If no annotations are present, the function arguments are assumed to be global. Mode inference can still function within the file as long as the interface file is present. Here's an example from the docs:

(* in a.mli *)
val f1 : foo:int option @ local -> unit
val f2 : int -> unit
(* in a.ml when a.mli is present *)
let f1 ~foo:_ = ()
let f2 x = f1 ~foo:(Some x) (* [Some x] is stack allocated *)
(* in a.ml when .mli is missing *)
let f1 ~foo:_ = ()
let f2 x = f1 ~foo:(Some x) (* [Some x] is heap allocated *)

When a.mli is present (irrespective of mode annotations), the foo parameter is inferred to be local, as it is not referenced outside of f1's function body and is local to the region. The function argument Some x is stack allocated. In the absence of a.mli, mode inference does not take place and foo is assumed to be global, resulting in a heap allocation of Some x.

A closure can also be stack allocated. When a closure captures a local variable, it cannot reference that variable outside of the region, requiring the closure itself to be stack allocated and local to the region.

Since a tail call requires clearing the current stack frame for constant-space function invocations and recursion, there are restrictions on the usage of local values. Values with the local mode cannot be passed to a tail call, and local closures cannot themselves be invoked in tail position, since they would be cleared from the frame before the call takes place. These restrictions can be circumvented by removing the tail call altogether, by either assigning the result to a variable which is then returned, or using the [@nontail] annotation.

let f () =
  let x = stack_ (Some 1) in
  let res = some_fn x in
  res
(* OR *)
let f () =
  let x = stack_ (Some 1) in
  some_fn x [@nontail]

exclave_ OxCaml allows terminating the current region and performing an allocation in the caller's stack frame. This is performed by means of the exclave_ keyword and is useful in a variety of cases. It could even be used in recursive functions, with all allocations taking place in the stack frame of the initial call site.

let f (x @ local) =
  let g @ local = some_fn x in
  if some_check g then exclave_ Some g
  else exclave_ None

Here, both None and Some g are allocated in the caller's stack frame, avoiding a heap allocation for these values being returned from the function and outliving the callee's stack frame.

As exclave_ terminates the current region, local values cannot be used in the exclave_ region. It must appear in the tail position of the current region.

The compiler typically infers the locality of a value as a whole and not separately for its parts. If a list is defined to be @ local, then the elements inside it are also assumed to be local.

let f x y =
  let p @ local = (x, y) in
  let (a, b) = p in
  a
(* Error: This value is local because it is an element of the tuple at ... which is local. However, the highlighted expression is expected to be local to the parent region or global because it is a function return value.
Hint: Use exclave_ to return a local value. *)

Records as well as arguments of a constructor (of a variant) are exceptions to this convention. Individual fields of the type can be annotated with global_ and must be passed heap-allocated values during construction.

type ('a, 'b) t = {global_ global: 'a; non_global: 'b}
type ('a, 'b) u = Foo of global_ 'a * 'b (* Foo here is a constructor *)
let f x y =
  let p @ local = {global = x; non_global = y} in
  let {global = a; non_global = b} = p in
  a

Since a is the value from the global_ field, it can be returned. Attempting to return b would fail to compile as it is local to the current region and is stack allocated. In this manner, a string t list could be defined, with the list and its elements of type t being local, but the string values in t being global and heap allocated. The base module provides this as the type Modes.Global.t for convenience.

type 'a t = { global_ global : 'a } [@@unboxed]

The locality of individual tuple values is only considered in a special case when the tuple is being matched. Consider matching on a tuple of three values:

match a, b, c with
  | p, q, r -> ...

a, b, c is syntactically a single value - a tuple with three elements. While matching them, the three elements are treated as three separate values with unique localities. This exception also applies during variable binding:

let a, b, c = ...

where (a, b, c) is syntactically a single tuple.

Mutable fields must point to heap-allocated values. For instance, stack_ refs and arrays can be local, but the values they point to (or are in them) must be global. (Arrays in OCaml are traditionally mutable, with immutable arrays available since OCaml 5.4.)

Currying introduces further subtleties around locality, since applying a curried function one argument at a time produces intermediate closures whose own locality needs to be tracked. Consider the following two function types:

let f ... : 'a -> 'b -> 'c
let g ... : 'a -> ('b -> 'c)

In vanilla OCaml, both of these function types are equivalent. On applying f to the first argument, the function is curried and returns a closure of the type 'b -> 'c, resulting in a function of the same type as g.

let f ... : 'a @ local -> 'b -> 'c
let g ... : 'a @ local -> ('b -> 'c)

In OxCaml, these function types are no longer equivalent due to currying. The first argument to f is local, requiring the closure to be of type ('b -> 'c) @ local. The closure is hence stack allocated and local, since it captures a reference to another stack-allocated value (the first argument). The second function g returns a global closure that cannot reference the first argument, which is stack allocated.

From the docs:

In general, in a curried function type ... -> ... -> ... (without parentheses), then after the first use of local, all arrow types except the last will implicitly be given local return types, enabling the expected partial application behaviour.

This is why the closure returned in g, the final value, is global. It is also why a type like

(a -> b -> c -> d) @ local -> e -> f -> g

is read as

(a -> (b -> (c -> d) @ local) @ local) @ local -> (e -> (f -> g) @ local) @ local

Except for the final values d and g, all the other function argument types are of mode local.

let f : ('a -> 'b -> 'c) @ local = ...
let g : 'a -> ('b -> 'c) @ local = ...

These functions f and g are of equivalent types, since the first type desugars to the second. Currying only affects the locality of the latter type arguments ('b -> 'c), which capture the first argument 'a. The first argument 'a itself remains global and accepts only heap-allocated values. To allow the first argument to also accept local values, the type is to be specified as follows:

'a @ local -> ('b -> 'c) @ local

In this case, the second local annotation on the returned closure can be avoided, thanks to the locality convention for curried function type arguments.

Stack allocations and locality significantly reduce heap allocations in an OxCaml program, enhancing performance with minimal garbage collection interference during runtime. They are aided by other language extensions such as unboxed types, modes and SIMD extensions.