(* profiler.ml

   This file is the backbone of a simple little profiler.  The
   profiler needs to allocate a global data structure for each
   profiled function, and an array of all these data structures.
   Constructing this list of functions is challenging.  gprof gets the
   list of profiled functions from the executable, but we cannot
   because the assembler does not put the right information in the
   symbol table (in the future this might be the way to go).

   Each profiled function imports a data structure of type prof_t
   whose name is a mangled version of the function's name.  To find
   the function I scale through the _i.tali files for the project for
   values of type prof_t, and demangle the name.

    Once we have this list of functions we emit a file containing
    global declarations of all the imported data structures, and a
    global array of these data structures.  The profiled code does not
    know what is in the data structure; it just gets a pointer.

   We then compile this temporary file into an object file that can be
   linked with the profiled code.  We also need to generate a .tali
   file that popcorn programs can use in conjunction with the object
   file so that they link verify. *)

(* WARNING: The size of prof_typ_decl must be kept in sync with the profiled
   code.  We also need to ensure doubleword alignment. *)

let prof_typ_decl = "typedef struct { int  c1,c2,c3,c4,c5,c6,c7,c8,c9,c10; int depth; char *name; } prof_t;\n"

let prof_typ = "prof_t"
let prof_array = "prof_data"
let prof_array_size = "prof_data_size"

(* We strip off the leading underscores inserted to match the C output types. *)
let collect tal_in fns = 
  let template = 
(*    Str.regexp ";[ \009\012]*FUNCTION:[ \009\012]*\([a-zA-Z0-9\?_]+\).*" in *)
    Str.regexp "[ \009\012]*VAL[ \009\012]*_\([a-zA-Z0-9\$_]+\),<`prof_t\?>.*" in
  try
    while true do
      let line = input_line tal_in in
      if Str.string_match template line 0 then
	begin 
	  let f = Str.matched_group 1 line in
	  fns := f :: !fns
	end
    done;
    ()
  with End_of_file -> ()

(*
let rewrite_name s = 
  let s = String.copy s in
  let s_len = String.length s in
  for i = 0 to s_len - 1 do
    if s.[i] = '?' then s.[i] <- '$' else ()
  done;
  ("prof_"^s^"_t")
;;*)
let rewrite_name s = s;;

let strip_name s =
  let templ = Str.regexp "prof_\(.*\)_t" in
  if Str.string_match templ s 0 then (Str.matched_group 1 s) else s
;;

let gen_decls f names =
  let gen_decl n =
    (Printf.fprintf f "%s %s = {0,0,0,0,0,0,0,0,0,0,0,\"%s\"};\n" 
       prof_typ (rewrite_name n) (strip_name n))
  in
  List.iter gen_decl names
;;

let gen_array f names =
  let pr_f f i s = output_string f s; (i + String.length s) in
  let pr i s = pr_f f i s in
  let pr' s = ignore (pr 0 s) in
  let head = prof_typ ^ " *"^prof_array^"[]={" in
  let pr_padding () = pr 0 (String.make ((String.length head)-2) ' ') in
  let end_line () = ignore (pr 0 "\n") in
  let rec aux i names = 
    match names with 
    | [] -> ()
    | hd::tl ->
	let i = 
	  if i=0 then pr_padding () else 
	  if i + (String.length hd) > 77 then (end_line (); pr_padding ())
	  else i in
	aux (pr (pr i " ,&") hd) tl
  in
  let i = pr 0 head in
  (match names with
  | [] -> ()
  | hd::tl -> aux (pr i ("&" ^hd)) tl);
  pr' "};\n"

let gen_tali f names =
  let pr s = output_string f s in
  pr "; TAL interface file\n";
  pr "; Generated by the profiler.\n";
  (* Add back underscore here. *)
  let gen_val n = pr"\tVAL\t_"; pr n; pr",<`prof_t?>\n" in
  List.iter gen_val names;
  pr "\n"
;;

type fileKind = Tali | Unrecognized 

let file_kind f =
  if Filename.check_suffix f ".tali" then Tali else
  Unrecognized
  
let main () = 
  let out_f = ref "" in
  let fns = ref ["prof_GC_malloc_t";"prof_GC_malloc_atomic_t"] in
  let set_output s = out_f := s in
  let new_fns fname = 
    if Sys.file_exists fname then
      (match file_kind fname with
      |	Tali -> 
	  let in_c = open_in fname in
	  collect in_c fns; close_in in_c
      |	_ -> Printf.eprintf "Unrecognized file extension: %s.\n" fname)
    else Printf.eprintf "Warning: Could not find file %s" fname
  in
  let cmd_options = ["-o",Arg.String set_output,"specify output file"] in
  Arg.parse cmd_options new_fns "usage: profiler -o <file> files ...\n";

  let names = Sort.list (<=) !fns in
  let rec dup x = 
    (* Although global variables cannot be duplicated, the same static function
       could exist in two different files.  If this turns up the fix is to
       mangle the names of static functions with the file in which they occur.
       *)
    match x with
    | [] | [_] -> false
    | hd1::hd2::tl -> 
	if hd1=hd2 then (Printf.eprintf "Duplicate function: %s\n" hd1; 
			 true)
	else dup (hd2::tl) in
  if dup names then () else 
  try 
    let out = if !out_f = "" then stdout else open_out (!out_f) in
    let tali_f = 
      if !out_f = "" then "" 
      else (Filename.chop_extension (!out_f))^".tali" in
    let tali_out = if tali_f = "" then stdout else open_out tali_f in
    output_string out (prof_typ_decl^"\n");
    gen_decls out names;
    output_string out "\n";
    let prof_names = List.map rewrite_name names in
    gen_array out prof_names;
    output_string out "\n";
    Printf.fprintf out "int %s = %d;\n" prof_array_size (List.length names);
    gen_tali tali_out prof_names;
    close_out out;
    close_out tali_out;
    ()
  with Sys_error s -> Printf.eprintf "Sys Error: %s \n" s
;;

main ()
