heddle-object-model 0.15.2

// SPDX-License-Identifier: Apache-2.0
//! Byte sources for streamable Tree reads.

use std::{
    fs::File,
    io::{Read, Seek, SeekFrom},
};

use bytes::Bytes;

use super::tree_stream::TreeStreamError;

/// How a tree body may be trusted for content-address verification.
///
/// Sequential reads from offset 0 recompute the typed tree hash. Ranged
/// resume skips the prefix and is allowed only after the caller has hashed
/// these exact bytes. A hash-shaped path or an embedded tree id is not that
/// proof. Object-store backends use [`Self::SequentialVerify`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TreeBodyIntegrity {
    /// Caller hashed these exact bytes; prefix skip is allowed.
    VerifiedPlacement,
    /// The reader must start at entry 0 and call `finish_and_verify`.
    SequentialVerify,
}

/// Random-access source of an uncompressed canonical tree body.
pub trait TreeByteSource {
    fn read_exact_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<(), TreeStreamError>;
    fn len(&self) -> u64;
    fn integrity(&self) -> TreeBodyIntegrity;
    fn bytes_read(&self) -> u64;

    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// In-memory tree body, used by tests and stores that already hold the bytes.
#[derive(Debug)]
pub struct BytesTreeSource {
    bytes: Bytes,
    integrity: TreeBodyIntegrity,
    bytes_read: u64,
}

impl BytesTreeSource {
    pub fn verified_placement(bytes: impl Into<Bytes>) -> Self {
        Self {
            bytes: bytes.into(),
            integrity: TreeBodyIntegrity::VerifiedPlacement,
            bytes_read: 0,
        }
    }

    pub fn sequential_verify(bytes: impl Into<Bytes>) -> Self {
        Self {
            bytes: bytes.into(),
            integrity: TreeBodyIntegrity::SequentialVerify,
            bytes_read: 0,
        }
    }
}

impl TreeByteSource for BytesTreeSource {
    fn read_exact_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<(), TreeStreamError> {
        let start =
            usize::try_from(offset).map_err(|_| TreeStreamError::TruncatedFrame { offset })?;
        let end = start
            .checked_add(buf.len())
            .ok_or(TreeStreamError::TruncatedFrame { offset })?;
        let slice = self
            .bytes
            .get(start..end)
            .ok_or(TreeStreamError::TruncatedFrame { offset })?;
        buf.copy_from_slice(slice);
        self.bytes_read += buf.len() as u64;
        Ok(())
    }

    fn len(&self) -> u64 {
        self.bytes.len() as u64
    }

    fn integrity(&self) -> TreeBodyIntegrity {
        self.integrity
    }

    fn bytes_read(&self) -> u64 {
        self.bytes_read
    }
}

/// File-backed tree body. Sequential reads seek per frame; resume at
/// ordinal > 0 is refused because a path is not a content-address proof.
#[derive(Debug)]
pub struct FileTreeSource {
    file: File,
    len: u64,
    integrity: TreeBodyIntegrity,
    bytes_read: u64,
}

impl FileTreeSource {
    pub fn sequential_verify(file: File, len: u64) -> Self {
        Self {
            file,
            len,
            integrity: TreeBodyIntegrity::SequentialVerify,
            bytes_read: 0,
        }
    }
}

impl TreeByteSource for FileTreeSource {
    fn read_exact_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<(), TreeStreamError> {
        if offset
            .checked_add(buf.len() as u64)
            .is_none_or(|end| end > self.len)
        {
            return Err(TreeStreamError::TruncatedFrame { offset });
        }
        self.file.seek(SeekFrom::Start(offset))?;
        self.file.read_exact(buf)?;
        self.bytes_read += buf.len() as u64;
        Ok(())
    }

    fn len(&self) -> u64 {
        self.len
    }

    fn integrity(&self) -> TreeBodyIntegrity {
        self.integrity
    }

    fn bytes_read(&self) -> u64 {
        self.bytes_read
    }
}

/// Store-facing handle for a seekable or lazily generated tree body.
pub enum OpenedTreeBody {
    Bytes(BytesTreeSource),
    File(FileTreeSource),
    Dynamic(Box<dyn TreeByteSource + Send>),
}

impl std::fmt::Debug for OpenedTreeBody {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Bytes(source) => formatter.debug_tuple("Bytes").field(source).finish(),
            Self::File(source) => formatter.debug_tuple("File").field(source).finish(),
            Self::Dynamic(_) => formatter.write_str("Dynamic(..)"),
        }
    }
}

impl TreeByteSource for OpenedTreeBody {
    fn read_exact_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<(), TreeStreamError> {
        match self {
            Self::Bytes(source) => source.read_exact_at(offset, buf),
            Self::File(source) => source.read_exact_at(offset, buf),
            Self::Dynamic(source) => source.read_exact_at(offset, buf),
        }
    }

    fn len(&self) -> u64 {
        match self {
            Self::Bytes(source) => source.len(),
            Self::File(source) => source.len(),
            Self::Dynamic(source) => source.len(),
        }
    }

    fn integrity(&self) -> TreeBodyIntegrity {
        match self {
            Self::Bytes(source) => source.integrity(),
            Self::File(source) => source.integrity(),
            Self::Dynamic(source) => source.integrity(),
        }
    }

    fn bytes_read(&self) -> u64 {
        match self {
            Self::Bytes(source) => source.bytes_read(),
            Self::File(source) => source.bytes_read(),
            Self::Dynamic(source) => source.bytes_read(),
        }
    }
}