(* camlp4r *)
(****************************************************************************)
(*                                                                          *)
(*                              Objective Caml                              *)
(*                                                                          *)
(*                            INRIA Rocquencourt                            *)
(*                                                                          *)
(*  Copyright   2006    Institut National de Recherche en Informatique et   *)
(*  en Automatique.  All rights reserved.  This file is distributed under   *)
(*  the terms of the GNU Library General Public License, with the special   *)
(*  exception on linking described in LICENSE at the top of the Objective   *)
(*  Caml source tree.                                                       *)
(*                                                                          *)
(****************************************************************************)

(* Authors:
 * - Nicolas Pouillard: initial version
 * - bluestorm : genericity improvement
 *)

(* 
   Input :
     let foo bar = __LOG__ blah 
   Output :
     let foo bar = __LOG_FUNC__ "foo" blah

  Compilation :
     camlfind ocamlc -pp camlp4rf -package camlp4 -c Camlp4GenericProfiler.ml
   Use :
     camlp4 Camlp4GenericProfiler.cmo test.ml
*)

open Camlp4;

module Id = struct
  value name    = "Camlp4Profiler";
  value version = "$Id: Camlp4Profiler.ml,v 1.1 2007/02/07 10:09:22 ertai Exp $";
end;

(* functor used to parametrize over the decorating function *)
module type Decorator = functor (Ast : Camlp4.Sig.Camlp4Ast) -> sig
  value with_func_name : string -> Ast.expr -> Ast.expr;
end;

module GenericMake (Decorator : Decorator) (AstFilters : Camlp4.Sig.AstFilters) = struct
  open AstFilters;
  open Ast;

  module Decorate = Decorator Ast;

  value decorate_binding decorate_fun = object
    inherit Ast.map as super;
    method binding =
      fun
      [ <:binding@_loc< $lid:id$ = $(<:expr< fun [ $_$ ] >> as e)$ >> ->
          <:binding< $lid:id$ = $decorate_fun id e$ >>
      | b -> super#binding b ];
  end#binding;

  value decorate decorate_fun = object (o)
    inherit Ast.map as super;
    method str_item =
      fun
      [ <:str_item@_loc< value $rec:r$ $b$ >> ->
          <:str_item< value $rec:r$ $decorate_binding decorate_fun b$ >>
      | st -> super#str_item st ];
    method expr =
      fun
      [ <:expr@_loc< let $rec:r$ $b$ in $e$ >> ->
          <:expr< let $rec:r$ $decorate_binding decorate_fun b$ in $o#expr e$ >>
      | <:expr@_loc< fun [ $_$ ] >> as e -> decorate_fun "<fun>" e
      | e -> super#expr e ];
  end;

  value rec decorate_fun id =
    let decorate = decorate decorate_fun in
    let decorate_expr = decorate#expr in
    let decorate_match_case = decorate#match_case in
    fun
    [ <:expr@_loc< fun $p$ -> $e$ >> ->
        <:expr< fun $p$ -> $decorate_fun id e$ >>
    | <:expr@_loc< fun [ $m$ ] >> ->
        Decorate.with_func_name id <:expr< fun [ $decorate_match_case m$ ] >>
    | e -> Decorate.with_func_name id (decorate_expr e) ];

  register_str_item_filter (decorate decorate_fun)#str_item;
end;

(* old Camlp4Profiler behavior *)
module ProfilingDecorator (Ast : Camlp4.Sig.Camlp4Ast) = struct
  open Ast;
  value with_func_name id e =
    let buf = Buffer.create 42 in
    let _loc = Ast.loc_of_expr e in
    let () = Format.bprintf buf "%s @@ %a@?" id Loc.dump _loc in
    let s = Buffer.contents buf in
    <:expr< let () = Camlp4prof.count $`str:s$ in $e$ >> ;
end;


(* Logging Behavior *)
module LoggingDecorator (Ast : Camlp4.Sig.Camlp4Ast) = struct
  value with_func_name func_name =
    object
      inherit Ast.map as super;
      method expr = fun
        [ <:expr@_loc< __LOG__ >> -> <:expr< __LOG_FUNC__  $`str:func_name$ >>
        | e -> super#expr e ];
    end#expr;
end;

let module Make = GenericMake LoggingDecorator in
let module M = Camlp4.Register.AstFilter Id Make in ();