liteboxfs 0.2.0

pub type Chunk<'a> = &'a [u8];

/// A chunker for slicing byte streams into chunks.
pub trait ChunkSlice {
    /// Pushes bytes into the chunker's internal buffer.
    fn push_bytes(&mut self, bytes: &[u8]);

    /// Returns the next chunk from the internal buffer, if one is available.
    fn next_chunk(&mut self) -> Option<Chunk<'_>>;

    /// Returns the number of bytes currently buffered in the chunker.
    fn buffer_size(&mut self) -> usize;

    /// Returns any remaining bytes in the internal buffer that do not form a complete chunk.
    fn remaining(&mut self) -> Chunk<'_>;

    /// Returns the next chunk if available, otherwise returns any remaining bytes in the internal
    /// buffer. If the internal buffer is empty, return `None`.
    ///
    /// You would think this could have a default implementation, but the borrow checker disagrees.
    fn next_chunk_or_remaining(&mut self) -> Option<Chunk<'_>>;
}

/// A chunker that does not reclaim space in its internal buffer unless told to do so.
pub trait VacuumChunker {
    /// Reclaims space in the internal buffer.
    fn vacuum(&mut self);
}

/// A chunker whose state can be reset back to a previous state.
pub trait RewindChunker {
    type State;

    fn dump_state(&self) -> Self::State;
    fn rewind_to(&mut self, state: Self::State);
}

pub trait Chunker: ChunkSlice + VacuumChunker + RewindChunker<State = ChunkerState> {}

/// A guard around a chunker that enforces rewind-on-error semantics.
///
/// Because the chunker state lives in memory, it doesn't roll back automatically when the database
/// transaction rolls back. This wrappers allows us to easily rewind the chunker state on error.
pub struct ChunkerGuard {
    inner: Box<dyn Chunker>,
}

impl ChunkerGuard {
    pub fn new(inner: Box<dyn Chunker>) -> Self {
        Self { inner }
    }

    pub fn with_rewind<R, E, F>(&mut self, func: F) -> Result<R, E>
    where
        F: FnOnce(&mut dyn ChunkSlice) -> Result<R, E>,
    {
        let state = self.inner.dump_state();

        match func(self.inner.as_mut()) {
            Ok(value) => {
                self.inner.vacuum();
                Ok(value)
            }
            Err(err) => {
                self.inner.rewind_to(state);
                Err(err)
            }
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ChunkerState {
    pub pos: usize,
    pub buf_len: usize,
}

#[cfg(test)]
#[derive(Debug, Default)]
pub struct NopChunker {
    buf: Vec<u8>,
    pos: usize,
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
impl NopChunker {
    pub fn new() -> Self {
        Self::default()
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
impl ChunkSlice for NopChunker {
    fn push_bytes(&mut self, bytes: &[u8]) {
        self.buf.extend_from_slice(bytes);
    }

    fn next_chunk(&mut self) -> Option<Chunk<'_>> {
        None
    }

    fn buffer_size(&mut self) -> usize {
        0
    }

    fn remaining(&mut self) -> Chunk<'_> {
        let chunk = &self
            .buf
            .get(self.pos..)
            .expect("Cursor unexpectedly out of bounds.");
        self.pos = self.buf.len();
        chunk
    }

    fn next_chunk_or_remaining(&mut self) -> Option<Chunk<'_>> {
        let remaining = self.remaining();

        if remaining.is_empty() {
            None
        } else {
            Some(remaining)
        }
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
impl VacuumChunker for NopChunker {
    fn vacuum(&mut self) {
        self.buf.drain(0..self.pos);
        self.pos = 0;
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
impl RewindChunker for NopChunker {
    type State = ChunkerState;

    fn dump_state(&self) -> Self::State {
        ChunkerState {
            pos: self.pos,
            buf_len: self.buf.len(),
        }
    }

    fn rewind_to(&mut self, state: Self::State) {
        self.pos = state.pos;
        self.buf.truncate(state.buf_len);
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
impl Chunker for NopChunker {}

#[derive(Debug)]
pub struct FixedSizeChunker {
    buf: Vec<u8>,
    pos: usize,
    chunk_size: u32,
}

impl FixedSizeChunker {
    pub fn new(chunk_size: u32) -> Self {
        Self {
            buf: Vec::with_capacity(chunk_size as usize),
            pos: 0,
            chunk_size,
        }
    }
}

impl ChunkSlice for FixedSizeChunker {
    fn push_bytes(&mut self, bytes: &[u8]) {
        self.buf.extend_from_slice(bytes);
    }

    fn next_chunk(&mut self) -> Option<Chunk<'_>> {
        let remaining = self.buf.len() - self.pos;

        if self.buf.is_empty() || (remaining < self.chunk_size as usize) {
            return None;
        }

        let chunk = self
            .buf
            .get(self.pos..(self.pos + self.chunk_size as usize))
            .expect("Cursor unexpectedly out of bounds.");

        self.pos += chunk.len();

        Some(chunk)
    }

    fn buffer_size(&mut self) -> usize {
        self.buf.len() - self.pos
    }

    fn remaining(&mut self) -> Chunk<'_> {
        let chunk = &self
            .buf
            .get(self.pos..)
            .expect("Cursor unexpectedly out of bounds.");
        self.pos = self.buf.len();
        chunk
    }

    fn next_chunk_or_remaining(&mut self) -> Option<Chunk<'_>> {
        let remaining_len = self.buf.len() - self.pos;

        if remaining_len == 0 {
            return None;
        } else if remaining_len <= self.chunk_size as usize {
            return Some(self.remaining());
        }

        let chunk = self
            .buf
            .get(self.pos..(self.pos + self.chunk_size as usize))
            .expect("Cursor unexpectedly out of bounds.");

        self.pos += chunk.len();

        Some(chunk)
    }
}

impl VacuumChunker for FixedSizeChunker {
    fn vacuum(&mut self) {
        self.buf.drain(0..self.pos);
        self.pos = 0;
    }
}

impl RewindChunker for FixedSizeChunker {
    type State = ChunkerState;

    fn dump_state(&self) -> Self::State {
        ChunkerState {
            pos: self.pos,
            buf_len: self.buf.len(),
        }
    }

    fn rewind_to(&mut self, state: Self::State) {
        self.pos = state.pos;
        self.buf.truncate(state.buf_len);
    }
}

impl Chunker for FixedSizeChunker {}