Experiment with Atomic_array

src_lockfree/ws_deque.ml

@@ -29,6 +29,7 @@
  *)
 
 module Atomic = Multicore_magic.Transparent_atomic
+module Atomic_array = Multicore_magic.Atomic_array
 
 module type S = sig
   type 'a t
@@ -49,47 +50,46 @@ module M : S = struct
     top : int Atomic.t;
     bottom : int Atomic.t;
     top_cache : int ref;
-    mutable tab : 'a ref array;
+    mutable tab : 'a Atomic_array.t;
   }
 
   let create () =
     let top = Atomic.make 0 |> Multicore_magic.copy_as_padded in
-    let tab = Array.make min_capacity (Obj.magic ()) in
+    let tab = Atomic_array.make min_capacity (Obj.magic ()) in
     let bottom = Atomic.make 0 |> Multicore_magic.copy_as_padded in
     let top_cache = ref 0 |> Multicore_magic.copy_as_padded in
     { top; bottom; top_cache; tab } |> Multicore_magic.copy_as_padded
 
   let realloc a t b sz new_sz =
-    let new_a = Array.make new_sz (Obj.magic ()) in
-    ArrayExtra.blit_circularly a
+    let new_a = Atomic_array.make new_sz (Obj.magic ()) in
+    ArrayExtra.blit_circularly (Obj.magic a)
       (t land (sz - 1))
-      new_a
+      (Obj.magic new_a)
       (t land (new_sz - 1))
       (b - t);
     new_a
 
   let push q v =
-    let v = ref v in
     (* Read of [bottom] by the owner simply does not require a fence as the
        [bottom] is only mutated by the owner. *)
     let b = Atomic.fenceless_get q.bottom in
     let t_cache = !(q.top_cache) in
     let a = q.tab in
     let size = b - t_cache in
-    let capacity = Array.length a in
+    let capacity = Atomic_array.length a in
     if
       size < capacity
       ||
       let t = Atomic.get q.top in
       q.top_cache := t;
       t != t_cache
     then begin
-      Array.unsafe_set a (b land (capacity - 1)) v;
+      Atomic_array.unsafe_fenceless_set a (b land (capacity - 1)) v;
       Atomic.incr q.bottom
     end
     else
       let a = realloc a t_cache b capacity (capacity lsl 1) in
-      Array.unsafe_set a (b land (Array.length a - 1)) v;
+      Atomic_array.unsafe_fenceless_set a (b land (Atomic_array.length a - 1)) v;
       q.tab <- a;
       Atomic.incr q.bottom
 
@@ -104,10 +104,10 @@ module M : S = struct
     let size = b - t in
     if 0 < size then begin
       let a = q.tab in
-      let capacity = Array.length a in
-      let out = Array.unsafe_get a (b land (capacity - 1)) in
-      let res = !out in
-      out := Obj.magic ();
+      let capacity = Atomic_array.length a in
+      let i = b land (capacity - 1) in
+      let res = Atomic_array.unsafe_fenceless_get a i in
+      Atomic_array.unsafe_fenceless_set a i (Obj.magic ());
       if size + size + size <= capacity - min_capacity then
         q.tab <- realloc a t b capacity (capacity lsr 1);
       match poly with Option -> Some res | Value -> res
@@ -116,15 +116,15 @@ module M : S = struct
       (* Whether or not the [compare_and_set] below succeeds, [top_cache] can be
          updated, because in either case [top] has been incremented. *)
       q.top_cache := t + 1;
+      let a = q.tab in
+      let i = b land (Atomic_array.length a - 1) in
+      let res = Atomic_array.unsafe_fenceless_get a i in
       let got = Atomic.compare_and_set q.top t (t + 1) in
       (* This write of [bottom] requires no fence.  The deque is empty and
          remains so until the next [push]. *)
       Atomic.fenceless_set q.bottom (b + 1);
       if got then begin
-        let a = q.tab in
-        let out = Array.unsafe_get a (b land (Array.length a - 1)) in
-        let res = !out in
-        out := Obj.magic ();
+        Atomic_array.unsafe_fenceless_set a i (Obj.magic ());
         match poly with Option -> Some res | Value -> res
       end
       else match poly with Option -> None | Value -> raise_notrace Exit
@@ -149,10 +149,10 @@ module M : S = struct
     let size = b - t in
     if 0 < size then
       let a = q.tab in
-      let out = Array.unsafe_get a (t land (Array.length a - 1)) in
+      let i = t land (Atomic_array.length a - 1) in
+      let res = Atomic_array.unsafe_fenceless_get a i in
       if Atomic.compare_and_set q.top t (t + 1) then begin
-        let res = !out in
-        out := Obj.magic ();
+        Atomic_array.unsafe_compare_and_set a i res (Obj.magic ()) |> ignore;
         match poly with Option -> Some res | Value -> res
       end
       else steal_as q (Backoff.once backoff) poly