zstd-rs/rust/src/mem.rs

#![allow(non_snake_case)]
use std::mem;
use std::os::raw::c_int;

pub type U8 = u8;
pub type S8 = i8;
pub type U16 = u16;
pub type S16 = i16;
pub type U32 = u32;
pub type S32 = i32;
pub type U64 = u64;
pub type S64 = i64;
pub type BYTE = u8;

#[inline]
pub fn MEM_32bits() -> bool {
    mem::size_of::<usize>() == 4
}

#[inline]
pub fn MEM_64bits() -> bool {
    mem::size_of::<usize>() == 8
}

#[inline]
pub fn MEM_isLittleEndian() -> bool {
    cfg!(target_endian = "little")
}

#[inline]
pub unsafe fn MEM_read16(mem_ptr: *const std::ffi::c_void) -> U16 {
    (mem_ptr as *const U16).read_unaligned()
}

#[inline]
pub unsafe fn MEM_read32(mem_ptr: *const std::ffi::c_void) -> U32 {
    (mem_ptr as *const U32).read_unaligned()
}

#[inline]
pub unsafe fn MEM_read64(mem_ptr: *const std::ffi::c_void) -> U64 {
    (mem_ptr as *const U64).read_unaligned()
}

#[inline]
pub unsafe fn MEM_readST(mem_ptr: *const std::ffi::c_void) -> usize {
    (mem_ptr as *const usize).read_unaligned()
}

#[inline]
pub unsafe fn MEM_write16(mem_ptr: *mut std::ffi::c_void, value: U16) {
    (mem_ptr as *mut U16).write_unaligned(value);
}

#[inline]
pub unsafe fn MEM_write32(mem_ptr: *mut std::ffi::c_void, value: U32) {
    (mem_ptr as *mut U32).write_unaligned(value);
}

#[inline]
pub unsafe fn MEM_write64(mem_ptr: *mut std::ffi::c_void, value: U64) {
    (mem_ptr as *mut U64).write_unaligned(value);
}

#[inline]
pub unsafe fn MEM_readLE16(mem_ptr: *const std::ffi::c_void) -> U16 {
    U16::from_le(MEM_read16(mem_ptr))
}

#[inline]
pub unsafe fn MEM_readLE24(mem_ptr: *const std::ffi::c_void) -> U32 {
    let low = MEM_readLE16(mem_ptr);
    let high = *mem_ptr.cast::<U8>().add(2) as U32;
    (low as U32) | (high << 16)
}

#[inline]
pub unsafe fn MEM_readLE32(mem_ptr: *const std::ffi::c_void) -> U32 {
    U32::from_le(MEM_read32(mem_ptr))
}

#[inline]
pub unsafe fn MEM_readLE64(mem_ptr: *const std::ffi::c_void) -> U64 {
    U64::from_le(MEM_read64(mem_ptr))
}

#[inline]
pub unsafe fn MEM_readLEST(mem_ptr: *const std::ffi::c_void) -> usize {
    if MEM_32bits() {
        MEM_readLE32(mem_ptr) as usize
    } else {
        MEM_readLE64(mem_ptr) as usize
    }
}

#[inline]
pub unsafe fn MEM_writeLE16(mem_ptr: *mut std::ffi::c_void, val: U16) {
    MEM_write16(mem_ptr, val.to_le());
}

#[inline]
pub unsafe fn MEM_writeLE24(mem_ptr: *mut std::ffi::c_void, val: U32) {
    MEM_writeLE16(mem_ptr, val as U16);
    *mem_ptr.cast::<U8>().add(2) = (val >> 16) as U8;
}

#[inline]
pub unsafe fn MEM_writeLE32(mem_ptr: *mut std::ffi::c_void, val: U32) {
    MEM_write32(mem_ptr, val.to_le());
}

#[inline]
pub unsafe fn MEM_writeLE64(mem_ptr: *mut std::ffi::c_void, val: U64) {
    MEM_write64(mem_ptr, val.to_le());
}

#[inline]
pub unsafe fn MEM_writeLEST(mem_ptr: *mut std::ffi::c_void, val: usize) {
    if MEM_32bits() {
        MEM_writeLE32(mem_ptr, val as U32);
    } else {
        MEM_writeLE64(mem_ptr, val as U64);
    }
}

#[inline]
pub unsafe fn MEM_readBE32(mem_ptr: *const std::ffi::c_void) -> U32 {
    U32::from_be(MEM_read32(mem_ptr))
}

#[inline]
pub unsafe fn MEM_writeBE32(mem_ptr: *mut std::ffi::c_void, val: U32) {
    MEM_write32(mem_ptr, val.to_be());
}

#[inline]
pub unsafe fn MEM_readBE64(mem_ptr: *const std::ffi::c_void) -> U64 {
    U64::from_be(MEM_read64(mem_ptr))
}

#[inline]
pub unsafe fn MEM_writeBE64(mem_ptr: *mut std::ffi::c_void, val: U64) {
    MEM_write64(mem_ptr, val.to_be());
}

#[inline]
pub unsafe fn MEM_readBEST(mem_ptr: *const std::ffi::c_void) -> usize {
    if MEM_32bits() {
        MEM_readBE32(mem_ptr) as usize
    } else {
        MEM_readBE64(mem_ptr) as usize
    }
}

#[inline]
pub unsafe fn MEM_writeBEST(mem_ptr: *mut std::ffi::c_void, val: usize) {
    if MEM_32bits() {
        MEM_writeBE32(mem_ptr, val as U32);
    } else {
        MEM_writeBE64(mem_ptr, val as U64);
    }
}

#[inline]
pub fn MEM_swap32(in_val: U32) -> U32 {
    in_val.swap_bytes()
}

#[inline]
pub fn MEM_swap64(in_val: U64) -> U64 {
    in_val.swap_bytes()
}

#[inline]
pub fn MEM_swapST(in_val: usize) -> usize {
    if MEM_32bits() {
        MEM_swap32(in_val as U32) as usize
    } else {
        MEM_swap64(in_val as U64) as usize
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ZstdOverlap {
    NoOverlap,
    OverlapSrcBeforeDst,
}

/// # Safety
/// `dst`/`src` must be valid for the wildcopy contract (may over-read/write up to 32 bytes).
#[inline]
pub unsafe fn ZSTD_wildcopy(
    dst: *mut std::ffi::c_void,
    src: *const std::ffi::c_void,
    length: isize,
    ovtype: ZstdOverlap,
) {
    const WILDCOPY_VECLEN: isize = 16;
    let mut op = dst as *mut U8;
    let mut ip = src as *const U8;
    let oend = op.wrapping_offset(length);
    let diff = (op as isize).wrapping_sub(ip as isize);

    if ovtype == ZstdOverlap::OverlapSrcBeforeDst && diff < WILDCOPY_VECLEN {
        // Handle short offset copies.
        loop {
            std::ptr::copy_nonoverlapping(ip, op, 8);
            op = op.add(8);
            ip = ip.add(8);
            if op >= oend {
                break;
            }
        }
    } else {
        assert!(diff >= WILDCOPY_VECLEN || diff <= -WILDCOPY_VECLEN);
        // Keep the first COPY16 separate because short copies are common.
        std::ptr::copy_nonoverlapping(ip, op, 16);
        if WILDCOPY_VECLEN >= length {
            return;
        }
        op = op.add(16);
        ip = ip.add(16);
        loop {
            std::ptr::copy_nonoverlapping(ip, op, 16);
            op = op.add(16);
            ip = ip.add(16);
            std::ptr::copy_nonoverlapping(ip, op, 16);
            op = op.add(16);
            ip = ip.add(16);
            if op >= oend {
                break;
            }
        }
    }
}

/// Rust implementation of the C `ZSTD_wildcopy()` ABI wrapper.
///
/// # Safety
/// `dst`/`src` must be valid for the wildcopy contract (may over-read/write up to 32 bytes),
/// and `ovtype` must be one of the two `ZSTD_overlap_e` values.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_wildcopy(
    dst: *mut std::ffi::c_void,
    src: *const std::ffi::c_void,
    length: isize,
    ovtype: c_int,
) {
    let ovtype = match ovtype {
        0 => ZstdOverlap::NoOverlap,
        1 => ZstdOverlap::OverlapSrcBeforeDst,
        _ => unreachable!("invalid ZSTD_overlap_e value"),
    };
    ZSTD_wildcopy(dst, src, length, ovtype);
}

/// # Safety
/// `dst` must be valid for `dst_capacity` bytes; `src` for `src_size` bytes.
#[inline]
pub unsafe fn ZSTD_limitCopy(
    dst: *mut std::ffi::c_void,
    dst_capacity: usize,
    src: *const std::ffi::c_void,
    src_size: usize,
) -> usize {
    let length = std::cmp::min(dst_capacity, src_size);
    if length > 0 {
        std::ptr::copy_nonoverlapping(src as *const U8, dst as *mut U8, length);
    }
    length
}

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

    #[test]
    fn unaligned_endian_io_round_trips() {
        let mut bytes = [0u8; 24];
        unsafe {
            let p = bytes.as_mut_ptr().add(1).cast::<c_void>();
            MEM_writeLE16(p, 0x1234);
            assert_eq!(&bytes[1..3], &[0x34, 0x12]);
            assert_eq!(MEM_readLE16(p), 0x1234);

            MEM_writeLE24(p, 0x00ab_cdef);
            assert_eq!(&bytes[1..4], &[0xef, 0xcd, 0xab]);
            assert_eq!(MEM_readLE24(p), 0x00ab_cdef);

            MEM_writeLE64(p, 0x0123_4567_89ab_cdef);
            assert_eq!(MEM_readLE64(p), 0x0123_4567_89ab_cdef);

            MEM_writeBE64(p, 0x0123_4567_89ab_cdef);
            assert_eq!(&bytes[1..9], &[1, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]);
            assert_eq!(MEM_readBE64(p), 0x0123_4567_89ab_cdef);
        }
    }

    const WILDCOPY_REDZONE: usize = 32;
    const WILDCOPY_LENGTHS: &[usize] = &[0, 1, 16, 17, 31, 32, 33, 63, 64];

    fn no_overlap_write_len(length: usize) -> usize {
        if length <= 16 {
            16
        } else {
            16 + (length - 16).div_ceil(32) * 32
        }
    }

    fn short_overlap_write_len(length: usize) -> usize {
        length.max(8).div_ceil(8) * 8
    }

    fn assert_filled(bytes: &[u8], value: u8) {
        assert!(bytes.iter().all(|byte| *byte == value));
    }

    #[test]
    fn wildcopy_abi_no_overlap_respects_contract_boundaries() {
        for &length in WILDCOPY_LENGTHS {
            let written = no_overlap_write_len(length);
            let source_start = WILDCOPY_REDZONE;
            let destination_start = source_start + written + 2 * WILDCOPY_REDZONE;
            let total_size = destination_start + written + WILDCOPY_REDZONE;
            let mut bytes = vec![0xa5; total_size];

            for (index, byte) in bytes[source_start..source_start + written]
                .iter_mut()
                .enumerate()
            {
                *byte = (index as u8).wrapping_mul(17).wrapping_add(3);
            }
            bytes[destination_start..destination_start + written].fill(0xcc);
            let source = bytes[source_start..source_start + written].to_vec();

            unsafe {
                ZSTD_rust_wildcopy(
                    bytes.as_mut_ptr().add(destination_start).cast(),
                    bytes.as_ptr().add(source_start).cast(),
                    length as isize,
                    0,
                );
            }

            assert_filled(&bytes[..source_start], 0xa5);
            assert_eq!(&bytes[source_start..source_start + written], &source);
            assert_filled(&bytes[source_start + written..destination_start], 0xa5);
            assert_eq!(
                &bytes[destination_start..destination_start + written],
                &source,
                "length={length}"
            );
            assert_filled(&bytes[destination_start + written..], 0xa5);
        }
    }

    #[test]
    fn wildcopy_abi_short_offsets_respect_contract_boundaries() {
        for &offset in &[8usize, 15] {
            for &length in WILDCOPY_LENGTHS {
                let written = short_overlap_write_len(length);
                let source_start = WILDCOPY_REDZONE;
                let destination_start = source_start + offset;
                let data_end = destination_start + written;
                let total_size = data_end + WILDCOPY_REDZONE;
                let mut bytes = vec![0xa5; total_size];

                for (index, byte) in bytes[source_start..data_end].iter_mut().enumerate() {
                    *byte = (index % offset) as u8 + 1;
                }

                unsafe {
                    ZSTD_rust_wildcopy(
                        bytes.as_mut_ptr().add(destination_start).cast(),
                        bytes.as_ptr().add(source_start).cast(),
                        length as isize,
                        1,
                    );
                }

                assert_filled(&bytes[..source_start], 0xa5);
                assert_eq!(
                    &bytes[destination_start..data_end],
                    &(1..=offset)
                        .cycle()
                        .take(written)
                        .map(|value| value as u8)
                        .collect::<Vec<_>>(),
                    "offset={offset}, length={length}"
                );
                assert_filled(&bytes[data_end..], 0xa5);
            }
        }
    }

    #[test]
    fn limit_copy_reports_the_bytes_written() {
        let src = *b"abcdef";
        let mut dst = [0u8; 4];
        let copied = unsafe {
            ZSTD_limitCopy(
                dst.as_mut_ptr().cast(),
                dst.len(),
                src.as_ptr().cast(),
                src.len(),
            )
        };
        assert_eq!(copied, 4);
        assert_eq!(&dst, b"abcd");
    }
}