(*
   Evaluate the performance of a stack-safe implementation of List.map.

   Run with:

     ocamlopt -o list_map unix.cmxa list_map.ml
     ./list_map
*)

open Printf

(*
   This is safe and simple but a little slow.
*)
let naive_safe_map f l =
  List.rev_map f l |> List.rev

(*
   Custom list type used to store intermediate lists, while minimizing
   the number of allocated blocks.
*)
type 'a list5 =
  | Elt of 'a * 'a list5
  | Tuple of 'a * 'a * 'a * 'a * 'a * 'a list5
  | Empty

let rev5 l =
  let rec aux acc l =
    match l with
    | Tuple (e, d, c, b, a, l) ->
        (* common case *)
        aux (a :: b :: c :: d :: e :: acc) l
    | Elt (a, l) ->
        aux (a :: acc) l
    | Empty -> acc
  in
  aux [] l

let rec slow_map acc f l =
  match l with
  | [] -> rev5 acc
  | [a] -> rev5 (Elt (f a, acc))
  | [a; b] ->
      let a = f a in
      let b = f b in
      rev5 (Elt (b, Elt (a, acc)))
  | [a; b; c] ->
      let a = f a in
      let b = f b in
      let c = f c in
      rev5 (Elt (c, (Elt (b, (Elt (a, acc))))))
  | [a; b; c; d] ->
      let a = f a in
      let b = f b in
      let c = f c in
      let d = f d in
      rev5 (Elt (d, (Elt (c, (Elt (b, (Elt (a, acc))))))))
  | [a; b; c; d; e] ->
      let a = f a in
      let b = f b in
      let c = f c in
      let d = f d in
      let e = f e in
      rev5 (Elt (e, (Elt (d, (Elt (c, (Elt (b, (Elt (a, acc))))))))))
  | a :: b :: c :: d :: e :: l ->
      let a = f a in
      let b = f b in
      let c = f c in
      let d = f d in
      let e = f e in
      slow_map (Tuple (e, d, c, b, a, acc)) f l

let rec fast_map rec_calls_remaining f l =
  if rec_calls_remaining <= 0 then
    slow_map Empty f l
  else
    match l with
    | [] -> []
    | [a] -> [f a]
    | [a; b] ->
        let a = f a in
        let b = f b in
        [a; b]
    | [a; b; c] ->
        let a = f a in
        let b = f b in
        let c = f c in
        [a; b; c]
    | [a; b; c; d] ->
        let a = f a in
        let b = f b in
        let c = f c in
        let d = f d in
        [a; b; c; d]
    | [a; b; c; d; e] ->
        let a = f a in
        let b = f b in
        let c = f c in
        let d = f d in
        let e = f e in
        [a; b; c; d; e]
    | a :: b :: c :: d :: e :: l ->
        let a = f a in
        let b = f b in
        let c = f c in
        let d = f d in
        let e = f e in
        a :: b :: c :: d :: e :: fast_map (rec_calls_remaining - 1) f l

(*
   This implementation of List.map makes at most 1000 non-tailrec calls
   before switching to a slower tailrec implementation.

   Additionally, this implementation guarantees left-to-right evaluation.
*)
let map f l = fast_map 1000 f l

let time name f =
  let t1 = Unix.gettimeofday () in
  let res = f () in
  let t2 = Unix.gettimeofday () in
  printf "%s: %.6f s\n%!" name (t2 -. t1);
  res

let repeat n f =
  let res = f () in
  for i = 2 to n do
    ignore (f ())
  done;
  res

let benchmark name repeats input_list =
  let conv = (fun i -> i + 1) in
  let res1 =
    time (name ^ ": List.map") (fun () ->
      repeat repeats (fun () ->
        List.map conv input_list
      )
    )
  in
  let res2 =
    time (name ^ ": naive_safe_map") (fun () ->
      repeat repeats (fun () ->
        naive_safe_map conv input_list
      )
    )
  in
  let res3 =
    time (name ^ ": map") (fun () ->
      repeat repeats (fun () ->
        map conv input_list
      )
    )
  in
  assert (res1 = res2);
  assert (res1 = res3)

let main () =
  let params = [
    3, 20_000_000;
    7, 20_000_000;
    50, 1_000_000;
    200, 300_000;
    100_000, 100;
  ] in
  List.iter (fun (len, repeats) ->
    let name = sprintf "list length = %i" len in
    let input_list = List.init len (fun i -> i) in
    benchmark name repeats input_list
  ) params

let () = main ()