Without sticky bit

Author: polytypic

lib/picos_std.sync/stack.ml

@@ -1,60 +1,55 @@
 open Picos_std_awaitable
 
 type 'a state =
-  | Nil of { mutable capacity : int }
-  | Cons of { mutable capacity : int; value : 'a; rest : 'a state }
+  | Nil of { capacity : int }
+  | Cons of { capacity : int; value : 'a; rest : 'a state }
+
+let[@inline] capacity_of = function Nil r -> r.capacity | Cons r -> r.capacity
 
 type 'a t = 'a state Awaitable.t
 
 exception Empty
 
-let busy_bit = 0b01
-let one = 0b10
-let max_capacity = Int.max_int / one
+let max_capacity = Int.max_int
 
 let create ?padded ?capacity () =
   let capacity =
     match capacity with
-    | None -> max_capacity * one
+    | None -> max_capacity
     | Some capacity ->
         if capacity < 1 || max_capacity < capacity then invalid_arg "capacity"
-        else capacity * one
+        else capacity
   in
   Awaitable.make ?padded (Nil { capacity })
 
+let rec push_await t value backoff =
+  let before = Awaitable.get t in
+  let capacity = capacity_of before - 1 in
+  if 0 <= capacity then
+    let after = Cons { capacity; value; rest = before } in
+    if Awaitable.compare_and_set t before after then Awaitable.signal t
+    else push_await t value (Backoff.once backoff)
+  else begin
+    Awaitable.await t before;
+    push_await t value Backoff.default
+  end
+
 let rec push t value backoff =
-  match Awaitable.get t with
-  | Nil r as before ->
-      let capacity = r.capacity land lnot busy_bit in
-      if
-        Awaitable.compare_and_set t before
-          (Cons { capacity = capacity - one; value; rest = Nil { capacity } })
-      then begin
-        if r.capacity land busy_bit <> 0 then Awaitable.broadcast t
-      end
-      else push t value (Backoff.once backoff)
-  | Cons r as before ->
-      let capacity = r.capacity in
-      if one <= capacity then begin
-        if
-          not
-            (Awaitable.compare_and_set t before
-               (Cons { capacity = capacity - one; value; rest = before }))
-        then push t value (Backoff.once backoff)
-      end
-      else begin
-        if capacity <> capacity lor busy_bit then
-          r.capacity <- capacity lor busy_bit;
-        Awaitable.await t before;
-        push t value Backoff.default
-      end
+  let before = Awaitable.get t in
+  let capacity = capacity_of before - 1 in
+  if 0 <= capacity then
+    let after = Cons { capacity; value; rest = before } in
+    if Awaitable.compare_and_set t before after then
+      match before with Nil _ -> Awaitable.signal t | Cons _ -> ()
+    else push t value (Backoff.once backoff)
+  else push_await t value backoff
 
 let rec pop_exn t backoff =
   match Awaitable.get t with
   | Nil _ -> raise_notrace Empty
   | Cons r as before ->
       if Awaitable.compare_and_set t before r.rest then begin
-        if r.capacity land busy_bit <> 0 then Awaitable.broadcast t;
+        if r.capacity = 0 then Awaitable.signal t;
         r.value
       end
       else pop_exn t (Backoff.once backoff)