Title

class: center, middle

# Stack allocation for OCaml

Stephen Dolan, Leo White

Jane Street

---

.left-column[
## Allocation
### Frequent
]

```ocaml
type node = {
  name: string;
  successors: string list
}

type graph = node list

let count_self_edges (g : graph) =
  let count = ref 0 in
  List.iter
    (fun node ->
      List.iter
        (fun succ ->
          if succ = node.name then incr count)
        node.successors)
    g;
  !count
```

--
- 16 bytes for the `ref`

--
- 32 bytes for the `fun node -> ...` closure

--
- 40×N bytes for the `fun succ -> ...` closure

---

.left-column[
## Allocation
### Frequent
### Cheap, not free
]

.vspace[]
**Short-lived allocations are cheap, but**:

- Space is not reused quickly  
    causing poor L1 cache usage

- GC advances towards the next minor GC  
    so other allocations are promoted unnecessarily

---

.left-column[
## Allocation
### Frequent
### Cheap, not free
### Stack allocation
]

.vspace[]
**Allocating values on a stack**:

- reuses space quickly

- does not cause any GC work

Hard to do safely, though!

---
class:center
.vspace[]
.big[
When is it safe to pass  
**stack-allocated values**  
to a function?
]

---

.left-column[
## Prior work
### Region variables
]

.vspace[]
**Region variables** attach lifetime information to types.

```rust
struct TwoBorrowedStrings<'a> {
    x: &'a str,
    y: &'a str,
}
```

- Extremely expressive

- Syntactically heavyweight

---

.left-column[
## Prior work
### Region variables
### Stack arguments
]

.vspace[]

Functions that can accept stack arguments are typed with **region polymorphism**:

$$
\forall \alpha. \; {\tt TwoBorrowedStrings}[\alpha] \rightarrow ()
$$

Higher-order functions can require higher-rank (non-inferrable) types.
---

.left-column[
## Modes, not types
### Local and global
]

.vspace[]
Instead, we mark variable bindings as `local` or `global`:

- `global` bindings never refer to stack-allocated values

- `local` bindings never escape their *region*  
    (function body or loop)

Less expressive than region variables, but much simpler.
---

.left-column[
## Modes, not types
### Local and global
### Modes are deep
]

.vspace[]
The same types are used at `local` and `global` mode:

```ocaml
type node = {
  name: string;
  successors: string list
}

type graph = node list
```

A local `graph` has local contents.

```ocaml
type part_global = {
  foo : string;
  global_ bar : string;
}
```

`(...).bar` is always global.

---

.left-column[
## Modes, not types
### Local and global
### Modes are deep
### Function types
]

.vspace[]
Our function types specify the mode of their argument:

```ocaml
type s = string -> unit
type t = local_ string -> unit
```

A function of type `local_ 'a -> 'b` cannot capture its argument,
so can be passed a stack-allocated value.

No lifetime variables or polymorphism, so inference works.

---

.left-column[
## Modes, not types
### Local and global
### Modes are deep
### Function types
### Local returns
]

.vpsace[]
Function types also have a mode on the return type:

```ocaml
module M : sig
  val f : 'a -> local_ 'a option
end = struct
  let f x = local_ (Some x)
end
```

Separating the data from the control stack means values can be
allocated in the caller's region.

---

.left-column[
## Modes, not types
### Local and global
### Modes are deep
### Function types
### Local returns
### Typing closures
]

Typing rule for closures:
$$
\frac{
\Gamma,\;\phantom{\square_i,}\; \phantom{j\;} x: A \vdash e: B \phantom{\;\;@\; k}
}{
\Gamma \vdash {\tt fun}\;x \rightarrow e: \phantom{j\; }A \rightarrow \phantom{k\; }B \phantom{\;\;@\; i}
}
$$

---

.left-column[
## Modes, not types
### Local and global
### Modes are deep
### Function types
### Local returns
### Typing closures
]

Typing rule for closures with modes:
$$
\def\phantom#1{#1}
\frac{
\Gamma,\;\phantom{\square_i,}\; \phantom{j\;} x: A \vdash e: B \phantom{\;\;@\; k}
}{
\Gamma \vdash {\tt fun}\;x \rightarrow e: \phantom{j\; }A \rightarrow \phantom{k\; }B \phantom{\;\;@\; i}
}
$$

- \$i\$: mode of the closure itself
  - \$j\$: mode of the argument
  - \$k\$: mode of the return

Variable access must agree with locking:
$$
\frac{
i \leq j\qquad i \leq k
}{
\Gamma,\;i x: A,\;\dots,\;\square_j\;,\dots\;\vdash x : A \;\;@\; k
}
$$

where `global` \$\leq\$ `local`.

---

.left-column[
## Examples
### Iteration
]

.vspace[]

```ocaml
val iter : local_ ('a -> unit) -> 'a list -> unit
```
--
.vspace[]
```ocaml
let count_self_edges (g : graph) =
  let count = ref 0 in
  List.iter
    (fun node ->
      List.iter
        (fun succ ->
          if succ = node.name then incr count)
        node.successors;
      ())
    g;
  !count
```

---
.left-column[
## Examples
### Iteration
### Currying
]

.vspace[]

```ocaml
val iter : local_ ('a -> unit) -> 'a list -> unit
```
--
is not:
```ocaml
val iter : local_ ('a -> unit) -> ('a list -> unit)
```
--
but instead:
```ocaml
val iter : local_ ('a -> unit) -> local_ ('a list -> unit)
```

--
.vspace[]
```ocaml
let f = List.iter g
```

---

.left-column[
## Examples
### Iteration
### Currying
### Local functions
]

.vspace[]
```ocaml
val iter : local_ ('a -> unit) -> 'a list -> unit
```

--
.svspace[]
```ocaml
val with_file :
  filename:string ->
  local_ (local_ filehandle -> 'a) ->
  'a
```

--
.svspace[]
```ocaml
val immut_array :
  length:int ->
  init:'a ->
  local_ (local_ 'a array -> 'b) ->
  'a immut_array * 'b
```

---

.left-column[
## Examples
### Iteration
### Currying
### Local functions
### More uses
]

.vspace[]
```ocaml
val borrow :
   unique_ 'a ->
   local_ (local_ 'a -> 'b) ->
   unique_ 'a * 'b
```

.vspace[]
```ocaml
val effectful :
   local_ 'a handler -> unit
```

---

.left-column[
## Conclusion
]

.vspace[]

.biggish[.center[Stack allocation is **efficient**...]]

.biggish[.center[... but locals are useful for **more than speed**.]]

.center[
.vspace[]
.vspace[]
Code & docs at:

<https://github.com/ocaml-flambda/ocaml-jst>

`{sdolan,lwhite}@janestreet.com`
]