Folds over trees, options and other structures

One function, fold, gave the size, the height, the sum and the sorted contents of a tree of five numbers: 5, 3, 21 and [1; 3; 4; 5; 8], each from one line that says what an empty tree is worth and how a node combines its parts. Written once for a type, a fold saves rewriting the same recursion for every question you ask of it.

You will write fold, map and iter for a tree type, express its measures and traversal orders as folds, replace nested matches on options with Option.bind and its friends, and prove that a law holds for every tree by structural induction.

A fold replaces each constructor with a function

The tree type of the variants lesson has two constructors, Leaf and Node (l, v, r). A fold over it takes one argument per constructor: ~leaf, the value to use for an empty tree, and ~node, a function that gets the folded left subtree, the node’s own value and the folded right subtree. fold ~leaf ~node then walks the tree once and returns what those two arguments build. The session folds the search tree made by inserting 5, 3, 8, 1 and 4:

from session import toplevel, write

# A fold over a tree replaces Leaf with a value and every Node with a function of the folded left tree, the value and
# the folded right tree. size, height, sum and the in-order list are then one line each. The tree is the search tree
# built by inserting 5, 3, 8, 1 and 4.
write("tree.ml", """type 'a tree = Leaf | Node of 'a tree * 'a * 'a tree

let rec insert x = function
  | Leaf -> Node (Leaf, x, Leaf)
  | Node (l, v, r) as t ->
      if x < v then Node (insert x l, v, r)
      else if x > v then Node (l, v, insert x r)
      else t

let rec fold ~leaf ~node = function
  | Leaf -> leaf
  | Node (l, v, r) -> node (fold ~leaf ~node l) v (fold ~leaf ~node r)
""")
toplevel("""
#use "tree.ml";;
let t = List.fold_left (fun t x -> insert x t) Leaf [5; 3; 8; 1; 4];;
let size t = fold ~leaf:0 ~node:(fun l _ r -> l + 1 + r) t;;
let height t = fold ~leaf:0 ~node:(fun l _ r -> 1 + max l r) t;;
let sum t = fold ~leaf:0 ~node:(fun l v r -> l + v + r) t;;
let to_list t = fold ~leaf:[] ~node:(fun l v r -> l @ [v] @ r) t;;
(size t, height t, sum t, to_list t);;
""")

$ ocaml
OCaml version 5.5.1
Enter #help;; for help.

# #use "tree.ml";;
type 'a tree = Leaf | Node of 'a tree * 'a * 'a tree
val insert : 'a -> 'a tree -> 'a tree = <fun>
val fold : leaf:'a -> node:('a -> 'b -> 'a -> 'a) -> 'b tree -> 'a = <fun>
# let t = List.fold_left (fun t x -> insert x t) Leaf [5; 3; 8; 1; 4];;
val t : int tree =
  Node (Node (Node (Leaf, 1, Leaf), 3, Node (Leaf, 4, Leaf)), 5,
   Node (Leaf, 8, Leaf))
# let size t = fold ~leaf:0 ~node:(fun l _ r -> l + 1 + r) t;;
val size : 'a tree -> int = <fun>
# let height t = fold ~leaf:0 ~node:(fun l _ r -> 1 + max l r) t;;
val height : 'a tree -> int = <fun>
# let sum t = fold ~leaf:0 ~node:(fun l v r -> l + v + r) t;;
val sum : int tree -> int = <fun>
# let to_list t = fold ~leaf:[] ~node:(fun l v r -> l @ [v] @ r) t;;
val to_list : 'a tree -> 'a list = <fun>
# (size t, height t, sum t, to_list t);;
- : int * int * int * int list = (5, 3, 21, [1; 3; 4; 5; 8])

size used ~leaf:0 and fun l _ r -> l + 1 + r, and gave 5. height replaced the addition with 1 + max l r and gave 3. sum added the values, 21, and to_list joined l @ [v] @ r into [1; 3; 4; 5; 8]. The type of fold that the toplevel printed, leaf:'a -> node:('a -> 'b -> 'a -> 'a) -> 'b tree -> 'a, shows why one function serves all four: the result type 'a can be anything, an int, a list or another tree.

A node can only be combined after both of its subtrees are, so a fold finishes a node’s children first. That order is the post-order traversal; step through it on the same five numbers:

Traversals and map are folds that place the value differently

The three depth-first orders differ in one thing: where the node’s own value goes between the two folded subtrees. Before them gives pre-order, between them in-order, after them post-order. A fold whose ~leaf is Leaf and whose ~node rebuilds a Node is map for trees:

from session import toplevel, write

# The three depth-first orders differ only in where the node's own value goes, and map for a tree is a fold as well.
# iter is written with ; so that the order of its effects is fixed: a fold computes both subtrees before it calls
# node, in an order OCaml does not specify.
write("tree.ml", """type 'a tree = Leaf | Node of 'a tree * 'a * 'a tree

let rec insert x = function
  | Leaf -> Node (Leaf, x, Leaf)
  | Node (l, v, r) as t ->
      if x < v then Node (insert x l, v, r)
      else if x > v then Node (l, v, insert x r)
      else t

let rec fold ~leaf ~node = function
  | Leaf -> leaf
  | Node (l, v, r) -> node (fold ~leaf ~node l) v (fold ~leaf ~node r)
""")
toplevel("""
#use "tree.ml";;
let t = List.fold_left (fun t x -> insert x t) Leaf [5; 3; 8; 1; 4];;
let preorder t = fold ~leaf:[] ~node:(fun l v r -> [v] @ l @ r) t;;
let inorder t = fold ~leaf:[] ~node:(fun l v r -> l @ [v] @ r) t;;
let postorder t = fold ~leaf:[] ~node:(fun l v r -> l @ r @ [v]) t;;
(preorder t, inorder t, postorder t);;
let map_tree f t = fold ~leaf:Leaf ~node:(fun l v r -> Node (l, f v, r)) t;;
map_tree (fun x -> x * 10) t;;
let rec iter f = function Leaf -> () | Node (l, v, r) -> iter f l; f v; iter f r;;
iter (fun x -> Printf.printf "%d " x) t;;
fold ~leaf:() ~node:(fun () v () -> Printf.printf "%d " v) t;;
""")

$ ocaml
OCaml version 5.5.1
Enter #help;; for help.

# #use "tree.ml";;
type 'a tree = Leaf | Node of 'a tree * 'a * 'a tree
val insert : 'a -> 'a tree -> 'a tree = <fun>
val fold : leaf:'a -> node:('a -> 'b -> 'a -> 'a) -> 'b tree -> 'a = <fun>
# let t = List.fold_left (fun t x -> insert x t) Leaf [5; 3; 8; 1; 4];;
val t : int tree =
  Node (Node (Node (Leaf, 1, Leaf), 3, Node (Leaf, 4, Leaf)), 5,
   Node (Leaf, 8, Leaf))
# let preorder t = fold ~leaf:[] ~node:(fun l v r -> [v] @ l @ r) t;;
val preorder : 'a tree -> 'a list = <fun>
# let inorder t = fold ~leaf:[] ~node:(fun l v r -> l @ [v] @ r) t;;
val inorder : 'a tree -> 'a list = <fun>
# let postorder t = fold ~leaf:[] ~node:(fun l v r -> l @ r @ [v]) t;;
val postorder : 'a tree -> 'a list = <fun>
# (preorder t, inorder t, postorder t);;
- : int list * int list * int list =
([5; 3; 1; 4; 8], [1; 3; 4; 5; 8], [1; 4; 3; 8; 5])
# let map_tree f t = fold ~leaf:Leaf ~node:(fun l v r -> Node (l, f v, r)) t;;
val map_tree : ('a -> 'b) -> 'a tree -> 'b tree = <fun>
# map_tree (fun x -> x * 10) t;;
- : int tree =
Node (Node (Node (Leaf, 10, Leaf), 30, Node (Leaf, 40, Leaf)), 50,
 Node (Leaf, 80, Leaf))
# let rec iter f = function Leaf -> () | Node (l, v, r) -> iter f l; f v; iter f r;;
val iter : ('a -> 'b) -> 'a tree -> unit = <fun>
# iter (fun x -> Printf.printf "%d " x) t;;
1 3 4 5 8 - : unit = ()
# fold ~leaf:() ~node:(fun () v () -> Printf.printf "%d " v) t;;
8 4 1 3 5 - : unit = ()

Each order is one node function, and the session printed all three lists for its tree:

map_tree (fun x -> x * 10) kept the shape and gave a tree of 10, 30, 40, 50 and 80. For iter, which exists only for its effects, the session wrote the recursion by hand with ;, which fixes the order: it printed 1 3 4 5 8. The fold version printed 8 4 1 3 5. A fold computes both subtrees as arguments of node, and OCaml leaves the order in which arguments are evaluated unspecified; this compiler evaluated the right one first.

A fold over a tree is not tail-recursive: it keeps a frame for each level it is inside, so its stack use is the tree’s depth. The program below folds two trees of 1,000,000 values under a stack limited to 100,000 words:

from session import run, write

# A fold over a tree is not tail-recursive: it keeps a frame for every level it is inside. A balanced tree of 1,000,000
# values is 20 levels deep; a chain of 1,000,000 nodes, the tree sorted input gives, is 1,000,000 levels deep.
write("depth.ml", """type 'a tree = Leaf | Node of 'a tree * 'a * 'a tree

let rec fold ~leaf ~node = function
  | Leaf -> leaf
  | Node (l, v, r) -> node (fold ~leaf ~node l) v (fold ~leaf ~node r)

let size t = fold ~leaf:0 ~node:(fun l _ r -> l + 1 + r) t

(* A balanced tree of the values lo .. hi: the middle value at the root. *)
let rec balanced lo hi =
  if lo > hi then Leaf
  else
    let mid = (lo + hi) / 2 in
    Node (balanced lo (mid - 1), mid, balanced (mid + 1) hi)

(* A chain: every node has only a right child, built with a loop. *)
let chain n =
  let t = ref Leaf in
  for i = n downto 1 do t := Node (Leaf, i, !t) done;
  !t

let rec depth = function Leaf -> 0 | Node (l, _, r) -> 1 + max (depth l) (depth r)

let () =
  let n = 1_000_000 in
  match Sys.argv.(1) with
  | "balanced" ->
      let t = balanced 1 n in
      Printf.printf "balanced: size %d, depth %d\\n" (size t) (depth t)
  | _ -> Printf.printf "chain: size %d\\n" (size (chain n))
""")
run("""
ocamlopt depth.ml -o depth
export OCAMLRUNPARAM=l=100000
./depth balanced
!./depth chain
""")

$ ocamlopt depth.ml -o depth
$ export OCAMLRUNPARAM=l=100000
$ ./depth balanced
balanced: size 1000000, depth 20
$ ./depth chain
Fatal error: exception Stack_overflow
[exit status 2]

The balanced tree was 20 levels deep, and its size came out as 1000000. The chain, a tree in which every node has only a right child, which is what sorted input builds, raised Stack_overflow: a million levels need a million frames.

Option.bind chains steps that may each fail

An option is a variant with two constructors, None and Some v, and the Option module provides its maps and its fold. The session reads a port from a list of settings, where each step can fail: the key may be missing, the text may not be a number, and the number may be out of range. Step through the four settings lists it tries:

from session import toplevel

# Reading a port from settings has three ways to fail: the key is missing, the text is not a number, or the number is
# out of range. Option.bind chains the steps, the binding operator let* writes the chain as a sequence, and
# Option.value, Option.fold and Option.map use the result.
toplevel("""
let settings = [("host", "localhost"); ("port", "8080"); ("retries", "three")];;
let in_range n = if n > 0 && n < 65536 then Some n else None;;
let port_nested s = match List.assoc_opt "port" s with None -> None | Some text -> (match int_of_string_opt text with None -> None | Some n -> in_range n);;
let port s = Option.bind (Option.bind (List.assoc_opt "port" s) int_of_string_opt) in_range;;
let ( let* ) = Option.bind;;
let port_seq s = let* text = List.assoc_opt "port" s in let* n = int_of_string_opt text in in_range n;;
let cases = [settings; []; [("port", "eighty")]; [("port", "70000")]];;
(List.map port_nested cases, List.map port cases, List.map port_seq cases);;
Option.value ~default:80 (port []);;
List.map (fun s -> Option.fold ~none:"no port" ~some:(Printf.sprintf "port %d") (port s)) cases;;
Option.map (fun p -> p + 1) (port settings);;
""")

$ ocaml
OCaml version 5.5.1
Enter #help;; for help.

# let settings = [("host", "localhost"); ("port", "8080"); ("retries", "three")];;
val settings : (string * string) list =
  [("host", "localhost"); ("port", "8080"); ("retries", "three")]
# let in_range n = if n > 0 && n < 65536 then Some n else None;;
val in_range : int -> int option = <fun>
# let port_nested s = match List.assoc_opt "port" s with None -> None | Some text -> (match int_of_string_opt text with None -> None | Some n -> in_range n);;
val port_nested : (string * string) list -> int option = <fun>
# let port s = Option.bind (Option.bind (List.assoc_opt "port" s) int_of_string_opt) in_range;;
val port : (string * string) list -> int option = <fun>
# let ( let* ) = Option.bind;;
val ( let* ) : 'a option -> ('a -> 'b option) -> 'b option = <fun>
# let port_seq s = let* text = List.assoc_opt "port" s in let* n = int_of_string_opt text in in_range n;;
val port_seq : (string * string) list -> int option = <fun>
# let cases = [settings; []; [("port", "eighty")]; [("port", "70000")]];;
val cases : (string * string) list list =
  [[("host", "localhost"); ("port", "8080"); ("retries", "three")]; [];
   [("port", "eighty")]; [("port", "70000")]]
# (List.map port_nested cases, List.map port cases, List.map port_seq cases);;
- : int option list * int option list * int option list =
([Some 8080; None; None; None], [Some 8080; None; None; None],
 [Some 8080; None; None; None])
# Option.value ~default:80 (port []);;
- : int = 80
# List.map (fun s -> Option.fold ~none:"no port" ~some:(Printf.sprintf "port %d") (port s)) cases;;
- : string list = ["port 8080"; "no port"; "no port"; "no port"]
# Option.map (fun p -> p + 1) (port settings);;
- : int option = Some 8081

All three versions gave [Some 8080; None; None; None]. The nested match needs a None -> None for every step. Option.bind o f is None when o is None and f v when o is Some v, so two calls of it chain the three steps. The binding operatorbinding operator: An operator such as let* that the programmer defines and OCaml uses to read let* x = e in body as a call of that operator on e and fun x -> body; defined as Option.bind, it chains steps that may give None. let*, defined as Option.bind, writes the same chain as a sequence of lets that stops at the first None.

The session used four functions of the Option module, bind among them, each for its own job:

Structural induction proves a law for every tree

A test checks the trees it tries; a proof covers all of them. Structural inductionstructural induction: A way to prove a claim about every value of a recursive type: prove it for each constructor, assuming it already holds for the smaller values of the same type inside that constructor. proves a claim about every value of a type by proving it for each constructor, assuming it already holds for the smaller values inside. Step through the proof that mirror, which swaps every left and right subtree, keeps the size of a tree:

The proof has two cases because the type has two constructors, and it uses the claim for l and r exactly where size and mirror call themselves. The session complements it by testing: it builds all 65 tree shapes with at most five nodes and checks claims about mirror on each one:

from session import toplevel

# A proof by structural induction covers every tree; a test can cover every tree up to a size. The session builds all
# 65 tree shapes with at most 5 nodes and checks three claims about mirror on each: two hold, and one is false.
toplevel("""
type tree = Leaf | Node of tree * int * tree;;
let rec size = function Leaf -> 0 | Node (l, _, r) -> size l + 1 + size r;;
let rec mirror = function Leaf -> Leaf | Node (l, v, r) -> Node (mirror r, v, mirror l);;
let rec to_list = function Leaf -> [] | Node (l, v, r) -> to_list l @ [v] @ to_list r;;
let rec shapes n = if n = 0 then [Leaf] else List.concat_map (fun k -> List.concat_map (fun l -> List.map (fun r -> Node (l, n, r)) (shapes (n - 1 - k))) (shapes k)) (List.init n Fun.id);;
let all () = List.concat_map shapes [0; 1; 2; 3; 4; 5];;
List.length (all ());;
List.for_all (fun t -> size (mirror t) = size t) (all ());;
List.for_all (fun t -> to_list (mirror t) = List.rev (to_list t)) (all ());;
List.find_opt (fun t -> to_list (mirror t) <> to_list t) (all ());;
""")

$ ocaml
OCaml version 5.5.1
Enter #help;; for help.

# type tree = Leaf | Node of tree * int * tree;;
type tree = Leaf | Node of tree * int * tree
# let rec size = function Leaf -> 0 | Node (l, _, r) -> size l + 1 + size r;;
val size : tree -> int = <fun>
# let rec mirror = function Leaf -> Leaf | Node (l, v, r) -> Node (mirror r, v, mirror l);;
val mirror : tree -> tree = <fun>
# let rec to_list = function Leaf -> [] | Node (l, v, r) -> to_list l @ [v] @ to_list r;;
val to_list : tree -> int list = <fun>
# let rec shapes n = if n = 0 then [Leaf] else List.concat_map (fun k -> List.concat_map (fun l -> List.map (fun r -> Node (l, n, r)) (shapes (n - 1 - k))) (shapes k)) (List.init n Fun.id);;
val shapes : int -> tree list = <fun>
# let all () = List.concat_map shapes [0; 1; 2; 3; 4; 5];;
val all : unit -> tree list = <fun>
# List.length (all ());;
- : int = 65
# List.for_all (fun t -> size (mirror t) = size t) (all ());;
- : bool = true
# List.for_all (fun t -> to_list (mirror t) = List.rev (to_list t)) (all ());;
- : bool = true
# List.find_opt (fun t -> to_list (mirror t) <> to_list t) (all ());;
- : tree option = Some (Node (Leaf, 2, Node (Leaf, 1, Leaf)))

size (mirror t) = size t held for all 65, as the proof says it must for every tree. So did to_list (mirror t) = List.rev (to_list t). The third claim, that mirroring keeps the in-order list, was false, and List.find_opt returned the first tree that breaks it, Node (Leaf, 2, Node (Leaf, 1, Leaf)), a tree of only two nodes whose list [2; 1] becomes [1; 2]. A search over small cases finds a false law at once; only a proof shows a true one holds beyond them.

Practise: folds, traversals and options

The questions use the session’s tree, its folds and the port settings with new values and new laws. Work out each result, or whether a claim holds, before you check:

Interview questions

Warm-up (2–5 years): what is a fold over a tree, and how does it relate to the functions on that tree?

A fold over a tree type takes one argument per constructor: a value for an empty tree, and a function that combines a node’s value with the already folded results of its subtrees. It walks the tree once, replacing every constructor with the matching argument. Most recursive functions on the tree are then special cases: size is a fold that adds 1 at each node, height one that takes the larger subtree plus 1, and the in-order list one that joins left, value and right. Like those functions, a tree fold is not tail-recursive, so its stack use is the depth of the tree, small for a balanced tree and as large as the tree for a chain.

Two harder questions follow, about an expression tree too deep for its fold and about a chain of optional steps that must say which step failed. For each, say what you would measure or print first, then open its first step:

Key takeaways

• A fold over a type takes one argument per constructor, so size, height, sum and to_list of a tree are each one line with the same fold.
• Pre-, in- and post-order differ only in where the node’s value goes between the folded subtrees, and map for a tree is a fold that rebuilds each node.
• A fold leaves the order of evaluating the two subtrees to the compiler, so an iter whose effects must come in order is written with ;, and a tree fold needs stack in proportion to the tree’s depth.
• Option.map, Option.bind, Option.value and Option.fold, or let* defined as Option.bind, replace nested matches on options and stop at the first None.
• Structural induction proves a claim for every tree by proving it for Leaf and for Node assuming it for the subtrees, while a test over all small trees finds a false claim’s counterexample.