From 8ce25c28150e2acf7a49e8e4e13725034ee05826 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 22 Dec 2023 21:24:53 -0500
Subject: [PATCH] add ppx_trace

---
 dune-project         | 12 +++++++
 ppx_trace.opam       | 35 +++++++++++++++++++
 src/ppx/dune         |  8 +++++
 src/ppx/ppx_trace.ml | 81 ++++++++++++++++++++++++++++++++++++++++++++
 test/dune            |  7 ++++
 5 files changed, 143 insertions(+)
 create mode 100644 ppx_trace.opam
 create mode 100644 src/ppx/dune
 create mode 100644 src/ppx/ppx_trace.ml

diff --git a/dune-project b/dune-project
index e69091a..e3505c9 100644
--- a/dune-project
+++ b/dune-project
@@ -21,6 +21,18 @@
  (tags
   (trace tracing observability profiling)))
 
+(package
+ (name ppx_trace)
+ (synopsis "ppx-based instrumentation for trace")
+ (depends
+   (ocaml (>= 4.08))
+   ppxlib
+   (trace (= :version))
+   (trace-tef (and (= :version) :with-test))
+   dune)
+ (tags
+  (trace tracing observability profiling ppx)))
+
 (package
  (name trace-tef)
  (synopsis "A simple backend for trace, emitting Catapult/TEF JSON into a file")
diff --git a/ppx_trace.opam b/ppx_trace.opam
new file mode 100644
index 0000000..14992f3
--- /dev/null
+++ b/ppx_trace.opam
@@ -0,0 +1,35 @@
+# This file is generated by dune, edit dune-project instead
+opam-version: "2.0"
+version: "0.5"
+synopsis: "ppx-based instrumentation for trace"
+maintainer: ["Simon Cruanes"]
+authors: ["Simon Cruanes"]
+license: "MIT"
+tags: ["trace" "tracing" "observability" "profiling" "ppx"]
+homepage: "https://github.com/c-cube/ocaml-trace"
+bug-reports: "https://github.com/c-cube/ocaml-trace/issues"
+depends: [
+  "ocaml" {>= "4.08"}
+  "ppxlib"
+  "trace" {= version}
+  "trace-tef" {= version & with-test}
+  "dune" {>= "2.9"}
+  "odoc" {with-doc}
+]
+build: [
+  ["dune" "subst"] {dev}
+  [
+    "dune"
+    "build"
+    "-p"
+    name
+    "-j"
+    jobs
+    "--promote-install-files=false"
+    "@install"
+    "@runtest" {with-test}
+    "@doc" {with-doc}
+  ]
+  ["dune" "install" "-p" name "--create-install-files" name]
+]
+dev-repo: "git+https://github.com/c-cube/ocaml-trace.git"
diff --git a/src/ppx/dune b/src/ppx/dune
new file mode 100644
index 0000000..451c5fc
--- /dev/null
+++ b/src/ppx/dune
@@ -0,0 +1,8 @@
+
+(library
+ (name ppx_trace)
+ (public_name ppx_trace)
+ (kind ppx_rewriter)
+ (preprocess (pps ppxlib.metaquot))
+ (ppx_runtime_libraries trace.core)
+ (libraries ppxlib))
diff --git a/src/ppx/ppx_trace.ml b/src/ppx/ppx_trace.ml
new file mode 100644
index 0000000..9601005
--- /dev/null
+++ b/src/ppx/ppx_trace.ml
@@ -0,0 +1,81 @@
+open Ppxlib
+
+let location_errorf ~loc fmt =
+  Format.kasprintf
+    (fun err ->
+      raise (Ocaml_common.Location.Error (Ocaml_common.Location.error ~loc err)))
+    fmt
+
+(** {2 let expression} *)
+
+let expand_let ~ctxt (name : string) body =
+  let loc = Expansion_context.Extension.extension_point_loc ctxt in
+  Ast_builder.Default.(
+    [%expr
+      let _trace_span =
+        Trace_core.enter_span ~__FILE__ ~__LINE__ [%e estring ~loc name]
+      in
+      try
+        let res = [%e body] in
+        Trace_core.exit_span _trace_span;
+        res
+      with exn ->
+        Trace_core.exit_span _trace_span;
+        raise exn])
+
+let extension_let =
+  Extension.V3.declare "trace" Extension.Context.expression
+    (let open! Ast_pattern in
+    single_expr_payload
+      (pexp_let nonrecursive
+         (value_binding
+            ~pat:(ppat_construct (lident (string "()")) none)
+            ~expr:(estring __)
+         ^:: nil)
+         __))
+    expand_let
+
+let rule_let = Ppxlib.Context_free.Rule.extension extension_let
+
+(** {2 Toplevel binding} *)
+
+let expand_top_let ~ctxt rec_flag (vbs : _ list) =
+  let loc = Expansion_context.Extension.extension_point_loc ctxt in
+  Ast_builder.Default.(
+    (* go in functions, and add tracing around the body *)
+    let rec push_into_fun (e : expression) : expression =
+      match e.pexp_desc with
+      | Pexp_fun (lbl, lbl_expr, pat, body) ->
+        pexp_fun ~loc:e.pexp_loc lbl lbl_expr pat @@ push_into_fun body
+      | _ ->
+        [%expr
+          let _trace_span =
+            Trace_core.enter_span ~__FILE__ ~__LINE__ __FUNCTION__
+          in
+          try
+            let res = [%e e] in
+            Trace_core.exit_span _trace_span;
+            res
+          with exn ->
+            Trace_core.exit_span _trace_span;
+            raise exn]
+    in
+
+    let tr_vb (vb : value_binding) : value_binding =
+      let expr = push_into_fun vb.pvb_expr in
+      { vb with pvb_expr = expr }
+    in
+
+    let vbs = List.map tr_vb vbs in
+    pstr_value ~loc rec_flag vbs)
+
+let extension_top_let =
+  Extension.V3.declare "trace" Extension.Context.structure_item
+    (let open! Ast_pattern in
+    pstr (pstr_value __ __ ^:: nil))
+    expand_top_let
+
+let rule_top_let = Ppxlib.Context_free.Rule.extension extension_top_let
+
+let () =
+  Driver.register_transformation ~rules:[ rule_let; rule_top_let ] "ppx_trace"
diff --git a/test/dune b/test/dune
index fdefb87..e52a276 100644
--- a/test/dune
+++ b/test/dune
@@ -2,5 +2,12 @@
 (test
   (name t1)
   (package trace-tef)
+  (modules t1)
   (libraries trace trace-tef))
 
+(test
+  (name t2)
+  (package ppx_trace)
+  (modules t2)
+  (preprocess (pps ppx_trace))
+  (libraries trace-tef))