From f2e9f99b3624a8f7ef9de00da1b97ec486f0c522 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 25 Oct 2023 00:16:54 -0400
Subject: [PATCH] perf ws_deque: implement shrinking and a push optim

---
 src/ws_deque_.ml | 40 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/ws_deque_.ml b/src/ws_deque_.ml
index 7ca39f86..be5e4b47 100644
--- a/src/ws_deque_.ml
+++ b/src/ws_deque_.ml
@@ -17,6 +17,7 @@ module CA : sig
   val get : 'a t -> int -> 'a
   val set : 'a t -> int -> 'a -> unit
   val grow : 'a t -> bottom:int -> top:int -> 'a t
+  val shrink : 'a t -> bottom:int -> top:int -> 'a t
 end = struct
   type 'a t = {
     log_size: int;
@@ -42,43 +43,69 @@ end = struct
       set new_arr i (get self i)
     done;
     new_arr
+
+  let shrink (self : _ t) ~bottom ~top : 'a t =
+    let new_arr = create ~log_size:(self.log_size - 1) () in
+    for i = top to bottom - 1 do
+      set new_arr i (get self i)
+    done;
+    new_arr
 end
 
 type 'a t = {
   top: int A.t;  (** Where we steal *)
   bottom: int A.t;  (** Where we push/pop from the owning thread *)
+  mutable top_cached: int;  (** Last read value of [top] *)
   mutable arr: 'a CA.t;  (** The circular array *)
 }
 
 let create () : _ t =
   let arr = CA.create ~log_size:4 () in
-  { top = A.make 0; bottom = A.make 0; arr }
+  { top = A.make 0; top_cached = 0; bottom = A.make 0; arr }
 
 let[@inline] size (self : _ t) : int = max 0 (A.get self.bottom - A.get self.top)
 
 let push (self : 'a t) (x : 'a) : unit =
   let b = A.get self.bottom in
-  let t = A.get self.top in
-  let size = b - t in
+  let t_approx = self.top_cached in
 
-  if size >= CA.size self.arr - 1 then
-    self.arr <- CA.grow self.arr ~top:t ~bottom:b;
+  (* Section 2.3: over-approximation of size.
+     Only if it seems too big do we actually read [t]. *)
+  let size_approx = b - t_approx in
+  if size_approx >= CA.size self.arr - 1 then (
+    let t = A.get self.top in
+    self.top_cached <- t;
+    let size = b - t in
+
+    if size >= CA.size self.arr - 1 then
+      self.arr <- CA.grow self.arr ~top:t ~bottom:b
+  );
 
   CA.set self.arr b x;
   A.set self.bottom (b + 1)
 
+let perhaps_shrink (self : _ t) ~top ~bottom : unit =
+  let size = bottom - top in
+  let ca_size = CA.size self.arr in
+  if ca_size >= 256 && size <= ca_size / 3 then
+    self.arr <- CA.shrink self.arr ~top ~bottom
+
 let pop (self : 'a t) : 'a option =
   let b = A.get self.bottom in
   let arr = self.arr in
   let b = b - 1 in
   A.set self.bottom b;
+
   let t = A.get self.top in
+  self.top_cached <- t;
+
   let size = b - t in
   if size < 0 then (
     A.set self.bottom t;
     None
   ) else if size > 0 then (
     let x = CA.get arr b in
+    perhaps_shrink self ~bottom:b ~top:t;
     Some x
   ) else if A.compare_and_set self.top t (t + 1) then (
     (* exactly one slot, so we might be racing against stealers
@@ -90,7 +117,10 @@ let pop (self : 'a t) : 'a option =
     None
 
 let steal (self : 'a t) : 'a option =
+  (* read [top], but do not update [top_cached]
+     as we're in another thread *)
   let t = A.get self.top in
+
   let b = A.get self.bottom in
   let arr = self.arr in