stack_size_effect.ml

[@@@warning "-32"]
(*  C stubs:

#define CAML_INTERNALS

#include <caml/mlvalues.h>
#include <caml/fiber.h>

value caml_used_stack_size(value cont) {
  value vstack = Field(cont, 0);
  struct stack_info* stack = Ptr_val(vstack);
  uint s = Stack_high(stack) - stack-> sp;
  return Val_int(s);
}


value caml_total_stack_size(value cont) {
  value vstack = Field(cont, 0);
  struct stack_info* stack = Ptr_val(vstack);
  uint s = Stack_high(stack) - Stack_base(stack);
  return Val_int(s);
}
*)

type _ Effect.t += Stack_size: (int*int) Effect.t


external caml_used_stack_size: ('a,'b) continuation -> int  =
  "caml_used_stack_size" [@@noalloc]

external caml_total_stack_size: ('a,'b) continuation -> int  =
  "caml_total_stack_size" [@@noalloc]

let suspended_stack_size k = caml_used_stack_size k, caml_total_stack_size k

let register = ref []

let record x = register := x :: !register
let delete x = register := List.filter ((!=) x) !register

let with_record x f =
  record x;
  Fun.protect f ~finally:(fun () -> delete x)

let run_with_stack_observer f = match f () with
  | x -> x
  | effect Stack_size, k ->
    let n = suspended_stack_size k in
    Effect.Deep.continue k n

let max_mem = ref 0
let alarm_size () =
  let stacks = List.fold_left (fun sum c -> sum + fst (suspended_stack_size c)) 0 !register in
  let current = 0 in
  let gc = Gc.quick_stat () in
  let total = gc.heap_words + current + stacks in
  if total > !max_mem then max_mem := total

let alarm = Gc.create_alarm alarm_size


let rec test k =
  if k = 0 then (
    Gc.full_major ();
    let us, fs = Effect.perform Stack_size in
    Format.eprintf "Stack size:%d/%d@." us fs;
    0.
  )
  else 1. +. test (k-1)

let () =
  run_with_stack_observer (fun () -> ignore @@ test 200;  Gc.delete_alarm alarm; Gc.full_major ());
  Format.eprintf "Max size = %d@." !max_mem