module StringMap = Map.Make (String)

let find_opt_back k t =
  (* avoid raising exceptions, they can be expensive *)
  if StringMap.mem k t then Some (StringMap.find k t) else None

let find_opt_raise k t = try Some (StringMap.find k t) with Not_found -> None

let update_raise k f t =
  let r = find_opt_raise k t in
  let r' = f r in
  match (r, r') with
  | None, None ->
      t
  | Some _, None ->
      StringMap.remove k t
  | Some r, Some r' when r == r' ->
      t
  | _, Some r' ->
      StringMap.add k r' t

let update k f t =
  let r = find_opt_back k t in
  let r' = f r in
  match (r, r') with
  | None, None ->
      t
  | Some _, None ->
      StringMap.remove k t
  | Some r, Some r' when r == r' ->
      t
  | _, Some r' ->
      StringMap.add k r' t

let do_raise () = raise Not_found

let do_raise_notrace () = raise_notrace Not_found

let dummy = ref 0

let wrap f () =
  try if Sys.opaque_identity !dummy = 0 then f () with Not_found -> ()

(* open here to make sure we use stdlib impl above *)
open! Core
open Core_bench

let pre = "/local/domain/0"

let test_map = List.init 1000 (fun i -> pre ^ string_of_int i)

let args = [10; 100; 1000]

let mk_map_bench ~name f =
  Bench.Test.create_indexed ~name ~args (fun len ->
      (* this runs before the benchmark *)
      let m =
        List.init len (fun i -> (pre ^ string_of_int i, "value"))
        |> Caml.List.to_seq |> StringMap.of_seq
      in
      Staged.stage (fun () ->
          (* this is the benchmarked function, the benchmark framework takes care of
             running it multiple times, and avoiding the compiler optimizing it away
             by "using" the result *)
          f m))

let () =
  Printexc.record_backtrace true ;
  let key = pre ^ "/nonexistent" in
  let keyf = pre ^ "5" in
  let f = function Some x -> None | None -> Some "value2" in
  Command.run
    (Command.group ~summary:"exception handling benchmarkls"
       [ ( "raise"
         , Bench.make_command
             [ Bench.Test.create ~name:"raise" (wrap do_raise)
             ; Bench.Test.create ~name:"raise_notrace" (wrap do_raise_notrace)
             ] )
       ; ( "find-opt"
         , Bench.make_command
             [ mk_map_bench ~name:"find_opt 4.06" (fun m ->
                   StringMap.find_opt key m)
             ; mk_map_bench ~name:"find_opt=mem+find" (fun m ->
                   find_opt_back key m)
             ; mk_map_bench ~name:"find_opt=find+catch" (fun m ->
                   find_opt_raise key m) ] )
       ; ( "find-opt-found"
         , Bench.make_command
             [ mk_map_bench ~name:"find_opt 4.06" (fun m ->
                   StringMap.find_opt keyf m)
             ; mk_map_bench ~name:"find_opt=mem+find" (fun m ->
                   find_opt_back keyf m)
             ; mk_map_bench ~name:"find_opt=find+catch" (fun m ->
                   find_opt_raise keyf m) ] )
       ; ( "update"
         , Bench.make_command
             [ mk_map_bench ~name:"update 4.06" (fun m ->
                   StringMap.update key f m)
             ; mk_map_bench ~name:"update=find+catch+add/remove" (fun m ->
                   update_raise key f m)
             ; mk_map_bench ~name:"update=mem+find+add/remove" (fun m ->
                   update key f m) ]

             )
       ; ( "update-found"
         , Bench.make_command
             [ mk_map_bench ~name:"update 4.06" (fun m ->
                   StringMap.update key f m)
             ; mk_map_bench ~name:"update=find+catch+add/remove" (fun m ->
                   update_raise keyf f m)
             ; mk_map_bench ~name:"update=mem+find+add/remove" (fun m ->
                   update keyf f m) ] ) ])