zest4096 0.1.6

use std::{
    intrinsics::transmute,
    io::{Read, Result},
    mem::{size_of, zeroed},
    num::Wrapping,
    ptr::addr_of,
};

use num::{Bounded, Integer};

use crate::zest_hash::zest4096;

// Notice: The seed of the PRNG also holds the counter, in the higher array indicies.
// Keep your key/seed in the low bits, and the nonce after that if you wish.
// The top 16 bytes for the CSPRNG will be "Bytes are fun :)" in big endian,
// The next highest bytes will be counted over once the current 512 byte buffer
// has been read from.

const MAGIC_BYTES: [u64; 2] = [
    u64::from_be_bytes(*b"Bytes ar"),
    u64::from_be_bytes(*b"e fun :)"),
];

#[derive(Clone, Copy)]
pub struct ZestRand {
    // Seed of the CSPRNG (THIS INCLUDES THE COUNTER) -- higher bits get changed first
    seed: [u64; 62],
    // Current output of the CSPRNG
    hash_output: Option<[u64; 64]>,
    // Current byte of the CSPRNG (in the output)
    byte_offset: usize,
}
impl ZestRand {
    // (Re)seed the CSPRNG, and start at beginning of block + offset
    #[inline(always)]
    pub fn seed(&mut self, seed: [u64; 62], byte_offset: usize) {
        if self.seed != seed {
            self.seed = seed;
            self.hash_output = None;
        }
        self.byte_offset = byte_offset & 0x1FF;
    }
    // Counts to the next block, (lower index ptr = higher place value)
    // and sets the hash output to None
    #[inline(always)]
    fn next_block(&mut self) {
        let mut carry: u64 = 1;
        for elem in self.seed.iter_mut().rev() {
            *elem = (Wrapping(*elem) + Wrapping(carry)).0;
            carry = (*elem < carry) as u64;
        }
        self.hash_output = None;
    }
    // Hash the current block, return it for use.
    #[inline(always)]
    pub fn compute_block(&mut self) -> [u64; 64] {
        zest4096(self.generate_input())
    }
    // Generate the input to the zest hash based on seed, index, and magic bytes
    #[inline(always)]
    fn generate_input(&self) -> [u64; 64] {
        let mut block = [0_u64; 64];
        let (seed_split, magic_split) = block.split_at_mut(62);
        // Copy the seed to the first 62 64 bit blocks
        seed_split.copy_from_slice(&self.seed);
        // Copy the magic bytes to the last 2 64 bit blocks
        magic_split.copy_from_slice(&MAGIC_BYTES);
        block
    }
}

// Functions which output random data:
impl ZestRand {
    // Fetch the next byte, this constructs everything.
    #[inline(always)]
    pub fn next_byte(&mut self) -> u8 {
        if self.hash_output.is_none() {
            self.hash_output = Some(self.compute_block());
        }
        // Safety: known size and type before and after conversion
        let cast_slice: &[u8; 512] = unsafe { transmute(&self.hash_output.unwrap()) };
        let byte_val = cast_slice[self.byte_offset & 0x1FF];
        self.byte_offset += 1;
        if self.byte_offset >= 0x200 {
            self.byte_offset = 0;
            self.next_block();
        }
        byte_val
    }
    // Fills the output array with random bytes --
    // useful for types with "from_be_bytes" and "from_le_bytes" options.
    #[inline(always)]
    pub fn next_bytes<const N: usize>(&mut self) -> [u8; N] {
        let mut output = [0; N];
        output.fill_with(|| self.next_byte());
        output
    }
    // Fill given type with random data.
    /// # Safety
    /// Do not use with objects containing pointers or references,
    /// as they are written over with random data, and can cause a seg fault.
    /// Do not use with enums, as not all values could map to a valid enum.
    /// Think before you use!
    #[inline(always)]
    pub unsafe fn fill_rand<A: Copy>(&mut self) -> A {
        let mut data = vec![];
        data.resize_with(size_of::<A>(), || self.next_byte());
        let data_address = addr_of!(data[..]) as *const A;
        *data_address
    }
    // Use to generate random integers.
    #[inline(always)]
    pub fn next_num<T: Integer + Copy>(&mut self) -> T {
        unsafe { self.fill_rand() }
    }
    // Function takes in any signed integer type, and the max and minimum value it could be.
    // Returns an unbiased result in that type.
    // The maximum range is **not** checked to be the actual highest of the two arguments,
    // the maximum range is also ***non-inclusive***.
    #[inline(always)]
    pub fn bounded<A: Integer + Copy + From<u8> + Bounded>(&mut self, min: A, max: A) -> A {
        let max_val: A = Bounded::max_value();
        let mut not_biased = self.next_num::<A>();
        let range = max - min;
        while not_biased > (max_val - (((max_val % range) + A::from(1)) % range)) {
            not_biased = self.next_num::<A>();
        }
        if not_biased < A::from(0) {
            // Calculate absolute value to make certain
            // the numbers are not symetrical over the positive/negatice range
            not_biased = not_biased * (A::from(0) - A::from(1));
        }
        not_biased % range + min
    }
}
impl Default for ZestRand {
    #[inline(always)]
    fn default() -> Self {
        // Safety: All fields are set to 0, and none
        // are pointers to any other object.
        unsafe { zeroed() }
    }
}

// Will always return Some. You can safely unwrap.
// This will copy the current state, so take care to
// not assume that this function updates the original instance.
impl Iterator for ZestRand {
    type Item = u8;
    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        Some(self.next_byte())
    }
}

// Simply fills the requested buffer.
// Always returns ok. You can safely unwrap.
impl Read for ZestRand {
    #[inline(always)]
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        buf.fill_with(|| self.next_byte());
        Ok(buf.len())
    }
    // The number of possible outputs is far greater than any ram stick array.
    #[inline(always)]
    fn read_to_end(&mut self, _: &mut Vec<u8>) -> Result<usize> {
        panic!();
    }
}

// Functions for shuffling slice types:
impl ZestRand {
    #[inline(always)]
    pub fn slice_shuffle<T: Sized + Copy>(&mut self, slice: &mut [T]) {
        for i in 0..slice.len() {
            let swap_index = self.bounded(i, slice.len());
            slice.swap(i, swap_index);
        }
    }
}