Reimplement queue using a new two stack representation

Author: polytypic

src/kcas_data/queue.ml

@@ -1,162 +1,219 @@
 open Kcas
 
-type 'a t = {
-  front : 'a Elems.t Loc.t;
-  middle : 'a Elems.t Loc.t;
-  back : 'a Elems.t Loc.t;
-}
-
-let alloc ~front ~middle ~back =
-  (* We allocate locations in specific order to make most efficient use of the
-     splay-tree based transaction log. *)
-  let front = Loc.make ~padded:true front
-  and middle = Loc.make ~padded:true middle
-  and back = Loc.make ~padded:true back in
-  Multicore_magic.copy_as_padded { back; middle; front }
+let unique = ref ()
+let null () = Obj.magic unique
 
-let create () = alloc ~front:Elems.empty ~middle:Elems.empty ~back:Elems.empty
+module Elems = struct
+  type 'a t = { value : 'a; tl : 'a t; length : int }
 
-let copy q =
-  let tx ~xt = (Xt.get ~xt q.front, Xt.get ~xt q.middle, Xt.get ~xt q.back) in
-  let front, middle, back = Xt.commit { tx } in
-  alloc ~front ~middle ~back
+  let rec empty = { value = null (); tl = empty; length = 0 }
+  let[@inline] length t = t.length lxor (t.length asr (Sys.int_size - 1))
 
-module Xt = struct
-  let is_empty ~xt t =
-    (* We access locations in order of allocation to make most efficient use of
-       the splay-tree based transaction log. *)
-    Xt.get ~xt t.front == Elems.empty
-    && Xt.get ~xt t.middle == Elems.empty
-    && Xt.get ~xt t.back == Elems.empty
-
-  let length ~xt { back; middle; front } =
-    Elems.length (Xt.get ~xt front)
-    + Elems.length (Xt.get ~xt middle)
-    + Elems.length (Xt.get ~xt back)
-
-  let add ~xt x q = Xt.unsafe_modify ~xt q.back @@ Elems.cons x
-  let push = add
+  let rec rev_append length t tl =
+    if length = 0 then tl
+    else rev_append (length - 1) t.tl { value = t.value; tl; length }
 
-  (** Cooperative helper to move elems from back to middle. *)
-  let back_to_middle ~middle ~back =
-    let tx ~xt =
-      let xs = Xt.exchange ~xt back Elems.empty in
-      if xs == Elems.empty || Xt.exchange ~xt middle xs != Elems.empty then
-        raise_notrace Exit
-    in
-    try Xt.commit { tx } with Exit -> ()
+  let tl_safe t = if -2 <= t.length then t.tl else t
+
+  let[@inline] tl res t =
+    let length = t.length in
+    if -2 <= length then begin
+      if length <> 0 then res := t.value;
+      t.tl
+    end
+    else
+      let length = lnot length in
+      let t =
+        rev_append (length - 1) t.tl { value = t.value; tl = empty; length }
+      in
+      res := t.value;
+      t.tl
+
+  let peek res t =
+    let length = t.length in
+    if -2 <= length then begin
+      if length <> 0 then res := t.value;
+      t
+    end
+    else
+      let length = lnot length in
+      let t =
+        rev_append (length - 1) t.tl { value = t.value; tl = empty; length }
+      in
+      res := t.value;
+      t
+
+  let rec prepend_to_seq t tl =
+    if t == empty then tl
+    else fun () -> Seq.Cons (t.value, prepend_to_seq t.tl tl)
+end
 
-  let take_opt_finish ~xt front elems =
-    let elems = Elems.rev elems in
-    Xt.set ~xt front (Elems.tl_safe elems);
-    Elems.hd_opt elems
+module Back = struct
+  type 'a t = { length : int; front : 'a; elems : 'a Elems.t }
 
-  let take_opt ~xt t =
-    let front = t.front in
-    let elems = Xt.unsafe_update ~xt front Elems.tl_safe in
-    if elems != Elems.empty then Elems.hd_opt elems
+  let empty = { length = -1; front = null (); elems = Elems.empty }
+  let[@inline] length t = lnot t.length
+
+  let[@inline] snoc x t =
+    let length = t.length in
+    if length = -1 then { length = length - 1; front = x; elems = Elems.empty }
     else
-      let middle = t.middle and back = t.back in
-      if not (Xt.is_in_log ~xt middle || Xt.is_in_log ~xt back) then
-        back_to_middle ~middle ~back;
-      let elems = Xt.exchange ~xt middle Elems.empty in
-      if elems != Elems.empty then take_opt_finish ~xt front elems
+      {
+        length = length - 1;
+        front = t.front;
+        elems = { value = x; tl = t.elems; length };
+      }
+
+  let rev_prepend_to_seq t tl =
+    let tl =
+      if t.length >= -2 then Elems.prepend_to_seq t.elems tl
       else
-        let elems = Xt.exchange ~xt back Elems.empty in
-        if elems != Elems.empty then take_opt_finish ~xt front elems else None
+        let t = ref (Either.Left t.elems) in
+        fun () ->
+          let t =
+            match !t with
+            | Left t' ->
+                (* This is parallelism safe as the result is always equivalent. *)
+                let t' = Elems.rev_append (lnot t'.length) t' Elems.empty in
+                t := Right t';
+                t'
+            | Right t' -> t'
+          in
+          Elems.prepend_to_seq t tl ()
+    in
+    if t.length <= -2 then Seq.cons t.front tl else tl
+end
+
+type 'a t = { front : 'a Elems.t Loc.t; back : 'a Back.t Loc.t }
+
+let alloc ~front ~back =
+  let front = Loc.make ~padded:true front
+  and back = Loc.make ~padded:true back in
+  Multicore_magic.copy_as_padded { front; back }
+
+let create () = alloc ~front:Elems.empty ~back:Back.empty
+
+let copy t =
+  let tx ~xt = (Xt.get ~xt t.front, Xt.get ~xt t.back) in
+  let front, back = Xt.commit { tx } in
+  alloc ~front ~back
 
-  let take_blocking ~xt q = Xt.to_blocking ~xt (take_opt q)
+module Xt = struct
+  let is_empty ~xt t =
+    Xt.get ~xt t.front == Elems.empty && Xt.get ~xt t.back == Back.empty
 
-  let peek_opt_finish ~xt front elems =
-    let elems = Elems.rev elems in
-    Xt.set ~xt front elems;
-    Elems.hd_opt elems
+  let length ~xt t =
+    Elems.length (Xt.get ~xt t.front) + Back.length (Xt.get ~xt t.back)
+
+  let add ~xt x t = Xt.unsafe_modify ~xt t.back @@ Back.snoc x
+  let push = add
 
   let peek_opt ~xt t =
-    let front = t.front in
-    let elems = Xt.get ~xt front in
-    if elems != Elems.empty then Elems.hd_opt elems
-    else
-      let middle = t.middle and back = t.back in
-      if not (Xt.is_in_log ~xt middle || Xt.is_in_log ~xt back) then
-        back_to_middle ~middle ~back;
-      let elems = Xt.exchange ~xt middle Elems.empty in
-      if elems != Elems.empty then peek_opt_finish ~xt front elems
-      else
-        let elems = Xt.exchange ~xt back Elems.empty in
-        if elems != Elems.empty then peek_opt_finish ~xt front elems else None
+    let res = ref (null ()) in
+    Xt.unsafe_modify ~xt t.front @@ Elems.peek res;
+    let res = !res in
+    if res == null () then
+      let back = Xt.get ~xt t.back in
+      if back.length = -1 then None else Some back.front
+    else Some res
+
+  let peek_blocking ~xt t =
+    let res = ref (null ()) in
+    Xt.unsafe_modify ~xt t.front @@ Elems.peek res;
+    let res = !res in
+    if res == null () then
+      let back = Xt.get ~xt t.back in
+      if back.length = -1 then Retry.later () else back.front
+    else res
 
-  let peek_blocking ~xt q = Xt.to_blocking ~xt (peek_opt q)
+  let take_opt ~xt t =
+    let res = ref (null ()) in
+    Xt.unsafe_modify ~xt t.front @@ Elems.tl res;
+    let res = !res in
+    if res == null () then
+      let back = Xt.exchange ~xt t.back Back.empty in
+      if back.length = -1 then None
+      else begin
+        Xt.set ~xt t.front back.elems;
+        Some back.front
+      end
+    else Some res
+
+  let take_blocking ~xt t =
+    let res = ref (null ()) in
+    Xt.unsafe_modify ~xt t.front @@ Elems.tl res;
+    let res = !res in
+    if res == null () then
+      let back = Xt.exchange ~xt t.back Back.empty in
+      if back.length = -1 then Retry.later ()
+      else begin
+        Xt.set ~xt t.front back.elems;
+        back.front
+      end
+    else res
 
   let clear ~xt t =
     Xt.set ~xt t.front Elems.empty;
-    Xt.set ~xt t.middle Elems.empty;
-    Xt.set ~xt t.back Elems.empty
-
-  let swap ~xt q1 q2 =
-    let front = Xt.get ~xt q1.front
-    and middle = Xt.get ~xt q1.middle
-    and back = Xt.get ~xt q1.back in
-    let front = Xt.exchange ~xt q2.front front
-    and middle = Xt.exchange ~xt q2.middle middle
-    and back = Xt.exchange ~xt q2.back back in
-    Xt.set ~xt q1.front front;
-    Xt.set ~xt q1.middle middle;
-    Xt.set ~xt q1.back back
-
-  let seq_of ~front ~middle ~back =
-    (* Sequence construction is lazy, so this function is O(1). *)
-    Seq.empty
-    |> Elems.rev_prepend_to_seq back
-    |> Elems.rev_prepend_to_seq middle
-    |> Elems.prepend_to_seq front
+    Xt.set ~xt t.back Back.empty
+
+  let swap ~xt t1 t2 =
+    let front = Xt.get ~xt t1.front and back = Xt.get ~xt t1.back in
+    let front = Xt.exchange ~xt t2.front front
+    and back = Xt.exchange ~xt t2.back back in
+    Xt.set ~xt t1.front front;
+    Xt.set ~xt t1.back back
+
+  let seq_of ~front ~back =
+    Seq.empty |> Back.rev_prepend_to_seq back |> Elems.prepend_to_seq front
 
   let to_seq ~xt t =
-    let front = Xt.get ~xt t.front
-    and middle = Xt.get ~xt t.middle
-    and back = Xt.get ~xt t.back in
-    seq_of ~front ~middle ~back
+    let front = Xt.get ~xt t.front and back = Xt.get ~xt t.back in
+    seq_of ~front ~back
 
   let take_all ~xt t =
     let front = Xt.exchange ~xt t.front Elems.empty
-    and middle = Xt.exchange ~xt t.middle Elems.empty
-    and back = Xt.exchange ~xt t.back Elems.empty in
-    seq_of ~front ~middle ~back
+    and back = Xt.exchange ~xt t.back Back.empty in
+    seq_of ~front ~back
 end
 
 let is_empty q = Kcas.Xt.commit { tx = Xt.is_empty q }
 let length q = Kcas.Xt.commit { tx = Xt.length q }
 
-let add x q =
+let add x t =
   (* Fenceless is safe as we always update. *)
-  Loc.fenceless_modify q.back @@ Elems.cons x
+  Loc.fenceless_modify t.back @@ Back.snoc x
 
 let push = add
 
-let take_opt q =
+let take_opt t =
   (* Fenceless is safe as we revert to a transaction in case we didn't update. *)
-  match Loc.fenceless_update q.front Elems.tl_safe |> Elems.hd_opt with
-  | None -> Kcas.Xt.commit { tx = Xt.take_opt q }
-  | some -> some
+  let front = Loc.fenceless_update t.front Elems.tl_safe in
+  let length = front.length in
+  if 0 < length || length = -2 then Some front.value
+  else Kcas.Xt.commit { tx = Xt.take_opt t }
 
-let take_blocking ?timeoutf q =
+let take_blocking ?timeoutf t =
   (* Fenceless is safe as we revert to a transaction in case we didn't update. *)
-  match Loc.fenceless_update q.front Elems.tl_safe |> Elems.hd_opt with
-  | None -> Kcas.Xt.commit ?timeoutf { tx = Xt.take_blocking q }
-  | Some elem -> elem
+  let front = Loc.fenceless_update t.front Elems.tl_safe in
+  let length = front.length in
+  if 0 < length || length = -2 then front.value
+  else Kcas.Xt.commit ?timeoutf { tx = Xt.take_blocking t }
+
+let peek_opt t =
+  let front = Loc.get t.front in
+  let length = front.length in
+  if 0 < length || length = -2 then Some front.value
+  else Kcas.Xt.commit { tx = Xt.peek_opt t }
+
+let peek_blocking ?timeoutf t =
+  let front = Loc.get t.front in
+  let length = front.length in
+  if 0 < length || length = -2 then front.value
+  else Kcas.Xt.commit ?timeoutf { tx = Xt.peek_blocking t }
 
 let take_all q = Kcas.Xt.commit { tx = Xt.take_all q }
-
-let peek_opt q =
-  match Loc.get q.front |> Elems.hd_opt with
-  | None -> Kcas.Xt.commit { tx = Xt.peek_opt q }
-  | some -> some
-
-let peek_blocking ?timeoutf q =
-  Kcas.Xt.commit ?timeoutf { tx = Xt.peek_blocking q }
-
-let clear q = Kcas.Xt.commit { tx = Xt.clear q }
-let swap q1 q2 = Kcas.Xt.commit { tx = Xt.swap q1 q2 }
+let clear t = Kcas.Xt.commit { tx = Xt.clear t }
+let swap t1 t2 = Kcas.Xt.commit { tx = Xt.swap t1 t2 }
 let to_seq q = Kcas.Xt.commit { tx = Xt.to_seq q }
 let iter f q = Seq.iter f @@ to_seq q
 let fold f a q = Seq.fold_left f a @@ to_seq q