rng-pack 0.2.5

use crate::rng64::Rng64;
use std::slice::from_raw_parts_mut;

/// A Philox 2x64 random number generator.
///
/// This is a counter-based RNG suitable for parallel applications.
///
/// # Examples
///
/// ```
/// use rng_pack::philox64::Philox64;
///
/// let mut rng = Philox64::new([123, 456]);
/// ```
#[repr(C)]
pub struct Philox64 {
    c: [u64; 2],
    k: [u64; 2],
}

impl Philox64 {
    const fn chunk_size() -> usize {
        2
    }

    const fn m0() -> u64 {
        0xD2B74407B1CE6E93
    }

    /// Creates a new `Philox64` instance.
    pub fn new(seed: [u64; 2]) -> Self {
        Self { c: [1, 0], k: seed }
    }

    /// Advances the generator state by `count` steps.
    pub fn warm(&mut self, count: usize) {
        for _i in 0..count {
            let _ = self.nextu();
        }
    }

    /// Generates the next block of random numbers.
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::philox64::Philox64;
    ///
    /// let mut rng = Philox64::new([123, 456]);
    /// let val = rng.nextu();
    /// ```
    #[inline]
    pub fn nextu(&mut self) -> [u64; 2] {
        let mut out = [0u64; 2];
        let prod = (self.c[0] as u128).wrapping_mul(Self::m0() as u128);

        out[0] = (prod as u64) ^ self.c[1] ^ self.k[0];
        out[1] = (prod >> 64) as u64;

        self.c[0] = self.c[0].wrapping_add(1);
        self.c[1] = self.c[1].wrapping_add(1);

        out
    }

    /// Generates the next random `f64` value in the range [0, 1).
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::philox64::Philox64;
    ///
    /// let mut rng = Philox64::new([123, 456]);
    /// let val = rng.nextf();
    /// assert!(val >= 0.0 && val < 1.0);
    /// ```
    #[inline]
    pub fn nextf(&mut self) -> f64 {
        self.nextu()[0] as f64 * (1.0 / (u64::MAX as f64 + 1.0))
    }

    /// Generates a random `i64` value in the range [min, max].
    ///
    /// # Arguments
    ///
    /// * `min` - The lower bound (inclusive).
    /// * `max` - The upper bound (inclusive).
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::philox64::Philox64;
    ///
    /// let mut rng = Philox64::new([123, 456]);
    /// let val = rng.randi(1, 10);
    /// assert!(val >= 1 && val <= 10);
    /// ```
    #[inline]
    pub fn randi(&mut self, min: i64, max: i64) -> i64 {
        let range = (max as i128 - min as i128 + 1) as u128;
        let x = self.nextu()[0];
        ((x as u128 * range) >> 64) as i64 + min
    }

    /// Generates a random `f64` value in the range [min, max).
    ///
    /// # Arguments
    ///
    /// * `min` - The lower bound (inclusive).
    /// * `max` - The upper bound (exclusive).
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::philox64::Philox64;
    ///
    /// let mut rng = Philox64::new([123, 456]);
    /// let val = rng.randf(1.0, 10.0);
    /// assert!(val >= 1.0 && val < 10.0);
    /// ```
    #[inline]
    pub fn randf(&mut self, min: f64, max: f64) -> f64 {
        self.nextf() * (max - min) + min
    }

    /// Returns a random element from a slice.
    ///
    /// # Arguments
    ///
    /// * `choices` - The slice to choose from.
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::philox64::Philox64;
    ///
    /// let mut rng = Philox64::new([123, 456]);
    /// let choices = [1, 2, 3, 4, 5];
    /// let val = rng.choice(&choices);
    /// assert!(choices.contains(val));
    /// ```
    #[inline]
    pub fn choice<'a, T>(&mut self, choices: &'a [T]) -> &'a T {
        let index = self.randi(0, choices.len() as i64 - 1) as usize;
        &choices[index]
    }
}

impl Rng64 for Philox64 {
    #[inline]
    fn randi(&mut self, min: i64, max: i64) -> i64 {
        self.randi(min, max)
    }

    #[inline]
    fn randf(&mut self, min: f64, max: f64) -> f64 {
        self.randf(min, max)
    }

    #[inline]
    fn choice<'a, T>(&mut self, choices: &'a [T]) -> &'a T {
        self.choice(choices)
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn philox64_new(seed1: u64, seed2: u64) -> *mut Philox64 {
    Box::into_raw(Box::new(Philox64::new([seed1, seed2])))
}
#[unsafe(no_mangle)]
pub extern "C" fn philox64_warm(ptr: *mut Philox64, count: usize) {
    unsafe {
        let rng = &mut *ptr;
        rng.warm(count);
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn philox64_free(ptr: *mut Philox64) {
    if !ptr.is_null() {
        unsafe { drop(Box::from_raw(ptr)) }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn philox64_next_u64s(ptr: *mut Philox64, out: *mut u64, count: usize) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        let mut i = 0;
        while i < count {
            let chunk = rng.nextu();
            let take = (count - i).min(Philox64::chunk_size());
            buffer[i..i + take].copy_from_slice(&chunk[..take]);
            i += take;
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn philox64_next_f64s(ptr: *mut Philox64, out: *mut f64, count: usize) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        let mut i = 0;
        while i < count {
            let chunk = rng.nextu();
            let take = (count - i).min(Philox64::chunk_size());
            for j in 0..take {
                buffer[i + j] = chunk[j] as f64 * (1.0 / (u64::MAX as f64 + 1.0));
            }
            i += take;
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn philox64_rand_i64s(
    ptr: *mut Philox64,
    out: *mut i64,
    count: usize,
    min: i64,
    max: i64,
) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        let mut i = 0;
        while i < count {
            let chunk = rng.nextu();
            let take = (count - i).min(Philox64::chunk_size());
            for j in 0..take {
                let range = (max as i128 - min as i128 + 1) as u128;
                buffer[i + j] = ((chunk[j] as u128 * range) >> 64) as i64 + min;
            }
            i += take;
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn philox64_rand_f64s(
    ptr: *mut Philox64,
    out: *mut f64,
    count: usize,
    min: f64,
    max: f64,
) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        let mut i = 0;
        while i < count {
            let chunk = rng.nextu();
            let take = (count - i).min(Philox64::chunk_size());
            for j in 0..take {
                let val_01 = chunk[j] as f64 * (1.0 / (u64::MAX as f64 + 1.0));
                buffer[i + j] = val_01 * (max - min) + min;
            }
            i += take;
        }
    }
}

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

    #[test]
    fn it_works() {
        let mut rng = Philox64::new([1, 2]);
        assert_eq!(rng.nextu(), [15183679468541472402, 0]);
        assert_eq!(rng.nextf(), 0.6462178266116214);
        assert_eq!(rng.randi(10, 20), 15);
        assert_eq!(rng.randf(10.0, 20.0), 12.924356532232428);
        assert_eq!(*rng.choice(&[0, 1, 2, 3, 4]), 0);
        assert_eq!(*rng.choice(&[1, 2, 3, 4, 5]), 5);
        assert_eq!(*rng.choice(&[2, 3, 4, 5, 6]), 5);
        assert_eq!(*rng.choice(&[3, 4, 5, 6, 7]), 5);
    }
}