Source file ivar.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
open! Base
open! Import

(* The underlying state machine of an ivar:

   {v
                                                      [create_full]
                                                            |
                                                            v
     [create]--> Empty --[read]--> Awaiters --[fill_*]--> Value
                   |                                        ^
                   |                                        |
                   +----------------------------------------+
   v}

   The [Value] state is terminal. *)

module State : sig @@ portable
  type (+'a : value_or_null) t : immutable_data with 'a @@ contended portable
  (*== | Empty
       | Awaiters
       | Value of 'a @@ contended portable
  *)

  (** Constructors *)

  val empty : ('a : value_or_null). 'a t
  val awaiters : ('a : value_or_null). 'a t
  val of_value : ('a : value_or_null). 'a @ contended portable -> 'a t

  (** Destructors *)

  val is_empty : ('a : value_or_null). 'a t @ contended local -> bool
  val is_awaiters : ('a : value_or_null). 'a t @ contended local -> bool
  val is_value : ('a : value_or_null). 'a t @ contended local -> bool
  val value_opt : 'a t @ contended -> 'a or_null @ contended portable
  val value_unsafe : ('a : value_or_null). 'a t @ contended -> 'a @ contended portable
end = struct
  type (+'a : value_or_null) t : immutable_data with 'a @@ contended portable

  let empty = (Obj.magic [@mode contended portable]) (ref 0)
  let awaiters = (Obj.magic [@mode contended portable]) (ref 1)
  let[@inline] of_value x = (Obj.magic [@mode contended portable]) x
  let[@inline] is_empty (x @ local) = phys_equal empty x
  let[@inline] is_awaiters (x @ local) = phys_equal awaiters x
  let[@inline] is_value x = not (is_empty x || is_awaiters x)

  let[@inline] value_opt x =
    if is_value x then This ((Obj.magic [@mode contended portable]) x) else Null
  ;;

  let[@inline] value_unsafe x =
    assert%debug (is_value x);
    (Obj.magic [@mode contended portable]) x
  ;;
end

type ('a : value_or_null) t = { t : 'a State.t Awaitable.t } [@@unboxed]

let create () = { t = Awaitable.make State.empty }
let create_full v = { t = Awaitable.make (State.of_value v) }

type on_full =
  | Drop
  | Raise

let fill_as (type a : value_or_null) ({ t } : a t) (v : a @ portable) on_empty =
  let v = State.of_value v in
  let before = Awaitable.get t in
  let[@inline] try_empty t =
    State.is_empty before
    &&
    match Awaitable.compare_and_set t ~if_phys_equal_to:State.empty ~replace_with:v with
    | Set_here -> true
    | Compare_failed -> false
  in
  let[@inline] try_awaiters t =
    match
      Awaitable.compare_and_set t ~if_phys_equal_to:State.awaiters ~replace_with:v
    with
    | Set_here -> Awaitable.broadcast t
    | Compare_failed ->
      (match on_empty with
       | Drop -> ()
       | Raise -> raise Already_full)
  in
  if not (try_empty t) then try_awaiters t
;;

let fill_if_empty t v = fill_as t v Drop
let fill_exn t v = fill_as t v Raise

type ('a : value_or_null, 'r : value_or_null) result =
  | Value : ('a : value_or_null). ('a, 'a) result
  | Or_canceled : ('a : value_or_null). ('a, 'a Or_canceled.t) result

let read_as
  (type (a : value_or_null) (r : value_or_null))
  w
  c
  ({ t } : a t)
  (r : (a, r) result)
  : r
  =
  let value_or_empty_or_awaiters = Awaitable.get t in
  let value_or_awaiters =
    if State.is_empty value_or_empty_or_awaiters
    then (
      match
        Awaitable.compare_and_set
          t
          ~if_phys_equal_to:State.empty
          ~replace_with:State.awaiters
      with
      | Set_here -> State.awaiters
      | Compare_failed -> Awaitable.get t)
    else value_or_empty_or_awaiters
  in
  if State.is_awaiters value_or_awaiters
  then (
    match r with
    | Value ->
      (match Awaitable.await w t ~until_phys_unequal_to:State.awaiters with
       | Signaled ->
         (* [Awaitable.await] will not return [Signaled] unless the state has been changed
            and the only state change from [awaiters] is to a value. *)
         State.value_unsafe (Awaitable.get t)
       | Terminated -> raise Await.Terminated)
    | Or_canceled ->
      (match Awaitable.await_or_cancel w c t ~until_phys_unequal_to:State.awaiters with
       | Signaled ->
         (* [Awaitable.await] will not return [Signaled] unless the state has been changed
            and the only state change from [awaiters] is to a value. *)
         Completed (State.value_unsafe (Awaitable.get t))
       | Canceled -> Canceled
       | Terminated -> raise Await.Terminated))
  else (
    (* At this point we know that the state we got was a value. *)
    match r with
    | Value -> State.value_unsafe value_or_awaiters
    | Or_canceled -> Completed (State.value_unsafe value_or_awaiters))
;;

let read w t = read_as w Cancellation.never t Value
let read_or_cancel w c t = read_as w c t Or_canceled

let peek { t } =
  let value_or_empty_or_awaiters = Awaitable.get t in
  State.value_opt value_or_empty_or_awaiters
;;

let is_full { t } = State.is_value (Awaitable.get t)
let is_empty t = not (is_full t)