fstd 0.1.2

use crate::hash::wyhash::WYP;
use std::cell::Cell;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::ops::{Bound, Div, RangeBounds};
use std::stringify;

#[inline(always)]
#[allow(unused)]
fn mul_u32(a: u32, b: u32) -> u32 {
    ((a as u64).wrapping_mul(b as u64) >> 32) as u32
}

#[inline(always)]
#[allow(unused)]
fn mul_u64(a: u64, b: u64) -> u64 {
    ((a as u128).wrapping_mul(b as u128) >> 64) as u64
}

macro_rules! mod_impl {
    ($fname:tt,$typ:tt, $next_random:tt,$mul:ident) => {
        #[inline(always)]
        #[allow(unused)]
        #[cfg(feature = "rand_nobias")]
        pub(crate) fn $fname(&self, n: $typ) -> $typ {
            // For implementation details, see https://arxiv.org/abs/1805.10941
            let mut v = self.$next_random();
            let mut res = $mul(v, n);
            let mut low = v.wrapping_mul(n);
            if low < n {
                let t = n.wrapping_neg() % n;
                while low < t {
                    v = self.$next_random();
                    res = $mul(v, n);
                    low = v.wrapping_mul(n);
                }
            }
            res
        }
        #[inline(always)]
        #[allow(unused)]
        #[cfg(not(feature = "rand_nobias"))]
        pub(crate) fn $fname(&self, n: $typ) -> $typ {
            // For implementation details, see https://arxiv.org/abs/1805.10941
            // This function introduces a slight bias.
            $mul(self.$next_random(), n)
        }
    };
}

macro_rules! size_impl {
    ($fname:tt,$typ:tt) => {
        #[inline(always)]
        #[allow(unused)]
        fn $fname(&self) -> $typ {
            self.next_u64() as $typ
        }
    };
}

// The API is inspired by fastrand.
macro_rules! range_impl {
    ($typ:tt, $unsigned_typ:tt, $next_random:tt, $mod:tt) => {
        #[doc = "Returns a random `"]
        #[doc = stringify!($typ)]
        #[doc = "` in the given range."]
        ///
        /// Panics if the range is invalid.
        #[inline(always)]
        pub fn $typ(&self, range: impl RangeBounds<$typ>) -> $typ {
            let panic_invalid_range = || {
                panic!(
                    "invalid range: {:?}..{:?}",
                    range.start_bound(),
                    range.end_bound()
                )
            };

            let lo = match range.start_bound() {
                Bound::Unbounded => std::$typ::MIN,
                Bound::Included(&x) => x,
                Bound::Excluded(&x) => x.checked_add(1).unwrap_or_else(panic_invalid_range),
            };

            let hi = match range.end_bound() {
                Bound::Unbounded => std::$typ::MAX,
                Bound::Included(&x) => x,
                Bound::Excluded(&x) => x.checked_sub(1).unwrap_or_else(panic_invalid_range),
            };

            if lo > hi {
                panic_invalid_range();
            }

            if lo == std::$typ::MIN && hi == std::$typ::MAX {
                self.$next_random() as $typ
            } else {
                let len = hi.wrapping_sub(lo).wrapping_add(1);
                lo.wrapping_add(self.$mod(len as $unsigned_typ as _) as $typ)
            }
        }
    };
}

/// A random number generator based on [wyrand](https://github.com/wangyi-fudan/wyhash).
/// It's **NOT** cryptographically secure.
#[derive(Debug, PartialEq, Eq)]
pub struct WyRng(Cell<u64>);

// TODO(zyh):
// - Support if the target_pointer_width is 128.
impl WyRng {
    // Private functions. The first parameter is the function name.
    size_impl!(next_u32, u32);
    size_impl!(next_i32, i32);
    size_impl!(next_i64, i64);
    size_impl!(next_usize, usize);
    size_impl!(next_isize, isize);
    mod_impl!(next_u32_mod, u32, next_u32, mul_u32);
    mod_impl!(next_u64_mod, u64, next_u64, mul_u64);

    // Public APIs.
    range_impl!(u32, u32, next_u64, next_u64_mod);
    range_impl!(i32, u32, next_u64, next_u64_mod);
    range_impl!(u64, u64, next_u64, next_u64_mod);
    range_impl!(i64, u64, next_u64, next_u64_mod);
    range_impl!(usize, usize, next_u64, next_u64_mod);
    range_impl!(isize, usize, next_u64, next_u64_mod);

    /// Returns a new random number generator with a random seed.
    #[inline]
    #[allow(unused)]
    pub fn new() -> Self {
        WyRng::from_seed_u64(
            THREAD_WYRAND
                .try_with(|rng| rng.next_u64())
                .unwrap_or(WYP[3]),
        )
    }

    pub fn seed(&self) -> u64 {
        self.0.get()
    }

    /// Returns a new random number generator using the given seed.
    #[inline]
    #[allow(unused)]
    pub fn from_seed_u64(seed: u64) -> Self {
        let rng = WyRng(Cell::new(seed));
        rng
    }

    /// Generates a random `u64`.
    /// All generated random numbers are based on this function.
    #[inline(always)]
    #[allow(unused)]
    fn next_u64(&self) -> u64 {
        let v = self.0.get().wrapping_add(WYP[0]);
        self.0.set(v);
        let t = (v as u128).wrapping_mul((v ^ WYP[1]) as u128);
        (t as u64) ^ (t >> 64) as u64
    }

    /// Generates a random `f64` in the `[0.0,1.0)`.
    #[inline(always)]
    pub fn f64(&self) -> f64 {
        ((self.next_u64() >> 11) as f64).div(((1 as u64) << 53) as f64)
    }

    /// Generates a random `f32` in the `[0.0,1.0)`.
    #[inline(always)]
    pub fn f32(&self) -> f32 {
        self.f64() as f32
    }

    /// Generates a random `u128`.
    #[inline(always)]
    #[allow(unused)]
    fn next_u128(&self) -> u128 {
        (self.next_u64() as u128) << 64 + (self.next_u64() as u128)
    }
}

thread_local! {
    pub(crate) static THREAD_WYRAND: WyRng = WyRng::from_seed_u64({
        let mut hasher = DefaultHasher::new();
        std::time::Instant::now().hash(&mut hasher);
        std::thread::current().id().hash(&mut hasher);
        let wyrand = WyRng::from_seed_u64(hasher.finish());
        wyrand.next_u64()
    });
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::ops::Mul;

    use super::WyRng;

    const TEST_SEED: u64 = 123;

    #[test]
    fn simple_test() {
        let r = WyRng::from_seed_u64(TEST_SEED);
        let mut m = HashMap::new();
        let times = 1000000;
        let mod_num = 100;
        for _ in 0..times {
            let rad = r.next_u64_mod(mod_num);
            let pv = m.get(&rad);
            match pv {
                Some(v) => {
                    m.insert(rad, *v + 1);
                }
                _ => {
                    m.insert(rad, 1);
                }
            }
        }
        let avg_count = times / mod_num;
        let max_thres = (avg_count as f64 * 0.1) as u64;
        let num_range = avg_count - max_thres..avg_count + max_thres;
        for (k, v) in m.iter() {
            assert!(
                num_range.contains(v),
                "({},{}) not in range({}..{}).",
                k,
                v,
                num_range.start,
                num_range.end
            )
        }
    }

    #[test]
    fn range_test() {
        let r = WyRng::from_seed_u64(TEST_SEED);
        let mut m = HashMap::new();
        let test_range = -30..70;
        for _ in 0..10000 {
            m.insert(r.i32(test_range.clone()), 1);
        }
        for i in test_range.clone() {
            match m.get(&i) {
                Some(_) => {}
                None => {
                    assert!(false, "{} doesn't exist.", i)
                }
            }
        }
    }

    #[test]
    fn float_test() {
        let r = WyRng::from_seed_u64(TEST_SEED);
        let mut m = HashMap::new();
        for _ in 0..100 {
            let random = r.f32().mul(10.0 as f32) as i32;
            m.insert(random, 1);
        }
        assert!(m.len() == 10)
    }
}