Work around OCaml 5's Atomic issues

src_lockfree/fixed_atomic.ml

@@ -0,0 +1,24 @@
+(** This "fixes" [Atomic] of OCaml 5 in two ways:
+
+    * [Atomic.get] is incorrectly subject to CSE optimization in OCaml 5.
+
+    * OCaml 5 generates inefficient accesses of ['a Atomic.t array]s assuming
+      that the array might be a double array. *)
+
+include Atomic
+
+type 'a t = private 'a ref
+
+open struct
+  external as_atomic : 'a t -> 'a Atomic.t = "%identity"
+  external of_atomic : 'a Atomic.t -> 'a t = "%identity"
+end
+
+let[@inline] make x = of_atomic (make x)
+let[@inline] get x = get (Sys.opaque_identity (as_atomic x))
+let[@inline] compare_and_set x b a = compare_and_set (as_atomic x) b a
+let[@inline] exchange x v = exchange (as_atomic x) v
+let[@inline] set x v = set (as_atomic x) v
+let[@inline] fetch_and_add x v = fetch_and_add (as_atomic x) v
+let[@inline] incr x = incr (as_atomic x)
+let[@inline] decr x = decr (as_atomic x)

src_lockfree/size.ml

@@ -1,53 +1,55 @@
 module Snapshot = struct
-  type t = int Atomic.t array
+  type t = int Fixed_atomic.t array
 
-  let zero = [| Atomic.make 0 |]
+  let zero = [| Fixed_atomic.make 0 |]
   let collecting = Int.min_int
   let computing = -1
 
   let[@inline] is_collecting (s : t) =
-    Atomic.get (Array.unsafe_get s 0) = collecting
+    Fixed_atomic.get (Array.unsafe_get s 0) = collecting
 
-  let create n = Array.init n @@ fun _ -> Atomic.make collecting
+  let create n = Array.init n @@ fun _ -> Fixed_atomic.make collecting
 
   let add s i after =
     let snap = Array.unsafe_get s i in
     while
-      let before = Atomic.get snap in
-      before < after && not (Atomic.compare_and_set snap before after)
+      let before = Fixed_atomic.get snap in
+      before < after && not (Fixed_atomic.compare_and_set snap before after)
     do
       ()
     done
 
-  let rec compute s sum i =
-    if i < Array.length s then
+  let rec compute s sum i n =
+    if i < n then
       compute s
         (sum
-        - Atomic.get (Array.unsafe_get s i)
-        + Atomic.get (Array.unsafe_get s (i + 1)))
-        (i + 2)
+        - Fixed_atomic.get (Array.unsafe_get s i)
+        + Fixed_atomic.get (Array.unsafe_get s (i + 1)))
+        (i + 2) n
     else sum
 
+  let compute s = compute s 0 1 (Array.length s)
+
   let get s =
     let status = Array.unsafe_get s 0 in
-    if Atomic.get status = collecting then
-      Atomic.compare_and_set status collecting computing |> ignore;
-    if Atomic.get status = computing then begin
-      let computed = compute s 0 1 in
-      if Atomic.get status = computing then
-        Atomic.compare_and_set status computing computed |> ignore
+    if Fixed_atomic.get status = collecting then
+      Fixed_atomic.compare_and_set status collecting computing |> ignore;
+    if Fixed_atomic.get status = computing then begin
+      let computed = compute s in
+      if Fixed_atomic.get status = computing then
+        Fixed_atomic.compare_and_set status computing computed |> ignore
     end;
-    Atomic.get status
+    Fixed_atomic.get status
 end
 
-type t = Obj.t Atomic.t array
+type t = Obj.t Fixed_atomic.t array
 
 let[@inline] get_current_snapshot (t : t) : Snapshot.t =
-  Obj.magic (Atomic.get (Array.unsafe_get t 0))
+  Obj.magic (Fixed_atomic.get (Array.unsafe_get t 0))
 
 let[@inline] cas_current_snapshot (t : t) (before : Snapshot.t)
     (after : Snapshot.t) =
-  Atomic.compare_and_set (Array.unsafe_get t 0) (Obj.repr before)
+  Fixed_atomic.compare_and_set (Array.unsafe_get t 0) (Obj.repr before)
     (Obj.repr after)
 
 (* *)
@@ -62,7 +64,7 @@ let create ?n_way () =
     | Some n_way -> n_way |> Int.min n_way_max |> Bits.ceil_pow_2
   in
   Array.init ((n_way * 2) + 1) @@ fun i ->
-  Atomic.make (if i = 0 then Obj.repr Snapshot.zero else Obj.repr 0)
+  Fixed_atomic.make (if i = 0 then Obj.repr Snapshot.zero else Obj.repr 0)
   |> Multicore_magic.copy_as_padded
 
 (* *)
@@ -91,7 +93,11 @@ let new_once t ~incr =
 
 (* *)
 
-type tx = { counter : Obj.t Atomic.t; value : int; once : [ `Open ] state }
+type tx = {
+  counter : Obj.t Fixed_atomic.t;
+  value : int;
+  once : [ `Open ] state;
+}
 
 let finish (t : t) tx =
   (* At this point the [tx.counter] already has the [tx].  Before updating the
@@ -101,7 +107,8 @@ let finish (t : t) tx =
   if index != used_index then use_index tx.once;
   (* At this point the [add_once] is essentially done.  To free memory, we do
      one more [compare_and_set]. *)
-  Atomic.compare_and_set tx.counter (Obj.repr tx) (Obj.repr tx.value) |> ignore;
+  Fixed_atomic.compare_and_set tx.counter (Obj.repr tx) (Obj.repr tx.value)
+  |> ignore;
   if index != used_index then
     let snapshot = get_current_snapshot t in
     if Snapshot.is_collecting snapshot then
@@ -113,12 +120,12 @@ let rec update_once (t : t) once =
   let index = get_index once in
   if index != used_index then
     let counter = Array.unsafe_get t index in
-    let counter_state = Atomic.get counter in
+    let counter_state = Fixed_atomic.get counter in
     if index = get_index once then
       if Obj.is_int counter_state then begin
         let value = Obj.magic counter_state + 1 in
         let tx = { counter; value; once } in
-        if Atomic.compare_and_set counter counter_state (Obj.repr tx) then
+        if Fixed_atomic.compare_and_set counter counter_state (Obj.repr tx) then
           finish t tx
         else update_once t once
       end
@@ -145,7 +152,7 @@ let get_collecting_snapshot t =
 let rec collect t snapshot i =
   if i < Array.length t then begin
     let after =
-      let counter_state = Atomic.get (Array.unsafe_get t i) in
+      let counter_state = Fixed_atomic.get (Array.unsafe_get t i) in
       if Obj.is_int counter_state then Obj.magic counter_state
       else (Obj.magic counter_state).value
     in

src_lockfree/skiplist.ml

@@ -51,7 +51,7 @@ type ('k, 'v, _) node =
    specifically, it is assumed that [Link node == Link node] meaning that the
    [Link] constructor does not allocate. *)
 and ('k, 'v) link = Link : ('k, 'v, _) node -> ('k, 'v) link [@@unboxed]
-and ('k, 'v) links = ('k, 'v) link Atomic.t array
+and ('k, 'v) links = ('k, 'v) link Fixed_atomic.t array
 
 type 'k compare = 'k -> 'k -> int
 (* Encoding the [compare] function using an algebraic type would allow the
@@ -98,13 +98,13 @@ let[@inline] is_marked = function
 let rec find_path t key prev preds succs level lowest =
   (* Breaking the sequence of dependent loads could improve performance. *)
   let prev_at_level = Array.unsafe_get prev level in
-  match Atomic.get prev_at_level with
+  match Fixed_atomic.get prev_at_level with
   | Link Null ->
       Array.unsafe_set preds level prev_at_level;
       Array.unsafe_set succs level Null;
       lowest < level && find_path t key prev preds succs (level - 1) lowest
   | Link (Node r as curr_node) as before -> begin
-      match Atomic.get (Array.unsafe_get r.next level) with
+      match Fixed_atomic.get (Array.unsafe_get r.next level) with
       | Link (Null | Node _) ->
           let c = t.compare key r.key in
           if 0 < c then find_path t key r.next preds succs level lowest
@@ -126,7 +126,8 @@ let rec find_path t key prev preds succs level lowest =
             Size.update_once t.size r.decr;
             r.decr <- Size.used_once
           end;
-          Atomic.compare_and_set prev_at_level before (Link r.node) |> ignore;
+          Fixed_atomic.compare_and_set prev_at_level before (Link r.node)
+          |> ignore;
           find_path t key prev preds succs level lowest
     end
   | Link (Mark _) ->
@@ -144,10 +145,10 @@ let[@inline] find_path t key preds succs lowest =
 let rec find_node t key prev level =
   (* Breaking the sequence of dependent loads could improve performance. *)
   let prev_at_level = Array.unsafe_get prev level in
-  match Atomic.get prev_at_level with
+  match Fixed_atomic.get prev_at_level with
   | Link Null -> if 0 < level then find_node t key prev (level - 1) else Null
   | Link (Node r as node) as before -> begin
-      match Atomic.get (Array.unsafe_get r.next level) with
+      match Fixed_atomic.get (Array.unsafe_get r.next level) with
       | Link (Null | Node _) ->
           let c = t.compare key r.key in
           if 0 < c then find_node t key r.next level
@@ -167,7 +168,8 @@ let rec find_node t key prev level =
             Size.update_once t.size r.decr;
             r.decr <- Size.used_once
           end;
-          Atomic.compare_and_set prev_at_level before (Link r.node) |> ignore;
+          Fixed_atomic.compare_and_set prev_at_level before (Link r.node)
+          |> ignore;
           find_node t key prev level
     end
   | Link (Mark _) -> find_node t key t.root (Array.length t.root - 1)
@@ -185,7 +187,7 @@ let create ?n_way ?(max_height = 10) ?(compare = Stdlib.compare) () =
      practice. *)
   if max_height < 1 || 30 < max_height then
     invalid_arg "Skiplist: max_height must be in the range [1, 30]";
-  let root = Array.init max_height @@ fun _ -> Atomic.make (Link Null) in
+  let root = Array.init max_height @@ fun _ -> Fixed_atomic.make (Link Null) in
   let size = Size.create ?n_way () in
   { compare; root; size }
 
@@ -205,15 +207,16 @@ let rec add t key value preds succs =
   &&
   let next =
     let height = get_random_height (Array.length t.root) in
-    Array.init height @@ fun i -> Atomic.make (Link (Array.unsafe_get succs i))
+    Array.init height @@ fun i ->
+    Fixed_atomic.make (Link (Array.unsafe_get succs i))
   in
   let incr = Size.new_once t.size ~incr:true in
   let (Node r as node : (_, _, [ `Node ]) node) =
     Node { key; value; incr; next }
   in
   if
     let succ = Link (Array.unsafe_get succs 0) in
-    Atomic.compare_and_set (Array.unsafe_get preds 0) succ (Link node)
+    Fixed_atomic.compare_and_set (Array.unsafe_get preds 0) succ (Link node)
   then begin
     if r.incr != Size.used_once then begin
       Size.update_once t.size r.incr;
@@ -222,7 +225,7 @@ let rec add t key value preds succs =
     (* The node is now considered as added to the skiplist. *)
     let rec update_levels level =
       if Array.length next = level then begin
-        if is_marked (Atomic.get (Array.unsafe_get next (level - 1))) then begin
+        if is_marked (Fixed_atomic.get (Array.unsafe_get next (level - 1))) then begin
           (* The node we finished adding has been removed concurrently.  To
              ensure that no references we added to the node remain, we call
              [find_node] which will remove nodes with marked references along
@@ -233,19 +236,21 @@ let rec add t key value preds succs =
       end
       else if
         let succ = Link (Array.unsafe_get succs level) in
-        Atomic.compare_and_set (Array.unsafe_get preds level) succ (Link node)
+        Fixed_atomic.compare_and_set
+          (Array.unsafe_get preds level)
+          succ (Link node)
       then update_levels (level + 1)
       else
         let _found = find_path t key preds succs level in
         let rec update_nexts level' =
           if level' < level then update_levels level
           else
             let next = Array.unsafe_get next level' in
-            match Atomic.get next with
+            match Fixed_atomic.get next with
             | Link (Null | Node _) as before ->
                 let succ = Link (Array.unsafe_get succs level') in
                 if before != succ then
-                  if Atomic.compare_and_set next before succ then
+                  if Fixed_atomic.compare_and_set next before succ then
                     update_nexts (level' - 1)
                   else update_levels level
                 else update_nexts (level' - 1)
@@ -277,17 +282,18 @@ let remove t key =
       let rec update_upper_levels level =
         if 0 < level then
           let link = Array.unsafe_get next level in
-          match Atomic.get link with
+          match Fixed_atomic.get link with
           | Link (Mark _) -> update_upper_levels (level - 1)
           | Link ((Null | Node _) as succ) ->
               let marked_succ = Mark { node = succ; decr = Size.used_once } in
-              if Atomic.compare_and_set link (Link succ) (Link marked_succ) then
-                update_upper_levels (level - 1)
+              if
+                Fixed_atomic.compare_and_set link (Link succ) (Link marked_succ)
+              then update_upper_levels (level - 1)
               else update_upper_levels level
       in
       update_upper_levels (Array.length next - 1);
       let rec try_update_bottom_level link =
-        match Atomic.get link with
+        match Fixed_atomic.get link with
         | Link (Mark r) ->
             if r.decr != Size.used_once then begin
               Size.update_once t.size r.decr;
@@ -297,7 +303,8 @@ let remove t key =
         | Link ((Null | Node _) as succ) ->
             let decr = Size.new_once t.size ~incr:false in
             let marked_succ = Mark { node = succ; decr } in
-            if Atomic.compare_and_set link (Link succ) (Link marked_succ) then
+            if Fixed_atomic.compare_and_set link (Link succ) (Link marked_succ)
+            then
               (* We have finished marking references on the node.  To ensure
                  that no references to the node remain, we call [find_node]
                  which will remove nodes with marked references along the

test/skiplist/dune

@@ -1,13 +1,14 @@
 (rule
  (progn
   (copy ../../src_lockfree/skiplist.ml skiplist.ml)
+  (copy ../../src_lockfree/fixed_atomic.ml fixed_atomic.ml)
   (copy ../../src_lockfree/bits.ml bits.ml)
   (copy ../../src_lockfree/size.ml size.ml)))
 
 (test
  (name skiplist_dscheck)
  (libraries atomic dscheck alcotest multicore-magic)
- (modules skiplist size bits skiplist_dscheck))
+ (modules skiplist size bits fixed_atomic skiplist_dscheck))
 
 (test
  (name qcheck_skiplist)