feat: add Task_local_storage.get_opt

This commit is contained in:
Simon Cruanes 2024-02-29 15:29:08 -05:00
parent 7b5ecffc8c
commit 4195d4d61c
No known key found for this signature in database
GPG key ID: EBFFF6F283F3A2B4
2 changed files with 14 additions and 2 deletions

View file

@ -35,8 +35,7 @@ let[@inline] get_cur_ () : ls_value array ref =
| Some r -> r
| None -> failwith "Task local storage must be accessed from within a runner."
let get (type a) ((module K) : a key) : a =
let cur = get_cur_ () in
let get_from_ (type a) cur ((module K) : a key) : a =
if K.offset >= Array.length !cur then resize_ cur (K.offset + 1);
match !cur.(K.offset) with
| K.V x -> (* common case first *) x
@ -47,6 +46,15 @@ let get (type a) ((module K) : a key) : a =
v
| _ -> assert false
let[@inline] get (key : 'a key) : 'a =
let cur = get_cur_ () in
get_from_ cur key
let[@inline] get_opt key =
match TLS.get k_ls_values with
| None -> None
| Some cur -> Some (get_from_ cur key)
let set (type a) ((module K) : a key) (v : a) : unit =
let cur = get_cur_ () in
if K.offset >= Array.length !cur then resize_ cur (K.offset + 1);

View file

@ -35,6 +35,10 @@ val get : 'a key -> 'a
Must be run from inside a task running on a runner.
@raise Failure otherwise *)
val get_opt : 'a key -> 'a option
(** [get_opt k] gets the current task's value for key [k],
or [None] if not run from inside the task. *)
val set : 'a key -> 'a -> unit
(** [set k v] sets the storage for [k] to [v].
Must be run from inside a task running on a runner.