jon.recoil.org

Introduction to Stack Allocation

Based on the OxCaml documentation at oxcaml.org.

In OCaml, values are normally allocated on the garbage-collected heap. OxCaml can instead allocate values on the stack, which offers two performance advantages: the same few hot cache lines are constantly reused (lower cache footprint), and stack allocations never trigger a GC -- making them safe for zero-alloc, low-latency code.

The stack_ keyword

Use stack_ before an allocation to force it onto the stack:

let () = let x = stack_ (1, 2, 3) in let a, b, c = x in Printf.printf "(%d, %d, %d)\n" a b c

The keyword works shallowly -- it only affects the immediately following allocation. Most types can be stack-allocated: tuples, records, variants, closures, and boxed numbers.

type point = { x : float; y : float } let () = let p = stack_ { x = 3.0; y = 4.0 } in Printf.printf "distance = %.2f\n" (Float.sqrt (p.x *. p.x +. p.y *. p.y))

Regions and escaping

Stack-allocated values live in a region (usually a function body) and must not escape it. The type-checker enforces this. Try uncommenting the last line to see the error:

let make_pair () = let p = stack_ (42, 99) in let a, b = p in Printf.printf "pair: (%d, %d)\n" a b (* p -- uncommenting would give: "This value escapes its region" *) let () = make_pair ()

Local parameters

To pass a stack-allocated value to a function, the function must promise not to let it escape. This is done with @ local:

let sum_pair (p @ local) = let a, b = p in a + b let () = let p = stack_ (42, 99) in Printf.printf "sum = %d\n" (sum_pair p)

A function with @ local parameters can be called with either stack- or heap-allocated values -- the annotation only constrains the function's implementation, not its callers. Here sum_pair computes a global int result from a local tuple, so the result can safely escape.

let mutable for loop variables

OxCaml's let mutable provides mutable local variables that are always stack-allocated, avoiding any heap allocation:

let sum_to n = let mutable total = 0 in for i = 1 to n do total <- total + i done; total let () = Printf.printf "sum to 100 = %d\n" (sum_to 100)

Compare this with the traditional approach using a ref (which allocates on the heap):

let sum_to_ref n = let total = ref 0 in for i = 1 to n do total := !total + i done; !total let () = Printf.printf "sum to 100 = %d\n" (sum_to_ref 100)

Returning local values with exclave_

The exclave_ keyword lets a function allocate in its caller's region, enabling the caller to stack-allocate the result:

module Counter : sig type t val make : unit -> t @ local val next : t @ local -> int end = struct type t = int ref let make () = exclave_ ref 0 let next c = let x = !c in incr c; x end let () = let c = Counter.make () in Printf.printf "next: %d\n" (Counter.next c); Printf.printf "next: %d\n" (Counter.next c); Printf.printf "next: %d\n" (Counter.next c)

Inference

In practice, you rarely need to write stack_ explicitly. The compiler infers stack allocation whenever possible:

let use_callback ~(f : _ @ local -> _) = let pair = (1, 2) in (* inferred as stack-allocated *) f pair let () = use_callback ~f:(fun p -> let a, b = p in Printf.printf "got: (%d, %d)\n" a b)

Without the @ local annotation on f, the pair would need to be heap- allocated since the compiler couldn't prove f won't capture it.

The global_ field annotation

When you need to extract a value from a stack-allocated record and return it, mark the field global_:

type 'a box = { global_ value : 'a; tag : string } let get_value (b @ local) = b.value (* always safe to return *) let () = let b = stack_ { value = 42; tag = "answer" } in Printf.printf "value = %d\n" (get_value b)

The global_ annotation means the field's contents are always on the heap, even when the record itself is stack-allocated.