Move universal maps from Fiber to Stdune (#892)

Signed-off-by: Jeremie Dimino <jeremie@dimino.org>

src/fiber/fiber.ml

@@ -1,50 +1,5 @@
 open Stdune
 
-module Eq = struct
-  type ('a, 'b) t = T : ('a, 'a) t
-
-  let cast (type a) (type b) (T : (a, b) t) (x : a) : b = x
-end
-
-module Var0 = struct
-  module Key = struct
-    type 'a t = ..
-  end
-
-  module type T = sig
-    type t
-    type 'a Key.t += T : t Key.t
-    val id : int
-  end
-
-  type 'a t = (module T with type t = 'a)
-
-  let next = ref 0
-
-  let create (type a) () =
-    let n = !next in
-    next := n + 1;
-    let module M = struct
-      type t = a
-      type 'a Key.t += T : t Key.t
-      let id = n
-    end in
-    (module M : T with type t = a)
-
-  let id (type a) (module M : T with type t = a) = M.id
-
-  let eq (type a) (type b)
-        (module A : T with type t = a)
-        (module B : T with type t = b) : (a, b) Eq.t =
-    match A.T with
-    | B.T -> Eq.T
-    | _ -> assert false
-end
-
-module Binding = struct
-  type t = T : 'a Var0.t * 'a -> t
-end
-
 module Execution_context : sig
   type t
 
@@ -66,14 +21,14 @@ module Execution_context : sig
     -> on_error:(exn -> unit)
     -> t
 
-  val vars : t -> Binding.t Int.Map.t
-  val set_vars : t -> Binding.t Int.Map.t -> t
+  val vars : t -> Univ_map.t
+  val set_vars : t -> Univ_map.t -> t
 end = struct
   type t =
     { on_error : exn -> unit (* This callback must never raise *)
     ; fibers   : int ref (* Number of fibers running in this execution
                             context *)
-    ; vars     : Binding.t Int.Map.t
+    ; vars     : Univ_map.t
     ; on_release : unit -> unit
     }
 
@@ -83,7 +38,7 @@ end = struct
   let create_initial () =
     { on_error   = reraise
     ; fibers     = ref 1
-    ; vars       = Int.Map.empty
+    ; vars       = Univ_map.empty
     ; on_release = ignore
     }
 
@@ -269,29 +224,13 @@ let parallel_iter l ~f ctx k =
         EC.forward_error ctx exn)
 
 module Var = struct
-  include Var0
+  include Univ_map.Key
 
-  let find ctx var =
-    match Int.Map.find (EC.vars ctx) (id var) with
-    | None -> None
-    | Some (Binding.T (var', v)) ->
-      let eq = eq var' var in
-      Some (Eq.cast eq v)
+  let get     var ctx k = k (Univ_map.find     (EC.vars ctx) var)
+  let get_exn var ctx k = k (Univ_map.find_exn (EC.vars ctx) var)
 
-  let find_exn ctx var =
-    match Int.Map.find (EC.vars ctx) (id var) with
-    | None -> failwith "Fiber.Var.find_exn"
-    | Some (Binding.T (var', v)) ->
-      let eq = eq var' var in
-      Eq.cast eq v
-
-  let get     var ctx k = k (find     ctx var)
-  let get_exn var ctx k = k (find_exn ctx var)
-
-  let set (type a) (var : a t) x fiber ctx k =
-    let (module M) = var in
-    let data = Binding.T (var, x) in
-    let ctx = EC.set_vars ctx (Int.Map.add (EC.vars ctx) M.id data) in
+  let set var x fiber ctx k =
+    let ctx = EC.set_vars ctx (Univ_map.add (EC.vars ctx) var x) in
     fiber ctx k
 end
 

src/stdune/stdune.ml

@@ -22,6 +22,7 @@ module Sexp       = Sexp
 module Path       = Path
 module Fmt        = Fmt
 module Interned   = Interned
+module Univ_map   = Univ_map
 
 external reraise : exn -> _ = "%reraise"
 

src/stdune/univ_map.ml

@@ -0,0 +1,73 @@
+module Eq = struct
+  type ('a, 'b) t = T : ('a, 'a) t
+
+  let cast (type a) (type b) (T : (a, b) t) (x : a) : b = x
+end
+
+module Key = struct
+  module Witness = struct
+    type 'a t = ..
+  end
+
+  module type T = sig
+    type t
+    type 'a Witness.t += T : t Witness.t
+    val id : int
+  end
+
+  type 'a t = (module T with type t = 'a)
+
+  let next = ref 0
+
+  let create (type a) () =
+    let n = !next in
+    next := n + 1;
+    let module M = struct
+      type t = a
+      type 'a Witness.t += T : t Witness.t
+      let id = n
+    end in
+    (module M : T with type t = a)
+
+  let id (type a) (module M : T with type t = a) = M.id
+
+  let eq (type a) (type b)
+        (module A : T with type t = a)
+        (module B : T with type t = b) : (a, b) Eq.t =
+    match A.T with
+    | B.T -> Eq.T
+    | _ -> assert false
+end
+
+module Binding = struct
+  type t = T : 'a Key.t * 'a -> t
+end
+
+type t = Binding.t Int.Map.t
+
+let empty = Int.Map.empty
+let is_empty = Int.Map.is_empty
+
+let add (type a) t (key : a Key.t) x =
+  let (module M) = key in
+  let data = Binding.T (key, x) in
+  Int.Map.add t M.id data
+
+let mem t key = Int.Map.mem t (Key.id key)
+
+let remove t key = Int.Map.remove t (Key.id key)
+
+let find t key =
+  match Int.Map.find t (Key.id key) with
+  | None -> None
+  | Some (Binding.T (key', v)) ->
+    let eq = Key.eq key' key in
+    Some (Eq.cast eq v)
+
+let find_exn t key =
+  match Int.Map.find t (Key.id key) with
+  | None -> failwith "Univ_map.find_exn"
+  | Some (Binding.T (key', v)) ->
+    let eq = Key.eq key' key in
+    Eq.cast eq v
+

src/stdune/univ_map.mli

@@ -0,0 +1,19 @@
+(** Universal maps *)
+
+(** A universal map is a map that can store values for arbitrary
+    keys. It is the the key that conveys the type of the data
+    associated to it. *)
+type t
+
+module Key : sig
+  type 'a t
+  val create : unit -> 'a t
+end
+
+val empty    : t
+val is_empty : t -> bool
+val mem      : t -> 'a Key.t -> bool
+val add      : t -> 'a Key.t -> 'a -> t
+val remove   : t -> 'a Key.t -> t
+val find     : t -> 'a Key.t -> 'a option
+val find_exn : t -> 'a Key.t -> 'a