(* ocamlfind ocamlopt -O3 -g -o encoding -linkpkg -package zarith,unix,lwt.unix,lwt,core_bench,data-encoding -thread encoding.ml *)
(* Make sure your installed data-encoding-0.2 *)

open Data_encoding

module Bench1 = struct

  type tree = (* Empty | *) Leaf of int | Binary of tree * tree

  let tree_encoding : tree encoding =
    mu "tree" @@ fun self ->
    union [
    (* case
     *   (Tag 0)
     *   ~title:"Empty"
     *   unit
     *   (function Empty -> Some () | _ -> None)
     *   (fun () -> Empty); *)
    case
      (Tag 1)
      ~title:"Leaf"
      int31
      (function Leaf x -> Some x | _ -> None)
      (fun x -> Leaf x);
    case
      (Tag 2)
      ~title:"Binary"
      (obj2 (req "left" self) (req "right" self))
      (function Binary (left, right) -> Some (left, right) | _ -> None)
      (fun (left, right) -> Binary (left, right));
  ]

  let rec exp_tree n =
    if n = 0 then Leaf 0 else Binary (exp_tree (n - 1), exp_tree (n - 1))

  let large_tree = exp_tree 20

  let eval_data_encoding () =
    Data_encoding.Binary.to_bytes tree_encoding large_tree |> ignore

  let eval_marshal () =
    Marshal.to_bytes large_tree [] |> ignore

  let run () =
    let open Core in
    let open Core_bench in
    Command.run
      (Bench.make_command
         [
           Bench.Test.create ~name:"Marshal" eval_marshal;
           Bench.Test.create ~name:"Data encoding 0.2" eval_data_encoding
         ])

end

module Bench2 = struct

  type tree = Empty | Leaf of int | Binary of tree * tree

  let tree_encoding : tree encoding =
    mu "tree" @@ fun self ->
    union [
     case
        (Tag 0)
        ~title:"Empty"
        unit
        (function Empty -> Some () | _ -> None)
        (fun () -> Empty);
    case
      (Tag 1)
      ~title:"Leaf"
      int31
      (function Leaf x -> Some x | _ -> None)
      (fun x -> Leaf x);
    case
      (Tag 2)
      ~title:"Binary"
      (obj2 (req "left" self) (req "right" self))
      (function Binary (left, right) -> Some (left, right) | _ -> None)
      (fun (left, right) -> Binary (left, right));
  ]

  let rec exp_tree n =
    if n = 0 then Leaf 0 else Binary (exp_tree (n - 1), exp_tree (n - 1))

  let large_tree = exp_tree 20

  let eval_data_encoding () =
    Data_encoding.Binary.to_bytes tree_encoding large_tree |> ignore

  let eval_marshal () =
    Marshal.to_bytes large_tree [] |> ignore

  let run () =
    let open Core in
    let open Core_bench in
    Command.run
      (Bench.make_command
         [
           Bench.Test.create ~name:"Marshal" eval_marshal;
           Bench.Test.create ~name:"Data encoding 0.2" eval_data_encoding
         ])

end

let () =
  Bench1.run ();
  Bench2.run ()