Use bit mixing for safety

CHANGES.md

@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 
 Next version:
 
+- Use bit mixing (@polytypic)
 - Change `find` to use `raise_notrace` for performance (@polytypic)
 - Change license to ISC from 0BSD (@tarides)
 

src/dune

@@ -1,3 +1,20 @@
 (library
  (name Thread_table)
  (public_name thread-table))
+
+(rule
+ (targets mix.ml)
+ (deps mix.64.ml)
+ (enabled_if %{arch_sixtyfour})
+ (action
+  (progn
+   (copy mix.64.ml mix.ml))))
+
+(rule
+ (targets mix.ml)
+ (deps mix.32.ml)
+ (enabled_if
+  (not %{arch_sixtyfour}))
+ (action
+  (progn
+   (copy mix.32.ml mix.ml))))

src/mix.32.ml

@@ -0,0 +1,15 @@
+(* Mixing function proposed by "TheIronBorn" in a Github issue
+
+     https://github.com/skeeto/hash-prospector/issues/19
+
+   in the repository of Hash Prospector by Chris Wellons.
+
+   Note that the mixing function was originally designed for 32-bit unsigned
+   integers. *)
+
+let[@inline] int x =
+  let x = x lxor (x lsr 16) in
+  let x = x * 0x21f0aaad in
+  let x = x lxor (x lsr 15) in
+  let x = x * 0x735a2d97 in
+  x lxor (x lsr 15)

src/mix.64.ml

@@ -0,0 +1,13 @@
+(* Mixing function proposed by Jon Maiga:
+
+     https://jonkagstrom.com/mx3/mx3_rev2.html
+
+   Note that the mixing function was originally designed for 64-bit unsigned
+   integers. *)
+
+let[@inline] int x =
+  let x = x lxor (x lsr 32) in
+  let x = x * 0xe9846af9b1a615d in
+  let x = x lxor (x lsr 32) in
+  let x = x * 0xe9846af9b1a615d in
+  x lxor (x lsr 28)

src/thread_table.ml

@@ -31,14 +31,14 @@ let rec find k' = function
 let[@tail_mod_cons] rec filter bit chk = function
   | Nil -> Nil
   | Cons (k, v, kvs) ->
-      if k land bit = chk then Cons (k, v, filter bit chk kvs)
+      if Mix.int k land bit = chk then Cons (k, v, filter bit chk kvs)
       else filter bit chk kvs
 
 let[@inline] filter bit chk = function
   | Nil -> Nil
-  | Cons (k, _, Nil) as kvs -> if k land bit = chk then kvs else Nil
+  | Cons (k, _, Nil) as kvs -> if Mix.int k land bit = chk then kvs else Nil
   | Cons (k, v, kvs) ->
-      if k land bit = chk then Cons (k, v, filter bit chk kvs)
+      if Mix.int k land bit = chk then Cons (k, v, filter bit chk kvs)
       else filter bit chk kvs
 
 let[@tail_mod_cons] rec append kvs tail =
@@ -54,9 +54,10 @@ let create () = { rehash = 0; buckets = Array.make min_buckets Nil; length = 0 }
 let length t = t.length
 
 let find t k' =
+  let h = Mix.int k' in
   let buckets = t.buckets in
   let n = Array.length buckets in
-  let i = k' land (n - 1) in
+  let i = h land (n - 1) in
   find k' (Array.unsafe_get buckets i)
 
 (* Below we use [@poll error] to ensure that there are no safe-points where
@@ -128,10 +129,11 @@ let[@poll error] add_atomically t buckets n i before after =
      end
 
 let rec add t k' v' =
+  let h = Mix.int k' in
   maybe_rehash t;
   let buckets = t.buckets in
   let n = Array.length buckets in
-  let i = k' land (n - 1) in
+  let i = h land (n - 1) in
   let before = Array.unsafe_get buckets i in
   let after = Cons (k', v', before) in
   if not (add_atomically t buckets n i before after) then add t k' v'
@@ -149,11 +151,12 @@ let[@poll error] remove_atomically t buckets n i before after removed =
         end)
 
 let rec remove t k' =
+  let h = Mix.int k' in
   let removed = ref false in
   maybe_rehash t;
   let buckets = t.buckets in
   let n = Array.length buckets in
-  let i = k' land (n - 1) in
+  let i = h land (n - 1) in
   let before = Array.unsafe_get buckets i in
   let after = remove_first removed k' before in
   if not (remove_atomically t buckets n i before after removed) then remove t k'