urng 0.12.1

#[cfg(feature = "simd")]
use std::arch::x86_64::*;

use wrapn::Wrap;

use crate::{
    _internal::{impl_seed, sm64_from_seed32},
    Rng32, Rng64,
};

// --- Pcg32 ---

/// A PCG (Permuted Congruential Generator) random number generator.
///
/// This implementation uses the PCG-XSH-RR algorithm with 64-bit state and 32-bit output.
///
/// # Examples
///
/// ```
/// use urng::*;
///
/// let mut rng = Pcg32::new(1);
/// let _ = rng.nextu();
/// ```
#[repr(C)]
pub struct Pcg32 {
    state: Wrap<u64>,
    inc: Wrap<u64>,
}

impl Pcg32 {
    /// Creates a new `Pcg32` instance seeded with the given value.
    pub fn new(seed: u32) -> Self {
        let mut seedgen = sm64_from_seed32!(seed);
        Pcg32 {
            state: seedgen.nextu().into(),
            inc: seedgen.nextu().into(),
        }
    }
}

impl_seed!(Pcg32, 32);

impl Rng32 for Pcg32 {
    #[inline]
    fn nextu(&mut self) -> u32 {
        let oldstate = self.state;
        self.state = oldstate * 6364136223846793005 + self.inc;
        let xorshifted = *(((oldstate >> 18) ^ oldstate) >> 27).raw() as u32;
        let rot = *(oldstate >> 59).raw() as u32;
        xorshifted.rotate_right(rot)
    }
}

// --- Pcg32x8 (AVX-512) ---

#[cfg(feature = "simd")]
pub const PCG32X8_LANE: usize = 8;
#[cfg(feature = "simd")]
pub const PCG32X8_PAR_CHUNK: usize = 131_072;
#[cfg(feature = "simd")]
pub const PCG32X8_PAR_CHUNK_BLOCKS: u64 = (PCG32X8_PAR_CHUNK / PCG32X8_LANE) as u64;
#[cfg(feature = "simd")]
pub const PCG32_MULT: u64 = 6364136223846793005;

/// AVX-512 implementation of PCG32 producing 8 values per step.
///
/// # Examples
///
/// ```no_run
/// use urng::rng32::Pcg32x8;
///
/// unsafe {
///     let mut rng = Pcg32x8::new(1);
///     let _ = rng.nextu();
/// }
/// ```
#[cfg(all(feature = "simd", target_arch = "x86_64"))]
#[repr(C, align(64))]
pub struct Pcg32x8 {
    pub(crate) state: __m512i,
    pub(crate) inc: __m512i,
}

#[cfg(all(feature = "simd", target_arch = "x86_64"))]
impl Pcg32x8 {
    /// Creates a new `Pcg32x8` instance with 8 independent PCG32 streams.
    /// Requires AVX-512F support.
    ///
    /// # Safety
    ///
    /// Must only be called on a CPU that supports AVX-512F.
    #[target_feature(enable = "avx512f")]
    pub unsafe fn new(seed: u32) -> Self {
        let mut seedgen = sm64_from_seed32!(seed);

        let mut state = [0u64; PCG32X8_LANE];
        state.iter_mut().for_each(|v| *v = seedgen.nextu());

        let mut inc = [0u64; PCG32X8_LANE];
        inc.iter_mut().for_each(|v| *v = seedgen.nextu());

        unsafe {
            Pcg32x8 {
                state: _mm512_loadu_si512(state.as_ptr() as *const _),
                inc: _mm512_loadu_si512(inc.as_ptr() as *const _),
            }
        }
    }

    /// Advances all 8 PCG32 streams and returns their outputs.
    ///
    /// # Safety
    ///
    /// Must only be called on a CPU that supports AVX-512F.
    #[allow(unsafe_op_in_unsafe_fn)]
    #[target_feature(enable = "avx512f")]
    pub unsafe fn nextu(&mut self) -> [u32; PCG32X8_LANE] {
        let mult_lo = _mm512_set1_epi64(0x4C957F2D_i64);
        let mult_hi = _mm512_set1_epi64(0x5851F42D_i64);
        let mask32 = _mm512_set1_epi64(0xFFFFFFFF_i64);
        let out256 = Self::step_u32(&mut self.state, self.inc, mult_lo, mult_hi, mask32);
        std::mem::transmute(out256)
    }

    #[cfg(all(feature = "simd", target_arch = "x86_64"))]
    #[inline]
    #[allow(unsafe_op_in_unsafe_fn)]
    #[target_feature(enable = "avx512f")]
    pub(crate) unsafe fn step_u32(
        state: &mut __m512i,
        inc: __m512i,
        mult_lo: __m512i,
        mult_hi: __m512i,
        mask32: __m512i,
    ) -> __m256i {
        let oldstate = *state;
        let state_hi = _mm512_srli_epi64(oldstate, 32);
        let prod_lo = _mm512_mul_epu32(oldstate, mult_lo);
        let cross = _mm512_add_epi64(
            _mm512_mul_epu32(state_hi, mult_lo),
            _mm512_mul_epu32(oldstate, mult_hi),
        );
        *state = _mm512_add_epi64(_mm512_add_epi64(prod_lo, _mm512_slli_epi64(cross, 32)), inc);

        let xs = _mm512_srli_epi64(
            _mm512_xor_si512(_mm512_srli_epi64(oldstate, 18), oldstate),
            27,
        );
        let rot = _mm512_srli_epi64(oldstate, 59);
        let rotated = _mm512_rorv_epi32(_mm512_and_si512(xs, mask32), rot);
        _mm512_cvtepi64_epi32(rotated)
    }
}

// -- Pcg32Simd --

/// Opaque handle for the Pcg32 RNG.
/// Dispatched at runtime to AVX-512 (`Pcg32x8`) or scalar (`Pcg32`) implementation.
///
/// # Examples
///
/// ```
/// use urng::rng32::Pcg32Simd;
///
/// let _ = core::mem::size_of::<Pcg32Simd>();
/// ```
#[cfg(feature = "simd")]
#[repr(C)]
pub struct Pcg32Simd([u8; 0]);

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

    crate::safe_test!(Pcg32);
}