Source file `once.ml`

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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
open! Import

type 'a t = 'a or_null Ref.t

let[@inline] make v = Ref.make (This v)
let[@inline] take_or_null t = Ref.exchange t Null

let already_accessed fname =
  match failwith (Format.sprintf "Once.%s failed: already accessed" fname) with
  | (_ : nothing) -> .
;;

let[@inline] take_exn t =
  match Ref.exchange t Null with
  | Null -> already_accessed "take_exn"
  | This v -> v
;;

(* We implement is_taken twice: once using mutation, to prove it can be done safely, and
   once using magic, to avoid the overhead of the extra [caml_modify] calls for the
   mutation *)

module type Is_taken = sig
  val is_taken : 'a t @ local -> bool
end

module Is_taken__slow : Is_taken = struct
  let is_taken t =
    match take_or_null t with
    | Null -> true
    | This a ->
      Ref.set t (This a);
      false
  ;;
end

module Is_taken__fast : Is_taken = struct
  let is_taken t =
    match t |> magic_unique |> Ref.get with
    | Null -> true
    | This _ -> false
  ;;
end

module _ = Is_taken__slow
include Is_taken__fast

let borrow_or_null (type a : value mod many) (t : a t) ~f =
  match take_or_null t with
  | Null -> Null
  | This a ->
    (match f (borrow_ a) with
     | res ->
       Ref.set t (This a);
       This res
     | exception exn ->
       let bt = Stdlib.Printexc.get_raw_backtrace () in
       Ref.set t (This a);
       Stdlib.Printexc.raise_with_backtrace exn bt)
;;

let borrow_exn (type a : value mod many) (t : a t) ~f =
  match take_or_null t with
  | Null -> already_accessed "borrow_exn"
  | This a ->
    (match f (borrow_ a) with
     | res ->
       Ref.set t (This a);
       res
     | exception exn ->
       let bt = Stdlib.Printexc.get_raw_backtrace () in
       Ref.set t (This a);
       (match Stdlib.Printexc.raise_with_backtrace exn bt with
        | (_ : nothing) -> .))
;;

module Atomic = struct
  type 'a t = 'a or_null Atomic.t

  let[@inline] make v = Atomic.make (This v)
  let[@inline] take_or_null t = Atomic.exchange t Null

  let[@inline] take_exn t =
    match Atomic.exchange t Null with
    | Null -> already_accessed "Atomic.take_exn"
    | This v -> v
  ;;

  (* We implement is_taken twice: once using mutation, to prove it can be done safely, and
     once using magic, to avoid the overhead of the extra [caml_modify] calls for the
     mutation *)

  module type Is_taken = sig
    val is_taken : 'a t @ local -> bool
  end

  module Is_taken__slow : Is_taken = struct
    let is_taken t =
      match take_or_null t with
      | Null -> true
      | This a ->
        Atomic.set t (This a);
        false
    ;;
  end

  module Is_taken__fast : Is_taken = struct
    let is_taken t =
      match t |> magic_unique |> Atomic.get with
      | Null -> true
      | This _ -> false
    ;;
  end

  module _ = Is_taken__slow
  include Is_taken__fast

  let borrow_or_null (type a : value mod many) (t : a t) ~f =
    match take_or_null t with
    | Null -> Null
    | This a ->
      (match f (borrow_ a) with
       | res ->
         Atomic.set t (This a);
         This res
       | exception exn ->
         let bt = Stdlib.Printexc.get_raw_backtrace () in
         Atomic.set t (This a);
         Stdlib.Printexc.raise_with_backtrace exn bt)
  ;;

  let borrow_exn (type a : value mod many) (t : a t) ~f =
    match take_or_null t with
    | Null -> already_accessed "Atomic.borrow_exn"
    | This a ->
      (match f (borrow_ a) with
       | res ->
         Atomic.set t (This a);
         res
       | exception exn ->
         let bt = Stdlib.Printexc.get_raw_backtrace () in
         Atomic.set t (This a);
         Stdlib.Printexc.raise_with_backtrace exn bt)
  ;;
end

module Local = struct
  type 'a t = 'a or_null Ref.Local.t

  let make v = exclave_ Ref.Local.make (This v)
  let take_or_null t = exclave_ Ref.Local.exchange_global t Null

  let take_exn t = exclave_
    match take_or_null t with
    | Null -> already_accessed "Local.take_exn"
    | This v -> v
  ;;
end