1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
open Common_
open Types
type 'a key = 'a Types.fls_key
let key_count_ = A.make 0
let new_key (type t) ~init () : t key =
let offset = A.fetch_and_add key_count_ 1 in
(module struct
type nonrec t = t
type fls_value += V of t
let offset = offset
let init = init
end : FLS_KEY
with type t = t)
type fls_value += Dummy
(** Resize array of TLS values *)
let[@inline never] resize_ (fib : _ Fiber.t) n =
let len = Array.length fib.fls in
let new_fls = Array.make (max n (len * 2)) Dummy in
Array.blit fib.fls 0 new_fls 0 len;
fib.fls <- new_fls
(** Access current fiber, or fail *)
let[@inline] cur_fiber_ () : any_fiber =
match !Fiber.get_current_ () with
| Some f -> f
| None -> failwith "FLS: must be run from inside a fiber"
let get (type a) ((module K) : a key) : a =
let (Any_fiber fib) = cur_fiber_ () in
if K.offset >= Array.length fib.fls then resize_ fib K.offset;
match fib.fls.(K.offset) with
| K.V x -> x
| Dummy ->
let v = K.init () in
fib.fls.(K.offset) <- K.V v;
v
| _ -> assert false
let set (type a) ((module K) : a key) (v : a) : unit =
let (Any_fiber fib) = cur_fiber_ () in
if K.offset >= Array.length fib.fls then resize_ fib K.offset;
fib.fls.(K.offset) <- K.V v
let with_value key x f =
let old = get key in
set key x;
Fun.protect ~finally:(fun () -> set key old) f