Add Picos_mpmcq.length and tweak for performance

polytypic · polytypic · commit 83bfb1e2aba3 · 2024-08-27T13:50:26.000+03:00
diff --git a/lib/picos_mpmcq/picos_mpmcq.ml b/lib/picos_mpmcq/picos_mpmcq.ml
@@ -48,32 +48,91 @@ let rev = function
         prefix
 
 let rec push t value backoff = function
-  | T (Snoc snoc_r as snoc) -> push_with t value backoff snoc_r.counter (T snoc)
-  | T (Tail tail_r as tail) -> begin
+  | T (Snoc snoc_r) as prefix ->
+      let after = Snoc { counter = snoc_r.counter + 1; prefix; value } in
+      if not (Atomic.compare_and_set t.tail prefix (T after)) then
+        let backoff = Backoff.once backoff in
+        push t value backoff (Atomic.fenceless_get t.tail)
+  | T (Tail tail_r) as prefix -> begin
       match tail_r.move with
-      | Used -> push_with t value backoff tail_r.counter (T tail)
-      | Snoc move_r as move -> begin
-          match Atomic.get t.head with
-          | H (Head head_r as head) when head_r.counter < move_r.counter ->
-              let after = rev move in
-              if
-                Atomic.fenceless_get t.head == H head
-                && Atomic.compare_and_set t.head (H head) (H after)
-              then tail_r.move <- Used;
-              let new_tail = Atomic.fenceless_get t.tail in
-              if new_tail != T tail then push t value backoff new_tail
-              else push_with t value backoff tail_r.counter (T tail)
-          | _ -> push_with t value backoff tail_r.counter (T tail)
-        end
+      | Used ->
+          let after = Snoc { counter = tail_r.counter + 1; prefix; value } in
+          if not (Atomic.compare_and_set t.tail prefix (T after)) then
+            let backoff = Backoff.once backoff in
+            push t value backoff (Atomic.fenceless_get t.tail)
+      | Snoc move_r as move ->
+          begin
+            match Atomic.get t.head with
+            | H (Head head_r as head) when head_r.counter < move_r.counter ->
+                let after = rev move in
+                if
+                  Atomic.fenceless_get t.head == H head
+                  && Atomic.compare_and_set t.head (H head) (H after)
+                then tail_r.move <- Used
+            | _ -> tail_r.move <- Used
+          end;
+          push t value backoff (Atomic.get t.tail)
     end
 
-and push_with t value backoff counter prefix =
-  let after = Snoc { counter = counter + 1; prefix; value } in
-  if not (Atomic.compare_and_set t.tail prefix (T after)) then
-    let backoff = Backoff.once backoff in
-    push t value backoff (Atomic.fenceless_get t.tail)
+exception Empty
 
-let push t value = push t value Backoff.default (Atomic.fenceless_get t.tail)
+let rec pop t backoff = function
+  | H (Cons cons_r as cons) ->
+      if Atomic.compare_and_set t.head (H cons) cons_r.suffix then cons_r.value
+      else
+        let backoff = Backoff.once backoff in
+        pop t backoff (Atomic.fenceless_get t.head)
+  | H (Head head_r as head) -> begin
+      match Atomic.get t.tail with
+      | T (Snoc snoc_r as move) ->
+          if head_r.counter = snoc_r.counter then
+            if Atomic.compare_and_set t.tail (T move) snoc_r.prefix then
+              snoc_r.value
+            else pop t backoff (Atomic.fenceless_get t.head)
+          else
+            let (Tail tail_r as tail : (_, [ `Tail ]) tdt) =
+              Tail { counter = snoc_r.counter; move }
+            in
+            let new_head = Atomic.get t.head in
+            if new_head != H head then pop t backoff new_head
+            else if Atomic.compare_and_set t.tail (T move) (T tail) then
+              let (Cons cons_r) = rev move in
+              let after = cons_r.suffix in
+              let new_head = Atomic.get t.head in
+              if new_head != H head then pop t backoff new_head
+              else if Atomic.compare_and_set t.head (H head) after then begin
+                tail_r.move <- Used;
+                cons_r.value
+              end
+              else
+                let backoff = Backoff.once backoff in
+                pop t backoff (Atomic.fenceless_get t.head)
+            else pop t backoff (Atomic.fenceless_get t.head)
+      | T (Tail tail_r) -> begin
+          match tail_r.move with
+          | Used ->
+              let new_head = Atomic.get t.head in
+              if new_head != H head then pop t backoff new_head
+              else raise_notrace Empty
+          | Snoc move_r as move ->
+              if head_r.counter < move_r.counter then
+                let (Cons cons_r) = rev move in
+                let after = cons_r.suffix in
+                let new_head = Atomic.get t.head in
+                if new_head != H head then pop t backoff new_head
+                else if Atomic.compare_and_set t.head (H head) after then begin
+                  tail_r.move <- Used;
+                  cons_r.value
+                end
+                else
+                  let backoff = Backoff.once backoff in
+                  pop t backoff (Atomic.fenceless_get t.head)
+              else
+                let new_head = Atomic.get t.head in
+                if new_head != H head then pop t backoff new_head
+                else raise_notrace Empty
+        end
+    end
 
 let rec push_head t value backoff =
   match Atomic.get t.head with
@@ -121,63 +180,29 @@ let rec push_head t value backoff =
                       Atomic.fenceless_get t.head == H head
                       && Atomic.compare_and_set t.head (H head) (H after)
                     then tail_r.move <- Used
-                | _ -> ()
+                | _ -> tail_r.move <- Used
               end;
               push_head t value backoff
         end
     end
 
-let push_head t value = push_head t value Backoff.default
-
-exception Empty
-
-let rec pop t backoff = function
-  | H (Cons cons_r as cons) ->
-      if Atomic.compare_and_set t.head (H cons) cons_r.suffix then cons_r.value
-      else
-        let backoff = Backoff.once backoff in
-        pop t backoff (Atomic.fenceless_get t.head)
-  | H (Head head_r as head) -> begin
-      match Atomic.get t.tail with
-      | T (Snoc snoc_r as move) ->
-          if head_r.counter = snoc_r.counter then
-            if Atomic.compare_and_set t.tail (T move) snoc_r.prefix then
-              snoc_r.value
-            else pop t backoff (Atomic.fenceless_get t.head)
-          else
-            let tail = Tail { counter = snoc_r.counter; move } in
-            let new_head = Atomic.get t.head in
-            if new_head != H head then pop t backoff new_head
-            else if Atomic.compare_and_set t.tail (T move) (T tail) then
-              pop_moving t backoff head move tail
-            else pop t backoff (Atomic.fenceless_get t.head)
-      | T (Tail tail_r as tail) -> begin
-          match tail_r.move with
-          | Used -> pop_emptyish t backoff head
-          | Snoc _ as move -> pop_moving t backoff head move tail
-        end
-    end
-
-and pop_moving t backoff (Head head_r as head : (_, [ `Head ]) tdt)
-    (Snoc move_r as move : (_, [ `Snoc ]) tdt)
-    (Tail tail_r : (_, [ `Tail ]) tdt) =
-  if head_r.counter < move_r.counter then
-    match rev move with
-    | Cons cons_r ->
-        let after = cons_r.suffix in
-        let new_head = Atomic.get t.head in
-        if new_head != H head then pop t backoff new_head
-        else if Atomic.compare_and_set t.head (H head) after then begin
-          tail_r.move <- Used;
-          cons_r.value
-        end
-        else
-          let backoff = Backoff.once backoff in
-          pop t backoff (Atomic.fenceless_get t.head)
-  else pop_emptyish t backoff head
-
-and pop_emptyish t backoff head =
-  let new_head = Atomic.get t.head in
-  if new_head == H head then raise_notrace Empty else pop t backoff new_head
-
-let pop_exn t = pop t Backoff.default (Atomic.fenceless_get t.head)
+let rec length t =
+  let head = Atomic.get t.head in
+  let tail = Atomic.fenceless_get t.tail in
+  if head != Atomic.get t.head then length t
+  else
+    let head_at =
+      match head with H (Cons r) -> r.counter | H (Head r) -> r.counter
+    in
+    let tail_at =
+      match tail with T (Snoc r) -> r.counter | T (Tail r) -> r.counter
+    in
+    tail_at - head_at + 1
+
+let[@inline] is_empty t = length t == 0
+let[@inline] pop_exn t = pop t Backoff.default (Atomic.fenceless_get t.head)
+
+let[@inline] push t value =
+  push t value Backoff.default (Atomic.fenceless_get t.tail)
+
+let[@inline] push_head t value = push_head t value Backoff.default
diff --git a/lib/picos_mpmcq/picos_mpmcq.mli b/lib/picos_mpmcq/picos_mpmcq.mli
@@ -31,6 +31,13 @@ val pop_exn : 'a t -> 'a
 
     @raise Empty in case the queue was empty. *)
 
+val length : 'a t -> int
+(** [length queue] returns the length or the number of values
+    in the [queue]. *)
+
+val is_empty : 'a t -> bool
+(** [is_empty queue] is equivalent to {{!length} [length queue = 0]}. *)
+
 (** {1 Examples}
 
     An example top-level session:
@@ -45,6 +52,9 @@ val pop_exn : 'a t -> 'a
       # Picos_mpmcq.push_head q 76
       - : unit = ()
 
+      # Picos_mpmcq.length q
+      - : int = 2
+
       # Picos_mpmcq.push q 101
       - : unit = ()
 
diff --git a/test/test_mpmcq.ml b/test/test_mpmcq.ml
@@ -13,12 +13,13 @@ let () =
   match Queue.pop_exn q with _ -> assert false | exception Queue.Empty -> ()
 
 module Spec = struct
-  type cmd = Push of int | Push_head of int | Pop_opt
+  type cmd = Push of int | Push_head of int | Pop_opt | Length
 
   let show_cmd = function
     | Push x -> "Push " ^ string_of_int x
     | Push_head x -> "Push_head " ^ string_of_int x
     | Pop_opt -> "Pop_opt"
+    | Length -> "Length"
 
   module State = struct
     type t = int list * int list
@@ -29,6 +30,8 @@ module Spec = struct
 
     let drop ((h, t) as s) =
       match h with [] -> s | [ _ ] -> (List.rev t, []) | _ :: h -> (h, t)
+
+    let length (h, t) = List.length h + List.length t
   end
 
   type state = State.t
@@ -39,6 +42,7 @@ module Spec = struct
       Gen.int_range 1 1000 |> Gen.map (fun x -> Push x);
       Gen.int_range 1 1000 |> Gen.map (fun x -> Push_head x);
       Gen.return Pop_opt;
+      Gen.return Length;
     ]
     |> Gen.oneof |> make ~print:show_cmd
 
@@ -51,6 +55,7 @@ module Spec = struct
     | Push x -> State.push x s
     | Push_head x -> State.push_head x s
     | Pop_opt -> State.drop s
+    | Length -> s
 
   let precond _ _ = true
 
@@ -64,12 +69,14 @@ module Spec = struct
             match Queue.pop_exn d with
             | v -> Some v
             | exception Queue.Empty -> None )
+    | Length -> Res (int, Queue.length d)
 
   let postcond c (s : state) res =
     match (c, res) with
     | Push _x, Res ((Unit, _), ()) -> true
     | Push_head _x, Res ((Unit, _), ()) -> true
     | Pop_opt, Res ((Option Int, _), res) -> res = State.peek_opt s
+    | Length, Res ((Int, _), res) -> res = State.length s
     | _, _ -> false
 end
 
