Philox counter-based random number generator

philox.rs

use std::{ops::{BitXor, Add}, num::Wrapping};

pub trait PhiloxWord
where
    Self: Sized + BitXor<Self, Output=Self> + Add<Self, Output=Self>,
{
    const W_0: Self;
    const W_1: Self;
    const M2: Self;
    const M4_0: Self;
    const M4_1: Self;

    fn mulhilo(self, other: Self) -> (Self, Self);

    fn round(self, multiplier: Self, pair: (Self, Self)) -> (Self, Self) {
        let (x0, x1) = pair;
        let (hi, lo) = multiplier.mulhilo(x0);
        (hi ^ self ^ x1, lo)
    }
}

impl PhiloxWord for Wrapping<u32> {
    const W_0: Self = Wrapping(0x9E3779B9);  // golden ratio
    const W_1: Self = Wrapping(0xBB67AE85);  // sqrt(3) - 1
    const M2: Self = Wrapping(0xD256D193);
    const M4_0: Self = Wrapping(0xCD9E8D57);
    const M4_1: Self = Wrapping(0xD2511F53);

    fn mulhilo(self, other: Self) -> (Self, Self) {
        let t = (self.0 as u64) * (other.0 as u64);
        (Wrapping((t >> 32) as _), Wrapping((t & ((1 << 32) - 1)) as _))
    }
}

impl PhiloxWord for Wrapping<u64> {
    const W_0: Self = Wrapping(0x9E3779B97F4A7C15);  // golden ratio
    const W_1: Self = Wrapping(0xBB67AE8584CAA73B);  // sqrt(3) - 1
    const M2: Self = Wrapping(0xD2B74407B1CE6E93);
    const M4_0: Self = Wrapping(0xCA5A826395121157);
    const M4_1: Self = Wrapping(0xD2E7470EE14C6C93);

    fn mulhilo(self, other: Self) -> (Self, Self) {
        let t = (self.0 as u128) * (other.0 as u128);
        (Wrapping((t >> 64) as _), Wrapping((t & ((1 << 64) - 1)) as _))
    }
}

pub trait Philox
where
    Self: Sized,
{
    type Word;

    fn round(self, key: Self::Word) -> (Self, Self::Word);

    fn rounds(self, key: Self::Word, number: usize) -> Self {
        (0..number)
            .fold((self, key), |(s, key), _| {
                s.round(key)
            })
            .0
    }
}

impl<T> Philox for (T, T) where T: PhiloxWord + Clone {
    type Word = T;

    fn round(self, key: Self::Word) -> (Self, Self::Word) {
        (key.clone().round(T::M2, self), key + T::W_0)
    }
}

impl<T> Philox for ((T, T), (T, T)) where T: PhiloxWord + Clone {
    type Word = (T, T);

    fn round(self, key: Self::Word) -> (Self, Self::Word) {
        let (k0, k1) = key;
        let ((x0, x1), (x2, x3)) = self;
        let (r0, k0) = (k0.clone().round(T::M4_0, (x2, x1)), k0 + T::W_0);
        let (r1, k1) = (k1.clone().round(T::M4_1, (x0, x3)), k1 + T::W_1);
        ((r0, r1), (k0, k1))
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn example() {
        let key = Wrapping(0x1234567u32);
        (0..8).for_each(|c| {
            let counter = (Wrapping(c), Wrapping(0));
            println!("{:?}", counter.rounds(key, 10));
        });
    }
}