Merge pull request #182 from akabe/feature/trace

Migrate trace directive from the ocaml compiler

jupyter/src/repl/compat.cppo.ml

@@ -74,3 +74,15 @@ let reset_fatal_warnings () =
 #else
   ()
 #endif
+
+#if OCAML_VERSION < (4,14,0)
+let types_get_desc t = t.Types.desc
+#else
+let types_get_desc = Types.get_desc
+#endif
+
+#if OCAML_VERSION < (4,13,0)
+let section_trace = "Tracing"
+#else
+let section_trace = Topdirs.section_trace
+#endif

jupyter/src/repl/dir_trace.ml

@@ -0,0 +1,129 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                                 OCaml                                  *)
+(*                                                                        *)
+(*             Xavier Leroy, projet Cristal, INRIA Rocquencourt           *)
+(*                                                                        *)
+(*   Copyright 1996 Institut National de Recherche en Informatique et     *)
+(*     en Automatique.                                                    *)
+(*                                                                        *)
+(*   All rights reserved.  This file is distributed under the terms of    *)
+(*   the GNU Lesser General Public License version 2.1, with the          *)
+(*   special exception on linking described in the file LICENSE.          *)
+(*                                                                        *)
+(**************************************************************************)
+
+
+(** The trace for OCaml 4.13.0 or above.
+
+    [#trace] directive implemented at
+
+    - [toplevel/topdirs.ml] on 4.12.0-, and
+    - [toplevel/byte/trace.ml] and [toplevel/byte/topmain.ml] on 4.13.0+.
+
+    This file is a part of [toplevel/byte/topmain.ml], migrated for
+    ocaml-jupyter. *)
+
+open Types
+open Trace
+open Toploop
+
+external current_environment: unit -> Obj.t = "caml_get_current_environment"
+
+let tracing_function_ptr =
+  get_code_pointer
+    (Obj.repr (fun arg -> Trace.print_trace (current_environment()) arg))
+
+let dir_trace ppf lid =
+  match Env.find_value_by_name lid !toplevel_env with
+  | (path, desc) -> begin
+      (* Check if this is a primitive *)
+      match desc.Types.val_kind with
+      | Types.Val_prim _ ->
+        Format.fprintf ppf
+          "%a is an external function and cannot be traced.@."
+          Printtyp.longident lid
+      | _ ->
+        let clos = Toploop.eval_value_path !toplevel_env path in
+        (* Nothing to do if it's not a closure *)
+        if Obj.is_block clos
+        && (Obj.tag clos = Obj.closure_tag || Obj.tag clos = Obj.infix_tag)
+        && (match
+              Compat.types_get_desc
+                (Ctype.expand_head !toplevel_env desc.val_type)
+            with Tarrow _ -> true | _ -> false)
+        then begin
+          match is_traced clos with
+          | Some opath ->
+            Format.fprintf ppf "%a is already traced (under the name %a).@."
+              Printtyp.path path
+              Printtyp.path opath
+          | None ->
+            (* Instrument the old closure *)
+            traced_functions :=
+              { path = path;
+                closure = clos;
+                actual_code = get_code_pointer clos;
+                instrumented_fun =
+                  instrument_closure
+                    !toplevel_env lid ppf desc.val_type }
+              :: !traced_functions;
+            (* Redirect the code field of the closure to point
+               to the instrumentation function *)
+            set_code_pointer clos tracing_function_ptr;
+            Format.fprintf ppf "%a is now traced.@." Printtyp.longident lid
+        end else
+          Format.fprintf ppf "%a is not a function.@." Printtyp.longident lid
+    end
+  | exception Not_found ->
+    Format.fprintf ppf "Unbound value %a.@." Printtyp.longident lid
+
+let dir_untrace ppf lid =
+  match Env.find_value_by_name lid !toplevel_env with
+  | (path, _desc) ->
+    let rec remove = function
+      | [] ->
+        Format.fprintf ppf "%a was not traced.@." Printtyp.longident lid;
+        []
+      | f :: rem ->
+        if Path.same f.path path then begin
+          set_code_pointer f.closure f.actual_code;
+          Format.fprintf ppf "%a is no longer traced.@."
+            Printtyp.longident lid;
+          rem
+        end else f :: remove rem in
+    traced_functions := remove !traced_functions
+  | exception Not_found ->
+    Format.fprintf ppf "Unbound value %a.@." Printtyp.longident lid
+
+let dir_untrace_all ppf () =
+  List.iter
+    (fun f ->
+       set_code_pointer f.closure f.actual_code;
+       Format.fprintf ppf "%a is no longer traced.@." Printtyp.path f.path)
+    !traced_functions;
+  traced_functions := []
+
+let add_directives ppf =
+  let _ = add_directive "trace"
+      (Directive_ident (dir_trace ppf))
+      {
+        section = Compat.section_trace;
+        doc = "All calls to the function \
+               named function-name will be traced.";
+      } in
+
+  let _ = add_directive "untrace"
+      (Directive_ident (dir_untrace ppf))
+      {
+        section = Compat.section_trace;
+        doc = "Stop tracing the given function.";
+      } in
+
+  let _ = add_directive "untrace_all"
+      (Directive_none (dir_untrace_all ppf))
+      {
+        section = Compat.section_trace;
+        doc = "Stop tracing all functions traced so far.";
+      } in
+  ()

jupyter/src/repl/dune

@@ -8,6 +8,7 @@
               Lwt_async_rewrite
               Caml_args
               Error
+              Dir_trace
               Compat)
  (flags       ((:include %{workspace_root}/config/ocaml_flags.sexp)))
  (preprocess  (pps lwt_ppx))

jupyter/src/repl/evaluation.cppo.ml

@@ -26,7 +26,7 @@ open Format
 open Jupyter
 
 let buffer = Buffer.create 256
-let ppf = formatter_of_buffer buffer
+let buffer_ppf = formatter_of_buffer buffer
 
 (** {2 Initialization} *)
 
@@ -45,7 +45,7 @@ let init_toploop () =
   try
     Toploop.initialize_toplevel_env ()
   with Env.Error _ | Typetexp.Error _ as exn ->
-    Location.report_exception ppf exn ;
+    Location.report_exception buffer_ppf exn ;
     exit 2
 
 let init ?(preinit = ignore) ?init_file () =
@@ -57,6 +57,7 @@ let init ?(preinit = ignore) ?init_file () =
   Compenv.readenv ppf Compenv.Before_link ;
   if not (Caml_args.prepare ppf) then exit 2 ;
   init_toploop () ;
+  Dir_trace.add_directives buffer_ppf ;
   preinit () ;
   begin match init_file with
     | None -> ()
@@ -88,7 +89,7 @@ let eval_phrase ~filename phrase =
   let phrase' = Compat.preprocess_phrase ~filename phrase in (* apply PPX *)
   let phrase' = Lwt_async_rewrite.rewrite phrase' in
   Env.reset_cache_toplevel () ;
-  let is_ok = Toploop.execute_phrase true ppf phrase' in
+  let is_ok = Toploop.execute_phrase true buffer_ppf phrase' in
   let message = Buffer.contents buffer in
   Buffer.clear buffer ;
   (is_ok, message)

test/repl/test_evaluation.ml

@@ -67,6 +67,19 @@ let test__directive ctxt =
   assert_equal ~ctxt ~printer:[%show: status] SHELL_OK status ;
   assert_equal ~ctxt ~printer:[%show: reply list] expected actual
 
+(* Implementation of [#trace] directive changes after OCaml 4.13.0. *)
+let test__trace_directive ctxt =
+  let status, actual = eval "let f x = x ;; #trace f ;; f 10 ;;" in
+  let expected = [
+    iopub_success ~count:0 "val f : 'a -> 'a = <fun>\n";
+    iopub_success ~count:0 "f is now traced.\n";
+    iopub_success ~count:0 "f <-- <poly>\n\
+                            f --> <poly>\n\
+                            - : int = 10\n";
+  ] in
+  assert_equal ~ctxt ~printer:[%show: status] SHELL_OK status ;
+  assert_equal ~ctxt ~printer:[%show: reply list] expected actual
+
 let test__external_command ctxt =
   let status, actual = eval "Sys.command \"ls -l >/dev/null 2>/dev/null\"" in
   let expected = [iopub_success ~count:0 "- : int = 0\n"] in
@@ -231,6 +244,7 @@ let suite =
       "simple_phrase" >:: test__simple_phrase;
       "multiple_phrases" >:: test__multiple_phrases;
       "directive" >:: test__directive;
+      "#trace directive" >:: test__trace_directive;
       "external_command" >:: test__external_command;
       "syntax_error" >:: test__syntax_error;
       "unbound_value" >:: test__unbound_value;