From dba8528aebaa3348e240243129400bcbe3d4a9f5 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 9 Jun 2023 19:08:11 -0400
Subject: [PATCH] thread-local: expose set/remove

---
 src/thread_local.ml  | 15 ++++++++++++---
 src/thread_local.mli |  4 ++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/thread_local.ml b/src/thread_local.ml
index a96ac414..f53fde6e 100644
--- a/src/thread_local.ml
+++ b/src/thread_local.ml
@@ -40,7 +40,7 @@ let remove_ref_ self key : unit =
     Thread.yield ()
   done
 
-let set_ self key (r : _ ref) : unit =
+let set_ref_ self key (r : _ ref) : unit =
   while
     let m = A.get self in
     let m' = Key_map_.add key r m in
@@ -59,9 +59,18 @@ let get_or_create_ref_ (self : _ t) key ~v : _ ref * _ option =
     r, Some old
   with Not_found ->
     let r = ref v in
-    set_ self key r;
+    set_ref_ self key r;
     r, None
 
+let set (self : _ t) v : unit =
+  let key = get_key_ () in
+  let _, _ = get_or_create_ref_ self key ~v in
+  ()
+
+let remove (self : _ t) : unit =
+  let key = get_key_ () in
+  remove_ref_ self key
+
 let get_or_create ~create (self : 'a t) : 'a =
   let key = get_key_ () in
   try
@@ -71,7 +80,7 @@ let get_or_create ~create (self : 'a t) : 'a =
     Gc.finalise (fun _ -> remove_ref_ self key) (Thread.self ());
     let v = create () in
     let r = ref v in
-    set_ self key r;
+    set_ref_ self key r;
     v
 
 let with_ self v f =
diff --git a/src/thread_local.mli b/src/thread_local.mli
index 1af43af9..7a33b709 100644
--- a/src/thread_local.mli
+++ b/src/thread_local.mli
@@ -15,6 +15,10 @@ val get_exn : 'a t -> 'a
 (** Like {!get} but fails with an exception
     @raise Not_found if no value was found *)
 
+val set : 'a t -> 'a -> unit
+
+val remove : _ t -> unit
+
 val get_or_create : create:(unit -> 'a) -> 'a t -> 'a
 
 val with_ : 'a t -> 'a -> ('a option -> 'b) -> 'b