rng-pack 0.2.5

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

/// A 64-bit Mersenne Twister (MT19937-64) random number generator.
///
/// # Examples
///
/// ```
/// use rng_pack::mt1993764::Mt1993764;
///
/// let mut rng = Mt1993764::new(12345, 1024);
/// let val = rng.nextu();
/// ```
#[repr(C)]
pub struct Mt1993764 {
    mt: [u64; N],
    mti: usize,
}

const N: usize = 312;
const M: usize = 156;
const MATRIX_A: u64 = 0xB5026F5AA96619E9;
const UPPER_MASK: u64 = 0xFFFFFFFF80000000;
const LOWER_MASK: u64 = 0x7FFFFFFF;

impl Mt1993764 {
    /// Creates a new `Mt1993764` instance.
    ///
    /// # Arguments
    ///
    /// * `seed` - The initial seed value.
    /// * `warm` - The number of initial iterations to skip (warm-up).
    pub fn new(seed: u64, warm: usize) -> Self {
        let mut mt = [0u64; N];
        mt[0] = seed;
        for i in 1..N {
            let prev = mt[i - 1];
            mt[i] = 6364136223846793005u64
                .wrapping_mul(prev ^ (prev >> 62))
                .wrapping_add(i as u64);
        }
        let mut rng = Self { mt, mti: N };
        (0..warm).into_iter().for_each(|_| {
            let _ = rng.nextu();
        });
        rng
    }

    /// Generates the next random `u64` value.
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::mt1993764::Mt1993764;
    ///
    /// let mut rng = Mt1993764::new(12345, 1024);
    /// let val = rng.nextu();
    /// ```
    #[inline]
    pub fn nextu(&mut self) -> u64 {
        if self.mti >= N {
            self.twist();
        }

        let mut y = self.mt[self.mti];
        self.mti += 1;

        y ^= (y >> 29) & 0x5555555555555555;
        y ^= (y << 17) & 0x71D67FFFEDA60000;
        y ^= (y << 37) & 0xFFF7EEE000000000;
        y ^= y >> 43;

        y
    }

    fn twist(&mut self) {
        for i in 0..N - M {
            let x = (self.mt[i] & UPPER_MASK) | (self.mt[i + 1] & LOWER_MASK);
            let mut x_a = x >> 1;
            if x & 1 != 0 {
                x_a ^= MATRIX_A;
            }
            self.mt[i] = self.mt[i + M] ^ x_a;
        }

        for i in N - M..N - 1 {
            let x = (self.mt[i] & UPPER_MASK) | (self.mt[i + 1] & LOWER_MASK);
            let mut x_a = x >> 1;
            if x & 1 != 0 {
                x_a ^= MATRIX_A;
            }
            self.mt[i] = self.mt[i + M - N] ^ x_a;
        }

        let x = (self.mt[N - 1] & UPPER_MASK) | (self.mt[0] & LOWER_MASK);
        let mut x_a = x >> 1;
        if x & 1 != 0 {
            x_a ^= MATRIX_A;
        }
        self.mt[N - 1] = self.mt[M - 1] ^ x_a;

        self.mti = 0;
    }

    /// Generates the next random `f64` value in the range [0, 1).
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::mt1993764::Mt1993764;
    ///
    /// let mut rng = Mt1993764::new(12345, 1024);
    /// let val = rng.nextf();
    /// assert!(val >= 0.0 && val < 1.0);
    /// ```
    #[inline]
    pub fn nextf(&mut self) -> f64 {
        self.nextu() 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::mt1993764::Mt1993764;
    ///
    /// let mut rng = Mt1993764::new(12345, 1024);
    /// 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();
        ((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::mt1993764::Mt1993764;
    ///
    /// let mut rng = Mt1993764::new(12345, 1024);
    /// 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 {
        let range = max - min;
        let scale = range * (1.0 / (u64::MAX as f64 + 1.0));
        (self.nextu() as f64 * scale) + min
    }

    /// Returns a random element from a slice.
    ///
    /// # Arguments
    ///
    /// * `choices` - The slice to choose from.
    ///
    /// # Examples
    ///
    /// ```
    /// use rng_pack::mt1993764::Mt1993764;
    ///
    /// let mut rng = Mt1993764::new(12345, 1024);
    /// 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);
        &choices[index as usize]
    }
}

impl Rng64 for Mt1993764 {
    #[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 mt1993764_new(seed: u64, warm: usize) -> *mut Mt1993764 {
    Box::into_raw(Box::new(Mt1993764::new(seed, warm)))
}

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

#[unsafe(no_mangle)]
pub extern "C" fn mt1993764_next_u64s(ptr: *mut Mt1993764, out: *mut u64, count: usize) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        for v in buffer {
            *v = rng.nextu();
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn mt1993764_next_f64s(ptr: *mut Mt1993764, out: *mut f64, count: usize) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        for v in buffer {
            *v = rng.nextf();
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn mt1993764_rand_i64s(
    ptr: *mut Mt1993764,
    out: *mut i64,
    count: usize,
    min: i64,
    max: i64,
) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        for v in buffer {
            *v = rng.randi(min, max);
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn mt1993764_rand_f64s(
    ptr: *mut Mt1993764,
    out: *mut f64,
    count: usize,
    min: f64,
    max: f64,
) {
    unsafe {
        let rng = &mut *ptr;
        let buffer = from_raw_parts_mut(out, count);
        for v in buffer {
            *v = rng.randf(min, max);
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::mt1993764::Mt1993764;

    #[test]
    fn it_works() {
        let mut rng = Mt1993764::new(1, 1024);

        assert_eq!(rng.nextu(), 4804558718269354394);
        assert_eq!(rng.nextf(), 0.6894973013138909);
        assert_eq!(rng.randi(10, 20), 16);
        assert_eq!(rng.randf(10.0, 20.0), 16.300278229970864);
        assert_eq!(*rng.choice(&[0, 1, 2, 3, 4]), 4);
        assert_eq!(*rng.choice(&[1, 2, 3, 4, 5]), 3);
        assert_eq!(*rng.choice(&[2, 3, 4, 5, 6]), 6);
        assert_eq!(*rng.choice(&[3, 4, 5, 6, 7]), 7);
    }
}