Compiler Explorer

Source code

unsafe extern "C" {
    fn generate_block(state: &mut [u32; 8], block: &mut [u32; 8]);
}

#[unsafe(no_mangle)]
pub fn fill_bytes(state: &mut [u32; 8], buf: &mut [u32; 8], dst: &mut [u8]) {
    fill_bytes_via_gen_block(dst, buf, |block| unsafe { generate_block(state, block) })
}

/// Implement `fill_bytes` using buffer and block generation closure.
fn fill_bytes_via_gen_block<W: Word, const N: usize>(
    mut dst: &mut [u8],
    buf: &mut [W; N],
    mut generate_block: impl FnMut(&mut [W; N]),
) {
    let word_size = size_of::<W>();

let pos = buf[0];
    let pos_usize = pos.into_usize();
    if pos_usize < buf.len() {
        let buf_tail = &buf[pos_usize..];
        let buf_rem = size_of_val(buf_tail);

if buf_rem >= dst.len() {
            let new_pos = read_bytes(buf, dst, pos);
            buf[0] = new_pos;
            return;
        }

let (l, r) = dst.split_at_mut(buf_rem);
        read_bytes(buf, l, pos);
        dst = r;
    }

let mut blocks = dst.chunks_exact_mut(N * word_size);
    let zero = W::from_usize(0);
    for block in &mut blocks {
        // We intentionally use the temporary buffer to prevent unnecessary writes
        // to the original `buf` and to enable potential optimization of writing
        // generated data directly into `block`.
        let mut buf = [zero; N];
        generate_block(&mut buf);
        read_bytes(&buf, block, zero);
    }

let rem = blocks.into_remainder();
    let new_pos = if rem.is_empty() {
        W::from_usize(N)
    } else {
        generate_block(buf);
        read_bytes::<W, N>(buf, rem, zero)
    };
    buf[0] = new_pos;
}

/// Reads bytes from `&block[pos..new_pos]` to `dst` using little endian byte order
/// ignoring the tail bytes if necessary and returns `new_pos`.
///
/// This function is written in a way which helps the compiler to compile it down
/// to one `memcpy`. The temporary buffer gets eliminated by the compiler, see:
/// https://rust.godbolt.org/z/xbo88cbsn
#[inline]
fn read_bytes<W: Word, const N: usize>(block: &[W; N], dst: &mut [u8], pos: W) -> W {
    let word_size = size_of::<W>();
    let pos = pos.into_usize();
    assert!(size_of_val(&block[pos..]) >= size_of_val(dst));

// TODO: replace with `[0u8; { size_of::<W>() * N }]` on
    // stabilization of `generic_const_exprs`
    let mut buf = [W::from_usize(0); N];
    // SAFETY: it's safe to reference `[u32/u64; N]` as `&mut [u8]`
    // with length equal to `size_of::<u32/u64>() * N`
    let buf: &mut [u8] = unsafe {
        let p: *mut u8 = buf.as_mut_ptr().cast();
        let len = word_size * N;
        core::slice::from_raw_parts_mut(p, len)
    };

for (src, dst) in block.iter().zip(buf.chunks_exact_mut(word_size)) {
        let val = src.to_le_bytes();
        dst.copy_from_slice(val.as_ref())
    }

let offset = pos * word_size;
    dst.copy_from_slice(&buf[offset..][..dst.len()]);
    let read_words = dst.len().div_ceil(word_size);
    W::from_usize(pos + read_words)
}

/// Sealed trait implemented for `u32` and `u64`.
trait Word: crate::sealed::Sealed {}

impl Word for u32 {}
impl Word for u64 {}

mod sealed {
    /// Sealed trait implemented for `u32` and `u64`.
    pub trait Sealed: Default + Copy + TryFrom<usize> {
        type Bytes: Sized + AsRef<[u8]> + for<'a> TryFrom<&'a [u8]>;

fn from_usize(val: usize) -> Self;
        fn into_usize(self) -> usize;
        fn to_le_bytes(self) -> Self::Bytes;
        fn from_le_bytes(bytes: Self::Bytes) -> Self;
        fn increment(&mut self);
    }

impl Sealed for u32 {
        type Bytes = [u8; 4];

#[inline(always)]
        fn from_usize(val: usize) -> Self {
            val.try_into().unwrap()
        }
        #[inline(always)]
        fn into_usize(self) -> usize {
            self.try_into().unwrap()
        }
        #[inline(always)]
        fn to_le_bytes(self) -> Self::Bytes {
            u32::to_le_bytes(self)
        }
        #[inline(always)]
        fn from_le_bytes(bytes: Self::Bytes) -> Self {
            u32::from_le_bytes(bytes)
        }
        #[inline(always)]
        fn increment(&mut self) {
            *self += 1;
        }
    }

impl Sealed for u64 {
        type Bytes = [u8; 8];

#[inline(always)]
        fn from_usize(val: usize) -> Self {
            val.try_into().unwrap()
        }
        #[inline(always)]
        fn into_usize(self) -> usize {
            self.try_into().unwrap()
        }
        #[inline(always)]
        fn to_le_bytes(self) -> Self::Bytes {
            u64::to_le_bytes(self)
        }
        #[inline(always)]
        fn from_le_bytes(bytes: Self::Bytes) -> Self {
            u64::from_le_bytes(bytes)
        }
        #[inline(always)]
        fn increment(&mut self) {
            *self += 1;
        }
    }
}