random_ascii.rs

#![feature(platform_intrinsics)]
#![feature(portable_simd)]
use core_simd::*;

extern "platform-intrinsic" {
    fn simd_cast<T,U>(x: T) -> U;
}

fn cast_16_32<const LANES: usize>(x: Simd<u16, LANES>) -> Simd<u32, LANES>
where LaneCount<LANES>: SupportedLaneCount
{
    unsafe { simd_cast(x) }
}

fn cast_32_8<const LANES: usize>(x: Simd<u32, LANES>) -> Simd<u8, LANES>
where LaneCount<LANES>: SupportedLaneCount
{
    unsafe { simd_cast(x) }
}

const LANES: usize = 4;

pub struct SRng {
    seed: Simd<u64, LANES>,
}

impl SRng {
    /// Construct RNG with a given seed
    pub fn new(seed: [u64; LANES]) -> Self {
        Self {
            seed: Simd::from_array(seed),
        }
    }
    /// Generate a random number using mxm from "The construct of a bit mixer"
    /// http://jonkagstrom.com/bit-mixer-construction/
    #[inline]
    pub fn next(&mut self) -> Simd<u64, LANES> {
        let mulc = Simd::splat(0x94d049bb133111eb);
        self.seed ^= self.seed >> 17;
        self.seed += mulc;
        // self.seed *= mulc;
        self.seed ^= self.seed << 43;
        self.seed
    }


    /// Generates 4*LANES random ASCII characters
    pub fn random_ascii(&mut self) -> [u8; 4 * LANES] {
        let max = Simd::splat(0x7f - 0x20);
        let space = Simd::splat(0x20);
        let random1 = Simd::<u16, {4*LANES}>::from_ne_bytes(self.next().to_ne_bytes());
        let result = cast_16_32(random1) * cast_16_32(max);
        (cast_32_8(result >> 16) + space).to_array()
    }

    /// Generates 4*LANES random alphanumeric characters
    pub fn random_alphanum(&mut self) -> [u8; 4 * LANES] {
        let max = Simd::splat(10 + 26 + 26);
        let random1 = Simd::<u16, {4*LANES}>::from_ne_bytes(self.next().to_ne_bytes());
        let result = cast_16_32(random1) * cast_16_32(max);
        let mut result = cast_32_8(result >> 16) + Simd::splat(b'0');
        let skip_ranges = [
            (b'9', b'A' - b'9' - 1),
            (b'Z', b'a' - b'Z' - 1),
        ];
        for (start, shift) in skip_ranges {
            let mask = result.lanes_gt(Simd::splat(start));
            let shifted = result + Simd::splat(shift);
            result = mask.select(shifted, result);
        }
        result.to_array()
    }

    /// Generates 8*LANES url-safe base64 digits
    pub fn random_ub64(&mut self) -> [u8; 8 * LANES] {
        let mask = Simd::splat(0x3f);
        let mut result = self.next().to_ne_bytes() & mask;
        result += Simd::splat(b'+');
        let skip_ranges = [
            (b'-', b'0' - b'-' - 1),
            (b'9', b'A' - b'9' - 1),
            (b'Z', b'_' - b'Z' - 1),
            (b'_', b'a' - b'_' - 1),
        ];
        for (start, shift) in skip_ranges {
            let mask = result.lanes_gt(Simd::splat(start));
            let shifted = result + Simd::splat(shift);
            result = mask.select(shifted, result);
        }
        result.to_array()
    }

    /// Generates 8*LANES base64 digits
    pub fn random_b64(&mut self) -> [u8; 8 * LANES] {
        let mask = Simd::splat(0x3f);
        let mut result = self.next().to_ne_bytes() & mask;
        result += Simd::splat(b'+');
        let skip_ranges = [
            (b'+', b'/' - b'+' - 1),
            (b'9', b'A' - b'9' - 1),
            (b'Z', b'a' - b'Z' - 1),
        ];
        for (start, shift) in skip_ranges {
            let mask = result.lanes_gt(Simd::splat(start));
            let shifted = result + Simd::splat(shift);
            result = mask.select(shifted, result);
        }
        result.to_array()
    }
}
fn main() {
    let mut rng = SRng::new([0x1, 0x2, 0x3, 0x4]);
    let iters = 2_000_000_000;
    let start = std::time::Instant::now();
    let mut sum = Simd::<u64, 2>::splat(0);
    let size = 16;
    for _ in 0..iters {
        let block = rng.random_ascii();
        assert_eq!(block.len(), size);
        sum ^= Simd::<u64, 2>::from_ne_bytes(Simd::<u8, 16>::from_array(block));
        // print!("{}", core::str::from_utf8(&block).unwrap());
        // sum += block[0];
    }
    let size = size as f64;
    let elapsed = start.elapsed();
    println!("sum={:?}", sum);
    println!("random ascii: {:4.3?} MB/s", (iters as f64) * size / 1_000_000.0 / elapsed.as_secs_f64())
}