diff --git a/src/core/task_local_storage.ml b/src/core/task_local_storage.ml index c6920261..56ff354f 100644 --- a/src/core/task_local_storage.ml +++ b/src/core/task_local_storage.ml @@ -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); diff --git a/src/core/task_local_storage.mli b/src/core/task_local_storage.mli index a1f96fab..4c10155c 100644 --- a/src/core/task_local_storage.mli +++ b/src/core/task_local_storage.mli @@ -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.