Tail-Recursive List Mapping Function in OCaml Example

Implement the function map : ('a -> 'b) -> 'a list -> 'b list such that map f [x1; x2; ...; xn] is [f x1; f x2; ...; f xn]. Make sure that f x1, f x2, ..., f xn are computed in that order, and that the function is tail-recursive so that large input lists are supported.

Solution

The function to implement is List.map from the standard library. We proceed recursively based on how the input list is constructed. To ensure that calls to f are made following the order of the input list, in the case where l = x :: xs, we first compute f x and then recursively compute map f xs.

let rec map : ('a -> 'b) -> 'a list -> 'b list =
  fun f l ->
    match l with
    | [] -> []
    | x :: xs ->
      let y = f x in
      let ys = map f xs in
      y :: ys

This previous implementation is not tail-recursive because the call to map in map f xs does not occur in a tail position. Indeed, we need to perform the computation y :: ys after the recursive call, so we need a new stack frame for map f xs. However, it is worth noting that this function is tail-recursive modulo cons.

To implement the function tail-recursively, we introduce a helper function map_tl : ('a -> 'b) -> 'a list -> ('b list -> 'c) -> 'c in continuation-passing style. This helper function is such that map_tl f [x1; x2; ...; xn] return is return [f x1; f x2; ...; f xn], with return : 'b list -> 'c being the continuation. We'll use this continuation to perform that previous y :: ys continuation.

let rec map_tl : ('a -> 'b) -> 'a list -> ('b list -> 'c) -> 'c =
  fun f l return ->
    match l with
    | [] -> return []
    | x :: xs ->
      let y = f x in
      map_tl f xs (fun ys -> return (y :: ys))

let map : ('a -> 'b) -> 'a list -> 'b list =
  fun f l -> map_tl f l (fun l' -> l')

What we have done is effectively replace memory allocations from the stack to the heap by way of the continuation-passing style. If n is the length of the input list, then we need O(n) stack memory in the first implementation. On the other hand, we need O(n) heap memory to construct the continuation closures in the tail-recursive implementation.