open Step

type thunk = unit -> result

type trace_entry = {
  pc : int;
  stack_len : int;
  cont_depth : int;
}

type config = {
  max_steps : int option;
  trace : bool;
  trace_entries : trace_entry list ref;
}

let default_config = {
  max_steps = None;
  trace = false;
  trace_entries = ref [];
}

let rec run ?(config = default_config) thunk =
  match config.max_steps with
  | Some max when max <= 0 -> None
  | Some max -> run_with_limit ~config max thunk
  | None -> run_unlimited ~config thunk

and run_unlimited ~config thunk =
  match thunk () with
  | Continue { state; cont } ->
    (if config.trace then
       config.trace_entries := {
         pc = state.pc;
         stack_len = List.length state.stack;
         cont_depth = Cont.depth cont;
       } :: !(config.trace_entries));
    run_unlimited ~config (fun () -> step state cont)
  | Halted v -> Some v
  | Error e -> None

and run_with_limit ~config max_steps thunk =
  let rec aux steps thunk =
    if steps >= max_steps then
      None
    else
      match thunk () with
      | Continue { state; cont } ->
        (if config.trace then
           config.trace_entries := {
             pc = state.pc;
             stack_len = List.length state.stack;
             cont_depth = Cont.depth cont;
           } :: !(config.trace_entries));
        aux (steps + 1) (fun () -> step state cont)
      | Halted v -> Some v
      | Error e -> None
  in
  aux 0 thunk

let get_trace config = List.rev !(config.trace_entries)

let clear_trace config = config.trace_entries := []