(**************************************************************************)
(*                                                                        *)
(*  Ocamlviz --- real-time profiling tools for Objective Caml             *)
(*  Copyright (C) by INRIA - CNRS - Universite Paris Sud                  *)
(*  Authors: Julien Robert                                                *)
(*           Guillaume Von Tokarski                                       *)
(*           Sylvain Conchon                                              *)
(*           Jean-Christophe Filliatre                                    *)
(*           Fabrice Le Fessant                                           *)
(*  GNU Library General Public License version 2                          *)
(*  See file LICENSE for details                                          *)
(*                                                                        *)
(**************************************************************************)

open Format
open Gc
open Unix
open Protocol


(* datas *)

let buffer_size = 1024

let old_decl = ref []
let new_decl = ref []
let bind_decl = ref []
let old_bind_decl = ref []
let killed_decl = ref []

let nb_client = ref 0

let new_tag = 
  let nb_tag = ref 4 in 
  fun () -> incr nb_tag; !nb_tag

let create_tag kind name =
  let t = new_tag () in
  new_decl := Declare (t,kind,name) :: !new_decl;
  t

let word_size_coeff = if Sys.word_size = 32 then Int64.of_int 4 else Int64.of_int 8




(* GC *)

let period = 
  try float_of_string (Sys.getenv "OCAMLVIZ_PERIOD")
  with Not_found -> 0.1

let gc_period =  
  try float_of_string (Sys.getenv "OCAMLVIZ_HEAP_PERIOD")
  with Not_found -> 1.0

let gc_freq = 
  let f = gc_period /. period in
  int_of_float f

let heap_total_size = ref 0
let heap_alive_size = ref 0
let gc_cpt = ref 0

let tag_heap_total_size = (new_decl := Declare (0,Special,"heap_total_size") :: !new_decl); 0
let tag_heap_alive_size = (new_decl := Declare (1,Special,"heap_alive_size") :: !new_decl); 1

let heap_stats () =
  let gc_stats = 
    if !gc_cpt = 0 then Gc.stat ()
    else Gc.quick_stat ()
  in
  let total_size = gc_stats.heap_words in
  let alive_size = if !gc_cpt = 0 then gc_stats.live_words else !heap_alive_size in
  gc_cpt := (!gc_cpt + 1) mod gc_freq;
  if (total_size <> !heap_total_size) || (alive_size <> !heap_alive_size)
  then 	begin
    heap_total_size := total_size;
    heap_alive_size := alive_size;
    true
  end
  else false


(* Log *)

let log_size = 200
let log_list = ref []
let log_hist = Array.make log_size (0.,"")
let log_index = ref 0
let log_hist_index = ref 0
let log_tag = (new_decl := Declare (3,Klog,"log") :: !new_decl); 3


let log f = 
  let buf = Buffer.create 512 in
  let fmt = Format.formatter_of_buffer buf in
  Format.kfprintf 
    (fun _ -> 
       Format.pp_print_flush fmt ();
       let s = Buffer.contents buf in
       Buffer.clear buf;
       let t = (times()).tms_utime in
       if !log_index < log_size then
         begin
           log_list := (t,s) :: !log_list;
	   incr log_index;
         end;
       if !log_hist_index < log_size then
	 begin
	   log_hist.(!log_hist_index) <- (t,s);
	   incr log_hist_index;
	 end	   
       else
	 begin
	   log_hist.(0) <- (t,s);
	   log_hist_index := 1;
	 end;
       )
    fmt f
    

let list_of_log_hist () = 
  let temp = ref [] in
  for i = !log_hist_index to log_size -1 do
    let t,s = log_hist.(i) in
    if s <> "" then 
      temp := (t,s):: !temp;
  done;
  for i = 0 to !log_hist_index - 1 do
    let t,s = log_hist.(i) in
    temp := (t,s):: !temp;
  done;
  !temp


(* Observes *)

module Point = struct

  type t = {tag:tag;
	    name:string;
	    mutable count:int}
      
  let occur_table = Hashtbl.create 17
  let killed_list = ref []

  let create name = 
    let tag = create_tag Point name in
    let t = {tag = tag; name = name; count = 0} in
    Hashtbl.add occur_table t.tag t;
    t
    
      
  let kill t = 
    Hashtbl.remove occur_table t.tag;
    killed_list := (Send(t.tag,Int t.count)) ::(Send (t.tag,Killed)) ::  !killed_list;
    killed_decl := (Send(t.tag,Int t.count))::(Send (t.tag,Killed)) ::  !killed_decl

  let observe point = point.count<-point.count+1
    

  let observe_calls name fct =
    let tag = create_tag Point name in
    let t = {tag = tag ; name = name; count = 0} in
    Hashtbl.replace occur_table t.tag t;
    fun x -> t.count<-t.count+1; fct x


  let observe_calls2 name fct =
    let tag = create_tag Point name in
    let t = {tag = tag ; name = name; count = 0} in
    Hashtbl.replace occur_table t.tag t;
    fun x y -> t.count<-t.count+1; fct x y
      
      
  let observe_calls_rec name fct =
    let tag = create_tag Point name in
    let t = {tag = tag ; name = name; count = 0} in
    Hashtbl.replace occur_table t.tag t;
    let rec f x = t.count<-t.count+1; fct f x in
    f

      
  let observe_calls_rec2 name fct =
    let tag = create_tag Point name in
    let t = {tag = tag ; name = name; count = 0} in
    Hashtbl.replace occur_table t.tag t;
    let rec f x y = t.count<-t.count+1; fct f x y in
    f
      

  let msg_of_occur_table () =
    let l = !killed_list 
    in
    killed_list := [];
    Hashtbl.fold
      (fun tag t l ->
	 Send (tag,Int(t.count))::l
      ) occur_table l
      

  let send buf = 
    Bproto.encode buf (msg_of_occur_table ())


end
  
(*************)



(* tags *)

module Tag = struct

  module WeakHash = Weak.Make
    (struct type t = Obj.t let hash = Hashtbl.hash let equal = (==) end)

  type t = {
    tag_count:tag;
    tag_size:tag;
    name : string;
    count : bool;
    size : bool;
    mutable period: int;
    mutable ticks : int;
    mutable data : WeakHash.t;
    mutable idx : int;
    mutable fct_size : WeakHash.t -> Int64.t
  }

  let all_tags = ref []
  let killed_list = ref []

  module H = Hashtbl.Make(
    struct 
      type t = Obj.t 
      let equal = (==) 
      let hash = Hashtbl.hash
    end)
    

  let size_table = H.create 5003

  let size_of_double = Int64.of_int (Obj.size (Obj.repr 1.0))
  let size_of_word = Int64.of_int (Sys.word_size / 8)

  let (++) = Int64.add

  
  let count_size data =
    let count = ref 0L in (* words *)
    let rec traverse t =
      if not (H.mem size_table t) then begin
	H.add size_table t ();
	if Obj.is_block t then begin
	  let n = Obj.size t in
	  let tag = Obj.tag t in
	  if tag < Obj.no_scan_tag then begin
	    count := !count ++ 1L ++ Int64.of_int n;
	    for i = 0 to n - 1 do
      	      let f = Obj.field t i in 
	      if Obj.is_block f then traverse f
	    done
	  end else if tag = Obj.string_tag then
	    count := !count ++ 1L ++ Int64.of_int n 
	  else if tag = Obj.double_tag then
	    count := !count ++ size_of_double
	  else if tag = Obj.double_array_tag then
	    count := !count ++ 1L ++ Int64.mul size_of_double (Int64.of_int n)
	  else
	    count := Int64.succ !count
	end
      end
    in
    WeakHash.iter traverse data;
    H.clear size_table;
    Int64.mul !count size_of_word

  let create ?(size=false) ?(count=true) ?(fct_size:(WeakHash.t -> Int64.t)=count_size) ?(period=100) name = 
    let t = {
      tag_count = new_tag ();
      tag_size = new_tag ();
      name = name;
      count = count;
      size = size;
      period = period;
      ticks = 0;
      data = WeakHash.create 17;
      idx = 0;
      fct_size = fct_size;
    } 
    in
    all_tags := t :: !all_tags;
    begin 
      match size,count with
	|true,false -> 
	   new_decl := Declare (t.tag_size,Tag_size,name)::!new_decl;
	|true,true -> 
	   new_decl := Declare (t.tag_size,Tag_size,name)::!new_decl;
	    new_decl := Declare (t.tag_count,Tag_count,name)::!new_decl;
	    bind_decl := Bind [t.tag_size;t.tag_count]::!bind_decl
	|false,true -> new_decl := Declare (t.tag_count,Tag_count,name)::!new_decl;
	|false,false -> ()
    end;
(*     if size then *)
(*       new_decl := Declare (t.tag_size,Tag_size,name)::!new_decl; *)
(*     if count then *)
(*       new_decl := Declare (t.tag_count,Tag_count,name)::!new_decl; *)
    t

    

  let mark t x =
    WeakHash.add t.data (Obj.repr x);
    x


  let kill t = 
    let all =  (List.filter (fun t' -> t'!=t) !all_tags) in
    all_tags := all;
    if t.count then begin
      let count = WeakHash.count t.data
      in
      killed_list := (Send (t.tag_count,Int count))::(Send (t.tag_count,Killed)) :: !killed_list;
      killed_decl := (Send (t.tag_count,Int count))::(Send (t.tag_count,Killed)) :: !killed_decl;
    end;
    if t.size then begin
      let size = t.fct_size t.data 
      in
      killed_list := (Send (t.tag_size,Int64 size))::(Send (t.tag_size,Killed)) :: !killed_list;
      killed_decl := (Send (t.tag_size,Int64 size))::(Send (t.tag_size,Killed)) :: !killed_decl
    end


  let set_period t i = 
    t.period <- i;
    t.ticks <- 0


  let string_of_one buf t =
    if (t.ticks = 0) then
      begin
	let _ = 
	  if t.count then
	    let count = WeakHash.count t.data in
	    Bproto.encode_one buf (Send (t.tag_count,Int (count)))
	in
	if t.size then
	  let size = t.fct_size t.data in
	  Bproto.encode_one buf (Send (t.tag_size,Int64 size))
      end;
    t.ticks <- (t.ticks + 100) mod t.period

      

  let send buf =
    List.iter (string_of_one buf) !all_tags;
    Bproto.encode buf !killed_list;
    killed_list := []

end


(****** Value ******)

module Value = struct

  type 'a t = {tag:tag;
	       name:string;
	       data:'a;
	       mutable period:int;
	       mutable ticks:int}
      
  let int_table = (Hashtbl.create 17) 

  let int_ref_table = (Hashtbl.create 17)

  let int_weak_table = (Hashtbl.create 17) 
    
  let float_table = (Hashtbl.create 17) 
  
  let float_ref_table = (Hashtbl.create 17) 
    
  let float_weak_table = (Hashtbl.create 17) 

  let string_table = (Hashtbl.create 17) 
    
  let string_ref_table = (Hashtbl.create 17) 
  
  let string_weak_table = (Hashtbl.create 17) 

  let bool_table = (Hashtbl.create 17)

  let bool_ref_table = (Hashtbl.create 17) 
  
  let bool_weak_table = (Hashtbl.create 17) 

  
  let observe_ref period tag name r tbl =
    let t = {tag = tag;
	     name = name;
	     data = Weak.create 1;
	     period = period;
	     ticks = 0}
    in
    Weak.set t.data 0 (Some (r));
    Hashtbl.add tbl t.tag t

  let observe period tag name fct tbl =
    let t = {tag = tag;
	     name = name;
	     data = fct;
	     period = period;
	     ticks = 0}
    in
    Hashtbl.add tbl t.tag t

  let observe_weak period tag name fct tbl =
    let t = {tag = tag;
	     name = name;
	     data = Weak.create 1;
	     period = period;
	     ticks = 0}
    in
    Weak.set t.data 0 (Some (fct));
    Hashtbl.add tbl t.tag t
      
  let observe_int_fct ?(weak=false) ?(period=100) name fct =
    let tag = create_tag Value_int name in
    if weak then
      observe_weak period tag name fct int_weak_table
    else
      observe period tag name fct int_table
	
  let observe_int ?(weak=false) ?(period=100) name i =
    observe_int_fct ~weak ~period name (fun () -> i)
	
  let observe_float_fct ?(weak=false) ?(period=100) name fct =
    let tag = create_tag Value_float name in
    if weak then
      observe_weak period tag name fct float_weak_table
    else
      observe period tag name fct float_table

  let observe_float ?(weak=false) ?(period=100) name f =
    observe_float_fct ~weak ~period name (fun () -> f)

  let observe_string_fct ?(weak=false) ?(period=100) name fct =
    let tag = create_tag Value_string name in
    if weak then
      observe_weak period tag name fct string_weak_table
    else
      observe period tag name fct string_table

  let observe_string ?(weak=false) ?(period=100) name s =
    observe_string_fct ~weak ~period name (fun () -> s)

  let observe_bool_fct ?(weak=false) ?(period=100) name fct =
    let tag = create_tag Value_bool name in
    if weak then
      observe_weak period tag name fct bool_weak_table
    else
      observe period tag name fct bool_table

  let observe_bool ?(weak=false) ?(period=100) name b =
    observe_bool_fct ~weak ~period name (fun () -> b)

  let observe_int_ref ?(period=100) name (r:int ref) =
    let tag = create_tag Value_int name in
    observe_ref period tag name r int_ref_table;
    r
      
  let observe_float_ref ?(period=100) name (r:float ref) =
    let tag = create_tag Value_float name in
    observe_ref period tag name r float_ref_table;
    r

  let observe_string_ref ?(period=100) name (r:string ref) =
    let tag = create_tag Value_string name in
    observe_ref period tag name r string_ref_table;
    r

  let observe_bool_ref ?(period=100) name (r:bool ref) =
    let tag = create_tag Value_bool name in
    observe_ref period tag name r bool_ref_table;
    r
      
  let msg_of_values ()=
    
    let msg_of_value_ref tbl typ l =
      Hashtbl.fold
      (fun tag t l  ->
	 if (t.ticks = 0) then
	   begin
	     t.ticks <- (t.ticks + 100) mod t.period;
	     match (Weak.get t.data 0) with
	       |None -> Hashtbl.remove tbl tag;
		   killed_decl := (Send (tag, Collected))::!killed_decl;
		   Send (tag, Collected)::l
	       |Some r ->
		  Send (tag,typ (!r))::l
	   end
	 else
	   begin
	     t.ticks <- (t.ticks + 100) mod t.period;
	     l
	   end
	 
      ) tbl l
    in

    let msg_of_value_weak tbl typ l =
      Hashtbl.fold
      (fun tag t l  ->
	 if (t.ticks = 0) then
	   begin
	     t.ticks <- (t.ticks + 100) mod t.period;
	     match (Weak.get t.data 0) with
	       |None -> Hashtbl.remove tbl tag;
		   killed_decl := (Send (tag, Collected))::!killed_decl;
		   Send (tag, Collected)::l
	       |Some f ->
		  let v = typ (f ()) in 
		  Send (tag,v) :: l
		  (*
		    if v <> t.old_value then begin
		    t.old_value <- v;
		    Send (tag,v) :: l
		  end else
		    l
		  *)
	   end
	 else
	   begin
	     t.ticks <- (t.ticks + 100) mod t.period;
	     l
	   end
	 
      ) tbl l
    in

    let msg_of_value tbl typ f l =
      Hashtbl.fold
	(fun tag t l  ->
	   if (t.ticks = 0) then
	     begin
	       t.ticks <- (t.ticks + 100) mod t.period;
	       Send (tag,typ (t.data ()))::l
	     end
	   else
	     begin
	       t.ticks <- (t.ticks + 100) mod t.period;
	       l
	     end
	       
	) tbl l
    in
    
    let l = msg_of_value int_table (fun x -> Int x) string_of_int [] in
    let l = msg_of_value float_table (fun x -> Float x) string_of_float l in
    let l = msg_of_value string_table (fun x -> String x) (fun x -> x) l in
    let l = msg_of_value bool_table (fun x -> Bool x) string_of_bool l in
    let l = msg_of_value_ref int_ref_table (fun x -> Int x) l in
    let l = msg_of_value_ref float_ref_table (fun x -> Float x) l in
    let l = msg_of_value_ref string_ref_table (fun x -> String x) l in
    let l = msg_of_value_ref bool_ref_table (fun x -> Bool x) l in
    let l = msg_of_value_weak int_weak_table (fun x -> Int x) l in
    let l = msg_of_value_weak float_weak_table (fun x -> Float x) l in
    let l = msg_of_value_weak string_weak_table (fun x -> String x) l in
    msg_of_value_weak bool_weak_table (fun x -> Bool x) l


  let send buf = 
    Bproto.encode buf (msg_of_values ())
      
      
end

(****** Times ******)


module Time = struct

  type t = {tag:tag;
	   name:string;
	   mutable time:float;
	   mutable started:bool}

  let start_table = Hashtbl.create 17
    
  let time_table = Hashtbl.create 17
    
  let killed_list = ref []

  let init_time = (times()).tms_utime (* Unix.gettimeofday () *)

  let etime = 
    (new_decl := Declare (2,Special,"etime") :: !new_decl);
    let tag = 2 in
    {tag = tag;
     name = "etime";
     time = init_time;
     started = true}

  exception Already_started
  exception Not_started
      
  let elapsed_time () = 
    (times()).tms_utime (*Unix.gettimeofday ()*) -. init_time
      
  let create name = 
    let tag = create_tag Time name in
    let t = {tag = tag;
	     name = name;
	     time = 0.;
	     started = false}
    in
    Hashtbl.add start_table t.tag t; 
    t

  let kill t = 
    Hashtbl.remove start_table t.tag;
    Hashtbl.remove time_table t.tag;
    killed_list := (Send(t.tag,Float t.time)):: (Send (t.tag,Killed)) :: !killed_list;
    killed_decl := (Send(t.tag,Float t.time)):: (Send (t.tag,Killed)) :: !killed_decl
 
  let start t = 
    let u = (times()).tms_utime in
    if t.started then raise Already_started;
    if Hashtbl.mem time_table t.tag then
      begin
	
	Hashtbl.remove time_table t.tag;
	t.time <- u -. t.time ;
	Hashtbl.add start_table t.tag t;
      end
    else 
      begin 
	t.time <- t.time +. u;   (*Unix.gettimeofday ()*)
      end;
    t.started <- true

    

  let stop t = 
    let u = (times()).tms_utime (*Unix.gettimeofday ()*) in
    try
      Hashtbl.remove start_table t.tag;
      t.time <- u -. t.time;
      t.started <- false;
      Hashtbl.add time_table t.tag t;
    with Not_found ->
      raise Not_started


  let time name fct x =     
    let tag = create_tag Time name in
    let t = {tag = tag;
	     name = name;
	     time = (times()).tms_utime (*Unix.gettimeofday ()*);
	     started = true;}
    in
    Hashtbl.add start_table t.tag t;
    let y = fct x in
    t.time <- (times()).tms_utime (*Unix.gettimeofday ()*) -. t.time;
    t.started <- false;
    Hashtbl.replace time_table t.tag t;
    Hashtbl.remove start_table t.tag;
    y


  let msg_of_time_tables ()=
    let u = (times()).tms_utime (* Unix.gettimeofday () *) in
    let l = !killed_list
    in
    let l = Hashtbl.fold
      (fun tag t l ->
	 Send (tag,Float (u -. t.time))::l
      ) start_table l 
    in
    Hashtbl.fold
      (fun tag t l ->
	 Send (tag,Float (t.time))::l
      ) time_table l


  let send buf =
    Bproto.encode_one buf (Send(etime.tag,Float (elapsed_time())));
    Bproto.encode buf (msg_of_time_tables ())

end
 


module Tree = struct

    
  type t = {tag:tag;
	    name:string;
	    mutable period:int;
	    mutable ticks:int;
	    mutable var:(unit -> Protocol.variant)}


  let tree_table = Hashtbl.create 17


  let observe ?(period=100) name fct =
    let tag = create_tag Ktree name in
    let t = {tag = tag;
	     name = name;
	     period = period;
	     ticks = 0;
	     var = fct}
    in
    Hashtbl.add tree_table t.tag t


  let msg_of_tree_table () =
    Hashtbl.fold
      (fun tag t l ->
	 if t.ticks = 0 then
	   begin
	     t.ticks <- (t.ticks + 100) mod t.period;
	     Send (tag,Tree (t.var ()))::l
	   end
	 else  
	   begin 
	     t.ticks <- (t.ticks + 100) mod t.period;
	     l
	   end
      ) tree_table [] 

  let send buf = 
    Bproto.encode buf (msg_of_tree_table ())


end


module Hashtable = struct
   
  type 'a t = {tag:tag;
		    name:string;
		    mutable period:int;
		    mutable ticks:int;
		    mutable data:'a}
      
  let hash_table = Hashtbl.create 17

  let observe ?(period=100) name h =
    let tag = create_tag Hash name in
    let t = {tag = tag;
	     name = name;
	     period = period;
	     ticks = 0;
	     data = Weak.create 1}
    in
    Weak.set t.data 0 (Some (h));
    Hashtbl.add hash_table t.tag ( Obj.repr t);
    h


  exception Invalid_hashtbl
      
  let bucket_length  b = 
    let rec fct acc b = 
      if Obj.is_block b then
	begin
	  if Obj.size b <> 3 then raise Invalid_hashtbl;
	  fct (acc + 1) (Obj.field b 2)
	end
      else acc
    in
    fct 0 b
      
    
  let stats h =
    if Obj.is_block h then begin
      if Obj.size h <> 2 then raise Invalid_hashtbl;
      let a = Obj.field h 1 in
      let len = Obj.size a in
      let empty_buckets = ref 0 in
      let nb_elements = ref 0 in
      let max_length_bucket = ref 0 in
      for i=0 to len-1 do
	let ei = bucket_length (Obj.field a i) in
	nb_elements := !nb_elements + ei;
	if ei = 0 then incr empty_buckets;
	if ei > !max_length_bucket then max_length_bucket := ei;
      done;
      len,!nb_elements,!empty_buckets,!max_length_bucket
    end
    else raise Invalid_hashtbl
    


  let msg_of_hash_table () =
    Hashtbl.fold
      (fun tag t l ->
	 let t = Obj.magic t in
	 if t.ticks = 0 then
	   begin
	     t.ticks <- (t.ticks + 100) mod t.period;
	     match (Weak.get t.data 0) with
	       |None -> Hashtbl.remove hash_table tag;
		   killed_decl := (Send (tag, Collected))::!killed_decl;
		   Send (tag, Collected)::l
	       |Some h ->
		  let len,nb_elt,empty_b,max_len_b = stats h in
		  Send (tag,Hashtable (Int len,Int nb_elt,Int empty_b,Int max_len_b))::l
	   end
	 else
	   begin
	     t.ticks <- (t.ticks + 100) mod t.period;
	     l
	   end
      ) hash_table []



  let send buf =
    Bproto.encode buf (msg_of_hash_table ())
    
end



let to_send () =
  let buf = Buffer.create buffer_size in
  Bproto.encode buf !new_decl; 
  Bproto.encode buf !bind_decl;
  old_decl := !new_decl @ !old_decl;
  old_bind_decl := !bind_decl @ !old_bind_decl;
  new_decl := [];
  bind_decl := [];

  Point.send buf;
  Time.send buf;
  Value.send buf;
  Tag.send buf;
  Tree.send buf;
  Hashtable.send buf;

  if (heap_stats ()) then
    begin
      let total_size = Int64.mul (Int64.of_int !heap_total_size) word_size_coeff in
      let alive_size = Int64.mul (Int64.of_int !heap_alive_size) word_size_coeff in
      Bproto.encode_one buf (Send (tag_heap_total_size,Int64 (total_size)));
      Bproto.encode_one buf (Send (tag_heap_alive_size,Int64 (alive_size)));
    end;

  if !log_index > 0 then 
    begin 
      Bproto.encode_one buf (Send (log_tag, Log (!log_list)));
      log_index := 0;
      log_list := [];
    end;
 
  Buffer.contents buf
 

let declare_tags () = 
  let buf = Buffer.create 1024 in
  Bproto.encode buf !old_decl;
  Bproto.encode buf !old_bind_decl;
  Bproto.encode buf !killed_decl;
  Bproto.encode_one buf (Send (log_tag, (Log (list_of_log_hist ()))));
  Buffer.contents buf
  


let wait_for_killed_clients () = 
  while !nb_client > 0 do Unix.sleep 1; ignore [Random.int 10] done

let wait_for_connected_clients n = while !nb_client < n do Unix.sleep 1 done

let set_nb_clients nb = nb_client := nb
 


(*
Local Variables: 
compile-command: "unset LANG; make -C .."
End: 
*)
