module Tree = struct
  type t = Leaf | Node of t * t

  let rec size = function
    | Leaf -> 0
    | Node (x, y) -> 1 + size x + size y
end

type instr = Gen | Build

let pop2 = function
  | x :: y :: l -> x, y, l
  | _ -> invalid_arg "pop2"

module type Boltz = sig
  val size : unit -> int
  val gen : unit -> Tree.t
end

module Specialised = struct
  let size =
    let rec aux s = function
      | 0 -> s
      | n ->
        let n = n - 1 in
        if Random.bool () then aux s n
        else aux (s + 1) (n + 2)
    in
    fun () -> aux 0 1

  let gen =
    let rec aux res = function
      | Gen :: stack ->
        if Random.bool () then aux (Tree.Leaf :: res) stack
        else aux res (Gen :: Gen :: Build :: stack)
      | Build :: stack ->
        let x, y, res = pop2 res in
        aux (Tree.Node (x, y) :: res) stack
      | [] -> res
    in
    fun () -> aux [] [Gen] |> List.hd
end

module Generic = struct
  let meta ~next ~finish ~do_leaf ~do_node =
    let rec aux state =
      let continue, state = next state in
      if continue then
        if Random.bool () then aux (do_leaf state)
        else aux (do_node state)
      else
        finish state
    in
    aux

  let size : unit -> int =
    let next (s, n) = (n <> 0), (s, n - 1) in
    let do_leaf (s, n) = s, n in
    let do_node (s, n) = s + 1, n + 2 in
    let finish (s, _) = s in
    fun () -> meta ~next ~finish ~do_leaf ~do_node (0, 1)

  let gen : unit -> Tree.t =
    let rec next (trees, instrs) = match instrs with
      | [] -> false, (trees, [])
      | Gen :: instrs -> true, (trees, instrs)
      | Build :: instrs ->
        let l, r, trees = pop2 trees in
        next (Tree.Node (l, r) :: trees, instrs)
    in
    let do_leaf (trees, instrs) = Tree.Leaf :: trees, instrs in
    let do_node (trees, instrs) = trees, Gen :: Gen :: Build :: instrs in
    let finish (trees, _) = List.hd trees in
    fun () -> meta ~next ~finish ~do_leaf ~do_node ([], [Gen])
end

let test ~name (module B: Boltz) =
  Format.printf " ---- %s ---- @." name;
  Random.init 4242424242;
  for _ = 0 to 20 do
    let rand_state = Random.get_state () in
    Format.printf "=> %d / @?" (B.size ());
    Random.set_state rand_state;
    let tree = B.gen () in
    Format.printf "%d@." (Tree.size tree)
  done

let () =
  test ~name:"Specialised" (module Specialised);
  test ~name:"Generic" (module Generic);