fuel-formats 0.20.71

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Safetensors support — transport convenience over the upstream
//! parser.
//!
//! Unlike the other modules in `fuel-formats`, the byte-level
//! safetensors parser is **not implemented here**. The upstream
//! [`safetensors`] crate already provides a complete, transport-
//! independent parser ([`SafeTensors::deserialize`] over `&[u8]`,
//! [`TensorView`], [`Dtype`], [`View`]). Re-implementing it would be
//! wasted effort.
//!
//! What this module *does* provide:
//!
//! - **Re-exports** of the parser surface so transport-only consumers
//!   can depend on `fuel-formats` instead of pulling in upstream
//!   `safetensors` directly.
//! - [`MmapedFile`] — a low-level memory-mapped file handle that
//!   produces a [`SafeTensors`] view via [`MmapedFile::deserialize`].
//!   Transport-independent in the same sense as the rest of
//!   `fuel-formats`: it owns the bytes, and the parser surface
//!   above sits on top.
//!
//! The mmap/view wrappers `fuel_loaders::safetensors::MmapedSafetensors`
//! and `fuel_loaders::safetensors::BufferedSafetensors` sit downstream in
//! `fuel-loaders`. They hand out `TensorView`s and nothing more — the
//! caller decodes the bytes.
//!
//! B6 deleted the eager tensor-construction half that used to live
//! alongside them (`load`, `load_buffer`, `save`, `SliceSafetensors`,
//! and the `Load` trait's `.load(device)`), because each called
//! `Tensor::from_*` or `Storage::*` on the retired eager `Tensor`.
//! Nothing replaced them: the lazy loaders read `view.data()` directly.

use std::path::{Path, PathBuf};

use fuel_ir::{Error, Result};

pub use safetensors::SafeTensors;
pub use safetensors::tensor::{Dtype, TensorView, View};

/// A low-level memory-mapped file handle for a safetensors file.
///
/// This does not eagerly deserialize the header. Call
/// [`MmapedFile::deserialize`] to obtain a [`SafeTensors`] view that
/// can be iterated or queried for individual tensors.
///
/// # Safety
///
/// Construction is `unsafe` because it relies on memory-mapped I/O.
/// The underlying file must not be modified while this handle is
/// alive.
pub struct MmapedFile {
    path: PathBuf,
    inner: memmap2::Mmap,
}

impl MmapedFile {
    /// Open `p` and memory-map its contents.
    ///
    /// # Safety
    ///
    /// Inherits the safety contract of [`memmap2::MmapOptions`]: the
    /// caller must ensure the file is not concurrently modified or
    /// truncated for the lifetime of the returned handle.
    pub unsafe fn new<P: AsRef<Path>>(p: P) -> Result<Self> {
        let p = p.as_ref();
        let file = std::fs::File::open(p).map_err(|e| Error::from(e).with_path(p))?;
        let inner = unsafe {
            memmap2::MmapOptions::new()
                .map(&file)
                .map_err(|e| Error::from(e).with_path(p))?
        };
        Ok(Self {
            inner,
            path: p.to_path_buf(),
        })
    }

    /// Deserialize the safetensors header and return a borrowed
    /// [`SafeTensors`] view.
    ///
    /// The returned view borrows from the memory-mapped region and
    /// can be used to iterate over tensor names or load individual
    /// tensor views.
    pub fn deserialize(&self) -> Result<SafeTensors<'_>> {
        let st = SafeTensors::deserialize(&self.inner)
            .map_err(|e| Error::from(e).with_path(&self.path))?;
        Ok(st)
    }

    /// Path this handle was opened from.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Borrow the underlying mmapped bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.inner
    }
}