hrx-rs 0.10.0

//! Indexed owned and memory-mapped SafeTensors files.

use crate::{Error, Result};
use memmap2::Mmap;
use safetensors::tensor::{Dtype as NativeDType, Metadata, SafeTensors};
use std::{
    collections::BTreeMap,
    fs::File,
    path::{Path, PathBuf},
    sync::Arc,
};

/// SafeTensors element type, independent of the parser crate's public API.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
#[allow(missing_docs)]
#[allow(non_camel_case_types)]
pub enum DType {
    BOOL,
    F4,
    F6_E2M3,
    F6_E3M2,
    U8,
    I8,
    F8_E5M2,
    F8_E4M3,
    F8_E8M0,
    F8_E4M3FNUZ,
    F8_E5M2FNUZ,
    I16,
    U16,
    F16,
    BF16,
    I32,
    U32,
    F32,
    C64,
    F64,
    I64,
    U64,
    /// A type introduced by a newer `safetensors` parser.
    Unknown,
}

impl From<NativeDType> for DType {
    fn from(value: NativeDType) -> Self {
        match value {
            NativeDType::BOOL => Self::BOOL,
            NativeDType::F4 => Self::F4,
            NativeDType::F6_E2M3 => Self::F6_E2M3,
            NativeDType::F6_E3M2 => Self::F6_E3M2,
            NativeDType::U8 => Self::U8,
            NativeDType::I8 => Self::I8,
            NativeDType::F8_E5M2 => Self::F8_E5M2,
            NativeDType::F8_E4M3 => Self::F8_E4M3,
            NativeDType::F8_E8M0 => Self::F8_E8M0,
            NativeDType::F8_E4M3FNUZ => Self::F8_E4M3FNUZ,
            NativeDType::F8_E5M2FNUZ => Self::F8_E5M2FNUZ,
            NativeDType::I16 => Self::I16,
            NativeDType::U16 => Self::U16,
            NativeDType::F16 => Self::F16,
            NativeDType::BF16 => Self::BF16,
            NativeDType::I32 => Self::I32,
            NativeDType::U32 => Self::U32,
            NativeDType::F32 => Self::F32,
            NativeDType::C64 => Self::C64,
            NativeDType::F64 => Self::F64,
            NativeDType::I64 => Self::I64,
            NativeDType::U64 => Self::U64,
            _ => Self::Unknown,
        }
    }
}

impl DType {
    /// Bytes per element for byte-addressable formats. Sub-byte formats return `None`.
    #[must_use]
    pub fn bytes(self) -> Option<usize> {
        match self {
            Self::F4 | Self::F6_E2M3 | Self::F6_E3M2 | Self::Unknown => None,
            Self::BOOL
            | Self::U8
            | Self::I8
            | Self::F8_E5M2
            | Self::F8_E4M3
            | Self::F8_E8M0
            | Self::F8_E4M3FNUZ
            | Self::F8_E5M2FNUZ => Some(1),
            Self::I16 | Self::U16 | Self::F16 | Self::BF16 => Some(2),
            Self::I32 | Self::U32 | Self::F32 => Some(4),
            Self::C64 | Self::F64 | Self::I64 | Self::U64 => Some(8),
        }
    }
}

/// One tensor's metadata and placement in the SafeTensors data block.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Entry {
    /// Element type.
    pub dtype: DType,
    /// Tensor dimensions.
    pub shape: Vec<usize>,
    /// Byte offset from the start of the data block.
    pub offset: usize,
    /// Tensor byte length.
    pub bytes: usize,
}

impl Entry {
    /// Number of elements, treating a scalar as one element.
    pub fn elements(&self) -> Result<usize> {
        self.shape.iter().try_fold(1usize, |count, &dimension| {
            count
                .checked_mul(dimension)
                .ok_or_else(|| Error::Message("SafeTensors element count overflow".into()))
        })
    }

    /// First dimension, treating a scalar as one row.
    #[must_use]
    pub fn rows(&self) -> usize {
        self.shape.first().copied().unwrap_or(1)
    }

    /// Bytes per first-dimension row.
    pub fn row_bytes(&self) -> Result<usize> {
        let rows = self.rows();
        if rows == 0 || !self.bytes.is_multiple_of(rows) {
            return Err(Error::Message(
                "tensor bytes are not a whole number of rows".into(),
            ));
        }
        Ok(self.bytes / rows)
    }
}

enum Storage {
    Owned(Vec<u8>),
    Mapped(Mmap),
}

impl Storage {
    fn bytes(&self) -> &[u8] {
        match self {
            Self::Owned(bytes) => bytes,
            Self::Mapped(map) => map,
        }
    }
}

/// An indexed SafeTensors file.
pub struct FileView {
    path: PathBuf,
    storage: Storage,
    file: Option<Arc<File>>,
    data: usize,
    entries: BTreeMap<String, Entry>,
}

/// Work performed by synchronous preparation of a tensor's backing pages.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PagePreparation {
    /// Empty range or owned bytes; no mapping preparation was needed.
    Resident,
    /// The operating system populated the readable mapping before returning.
    Populated,
    /// Explicit population is unsupported; advisory prefetch was requested.
    Advised,
}

/// A tensor borrowing its file storage.
#[derive(Clone, Copy, Debug)]
pub struct Tensor<'a> {
    /// Element type.
    pub dtype: DType,
    /// Dimensions.
    pub shape: &'a [usize],
    /// Original encoded bytes.
    pub bytes: &'a [u8],
}

impl FileView {
    /// Read a file into owned memory and index its tensors.
    pub fn read(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref().to_path_buf();
        let bytes = std::fs::read(&path)
            .map_err(|error| Error::Io(error).context(format!("reading {}", path.display())))?;
        Self::from_storage(path, Storage::Owned(bytes))
    }

    /// Map and index a file without copying its tensor data.
    ///
    /// # Safety
    ///
    /// The file must not be modified or truncated until the returned value is dropped.
    pub unsafe fn map(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref().to_path_buf();
        let file = File::open(&path)
            .map_err(|error| Error::Io(error).context(format!("opening {}", path.display())))?;
        let map = unsafe { Mmap::map(&file) }
            .map_err(|error| Error::Io(error).context(format!("mapping {}", path.display())))?;
        let mut view = Self::from_storage(path, Storage::Mapped(map))?;
        view.file = Some(Arc::new(file));
        Ok(view)
    }

    fn from_storage(path: PathBuf, storage: Storage) -> Result<Self> {
        let bytes = storage.bytes();
        let (header, metadata): (usize, Metadata) =
            SafeTensors::read_metadata(bytes).map_err(|error| {
                Error::Message(format!("{} is not SafeTensors: {error}", path.display()))
            })?;
        let data = 8usize
            .checked_add(header)
            .ok_or_else(|| Error::Message("SafeTensors header size overflow".into()))?;
        let data_bytes = bytes
            .len()
            .checked_sub(data)
            .ok_or_else(|| Error::Message("SafeTensors header exceeds the file".into()))?;
        let mut entries = BTreeMap::new();
        for (name, info) in metadata.tensors() {
            let (start, end) = info.data_offsets;
            if end < start || end > data_bytes {
                return Err(Error::Message(format!(
                    "tensor {name} exceeds the SafeTensors data block"
                )));
            }
            let entry = Entry {
                dtype: info.dtype.into(),
                shape: info.shape.clone(),
                offset: start,
                bytes: end - start,
            };
            if entries.insert(name.clone(), entry).is_some() {
                return Err(Error::Message(format!("duplicate tensor {name}")));
            }
        }
        Ok(Self {
            path,
            storage,
            file: None,
            data,
            entries,
        })
    }

    /// Source path.
    #[must_use]
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Indexed entries in lexical order.
    #[must_use]
    pub fn entries(&self) -> &BTreeMap<String, Entry> {
        &self.entries
    }

    /// Whether a tensor exists.
    #[must_use]
    pub fn contains(&self, name: &str) -> bool {
        self.entries.contains_key(name)
    }

    /// Tensor names in lexical order.
    pub fn names(&self) -> impl Iterator<Item = &str> {
        self.entries.keys().map(String::as_str)
    }

    /// Look up a tensor.
    pub fn get(&self, name: &str) -> Result<Tensor<'_>> {
        let entry = self.entries.get(name).ok_or_else(|| {
            Error::Message(format!("missing tensor {name} in {}", self.path.display()))
        })?;
        let start = self.data + entry.offset;
        let end = start + entry.bytes;
        Ok(Tensor {
            dtype: entry.dtype,
            shape: &entry.shape,
            bytes: &self.storage.bytes()[start..end],
        })
    }

    /// Borrow the data range described by an indexed entry.
    pub fn bytes(&self, entry: &Entry) -> Result<&[u8]> {
        let start = self
            .data
            .checked_add(entry.offset)
            .ok_or_else(|| Error::Message("SafeTensors offset overflow".into()))?;
        let end = start
            .checked_add(entry.bytes)
            .ok_or_else(|| Error::Message("SafeTensors range overflow".into()))?;
        self.storage
            .bytes()
            .get(start..end)
            .ok_or_else(|| Error::Message("SafeTensors entry exceeds the file".into()))
    }

    /// Descriptor of the mapped file, retained across pathname replacement.
    /// Owned snapshots have no backing descriptor.
    pub fn backing_file(&self) -> Option<&Arc<File>> {
        self.file.as_ref()
    }

    /// Absolute file offset of a checked tensor subrange, including its header.
    /// This validates metadata without touching mapped tensor pages. The returned
    /// descriptor identifies the original mapped inode, even after a rename.
    pub fn file_range(
        &self,
        entry: &Entry,
        offset: usize,
        bytes: usize,
    ) -> Result<(Arc<File>, u64)> {
        if offset
            .checked_add(bytes)
            .is_none_or(|end| end > entry.bytes)
        {
            return Err(Error::Message("tensor subrange exceeds entry".into()));
        }
        let start = self
            .data
            .checked_add(entry.offset)
            .and_then(|start| start.checked_add(offset))
            .ok_or_else(|| Error::Message("SafeTensors offset overflow".into()))?;
        if start
            .checked_add(bytes)
            .is_none_or(|end| end > self.storage.bytes().len())
        {
            return Err(Error::Message("SafeTensors entry exceeds file".into()));
        }
        let file = self.file.as_ref().ok_or_else(|| {
            Error::Unsupported("owned SafeTensors snapshots have no backing file".into())
        })?;
        Ok((file.clone(), start as u64))
    }

    /// Look up and validate a tensor. Negative expected dimensions are wildcards.
    pub fn get_checked(&self, name: &str, dtype: DType, shape: &[i64]) -> Result<Tensor<'_>> {
        let tensor = self.get(name)?;
        let shape_matches = tensor.shape.len() == shape.len()
            && tensor.shape.iter().zip(shape).all(|(&actual, &expected)| {
                expected < 0 || usize::try_from(expected) == Ok(actual)
            });
        if tensor.dtype != dtype || !shape_matches {
            return Err(Error::Message(format!(
                "{name} is {:?} {:?}, expected {dtype:?} {shape:?} in {}",
                tensor.dtype,
                tensor.shape,
                self.path.display()
            )));
        }
        Ok(tensor)
    }

    /// Ask the operating system to prefetch the mapped pages backing a tensor.
    pub fn will_need(&self, tensor: Tensor<'_>) {
        self.advise(tensor.bytes, libc::MADV_WILLNEED, true);
    }

    /// Release complete mapped pages after a tensor has been consumed.
    pub fn done_with(&self, tensor: Tensor<'_>) {
        self.advise(tensor.bytes, libc::MADV_DONTNEED, false);
    }

    /// Ask the operating system to prefetch a checked range borrowed from this file.
    pub fn will_need_bytes(&self, bytes: &[u8]) {
        self.advise(bytes, libc::MADV_WILLNEED, true);
    }

    /// Prepare a checked range before registration, conversion, or upload.
    ///
    /// Reads only this range (rounded to pages), with at most 16 workers and
    /// 16 MiB chunks. All workers finish before return, including on failure.
    /// Unsupported explicit population falls back to advisory prefetch. Other
    /// OS errors propagate; this does not copy weights or change their bytes.
    /// Preparation does not pin pages against later reclamation by the OS.
    pub fn prepare_bytes(&self, bytes: &[u8]) -> Result<PagePreparation> {
        if bytes.is_empty() {
            return Ok(PagePreparation::Resident);
        }
        let storage = self.storage.bytes();
        let begin = bytes.as_ptr() as usize;
        let base = storage.as_ptr() as usize;
        if begin < base
            || begin
                .checked_add(bytes.len())
                .is_none_or(|end| end > base + storage.len())
        {
            return Err(Error::Message(
                "page preparation range is outside its file".into(),
            ));
        }
        let Storage::Mapped(_) = &self.storage else {
            return Ok(PagePreparation::Resident);
        };
        let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
        if page <= 0 {
            return Err(Error::Io(std::io::Error::other("page size unavailable")));
        }
        let (begin, end) = advise_pages(
            (base, storage.len()),
            (begin, bytes.len()),
            page as usize,
            true,
        )
        .ok_or_else(|| Error::Message("invalid page preparation range".into()))?;
        Ok(if super::mapped_pages::populate(begin, end - begin)? {
            PagePreparation::Populated
        } else {
            PagePreparation::Advised
        })
    }

    /// Release complete mapped pages in a checked range borrowed from this file.
    pub fn done_with_bytes(&self, bytes: &[u8]) {
        self.advise(bytes, libc::MADV_DONTNEED, false);
    }

    fn advise(&self, range: &[u8], how: i32, outward: bool) {
        let Storage::Mapped(map) = &self.storage else {
            return;
        };
        if range.is_empty() {
            return;
        }
        let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
        if page <= 0 {
            return;
        }
        let Some((begin, end)) = advise_pages(
            (map.as_ptr() as usize, map.len()),
            (range.as_ptr() as usize, range.len()),
            page as usize,
            outward,
        ) else {
            return;
        };
        let _ = unsafe { libc::madvise(begin as *mut libc::c_void, end - begin, how) };
    }
}

fn advise_pages(
    map: (usize, usize),
    range: (usize, usize),
    page: usize,
    outward: bool,
) -> Option<(usize, usize)> {
    let map_end = map.0.checked_add(map.1)?;
    let range_end = range.0.checked_add(range.1)?;
    if page == 0 || !page.is_power_of_two() || range.0 < map.0 || range_end > map_end {
        return None;
    }
    let mask = page - 1;
    let (begin, end) = if outward {
        (range.0 & !mask, range_end.checked_add(mask)? & !mask)
    } else {
        (range.0.checked_add(mask)? & !mask, range_end & !mask)
    };
    (end > begin).then_some((begin, end))
}

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

    #[test]
    fn file_ranges_retain_original_identity_without_faulting_data() {
        use std::{io::Write, os::unix::fs::FileExt};
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("model.safetensors");
        let header = br#"{"x":{"dtype":"U8","shape":[4],"data_offsets":[0,4]}}"#;
        let mut file = File::create(&path).unwrap();
        file.write_all(&(header.len() as u64).to_le_bytes())
            .unwrap();
        file.write_all(header).unwrap();
        file.write_all(&[1, 2, 3, 4]).unwrap();
        drop(file);
        // SAFETY: the original inode remains immutable, including after rename.
        let mapped = unsafe { FileView::map(&path) }.unwrap();
        let owned = FileView::read(&path).unwrap();
        let entry = &mapped.entries()["x"];
        let (file, offset) = mapped.file_range(entry, 1, 2).unwrap();
        assert_eq!(offset, header.len() as u64 + 9);
        assert!(mapped.file_range(entry, 3, 2).is_err());
        assert!(mapped.file_range(entry, usize::MAX, 1).is_err());
        assert!(owned.file_range(entry, 0, 4).is_err());
        std::fs::rename(&path, dir.path().join("old")).unwrap();
        std::fs::write(&path, b"replacement").unwrap();
        drop(mapped);
        let mut bytes = [0; 2];
        file.read_exact_at(&mut bytes, offset).unwrap();
        assert_eq!(bytes, [2, 3]);
    }

    #[test]
    fn preparation_checks_ownership_and_preserves_mapped_bytes() {
        use std::io::Write;
        let mut file = tempfile::NamedTempFile::new().unwrap();
        let header = br#"{"x":{"dtype":"U8","shape":[8193],"data_offsets":[0,8193]}}"#;
        file.write_all(&(header.len() as u64).to_le_bytes())
            .unwrap();
        file.write_all(header).unwrap();
        let expected: Vec<u8> = (0..8193).map(|i| (i % 251) as u8).collect();
        file.write_all(&expected).unwrap();
        // This test exclusively owns the file and never changes it after mapping.
        let mapped = unsafe { FileView::map(file.path()) }.unwrap();
        let bytes = mapped.get("x").unwrap().bytes;
        assert!(matches!(
            mapped.prepare_bytes(&bytes[1..]),
            Ok(PagePreparation::Populated | PagePreparation::Advised)
        ));
        assert_eq!(bytes, expected);
        assert_eq!(
            mapped.prepare_bytes(&[]).unwrap(),
            PagePreparation::Resident
        );
        assert!(mapped.prepare_bytes(&expected).is_err());
        let owned = FileView::read(file.path()).unwrap();
        assert_eq!(
            owned.prepare_bytes(owned.get("x").unwrap().bytes).unwrap(),
            PagePreparation::Resident
        );
        assert!(owned.prepare_bytes(&expected).is_err());
    }

    #[test]
    fn page_advice_rounding_stays_inside_for_discard() {
        assert_eq!(
            advise_pages((0x1000, 0x4000), (0x1800, 0x2100), 0x1000, false),
            Some((0x2000, 0x3000))
        );
        assert!(advise_pages((0x1000, 0x1000), (0, 16), 0x1000, false).is_none());
    }
}