open Core
open Core_bench

let write_bin_prot writer buf ~pos a =
  let len = writer.Bin_prot.Type_class.size a in
  assert (writer.Bin_prot.Type_class.write buf ~pos a = pos + len)
;;

let read_bin_prot reader buf ~pos_ref = reader.Bin_prot.Type_class.read buf ~pos_ref
let pos_ref = ref 0
let one = if Random.bool () then 1.0 else 1.0
let pi = if Random.bool () then 3.141597 else 3.141597

let read_float x =
  let buf = Bigstring.create 100 in
  write_bin_prot Float.bin_writer_t buf x ~pos:0;
  fun () ->
    pos_ref := 0;
    ignore (read_bin_prot Float.bin_reader_t buf ~pos_ref : float)
;;

let write_float x =
  let buf = Bigstring.create 100 in
  fun () -> write_bin_prot Float.bin_writer_t buf x ~pos:0
;;

module R = struct
  type t =
    { x : float
    ; y : float
    }
  [@@deriving bin_io]
end

let read_r x =
  let buf = Bigstring.create 100 in
  write_bin_prot R.bin_writer_t buf x ~pos:0;
  fun () ->
    pos_ref := 0;
    ignore (read_bin_prot R.bin_reader_t buf ~pos_ref : R.t)
;;

let write_r x =
  let buf = Bigstring.create 100 in
  fun () -> write_bin_prot R.bin_writer_t buf x ~pos:0
;;

let r_one : R.t = if Random.bool () then { x = 1.0; y = 1.0 } else { x = 1.0; y = 1.0 }
let r_pi : R.t = if Random.bool () then { x = pi; y = pi } else { x = pi; y = pi }

module IR = struct
  type t =
    | A of
        { x : float
        ; y : float
        }
    | B of
        { u : int
        ; v : int
        }
  [@@deriving bin_io]
end

let read_ir x =
  let buf = Bigstring.create 100 in
  write_bin_prot IR.bin_writer_t buf x ~pos:0;
  fun () ->
    pos_ref := 0;
    ignore (read_bin_prot IR.bin_reader_t buf ~pos_ref : IR.t)
;;

let write_ir x =
  let buf = Bigstring.create 100 in
  fun () -> write_bin_prot IR.bin_writer_t buf x ~pos:0
;;

let ir_one : IR.t = if Random.bool () then B { u = 1; v = 1 } else B { u = 1; v = 1 }
let ir_pi : IR.t = if Random.bool () then A { x = pi; y = pi } else A { x = pi; y = pi }
let lengths = [ 0; 1; 10; 100; 1000; 10_000 ]

let write_float_array len =
  let arr = Array.create ~len 1.0 in
  let buf = Bigstring.create 1_000_000 in
  Staged.stage (fun () -> write_bin_prot bin_writer_float_array buf arr ~pos:0)
;;

let read_float_array len =
  let arr = Array.create ~len 1.0 in
  let buf = Bigstring.create 1_000_000 in
  write_bin_prot bin_writer_float_array buf arr ~pos:0;
  Staged.stage (fun () ->
    pos_ref := 0;
    let arr = read_bin_prot bin_reader_float_array buf ~pos_ref in
    if not (Array.length arr = len)
    then failwithf "got len %d, expected %d" (Array.length arr) len ())
;;

let benchs =
  [ Bench.Test.create ~name:"write float one" (write_float one)
  ; Bench.Test.create ~name:"read float one" (read_float one)
  ; Bench.Test.create ~name:"write float pi" (write_float pi)
  ; Bench.Test.create ~name:"read float pi" (read_float pi)
  ; Bench.Test.create ~name:"write record one" (write_r r_one)
  ; Bench.Test.create ~name:"read record one" (read_r r_one)
  ; Bench.Test.create ~name:"write record pi" (write_r r_pi)
  ; Bench.Test.create ~name:"read record pi" (read_r r_pi)
  ; Bench.Test.create ~name:"write inline record one" (write_ir ir_one)
  ; Bench.Test.create ~name:"read inline record one" (read_ir ir_one)
  ; Bench.Test.create ~name:"write inline record pi" (write_ir ir_pi)
  ; Bench.Test.create ~name:"read inline record pi" (read_ir ir_pi)
  ; Bench.Test.create_indexed ~name:"write float array" ~args:lengths write_float_array
  ; Bench.Test.create_indexed ~name:"read float array" ~args:lengths read_float_array
  ]
;;

let () = Command_unix.run (Bench.make_command benchs)
