weavatrix-search-vector 0.2.0

use std::fs::File;
use std::io;
use std::mem::{align_of, size_of};
use std::path::Path;
use std::ptr::NonNull;

#[derive(Debug)]
pub(crate) struct Mapping {
    ptr: NonNull<u8>,
    len: usize,
    _file: File,
}

impl Mapping {
    pub(crate) fn open(path: &Path) -> io::Result<Self> {
        let file = File::open(path)?;
        let len = usize::try_from(file.metadata()?.len())
            .map_err(|_| io::Error::new(io::ErrorKind::FileTooLarge, "mapping is too large"))?;
        if len == 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "cannot map an empty file",
            ));
        }
        let ptr = map_file(&file, len)?;
        Ok(Self {
            ptr,
            len,
            _file: file,
        })
    }

    pub(crate) fn len(&self) -> usize {
        self.len
    }

    pub(crate) fn bytes(&self) -> &[u8] {
        // SAFETY: `ptr` describes an immutable mapping of exactly `len` bytes
        // and remains valid until `Drop`.
        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
    }

    pub(crate) fn u16_slice(&self, offset: usize, len: usize) -> Result<&[u16], &'static str> {
        self.typed_slice(offset, len)
    }

    pub(crate) fn u32_slice(&self, offset: usize, len: usize) -> Result<&[u32], &'static str> {
        self.typed_slice(offset, len)
    }

    pub(crate) fn u64_slice(&self, offset: usize, len: usize) -> Result<&[u64], &'static str> {
        self.typed_slice(offset, len)
    }

    pub(crate) fn f32_slice(&self, offset: usize, len: usize) -> Result<&[f32], &'static str> {
        self.typed_slice(offset, len)
    }

    fn typed_slice<T>(&self, offset: usize, len: usize) -> Result<&[T], &'static str> {
        if !cfg!(target_endian = "little") {
            return Err("memory-mapped snapshots require a little-endian target");
        }
        let bytes = len
            .checked_mul(size_of::<T>())
            .ok_or("mapped slice length overflowed")?;
        let end = offset
            .checked_add(bytes)
            .ok_or("mapped slice range overflowed")?;
        if end > self.len {
            return Err("mapped slice is outside the snapshot");
        }
        if !offset.is_multiple_of(align_of::<T>()) {
            return Err("mapped slice has invalid alignment");
        }
        let ptr = self.ptr.as_ptr().wrapping_add(offset).cast::<T>();
        // SAFETY: range, alignment, length, lifetime, and little-endian
        // representation were checked above. Callers only request integer and
        // finite-validated f32 sections from an immutable mapping.
        Ok(unsafe { std::slice::from_raw_parts(ptr, len) })
    }
}

// SAFETY: the mapping is immutable and its lifetime is owned by `Mapping`.
unsafe impl Send for Mapping {}
// SAFETY: shared reads do not mutate the mapping.
unsafe impl Sync for Mapping {}

impl Drop for Mapping {
    fn drop(&mut self) {
        unmap_file(self.ptr, self.len);
    }
}

#[cfg(unix)]
fn map_file(file: &File, len: usize) -> io::Result<NonNull<u8>> {
    use std::ffi::c_void;
    use std::os::fd::AsRawFd;

    const PROT_READ: i32 = 0x1;
    const MAP_PRIVATE: i32 = 0x2;

    unsafe extern "C" {
        fn mmap(
            address: *mut c_void,
            length: usize,
            protection: i32,
            flags: i32,
            descriptor: i32,
            offset: isize,
        ) -> *mut c_void;
    }

    // SAFETY: the descriptor is open for reading, the requested range is the
    // current file length, and no writable alias is created.
    let ptr = unsafe {
        mmap(
            std::ptr::null_mut(),
            len,
            PROT_READ,
            MAP_PRIVATE,
            file.as_raw_fd(),
            0,
        )
    };
    if ptr as isize == -1 {
        return Err(io::Error::last_os_error());
    }
    NonNull::new(ptr.cast()).ok_or_else(io::Error::last_os_error)
}

#[cfg(unix)]
fn unmap_file(ptr: NonNull<u8>, len: usize) {
    use std::ffi::c_void;

    unsafe extern "C" {
        fn munmap(address: *mut c_void, length: usize) -> i32;
    }

    // SAFETY: this is the exact address and length returned by `mmap`, and
    // `Drop` runs once.
    let _ = unsafe { munmap(ptr.as_ptr().cast(), len) };
}

#[cfg(windows)]
fn map_file(file: &File, len: usize) -> io::Result<NonNull<u8>> {
    use std::ffi::c_void;
    use std::os::windows::io::AsRawHandle;

    type Handle = *mut c_void;
    const PAGE_READONLY: u32 = 0x02;
    const FILE_MAP_READ: u32 = 0x0004;

    #[link(name = "kernel32")]
    unsafe extern "system" {
        fn CreateFileMappingW(
            file: Handle,
            attributes: *mut c_void,
            protection: u32,
            maximum_size_high: u32,
            maximum_size_low: u32,
            name: *const u16,
        ) -> Handle;
        fn MapViewOfFile(
            mapping: Handle,
            desired_access: u32,
            offset_high: u32,
            offset_low: u32,
            bytes_to_map: usize,
        ) -> *mut c_void;
        fn CloseHandle(object: Handle) -> i32;
    }

    // SAFETY: the file handle is valid for the duration of these calls and the
    // mapping is read-only.
    let mapping = unsafe {
        CreateFileMappingW(
            file.as_raw_handle().cast(),
            std::ptr::null_mut(),
            PAGE_READONLY,
            0,
            0,
            std::ptr::null(),
        )
    };
    if mapping.is_null() {
        return Err(io::Error::last_os_error());
    }
    // SAFETY: `mapping` is a valid read-only mapping object.
    let view = unsafe { MapViewOfFile(mapping, FILE_MAP_READ, 0, 0, len) };
    let view_error = if view.is_null() {
        Some(io::Error::last_os_error())
    } else {
        None
    };
    // SAFETY: the view retains its own reference; the mapping handle can be
    // closed after `MapViewOfFile`.
    let _ = unsafe { CloseHandle(mapping) };
    if let Some(error) = view_error {
        return Err(error);
    }
    NonNull::new(view.cast()).ok_or_else(io::Error::last_os_error)
}

#[cfg(windows)]
fn unmap_file(ptr: NonNull<u8>, _len: usize) {
    use std::ffi::c_void;

    #[link(name = "kernel32")]
    unsafe extern "system" {
        fn UnmapViewOfFile(base_address: *const c_void) -> i32;
    }

    // SAFETY: this is the exact base address returned by `MapViewOfFile`, and
    // `Drop` runs once.
    let _ = unsafe { UnmapViewOfFile(ptr.as_ptr().cast()) };
}