1(**************************************************************************)
2(*                                                                        *)
3(*                                 OCaml                                  *)
4(*                                                                        *)
5(*       Damien Doligez and Francois Rouaix, INRIA Rocquencourt           *)
6(*           Ported to Caml Special Light by John Malecki                 *)
7(*                                                                        *)
8(*   Copyright 1996 Institut National de Recherche en Informatique et     *)
9(*     en Automatique.                                                    *)
10(*                                                                        *)
11(*   All rights reserved.  This file is distributed under the terms of    *)
12(*   the GNU Lesser General Public License version 2.1, with the          *)
13(*   special exception on linking described in the file LICENSE.          *)
14(*                                                                        *)
15(**************************************************************************)
16
17open Printf
18
19open Location
20open Parsetree
21
22(* User programs must not use identifiers that start with these prefixes. *)
23let idprefix = "__ocaml_prof_";;
24let modprefix = "OCAML__prof_";;
25
26(* Errors specific to the profiler *)
27exception Profiler of string
28
29(* Modes *)
30let instr_fun    = ref false
31and instr_match  = ref false
32and instr_if     = ref false
33and instr_loops  = ref false
34and instr_try    = ref false
35
36let cur_point = ref 0
37and inchan = ref stdin
38and outchan = ref stdout
39
40(* To copy source fragments *)
41let copy_buffer = Bytes.create 256
42
43let copy_chars_unix nchars =
44  let n = ref nchars in
45  while !n > 0 do
46    let m = input !inchan copy_buffer 0 (min !n 256) in
47    if m = 0 then raise End_of_file;
48    output !outchan copy_buffer 0 m;
49    n := !n - m
50  done
51
52let copy_chars_win32 nchars =
53  for _i = 1 to nchars do
54    let c = input_char !inchan in
55    if c <> '\r' then output_char !outchan c
56  done
57
58let copy_chars =
59  match Sys.os_type with
60    "Win32" | "Cygwin" -> copy_chars_win32
61  | _       -> copy_chars_unix
62
63let copy next =
64  assert (next >= !cur_point);
65  seek_in !inchan !cur_point;
66  copy_chars (next - !cur_point);
67  cur_point := next;
68;;
69
70let prof_counter = ref 0;;
71
72let instr_mode = ref false
73
74type insert = Open | Close;;
75let to_insert = ref ([] : (insert * int) list);;
76
77let insert_action st en =
78  to_insert := (Open, st) :: (Close, en) :: !to_insert
79;;
80
81(* Producing instrumented code *)
82let add_incr_counter modul (kind,pos) =
83   copy pos;
84   match kind with
85   | Open ->
86         fprintf !outchan "(%sProfiling.incr %s%s_cnt %d; "
87                 modprefix idprefix modul !prof_counter;
88         incr prof_counter;
89   | Close -> fprintf !outchan ")";
90;;
91
92let counters = ref (Array.make 0 0)
93
94(* User defined marker *)
95let special_id = ref ""
96
97(* Producing results of profile run *)
98let add_val_counter (kind,pos) =
99  if kind = Open then begin
100    copy pos;
101    fprintf !outchan "(* %s%d *) " !special_id !counters.(!prof_counter);
102    incr prof_counter;
103  end
104;;
105
106(* ************* rewrite ************* *)
107
108let insert_profile rw_exp ex =
109  let st = ex.pexp_loc.loc_start.Lexing.pos_cnum
110  and en = ex.pexp_loc.loc_end.Lexing.pos_cnum
111  and gh = ex.pexp_loc.loc_ghost
112  in
113  if gh || st = en then
114    rw_exp true ex
115  else begin
116    insert_action st en;
117    rw_exp false ex;
118  end
119;;
120
121
122let pos_len = ref 0
123
124let init_rewrite modes mod_name =
125  cur_point := 0;
126  if !instr_mode then begin
127    fprintf !outchan "module %sProfiling = Profiling;; " modprefix;
128    fprintf !outchan "let %s%s_cnt = Array.make 000000000" idprefix mod_name;
129    pos_len := pos_out !outchan;
130    fprintf !outchan
131            " 0;; Profiling.counters := \
132              (\"%s\", (\"%s\", %s%s_cnt)) :: !Profiling.counters;; "
133            mod_name modes idprefix mod_name;
134  end
135
136let final_rewrite add_function =
137  to_insert := List.sort (fun x y -> compare (snd x) (snd y)) !to_insert;
138  prof_counter := 0;
139  List.iter add_function !to_insert;
140  copy (in_channel_length !inchan);
141  if !instr_mode then begin
142    let len = string_of_int !prof_counter in
143    if String.length len > 9 then raise (Profiler "too many counters");
144    seek_out !outchan (!pos_len - String.length len);
145    output_string !outchan len
146  end;
147  (* Cannot close because outchan is stdout and Format doesn't like
148     a closed stdout.
149    close_out !outchan;
150  *)
151;;
152
153let rec rewrite_patexp_list iflag l =
154  rewrite_exp_list iflag (List.map (fun x -> x.pvb_expr) l)
155
156and rewrite_cases iflag l =
157  List.iter
158    (fun pc ->
159      begin match pc.pc_guard with
160      | None -> ()
161      | Some g -> rewrite_exp iflag g
162      end;
163      rewrite_exp iflag pc.pc_rhs
164    )
165    l
166
167and rewrite_labelexp_list iflag l =
168  rewrite_exp_list iflag (List.map snd l)
169
170and rewrite_exp_list iflag l =
171  List.iter (rewrite_exp iflag) l
172
173and rewrite_exp iflag sexp =
174  if iflag then insert_profile rw_exp sexp
175           else rw_exp false sexp
176
177and rw_exp iflag sexp =
178  match sexp.pexp_desc with
179    Pexp_ident _lid -> ()
180  | Pexp_constant _cst -> ()
181
182  | Pexp_let(_, spat_sexp_list, sbody) ->
183    rewrite_patexp_list iflag spat_sexp_list;
184    rewrite_exp iflag sbody
185
186  | Pexp_function caselist ->
187    if !instr_fun then
188      rewrite_function iflag caselist
189    else
190      rewrite_cases iflag caselist
191
192  | Pexp_fun (_, _, p, e) ->
193      let l = [{pc_lhs=p; pc_guard=None; pc_rhs=e}] in
194      if !instr_fun then
195        rewrite_function iflag l
196      else
197        rewrite_cases iflag l
198
199  | Pexp_match(sarg, caselist) ->
200    rewrite_exp iflag sarg;
201    if !instr_match && not sexp.pexp_loc.loc_ghost then
202      rewrite_funmatching caselist
203    else
204      rewrite_cases iflag caselist
205
206  | Pexp_try(sbody, caselist) ->
207    rewrite_exp iflag sbody;
208    if !instr_try && not sexp.pexp_loc.loc_ghost then
209      rewrite_trymatching caselist
210    else
211      rewrite_cases iflag caselist
212
213  | Pexp_apply(sfunct, sargs) ->
214    rewrite_exp iflag sfunct;
215    rewrite_exp_list iflag (List.map snd sargs)
216
217  | Pexp_tuple sexpl ->
218    rewrite_exp_list iflag sexpl
219
220  | Pexp_construct(_, None) -> ()
221  | Pexp_construct(_, Some sarg) ->
222    rewrite_exp iflag sarg
223
224  | Pexp_variant(_, None) -> ()
225  | Pexp_variant(_, Some sarg) ->
226    rewrite_exp iflag sarg
227
228  | Pexp_record(lid_sexp_list, None) ->
229    rewrite_labelexp_list iflag lid_sexp_list
230  | Pexp_record(lid_sexp_list, Some sexp) ->
231    rewrite_exp iflag sexp;
232    rewrite_labelexp_list iflag lid_sexp_list
233
234  | Pexp_field(sarg, _) ->
235    rewrite_exp iflag sarg
236
237  | Pexp_setfield(srecord, _, snewval) ->
238    rewrite_exp iflag srecord;
239    rewrite_exp iflag snewval
240
241  | Pexp_array(sargl) ->
242    rewrite_exp_list iflag sargl
243
244  | Pexp_ifthenelse(scond, sifso, None) ->
245      rewrite_exp iflag scond;
246      rewrite_ifbody iflag sexp.pexp_loc.loc_ghost sifso
247  | Pexp_ifthenelse(scond, sifso, Some sifnot) ->
248      rewrite_exp iflag scond;
249      rewrite_ifbody iflag sexp.pexp_loc.loc_ghost sifso;
250      rewrite_ifbody iflag sexp.pexp_loc.loc_ghost sifnot
251
252  | Pexp_sequence(sexp1, sexp2) ->
253    rewrite_exp iflag sexp1;
254    rewrite_exp iflag sexp2
255
256  | Pexp_while(scond, sbody) ->
257    rewrite_exp iflag scond;
258    if !instr_loops && not sexp.pexp_loc.loc_ghost
259    then insert_profile rw_exp sbody
260    else rewrite_exp iflag sbody
261
262  | Pexp_for(_, slow, shigh, _, sbody) ->
263    rewrite_exp iflag slow;
264    rewrite_exp iflag shigh;
265    if !instr_loops && not sexp.pexp_loc.loc_ghost
266    then insert_profile rw_exp sbody
267    else rewrite_exp iflag sbody
268
269  | Pexp_constraint(sarg, _) | Pexp_coerce(sarg, _, _) ->
270    rewrite_exp iflag sarg
271
272  | Pexp_send (sobj, _) ->
273    rewrite_exp iflag sobj
274
275  | Pexp_new _ -> ()
276
277  | Pexp_setinstvar (_, sarg) ->
278    rewrite_exp iflag sarg
279
280  | Pexp_override l ->
281      List.iter (fun (_, sexp) -> rewrite_exp iflag sexp) l
282
283  | Pexp_letmodule (_, smod, sexp) ->
284      rewrite_mod iflag smod;
285      rewrite_exp iflag sexp
286
287  | Pexp_letexception (_cd, exp) ->
288      rewrite_exp iflag exp
289
290  | Pexp_assert (cond) -> rewrite_exp iflag cond
291
292  | Pexp_lazy (expr) -> rewrite_exp iflag expr
293
294  | Pexp_poly (sexp, _) -> rewrite_exp iflag sexp
295
296  | Pexp_object cl ->
297      List.iter (rewrite_class_field iflag) cl.pcstr_fields
298
299  | Pexp_newtype (_, sexp) -> rewrite_exp iflag sexp
300  | Pexp_open (_ovf, _, e) -> rewrite_exp iflag e
301  | Pexp_pack (smod) -> rewrite_mod iflag smod
302  | Pexp_extension _ -> ()
303  | Pexp_unreachable -> ()
304
305and rewrite_ifbody iflag ghost sifbody =
306  if !instr_if && not ghost then
307    insert_profile rw_exp sifbody
308  else
309    rewrite_exp iflag sifbody
310
311(* called only when !instr_fun *)
312and rewrite_annotate_exp_list l =
313  List.iter
314    (function
315     | {pc_guard=Some scond; pc_rhs=sbody} ->
316         insert_profile rw_exp scond;
317         insert_profile rw_exp sbody;
318     | {pc_rhs={pexp_desc = Pexp_constraint(sbody, _)}} (* let f x : t = e *)
319        -> insert_profile rw_exp sbody
320     | {pc_rhs=sexp} -> insert_profile rw_exp sexp)
321    l
322
323and rewrite_function iflag = function
324  | [{pc_lhs=_; pc_guard=None;
325      pc_rhs={pexp_desc = (Pexp_function _|Pexp_fun _)} as sexp}] ->
326        rewrite_exp iflag sexp
327  | l -> rewrite_funmatching l
328
329and rewrite_funmatching l =
330  rewrite_annotate_exp_list l
331
332and rewrite_trymatching l =
333  rewrite_annotate_exp_list l
334
335(* Rewrite a class definition *)
336
337and rewrite_class_field iflag cf =
338  match cf.pcf_desc with
339    Pcf_inherit (_, cexpr, _)     -> rewrite_class_expr iflag cexpr
340  | Pcf_val (_, _, Cfk_concrete (_, sexp))  -> rewrite_exp iflag sexp
341  | Pcf_method (_, _,
342                Cfk_concrete (_, ({pexp_desc = (Pexp_function _|Pexp_fun _)}
343                                    as sexp))) ->
344      rewrite_exp iflag sexp
345  | Pcf_method (_, _, Cfk_concrete(_, sexp)) ->
346      let loc = cf.pcf_loc in
347      if !instr_fun && not loc.loc_ghost then insert_profile rw_exp sexp
348      else rewrite_exp iflag sexp
349  | Pcf_initializer sexp ->
350      rewrite_exp iflag sexp
351  | Pcf_method (_, _, Cfk_virtual _)
352  | Pcf_val (_, _, Cfk_virtual _)
353  | Pcf_constraint _  -> ()
354  | Pcf_attribute _ -> ()
355  | Pcf_extension _ -> ()
356
357and rewrite_class_expr iflag cexpr =
358  match cexpr.pcl_desc with
359    Pcl_constr _ -> ()
360  | Pcl_structure st ->
361      List.iter (rewrite_class_field iflag) st.pcstr_fields
362  | Pcl_fun (_, _, _, cexpr) ->
363      rewrite_class_expr iflag cexpr
364  | Pcl_apply (cexpr, exprs) ->
365      rewrite_class_expr iflag cexpr;
366      List.iter (rewrite_exp iflag) (List.map snd exprs)
367  | Pcl_let (_, spat_sexp_list, cexpr) ->
368      rewrite_patexp_list iflag spat_sexp_list;
369      rewrite_class_expr iflag cexpr
370  | Pcl_constraint (cexpr, _) ->
371      rewrite_class_expr iflag cexpr
372  | Pcl_extension _ -> ()
373
374and rewrite_class_declaration iflag cl =
375  rewrite_class_expr iflag cl.pci_expr
376
377(* Rewrite a module expression or structure expression *)
378
379and rewrite_mod iflag smod =
380  match smod.pmod_desc with
381    Pmod_ident _ -> ()
382  | Pmod_structure sstr -> List.iter (rewrite_str_item iflag) sstr
383  | Pmod_functor(_param, _smty, sbody) -> rewrite_mod iflag sbody
384  | Pmod_apply(smod1, smod2) -> rewrite_mod iflag smod1; rewrite_mod iflag smod2
385  | Pmod_constraint(smod, _smty) -> rewrite_mod iflag smod
386  | Pmod_unpack(sexp) -> rewrite_exp iflag sexp
387  | Pmod_extension _ -> ()
388
389and rewrite_str_item iflag item =
390  match item.pstr_desc with
391    Pstr_eval (exp, _attrs) -> rewrite_exp iflag exp
392  | Pstr_value(_, exps)
393     -> List.iter (fun x -> rewrite_exp iflag x.pvb_expr) exps
394  | Pstr_module x -> rewrite_mod iflag x.pmb_expr
395        (* todo: Pstr_recmodule?? *)
396  | Pstr_class classes -> List.iter (rewrite_class_declaration iflag) classes
397  | _ -> ()
398
399(* Rewrite a .ml file *)
400let rewrite_file srcfile add_function =
401  inchan := open_in_bin srcfile;
402  let lb = Lexing.from_channel !inchan in
403  Location.input_name := srcfile;
404  Location.init lb srcfile;
405  List.iter (rewrite_str_item false) (Parse.implementation lb);
406  final_rewrite add_function;
407  close_in !inchan
408
409(* Copy a non-.ml file without change *)
410let null_rewrite srcfile =
411  inchan := open_in_bin srcfile;
412  copy (in_channel_length !inchan);
413  close_in !inchan
414;;
415
416(* Setting flags from saved config *)
417let set_flags s =
418  for i = 0 to String.length s - 1 do
419    match String.get s i with
420      'f' -> instr_fun := true
421    | 'm' -> instr_match := true
422    | 'i' -> instr_if := true
423    | 'l' -> instr_loops := true
424    | 't' -> instr_try := true
425    | 'a' -> instr_fun := true; instr_match := true;
426             instr_if := true; instr_loops := true;
427             instr_try := true
428    | _ -> ()
429    done
430
431(* Command-line options *)
432
433let modes = ref "fm"
434let dumpfile = ref "ocamlprof.dump"
435
436(* Process a file *)
437
438let process_intf_file filename = null_rewrite filename;;
439
440let process_impl_file filename =
441   let modname = Filename.basename(Filename.chop_extension filename) in
442       (* FIXME should let modname = String.capitalize modname *)
443   if !instr_mode then begin
444     (* Instrumentation mode *)
445     set_flags !modes;
446     init_rewrite !modes modname;
447     rewrite_file filename (add_incr_counter modname);
448   end else begin
449     (* Results mode *)
450     let ic = open_in_bin !dumpfile in
451     let allcounters =
452       (input_value ic : (string * (string * int array)) list) in
453     close_in ic;
454     let (modes, cv) =
455       try
456         List.assoc modname allcounters
457       with Not_found ->
458         raise(Profiler("Module " ^ modname ^ " not used in this profile."))
459     in
460     counters := cv;
461     set_flags modes;
462     init_rewrite modes modname;
463     rewrite_file filename add_val_counter;
464   end
465;;
466
467let process_anon_file filename =
468  if Filename.check_suffix filename ".ml" then
469    process_impl_file filename
470  else
471    process_intf_file filename
472;;
473
474(* Main function *)
475
476open Format
477
478let usage = "Usage: ocamlprof <options> <files>\noptions are:"
479
480let print_version () =
481  printf "ocamlprof, version %s@." Sys.ocaml_version;
482  exit 0;
483;;
484
485let print_version_num () =
486  printf "%s@." Sys.ocaml_version;
487  exit 0;
488;;
489
490let main () =
491  try
492    Warnings.parse_options false "a";
493    Arg.parse_expand [
494       "-f", Arg.String (fun s -> dumpfile := s),
495             "<file>     Use <file> as dump file (default ocamlprof.dump)";
496       "-F", Arg.String (fun s -> special_id := s),
497             "<s>        Insert string <s> with the counts";
498       "-impl", Arg.String process_impl_file,
499                "<file>  Process <file> as a .ml file";
500       "-instrument", Arg.Set instr_mode, "  (undocumented)";
501       "-intf", Arg.String process_intf_file,
502                "<file>  Process <file> as a .mli file";
503       "-m", Arg.String (fun s -> modes := s), "<flags>    (undocumented)";
504       "-version", Arg.Unit print_version,
505                   "     Print version and exit";
506       "-vnum", Arg.Unit print_version_num,
507                "        Print version number and exit";
508        "-args", Arg.Expand Arg.read_arg,
509                "<file> Read additional newline separated command line arguments \n\
510                \      from <file>";
511       "-args0", Arg.Expand Arg.read_arg0,
512               "<file> Read additional NUL separated command line arguments from \n\
513               \      <file>"
514    ] process_anon_file usage;
515    exit 0
516  with
517  | Profiler msg ->
518      fprintf Format.err_formatter "@[%s@]@." msg;
519      exit 2
520  | exn ->
521      Location.report_exception Format.err_formatter exn
522
523let _ = main ()
524