[ppx_trace] Add trace function and extension to support unhandled exceptions

Summary:
Add support to ppx_trace to rewrite `[%trace]` into a call to the
corresponding `Trace` function:

```lang=ocaml
val trace :
     ?call:(pf -> unit)
  -> ?retn:(pf -> 'a -> unit)
  -> ?rais:(pf -> exn -> Printexc.raw_backtrace -> unit)
  -> string
  -> string
  -> (unit -> 'a)
  -> 'a
(** [trace ~call ~retn ~rais mod_name fun_name k] either simply invokes
    [k ()], when not enabled, or else increases the indentation level and
    emits the [call] message, then invokes [k ()], then decreases the
    indentation level and either emits the [retn] or [rais] message,
    depending on whether [k ()] returned normally or exceptionally. *)
```

The main motivation over the existing `Trace.call` and `Trace.retn` is
that by packaging them together, unhandled exceptions can be treated
better.

Reviewed By: jvillard

Differential Revision: D23636200

fbshipit-source-id: f61c267fd

sledge/ppx_trace/ppx_trace.ml

@@ -11,22 +11,37 @@
     "debug" mode or not. To enable "debug" mode, pass
     [--cookie 'ppx_trace_enabled="1"'] (or with [true] instead or [1]).
 
-    It rewrites [\[%Trace.info f\]] to a call
-    [\[Trace.info mod_name fun_name f\]] where [mod_name] and [fun_name] are
-    the enclosing module and function names in the parsetree. This is only
-    done in debug mode, otherwise [\[%Trace.info f\]] is rewritten to [()].
-
-    Similarly, [\[%Trace.call\]] is rewritten to a call to [Trace.call] or
-    [()], and [\[%Trace.retn\]] to a call to [Trace.retn] or [Fun.id].
+    It rewrites [\[%trace\] ~call ~retn ~rais] to a call
+    [Trace.trace ~call ~retn ~rais mod_name fun_name] where [mod_name] and
+    [fun_name] are the enclosing module and function names in the parsetree.
+    This is only done in debug mode, otherwise
+    [\[%trace\] ~call ~retn ~rais] is rewritten to [(fun k -> k ())].
+
+    Similarly, [\[%Trace.info\]], [\[%Trace.infok\]], [\[%Trace.printf\]],
+    [\[%Trace.fprintf\]], [\[%Trace.kprintf\]], and [\[%Trace.call\]] are
+    rewritten to their analogues in the [Trace] module, or [()]; and
+    [\[%Trace.retn\]] is rewritten to a call to [Trace.retn] or
+    [(fun x -> x)].
 
     For example, this enables writing
 
     {[
-      [%Trace.call fun {pf} -> pf "%a" pp_arg_type arg]
-      ;
-      func arg
-      |>
-      [%Trace.retn fun {pf} -> pf "%a" pp_result_type]
+      let func arg =
+        [%trace]
+          ~call:(fun {pf} -> pf "%a" pp_arg_type arg)
+          ~retn:(fun {pf} -> pf "%a" pp_result_type)
+        @@ fun () -> func arg
+    ]}
+
+    or
+
+    {[
+      let func arg =
+        [%Trace.call fun {pf} -> pf "%a" pp_arg_type arg]
+        ;
+        func arg
+        |>
+        [%Trace.retn fun {pf} -> pf "%a" pp_result_type]
     ]}
 
     to trace calls to [func] in debug mode while completely compiling out
@@ -90,6 +105,11 @@ let vb_stack_with, vb_stack_top =
 (* (fun x -> x) *)
 let fun_id loc = pexp_fun ~loc Nolabel None (pvar ~loc "x") (evar ~loc "x")
 
+(* (fun k -> k ()) *)
+let fun_go loc =
+  pexp_fun ~loc Nolabel None (pvar ~loc "k")
+    (eapply ~loc (evar ~loc "k") [eunit ~loc])
+
 let mapper =
   object
     inherit Ast_traverse.map as super
@@ -106,6 +126,14 @@ let mapper =
         (Nolabel, mod_name) :: (Nolabel, fun_name) :: args
       in
       match exp.pexp_desc with
+      | Pexp_apply
+          ( { pexp_desc= Pexp_extension ({txt= "trace"; loc}, PStr [])
+            ; pexp_loc }
+          , args ) ->
+          if not !debug then fun_go pexp_loc
+          else
+            pexp_apply ~loc:exp.pexp_loc (evar ~loc "Trace.trace")
+              (append_here_args args)
       | Pexp_extension
           ( { txt=
                 ( "Trace.info" | "Trace.infok" | "Trace.printf"

sledge/ppx_trace/trace/trace.ml

@@ -173,16 +173,20 @@ let info mod_name fun_name fmt =
 let infok mod_name fun_name k =
   k {pf= (fun fmt -> info mod_name fun_name fmt)}
 
+let incf_ _mod_name fun_name fmt =
+  Format.fprintf fs "@\n@[<2>@[<hv 2>( %s:@ " fun_name ;
+  Format.kfprintf (fun fs -> Format.fprintf fs "@]") fs fmt
+
 let incf mod_name fun_name fmt =
-  if enabled mod_name fun_name then (
-    Format.fprintf fs "@\n@[<2>@[<hv 2>( %s:@ " fun_name ;
-    Format.kfprintf (fun fs -> Format.fprintf fs "@]") fs fmt )
+  if enabled mod_name fun_name then incf_ mod_name fun_name fmt
   else Format.ifprintf fs fmt
 
+let decf_ _mod_name fun_name fmt =
+  Format.fprintf fs "@]@\n@[<2>) %s:@ " fun_name ;
+  Format.kfprintf (fun fs -> Format.fprintf fs "@]") fs fmt
+
 let decf mod_name fun_name fmt =
-  if enabled mod_name fun_name then (
-    Format.fprintf fs "@]@\n@[<2>) %s:@ " fun_name ;
-    Format.kfprintf (fun fs -> Format.fprintf fs "@]") fs fmt )
+  if enabled mod_name fun_name then decf_ mod_name fun_name fmt
   else Format.ifprintf fs fmt
 
 let call mod_name fun_name k =
@@ -192,6 +196,33 @@ let retn mod_name fun_name k result =
   k {pf= (fun fmt -> decf mod_name fun_name fmt)} result ;
   result
 
+let trace :
+       ?call:(pf -> unit)
+    -> ?retn:(pf -> 'a -> unit)
+    -> ?rais:(pf -> exn -> Printexc.raw_backtrace -> unit)
+    -> string
+    -> string
+    -> (unit -> 'a)
+    -> 'a =
+ fun ?call ?retn ?rais mod_name fun_name k ->
+  let call = Option.value call ~default:(fun {pf} -> pf "") in
+  let retn = Option.value retn ~default:(fun {pf} _ -> pf "") in
+  let rais =
+    Option.value rais ~default:(fun {pf} exc _ ->
+        pf "%s" (Printexc.to_string exc) )
+  in
+  if not (enabled mod_name fun_name) then k ()
+  else (
+    call {pf= (fun fmt -> incf_ mod_name fun_name fmt)} ;
+    match k () with
+    | result ->
+        retn {pf= (fun fmt -> decf_ mod_name fun_name fmt)} result ;
+        result
+    | exception exc ->
+        let bt = Printexc.get_raw_backtrace () in
+        rais {pf= (fun fmt -> decf_ mod_name fun_name fmt)} exc bt ;
+        Printexc.raise_with_backtrace exc bt )
+
 let raisef ?margin exn fmt =
   let bt = Printexc.get_raw_backtrace () in
   let fs = Format.str_formatter in

sledge/ppx_trace/trace/trace.mli

@@ -45,6 +45,20 @@ val call : string -> string -> (pf -> 'a) -> 'a
 val retn : string -> string -> (pf -> 'a -> unit) -> 'a -> 'a
 (** Decrease indentation level and emit a message, if enabled. *)
 
+val trace :
+     ?call:(pf -> unit)
+  -> ?retn:(pf -> 'a -> unit)
+  -> ?rais:(pf -> exn -> Printexc.raw_backtrace -> unit)
+  -> string
+  -> string
+  -> (unit -> 'a)
+  -> 'a
+(** [trace ~call ~retn ~rais mod_name fun_name k] either simply invokes
+    [k ()], when not enabled, or else increases the indentation level and
+    emits the [call] message, then invokes [k ()], then decreases the
+    indentation level and either emits the [retn] or [rais] message,
+    depending on whether [k ()] returned normally or exceptionally. *)
+
 val flush : unit -> unit
 (** Flush the internal buffers. *)