Rearrange conditionals to mark atomic sections with comments

src/thread_table.ml

@@ -68,35 +68,39 @@ let rec maybe_rehash t =
     if old_n < new_n then
       let new_bit = new_n lsr 1 in
       let rec loop i =
-        if t.buckets != old_buckets then maybe_rehash t
-        else if i < old_n then begin
-          let kvs = Array.unsafe_get old_buckets i in
-          Array.unsafe_set new_buckets i (filter new_bit 0 kvs);
-          Array.unsafe_set new_buckets (i lor new_bit)
-            (filter new_bit new_bit kvs);
-          loop (i + 1)
-        end
-        else begin
-          t.buckets <- new_buckets;
-          t.rehash <- 0
-        end
+        (* BEGIN ATOMIC *)
+        if t.buckets == old_buckets then
+          if old_n <= i then begin
+            t.buckets <- new_buckets;
+            t.rehash <- 0 (* END ATOMIC *)
+          end
+          else begin
+            let kvs = Array.unsafe_get old_buckets i in
+            Array.unsafe_set new_buckets i (filter new_bit 0 kvs);
+            Array.unsafe_set new_buckets (i lor new_bit)
+              (filter new_bit new_bit kvs);
+            loop (i + 1)
+          end
+        else maybe_rehash t
       in
       loop 0
     else
       let old_bit = old_n lsr 1 in
       let rec loop i =
-        if t.buckets != old_buckets then maybe_rehash t
-        else if i < new_n then begin
-          Array.unsafe_set new_buckets i
-            (append
-               (Array.unsafe_get old_buckets (i + old_bit))
-               (Array.unsafe_get old_buckets i));
-          loop (i + 1)
-        end
-        else begin
-          t.buckets <- new_buckets;
-          t.rehash <- 0
-        end
+        (* BEGIN ATOMIC *)
+        if t.buckets == old_buckets then
+          if new_n <= i then begin
+            t.buckets <- new_buckets;
+            t.rehash <- 0 (* END ATOMIC *)
+          end
+          else begin
+            Array.unsafe_set new_buckets i
+              (append
+                 (Array.unsafe_get old_buckets (i + old_bit))
+                 (Array.unsafe_get old_buckets i));
+            loop (i + 1)
+          end
+        else maybe_rehash t
       in
       loop 0
 
@@ -112,16 +116,17 @@ let[@inline never] rec add t k' v' =
   let i = k' land (n - 1) in
   let before = Array.unsafe_get buckets i in
   let after = Cons (k', v', before) in
+  (* BEGIN ATOMIC *)
   if
-    t.rehash <> 0 || buckets != t.buckets
-    || before != Array.unsafe_get buckets i
-  then add t k' v'
-  else begin
+    t.rehash = 0 && buckets == t.buckets && before == Array.unsafe_get buckets i
+  then begin
     Array.unsafe_set buckets i after;
     let length = t.length + 1 in
     t.length <- length;
     if n < length && n < max_buckets_div_2 then t.rehash <- n * 2
+    (* END ATOMIC *)
   end
+  else add t k' v'
 
 let[@inline never] rec remove t k' =
   let removed = ref false in
@@ -131,13 +136,16 @@ let[@inline never] rec remove t k' =
   let i = k' land (n - 1) in
   let before = Array.unsafe_get buckets i in
   let after = remove_first removed k' before in
+  (* BEGIN ATOMIC *)
   if
-    t.rehash <> 0 || buckets != t.buckets
-    || before != Array.unsafe_get buckets i
-  then remove t k'
-  else if !removed then begin
-    Array.unsafe_set buckets i after;
-    let length = t.length - 1 in
-    t.length <- length;
-    if length * 4 < n && min_buckets < n then t.rehash <- n asr 1
+    t.rehash = 0 && buckets == t.buckets && before == Array.unsafe_get buckets i
+  then begin
+    if !removed then begin
+      Array.unsafe_set buckets i after;
+      let length = t.length - 1 in
+      t.length <- length;
+      if length * 4 < n && min_buckets < n then t.rehash <- n asr 1
+      (* END ATOMIC *)
+    end
   end
+  else remove t k'