Add push_head

src_lockfree/two_stack_queue.ml

@@ -1,4 +1,4 @@
-(* Copyright (c) 2023, Vesa Karvonen <[email protected]>
+(* Copyright (c) 2023-2024, Vesa Karvonen <[email protected]>
 
    Permission to use, copy, modify, and/or distribute this software for any
    purpose with or without fee is hereby granted, provided that the above
@@ -92,6 +92,71 @@ and push_with t value backoff counter prefix =
 
 let push t value = push t value Backoff.default (Atomic.fenceless_get t.tail)
 
+let rec push_head t value backoff =
+  match Atomic.get t.head with
+  | H (Cons cons_r) as suffix ->
+      let after = Cons { counter = cons_r.counter - 1; value; suffix } in
+      if not (Atomic.compare_and_set t.head suffix (H after)) then
+        push_head t value (Backoff.once backoff)
+  | H (Head head_r) as head -> begin
+      match Atomic.get t.tail with
+      | T (Snoc snoc_r as move) ->
+          if Atomic.get t.head != head then push_head t value backoff
+          else if head_r.counter = snoc_r.counter then begin
+            let after =
+              Snoc
+                {
+                  counter = snoc_r.counter + 1;
+                  value = snoc_r.value;
+                  prefix =
+                    T
+                      (Snoc
+                         {
+                           counter = snoc_r.counter;
+                           value;
+                           prefix = snoc_r.prefix;
+                         });
+                }
+            in
+            if not (Atomic.compare_and_set t.tail (T move) (T after)) then
+              push_head t value (Backoff.once backoff)
+          end
+          else
+            let tail = Tail { counter = snoc_r.counter; move } in
+            let backoff =
+              if Atomic.compare_and_set t.tail (T move) (T tail) then backoff
+              else Backoff.once backoff
+            in
+            push_head t value backoff
+      | T (Tail tail_r) as prefix -> begin
+          match tail_r.move with
+          | Used ->
+              if Atomic.get t.head == head then begin
+                let tail =
+                  Snoc { counter = tail_r.counter + 1; value; prefix }
+                in
+                if not (Atomic.compare_and_set t.tail prefix (T tail)) then
+                  push_head t value (Backoff.once backoff)
+              end
+              else push_head t value backoff
+          | 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
+                | _ -> ()
+              end;
+              push_head t value backoff
+        end
+    end
+
+let push_head t value = push_head t value Backoff.default
+
 type ('a, _) poly = Option : ('a, 'a option) poly | Value : ('a, 'a) poly
 
 exception Empty

src_lockfree/two_stack_queue.mli

@@ -7,6 +7,9 @@ val create : unit -> 'a t
 val push : 'a t -> 'a -> unit
 (** *)
 
+val push_head : 'a t -> 'a -> unit
+(** *)
+
 exception Empty
 (** Raised by {!pop_exn} in case the queue is empty. *)
 

test/two_stack_queue/stm_two_stack_queue.ml

@@ -11,17 +11,19 @@ let () =
   assert (Queue.pop_opt q = None)
 
 module Spec = struct
-  type cmd = Push of int | Pop_opt | Length
+  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
 
     let push x (h, t) = if h == [] then ([ x ], []) else (h, x :: t)
+    let push_head x (h, t) = (x :: h, t)
     let peek_opt (h, _) = match h with x :: _ -> Some x | [] -> None
     let length (h, t) = List.length h + List.length t
 
@@ -36,6 +38,7 @@ module Spec = struct
     let open QCheck in
     [
       Gen.int_range 1 10000 |> Gen.map (fun x -> Push x);
+      Gen.int_range 1 10000 |> Gen.map (fun x -> Push_head x);
       Gen.return Pop_opt;
       Gen.return Length;
     ]
@@ -48,6 +51,7 @@ module Spec = struct
   let next_state c s =
     match c with
     | Push x -> State.push x s
+    | Push_head x -> State.push_head x s
     | Pop_opt -> State.drop s
     | Length -> s
 
@@ -57,13 +61,15 @@ module Spec = struct
     let open STM in
     match c with
     | Push x -> Res (unit, Queue.push d x)
+    | Push_head x -> Res (unit, Queue.push_head d x)
     | Pop_opt -> Res (option int, Queue.pop_opt d)
     | Length -> Res (int, Queue.length d)
 
   let postcond c (s : state) res =
     let open STM in
     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