From e53986d7b46057027d5d9c10c79d452bfd136f61 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 29 Aug 2024 15:55:31 -0400
Subject: [PATCH] feat: optional dep on hmap, for inheritable FLS data

---
 .github/workflows/main.yml |  2 +-
 dune-project               |  1 +
 moonpool.opam              |  1 +
 src/fib/dune               |  6 ++++-
 src/fib/fiber.ml           |  8 +++++--
 src/fib/fiber.mli          |  6 +++++
 src/fib/hmap_fls.dummy.ml  |  7 ++++++
 src/fib/hmap_fls.real.ml   | 47 ++++++++++++++++++++++++++++++++++++++
 8 files changed, 74 insertions(+), 4 deletions(-)
 create mode 100644 src/fib/hmap_fls.dummy.ml
 create mode 100644 src/fib/hmap_fls.real.ml

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 7a3669cc..e5bb01a8 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -37,7 +37,7 @@ jobs:
       - run: opam exec -- dune build @install
 
       # install some depopts
-      - run: opam install thread-local-storage trace
+      - run: opam install thread-local-storage trace hmap
         if: matrix.ocaml-compiler == '5.2'
 
       - run: opam exec -- dune build --profile=release --force @install @runtest
diff --git a/dune-project b/dune-project
index eb20a706..21e85084 100644
--- a/dune-project
+++ b/dune-project
@@ -29,6 +29,7 @@
     (>= 1.9.0)
     :with-test)))
  (depopts
+   hmap
    (trace (>= 0.6)))
  (tags
   (thread pool domain futures fork-join)))
diff --git a/moonpool.opam b/moonpool.opam
index 101d0881..e7d7198d 100644
--- a/moonpool.opam
+++ b/moonpool.opam
@@ -20,6 +20,7 @@ depends: [
   "mdx" {>= "1.9.0" & with-test}
 ]
 depopts: [
+  "hmap"
   "trace" {>= "0.6"}
 ]
 build: [
diff --git a/src/fib/dune b/src/fib/dune
index 4c6594f8..64d4c080 100644
--- a/src/fib/dune
+++ b/src/fib/dune
@@ -2,7 +2,11 @@
  (name moonpool_fib)
  (public_name moonpool.fib)
  (synopsis "Fibers and structured concurrency for Moonpool")
- (libraries moonpool picos)
+ (libraries moonpool picos
+            (select hmap_fls.ml from
+                    (hmap -> hmap_fls.real.ml)
+                    (-> hmap_fls.dummy.ml)
+                    ))
  (enabled_if
   (>= %{ocaml_version} 5.0))
  (flags :standard -open Moonpool_private -open Moonpool)
diff --git a/src/fib/fiber.ml b/src/fib/fiber.ml
index 6fe034c0..b5daff35 100644
--- a/src/fib/fiber.ml
+++ b/src/fib/fiber.ml
@@ -257,8 +257,10 @@ let spawn_ ~parent ~runner (f : unit -> 'a) : 'a t =
   let comp = Picos.Computation.create () in
   let pfiber = PF.create ~forbid:false comp in
 
-  (* inherit FLS from parent, if present *)
-  Option.iter (fun (p : _ t) -> PF.copy_fls p.pfiber pfiber) parent;
+  (* copy local hmap from parent, if present *)
+  Option.iter
+    (fun (p : _ t) -> Hmap_fls.Private_hmap_fls_.copy_fls p.pfiber pfiber)
+    parent;
 
   (match parent with
   | Some p when is_closed p -> failwith "spawn: nursery is closed"
@@ -328,3 +330,5 @@ let yield () : unit =
     check_if_cancelled_ self;
     PF.yield ();
     check_if_cancelled_ self
+
+include Hmap_fls
diff --git a/src/fib/fiber.mli b/src/fib/fiber.mli
index 0da300e7..f83dd2ad 100644
--- a/src/fib/fiber.mli
+++ b/src/fib/fiber.mli
@@ -147,3 +147,9 @@ val spawn_ignore : ?protect:bool -> (unit -> _) -> unit
 (** [spawn_ignore f] is [ignore (spawn f)].
   The fiber will still affect termination of the parent, ie. the
   parent will exit only after this new fiber exits. *)
+
+(** {2 Local [Hmap.t]} *)
+
+include module type of struct
+  include Hmap_fls
+end
diff --git a/src/fib/hmap_fls.dummy.ml b/src/fib/hmap_fls.dummy.ml
new file mode 100644
index 00000000..af953380
--- /dev/null
+++ b/src/fib/hmap_fls.dummy.ml
@@ -0,0 +1,7 @@
+(**/**)
+
+module Private_hmap_fls_ = struct
+  let copy_fls _ _ = ()
+end
+
+(**/**)
diff --git a/src/fib/hmap_fls.real.ml b/src/fib/hmap_fls.real.ml
new file mode 100644
index 00000000..6ed269ac
--- /dev/null
+++ b/src/fib/hmap_fls.real.ml
@@ -0,0 +1,47 @@
+open Moonpool.Private.Types_
+
+open struct
+  module FLS = Picos.Fiber.FLS
+end
+
+(** A local hmap, inherited in children fibers *)
+let k_local_hmap : Hmap.t FLS.t = FLS.create ()
+
+(** Access the local [hmap], or an empty one if not set *)
+let[@inline] get_local_hmap () : Hmap.t =
+  let fiber = get_current_fiber_exn () in
+  FLS.get fiber ~default:Hmap.empty k_local_hmap
+
+let[@inline] set_local_hmap (h : Hmap.t) : unit =
+  let fiber = get_current_fiber_exn () in
+  FLS.set fiber k_local_hmap h
+
+let[@inline] update_local_hmap (f : Hmap.t -> Hmap.t) : unit =
+  let fiber = get_current_fiber_exn () in
+  let h = FLS.get fiber ~default:Hmap.empty k_local_hmap in
+  let h = f h in
+  FLS.set fiber k_local_hmap h
+
+(** @raise Invalid_argument if not present *)
+let get_exn (k : 'a Hmap.key) : 'a =
+  let h = get_local_hmap () in
+  Hmap.get k h
+
+let get_opt (k : 'a Hmap.key) : 'a option =
+  let h = get_local_hmap () in
+  Hmap.find k h
+
+let[@inline] set (k : 'a Hmap.key) (v : 'a) : unit =
+  update_local_hmap (Hmap.add k v)
+
+(**/**)
+
+module Private_hmap_fls_ = struct
+  (** Copy the hmap from f1.fls to f2.fls *)
+  let copy_fls (f1 : Picos.Fiber.t) (f2 : Picos.Fiber.t) : unit =
+    match FLS.get_exn f1 k_local_hmap with
+    | exception FLS.Not_set -> ()
+    | hmap -> FLS.set f2 k_local_hmap hmap
+end
+
+(**/**)