Add bitgenerator for the ChaCha family of PRNG's.

bin/bench.ml

@@ -21,6 +21,7 @@ let pairs = [
     "SFC64", (module SFC64: S);
     "Xoshiro256", (module Xoshiro256: S);
     "Philox64", (module Philox64: S);
+    "ChaCha", (module ChaCha: S);
 ]
 
 

bin/crush.ml

@@ -51,6 +51,7 @@ let to_module = function
     | "pcg64" -> (module PCG64 : S)
     | "philox64" -> (module Philox64 : S)
     | "sfc64" -> (module SFC64 : S)
+    | "chacha" -> (module ChaCha : S)
     | _ -> failwith "Unknown PRNG"
 
 

lib/bitgen.ml

@@ -36,3 +36,4 @@ module SFC64 = Sfc.SFC64
 module PCG64 = Pcg.PCG64
 module Xoshiro256 = Xoshiro.Xoshiro256StarStar
 module Philox64 = Philox.Philox
+module ChaCha = Chacha.ChaCha128Counter

lib/chacha.ml

@@ -0,0 +1,112 @@
+open Stdint
+
+
+module ChaCha128Counter : sig
+
+    include Common.BITGEN
+
+    val initialize_full : key:uint64 array -> counter:uint64 * uint64 -> rounds:int -> t
+end = struct
+    type t = {rounds: int; block : uint32 array; keysetup : uint32 array; ctr : uint64 * uint64}
+
+
+    let rotl32 x n =
+        let y = 32 - n in
+        let open Uint32 in logor (shift_left x n) (shift_right x y)
+
+
+    let qround (a, b, c, d) =
+        let open Uint32 in
+        (* round 1 *)
+        let a' = a + b in let (a, b, c, d) = a', b, c, rotl32 (logxor a' d) 16 in
+        (* round 2 *)
+        let c' = c + d in let (a, b, c, d) = a, rotl32 (logxor c' b) 12, c', d in
+        (* round 3 *)
+        let a' = a + b in let (a, b, c, d) = a', b, c, rotl32 (logxor a' d) 8 in
+        (* round 4 *)
+        let c' = c + d in a, rotl32 (logxor c' b) 7, c', d
+
+
+    let full_quarter state (w, x, y, z) =
+        let w', x', y', z' = qround (state.(w), state.(x), state.(y), state.(z)) in
+        Array.mapi
+            (fun i v -> match i with
+            | k when k = w -> w' | k when k = x -> x'
+            | k when k = y -> y' | k when k = z -> z' | _ -> v) state
+
+
+    let indices = [(0, 4, 8, 12); (1, 5, 9, 13); (2, 6, 10, 14); (3, 7, 11, 15);
+                    (0, 5, 10, 15); (1, 6, 11, 12); (2, 7, 8, 13); (3, 4, 9, 14)]
+
+
+    let core state n =
+        let f s = List.fold_left (fun acc idx -> full_quarter acc idx) s indices in
+        let rec loop state0 = function
+            | 0 -> state0
+            | i -> loop (f state0) (i - 1)
+        in
+        loop state n
+
+
+    let sixteen32 = Uint32.of_int 16
+    let sixteen64 = Uint64.of_int 16
+    let mask = Uint32.(max_int |> to_uint64)
+    let constants = Uint32.[| of_int 0x61707865; of_int 0x3320646e;
+                              of_int 0x79622d32; of_int 0x6b206574 |]
+
+
+    let generate_block ctr keysetup rounds =
+        let open Uint64 in
+        let ctr0, ctr1 = ctr in
+        let f x = shift_right x 4 |> logand mask |> to_uint32 in
+        let g x = shift_right (shift_right x 4) 32 |> to_uint32 in
+        let h x = Uint32.(shift_left (sixteen32 |> rem (of_uint64 x)) 28) in
+        let state = [|
+            constants.(0); constants.(1); constants.(2); constants.(3);
+            keysetup.(0); keysetup.(1); keysetup.(2); keysetup.(3);
+            keysetup.(4); keysetup.(5); keysetup.(6); keysetup.(7);
+            f ctr0; Uint32.logor (g ctr0) (h ctr1); f ctr1; g ctr1 |]
+        in
+        rounds lsr 1 |> core state |> Array.map2 Uint32.add state
+
+
+    let next_uint32 t =
+        let open Uint64 in
+        let idx, t' = match rem (fst t.ctr) sixteen64 with
+            (* this branch is unlikely *)
+            | i when i = zero -> i, {t with block = generate_block t.ctr t.keysetup t.rounds}
+            | i -> i, t
+        in
+        match t'.block.(to_int idx), fst t.ctr |> add one with
+        | u, v when v = zero -> u, {t' with ctr = v, snd t.ctr + one}
+        | u, v -> u, {t' with ctr = v, snd t.ctr}
+
+
+    let next_uint64 t =
+        let u, t1 = next_uint32 t in
+        let v, t2 = next_uint32 t1 in
+        Uint64.(logor (shift_left (of_uint32 v) 32) (of_uint32 u)), t2
+
+
+    let next_double t = Common.next_double ~nextu64:next_uint64 t
+
+
+    let set_seed seed stream ctr rounds =
+        let open Uint64 in
+        let f x = logand x mask |> to_uint32
+        and g x = shift_right x 32 |> to_uint32 in
+        let keysetup = [| f seed.(0); g seed.(0); f seed.(1); g seed.(1);
+                        f stream.(0); g stream.(0); f stream.(1); g stream.(1) |] in
+        let ctr' = shift_left (shift_right (fst ctr) 4) 4, snd ctr in
+        {block = generate_block ctr' keysetup rounds; ctr; keysetup; rounds}
+
+
+    let initialize_full ~key ~counter ~rounds =
+        set_seed (Array.sub key 0 2) (Array.sub key 2 2) counter rounds
+
+
+    let initialize seed = 
+        let istate = Seed.SeedSequence.generate_64bit_state 4 seed in
+        initialize_full ~key:istate ~counter:Uint64.(zero, zero) ~rounds:4
+end
+