From 3c2f804716fed857b546f443cd6b6e2169acb963 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sat, 24 Feb 2024 11:33:48 -0500
Subject: [PATCH] add hmap as a depopt (#28)

if present, we use `Hmap.t` as the meta map for manual spans
---
 .github/workflows/main.yml  |  4 ++
 dune-project                |  1 +
 src/core/dune               |  3 ++
 src/core/meta_map.hmap.ml   |  3 ++
 src/core/meta_map.ml        | 83 ---------------------------------
 src/core/meta_map.mli       | 37 ---------------
 src/core/meta_map.ourown.ml | 91 +++++++++++++++++++++++++++++++++++++
 src/core/types.ml           |  5 +-
 src/fuchsia/fcollector.ml   |  2 +-
 src/tef/trace_tef.ml        |  6 +--
 trace.opam                  |  1 +
 11 files changed, 111 insertions(+), 125 deletions(-)
 create mode 100644 src/core/meta_map.hmap.ml
 delete mode 100644 src/core/meta_map.ml
 delete mode 100644 src/core/meta_map.mli
 create mode 100644 src/core/meta_map.ourown.ml

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 30efe46..d7fc0c5 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -42,3 +42,7 @@ jobs:
       - run: opam install trace
       - run: opam exec -- dune runtest -p trace-tef,trace-fuchsia
 
+      # with depopts
+      - run: opam install hmap
+      - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
+
diff --git a/dune-project b/dune-project
index 083ff54..74d4fd4 100644
--- a/dune-project
+++ b/dune-project
@@ -19,6 +19,7 @@
    (ocaml (>= 4.08))
    dune)
  (depopts
+   hmap
    (mtime (>= 2.0)))
  (tags
   (trace tracing observability profiling)))
diff --git a/src/core/dune b/src/core/dune
index 2240246..b872cef 100644
--- a/src/core/dune
+++ b/src/core/dune
@@ -2,6 +2,9 @@
 (library
  (name trace_core)
  (public_name trace.core)
+ (libraries (select meta_map.ml from
+                    (hmap -> meta_map.hmap.ml)
+                    (-> meta_map.ourown.ml)))
  (synopsis "Lightweight stub for tracing")
 )
 
diff --git a/src/core/meta_map.hmap.ml b/src/core/meta_map.hmap.ml
new file mode 100644
index 0000000..0aee4db
--- /dev/null
+++ b/src/core/meta_map.hmap.ml
@@ -0,0 +1,3 @@
+include Hmap
+
+let find_exn = get
diff --git a/src/core/meta_map.ml b/src/core/meta_map.ml
deleted file mode 100644
index b4564c0..0000000
--- a/src/core/meta_map.ml
+++ /dev/null
@@ -1,83 +0,0 @@
-module type KEY_IMPL = sig
-  type t
-
-  exception Store of t
-
-  val id : int
-end
-
-module Key = struct
-  type 'a t = (module KEY_IMPL with type t = 'a)
-
-  let _n = ref 0
-
-  let create (type k) () =
-    incr _n;
-    let id = !_n in
-    let module K = struct
-      type t = k
-
-      let id = id
-
-      exception Store of k
-    end in
-    (module K : KEY_IMPL with type t = k)
-
-  let id (type k) (module K : KEY_IMPL with type t = k) = K.id
-
-  let equal : type a b. a t -> b t -> bool =
-   fun (module K1) (module K2) -> K1.id = K2.id
-end
-
-type pair = Pair : 'a Key.t * 'a -> pair
-type exn_pair = E_pair : 'a Key.t * exn -> exn_pair
-
-let pair_of_e_pair (E_pair (k, e)) =
-  let module K = (val k) in
-  match e with
-  | K.Store v -> Pair (k, v)
-  | _ -> assert false
-
-module M = Map.Make (struct
-  type t = int
-
-  let compare (i : int) j = Stdlib.compare i j
-end)
-
-type t = exn_pair M.t
-
-let empty = M.empty
-let mem k t = M.mem (Key.id k) t
-
-let find_exn (type a) (k : a Key.t) t : a =
-  let module K = (val k) in
-  let (E_pair (_, e)) = M.find K.id t in
-  match e with
-  | K.Store v -> v
-  | _ -> assert false
-
-let find k t = try Some (find_exn k t) with Not_found -> None
-
-let add_e_pair_ p t =
-  let (E_pair ((module K), _)) = p in
-  M.add K.id p t
-
-let add_pair_ p t =
-  let (Pair (((module K) as k), v)) = p in
-  let p = E_pair (k, K.Store v) in
-  M.add K.id p t
-
-let add (type a) (k : a Key.t) v t =
-  let module K = (val k) in
-  add_e_pair_ (E_pair (k, K.Store v)) t
-
-let remove (type a) (k : a Key.t) t =
-  let module K = (val k) in
-  M.remove K.id t
-
-let cardinal t = M.cardinal t
-let length = cardinal
-let iter f t = M.iter (fun _ p -> f (pair_of_e_pair p)) t
-let to_list t = M.fold (fun _ p l -> pair_of_e_pair p :: l) t []
-let add_list t l = List.fold_right add_pair_ l t
-let of_list l = add_list empty l
diff --git a/src/core/meta_map.mli b/src/core/meta_map.mli
deleted file mode 100644
index 94f9317..0000000
--- a/src/core/meta_map.mli
+++ /dev/null
@@ -1,37 +0,0 @@
-(** Associative containers with Heterogeneous Values *)
-
-(** Keys with a type witness. *)
-module Key : sig
-  type 'a t
-  (** A key of type ['a t] is used to access the portion of the
-      map or table that associates keys of type ['a] to values. *)
-
-  val create : unit -> 'a t
-  (** Make a new key. This is generative, so calling [create ()] twice with the
-      same return type will produce incompatible keys that cannot see each
-      other's bindings. *)
-
-  val equal : 'a t -> 'a t -> bool
-  (** Compare two keys that have compatible types. *)
-end
-
-type pair = Pair : 'a Key.t * 'a -> pair
-
-type t
-(** Immutable map from {!Key.t} to values *)
-
-val empty : t
-val mem : _ Key.t -> t -> bool
-val add : 'a Key.t -> 'a -> t -> t
-val remove : _ Key.t -> t -> t
-val length : t -> int
-val cardinal : t -> int
-val find : 'a Key.t -> t -> 'a option
-
-val find_exn : 'a Key.t -> t -> 'a
-(** @raise Not_found if the key is not in the table. *)
-
-val iter : (pair -> unit) -> t -> unit
-val add_list : t -> pair list -> t
-val of_list : pair list -> t
-val to_list : t -> pair list
diff --git a/src/core/meta_map.ourown.ml b/src/core/meta_map.ourown.ml
new file mode 100644
index 0000000..dadadd1
--- /dev/null
+++ b/src/core/meta_map.ourown.ml
@@ -0,0 +1,91 @@
+module type KEY_IMPL = sig
+  type t
+
+  exception Store of t
+
+  val id : int
+end
+
+module Key = struct
+  type 'a t = (module KEY_IMPL with type t = 'a)
+
+  let _n = ref 0
+
+  let create (type k) () =
+    incr _n;
+    let id = !_n in
+    let module K = struct
+      type t = k
+
+      let id = id
+
+      exception Store of k
+    end in
+    (module K : KEY_IMPL with type t = k)
+
+  let[@inline] id (type k) (module K : KEY_IMPL with type t = k) = K.id
+
+  let equal : type a b. a t -> b t -> bool =
+   fun (module K1) (module K2) -> K1.id = K2.id
+end
+
+type 'a key = 'a Key.t
+type binding = B : 'a Key.t * 'a -> binding
+
+open struct
+  type exn_pair = E_pair : 'a Key.t * exn -> exn_pair
+
+  let pair_of_e_pair (E_pair (k, e)) =
+    let module K = (val k) in
+    match e with
+    | K.Store v -> B (k, v)
+    | _ -> assert false
+end
+
+module M = Map.Make (struct
+  type t = int
+
+  let compare (i : int) j = Stdlib.compare i j
+end)
+
+type t = { m: exn_pair M.t } [@@unboxed]
+
+let empty : t = { m = M.empty }
+let[@inline] mem k (self : t) = M.mem (Key.id k) self.m
+
+let find_exn (type a) (k : a Key.t) (self : t) : a =
+  let module K = (val k) in
+  let (E_pair (_, e)) = M.find K.id self.m in
+  match e with
+  | K.Store v -> v
+  | _ -> assert false
+
+let find k (self : t) = try Some (find_exn k self) with Not_found -> None
+
+open struct
+  let add_e_pair_ p self =
+    let (E_pair ((module K), _)) = p in
+    { m = M.add K.id p self.m }
+
+  let add_pair_ p (self : t) : t =
+    let (B (((module K) as k), v)) = p in
+    let p = E_pair (k, K.Store v) in
+    { m = M.add K.id p self.m }
+end
+
+let add (type a) (k : a Key.t) v (self : t) : t =
+  let module K = (val k) in
+  add_e_pair_ (E_pair (k, K.Store v)) self
+
+let remove (type a) (k : a Key.t) (self : t) : t =
+  let module K = (val k) in
+  { m = M.remove K.id self.m }
+
+let[@inline] cardinal (self : t) = M.cardinal self.m
+let length = cardinal
+let iter f (self : t) = M.iter (fun _ p -> f (pair_of_e_pair p)) self.m
+
+let to_list (self : t) : binding list =
+  M.fold (fun _ p l -> pair_of_e_pair p :: l) self.m []
+
+let add_list (self : t) l = List.fold_right add_pair_ l self
diff --git a/src/core/types.ml b/src/core/types.ml
index 9472385..4e7f6f2 100644
--- a/src/core/types.ml
+++ b/src/core/types.ml
@@ -17,6 +17,9 @@ type explicit_span = {
   span: span;
       (** Identifier for this span. Several explicit spans might share the same
       identifier since we can differentiate between them via [meta]. *)
-  mutable meta: Meta_map.t;  (** Metadata for this span (and its context) *)
+  mutable meta: Meta_map.t;
+      (** Metadata for this span (and its context). This can be used by collectors to
+        carry collector-specific information from the beginning
+        of the span, to the end of the span. *)
 }
 (** Explicit span, with collector-specific metadata *)
diff --git a/src/fuchsia/fcollector.ml b/src/fuchsia/fcollector.ml
index 0f5b486..a6963e7 100644
--- a/src/fuchsia/fcollector.ml
+++ b/src/fuchsia/fcollector.ml
@@ -110,7 +110,7 @@ type async_span_info = {
   mutable data: (string * user_data) list;
 }
 
-let key_async_data : async_span_info Meta_map.Key.t = Meta_map.Key.create ()
+let key_async_data : async_span_info Meta_map.key = Meta_map.Key.create ()
 
 open struct
   let state_id_ = A.make 0
diff --git a/src/tef/trace_tef.ml b/src/tef/trace_tef.ml
index bd36ecb..bbbee69 100644
--- a/src/tef/trace_tef.ml
+++ b/src/tef/trace_tef.ml
@@ -92,12 +92,12 @@ type span_info = {
 }
 
 (** key used to carry a unique "id" for all spans in an async context *)
-let key_async_id : int Meta_map.Key.t = Meta_map.Key.create ()
+let key_async_id : int Meta_map.key = Meta_map.Key.create ()
 
-let key_async_data : (string * [ `Sync | `Async ] option) Meta_map.Key.t =
+let key_async_data : (string * [ `Sync | `Async ] option) Meta_map.key =
   Meta_map.Key.create ()
 
-let key_data : (string * user_data) list ref Meta_map.Key.t =
+let key_data : (string * user_data) list ref Meta_map.key =
   Meta_map.Key.create ()
 
 (** Writer: knows how to write entries to a file in TEF format *)
diff --git a/trace.opam b/trace.opam
index bd0315a..23ebddf 100644
--- a/trace.opam
+++ b/trace.opam
@@ -15,6 +15,7 @@ depends: [
   "odoc" {with-doc}
 ]
 depopts: [
+  "hmap"
   "mtime" {>= "2.0"}
 ]
 build: [