Source file benchmark_accumulator.ml

let unique_id =
  let r = ref 0 in
  fun () ->
    incr r;
    !r
;;

(* Used to track the current libname in such a way that for functor applications, it is
   the calling libraries name that gets registered. *)
module Current_libname = struct
  let null = "<unknown>"
  let libname_ref = ref null
  let set str = libname_ref := str
  let unset () = libname_ref := null
  let get () = !libname_ref
end

module Current_bench_module_stack = struct
  let t = ref []
  let push s = t := s :: !t
  let pop_exn () = t := List.tl !t

  let to_name () =
    match !t with
    | [] -> None
    | ms -> Some (String.concat "." (List.rev ms))
  ;;
end

(* This is the main data structure of this module. An [Entry.t] represents a benchmark
   along with some metadata about is position, arguments etc. *)
module Entry = struct
  type ('param, 'a) parameterised_spec =
    { arg_name : string
    ; params : (string * 'param) list
    ; thunk : 'param -> unit -> 'a
    }

  type test_spec =
    | Regular_thunk : ([ `init ] -> unit -> 'a) -> test_spec
    | Parameterised_thunk : ('param, 'a) parameterised_spec -> test_spec

  type t =
    { unique_id : int
    ; code : string
    ; type_conv_path : string
    ; name : string
    ; filename : string
    ; line : int
    ; startpos : int
    ; endpos : int
    ; test_spec : test_spec
    ; bench_module_name : string option
    }

  let compare t1 t2 = compare t1.unique_id t2.unique_id

  (* Extracts module name from ["filename.ml.Module"], which is the format of [ext_name]
     as set by [typeconv]. *)
  let get_module_name_opt t =
    let str = t.type_conv_path in
    let len = String.length str in
    let rec loop i =
      if i + 4 <= len
      then
        if String.sub str i 4 = ".ml."
        then Some (String.sub str (i + 4) (len - i - 4))
        else loop (i + 1)
      else None
    in
    loop 0
  ;;

  let with_test_spec t test_spec = { t with test_spec }
end

(* Inspect system environment variables to decide if benchmarks are being run. This is
   called by the code generated by the [pa_bench] syntax to decide if the global hashtable
   should be populated. *)
let add_environment_var =
  let v =
    try Sys.getenv "BENCHMARKS_RUNNER" with
    | Not_found -> ""
  in
  v = "TRUE"
;;

(* This hashtable contains all the benchmarks from all the of libraries that have been
   loaded. At the time the benchmarks are registering themselves with [ppx_bench_lib] we
   don't yet know which libraries will need to be run.  *)
let libs_to_entries : (string, Entry.t list) Hashtbl.t = Hashtbl.create 10

let lookup_rev_lib ~libname =
  try Hashtbl.find libs_to_entries libname with
  | Not_found -> []
;;

let lookup_lib ~libname = List.rev (lookup_rev_lib ~libname)

let force_drop =
  (* Useful for js_of_ocaml to perform deadcode elimination.
     see ppx/ppx_inline_test/runtime-lib/runtime.ml [Action.get] for more details *)
  try
    ignore (Sys.getenv "FORCE_DROP_BENCH" : string);
    true
  with
  | Not_found -> false
;;

let get_mode () = if force_drop then `Ignore else `Collect

let[@inline never] add_bench
  ~name
  ~code
  ~filename
  ~type_conv_path
  ~line
  ~startpos
  ~endpos
  test_spec
  =
  match get_mode () with
  | `Ignore -> ()
  | `Collect ->
    let libname = Current_libname.get () in
    let entry =
      { Entry.code
      ; unique_id = unique_id ()
      ; type_conv_path
      ; bench_module_name = Current_bench_module_stack.to_name ()
      ; name
      ; filename
      ; line
      ; startpos
      ; endpos
      ; test_spec
      }
    in
    Hashtbl.add libs_to_entries libname (entry :: lookup_rev_lib ~libname)
;;

let[@inline never] add_bench_module
  ~name
  ~code:_
  ~type_conv_path:_
  ~filename:_
  ~line:_
  ~startpos:_
  ~endpos:_
  f
  =
  match get_mode () with
  | `Ignore -> ()
  | `Collect ->
    (* Running f registers the benchmarks using BENCH *)
    Current_bench_module_stack.push name;
    (try
       f ();
       Current_bench_module_stack.pop_exn ()
     with
     | ex ->
       Current_bench_module_stack.pop_exn ();
       raise ex)
;;