Iterators in OCaml (chapter 4)

(* Section 4.1: Iterators, in OCaml *)

open Printf

(* Internal iterators over lists *)

let rec iter : ('a -> unit) -> 'a list -> unit =
  fun f l ->
    match l with [] -> () | h :: t -> f h; iter f t

let rec map : ('a -> 'b) -> 'a list -> 'b list =
  fun f l ->
    match l with [] -> [] | h :: t -> f h :: map f t

let rec fold_left : ('res -> 'a -> 'res) -> 'res -> 'a list -> 'res =
  fun f accu l ->
    match l with [] -> accu | h :: t -> fold_left f (f accu h) l

(* Imperative external iterator over arrays *)

let array_iterator (arr: 'a array) : unit -> 'a option =
    let i = ref 0 in
    fun () ->
        if !i >= Array.length arr
        then None
        else (let res = arr.(!i) in incr i; Some res)

exception StopIteration

let next (iter: unit -> 'a option) : 'a =
  match iter() with Some x -> x | None -> raise StopIteration

let rec same_iter (i1: unit -> 'a option) (i2: unit -> 'a option) : bool =
  match i1(), i2() with
  | None, None -> true
  | Some x1, Some x2 -> x1 = x2 && same_iter i1 i2
  | _, _ -> false

let same_array a1 a2 =
  same_iter (array_iterator a1) (array_iterator a2)

let _ =
  let a1 = [|1;2;3|] and a2 = [|1;3;5|] in
  printf "same_array a1 a1: %B\n" (same_array a1 a1);
  printf "same_array a1 a2: %B\n" (same_array a1 a2)

(* Functional external iterators *)

type 'a iter = unit -> 'a enum
and 'a enum = Done | More of 'a * 'a iter

let array_iterator (arr: 'a array) : 'a iter =
  let rec iter i () =
    if i >= Array.length arr then Done else More(arr.(i), iter (i+1))
  in iter 0

type 'a tree = Leaf | Node of 'a tree * 'a * 'a tree

let my_tree1 = Node(Node(Leaf, 1, Node(Leaf, 2, Leaf)), 3, Node(Leaf, 4, Leaf))
let my_tree2 = Node(Node(Leaf, 1, Leaf), 2, Node(Leaf, 3, Node(Leaf, 4, Leaf)))
let my_tree3 = Node(Node(Leaf, 1, Leaf), 2, Node(Leaf, 4, Node(Leaf, 6, Leaf)))

let tree_iterator (t: 'a tree) : 'a iter =
  let rec iter t k () =
    match t with
    | Leaf -> k
    | Node(l, x, r) -> iter l (More(x, iter r k)) ()
  in iter t Done

let rec same_iter (i1: 'a iter) (i2: 'a iter) : bool =
  match i1(), i2() with
  | Done, Done -> true
  | More(h1, t1), More(h2, t2) -> h1 = h2 && same_iter t1 t2
  | _, _ -> false

let same_fringe (t1: 'a tree) (t2: 'a tree): bool =
    same_iter (tree_iterator t1) (tree_iterator t2)

let _ =
  printf "same_fringe t1 t2: %B\n" (same_fringe my_tree1 my_tree2);
  printf "same_fringe t1 t3: %B\n" (same_fringe my_tree1 my_tree3)