liteboxfs 0.2.0

use std::fmt::{self, Debug};

use crate::chunker::ChunkerGuard;

use super::{
    buffer::{EjectBuf, FullSliceBuf, PartialSliceBuf, SliceBuf},
    operation::WriteOperation,
    store::{BlockStore, BlockStoreInput},
    types::{BlockIndex, BlockOffset, BlockSignature, HoleLen, IndexedBlock},
};

/// An owned fragment produced by splitting a block at a splice boundary in the partial path.
///
/// Tracking the fragment's kind (data vs hole) lets us preserve sparsity when splicing through
/// existing hole blocks instead of materializing their implicit zeros.
enum Fragment {
    Data(Vec<u8>),
    Hole(HoleLen),
}

impl Fragment {
    fn is_empty(&self) -> bool {
        match self {
            Fragment::Data(bytes) => bytes.is_empty(),
            Fragment::Hole(len) => *len == 0,
        }
    }

    fn as_input(&self) -> BlockStoreInput<'_> {
        match self {
            Fragment::Data(bytes) => bytes.as_slice().into(),
            Fragment::Hole(len) => BlockStoreInput::Hole { len: *len },
        }
    }
}

/// Build the `before` and `after` fragments produced by splitting a block at the given offset.
///
/// Hole blocks are split logically into smaller holes; data blocks are split as byte slices read
/// from the partial slice buffer.
fn split_block(
    signature: &BlockSignature,
    data: Option<&[u8]>,
    split_offset: BlockOffset,
) -> (Fragment, Fragment) {
    match signature {
        BlockSignature::Data { .. } => {
            let bytes = data.expect("Partial slice buffer unexpectedly missing data bytes.");
            let before = bytes
                .get(0..split_offset)
                .expect("Unexpected out of bounds when splitting a data block.")
                .to_vec();
            let after = bytes
                .get(split_offset..)
                .expect("Unexpected out of bounds when splitting a data block.")
                .to_vec();
            (Fragment::Data(before), Fragment::Data(after))
        }
        BlockSignature::Hole { len } => (
            Fragment::Hole(split_offset as HoleLen),
            Fragment::Hole((*len as usize - split_offset) as HoleLen),
        ),
    }
}

/// An abstraction representing mutations that can happen on a slice of blocks loaded in memory.
pub struct LoadedSlice<'a, Store, Buf> {
    op: WriteOperation<'a, Store>,
    chunker: &'a mut ChunkerGuard,
    blocks: Vec<IndexedBlock>,
    buf: Buf,
}

#[cfg_attr(coverage_nightly, coverage(off))]
impl<'a, Store, Bug> Debug for LoadedSlice<'a, Store, Bug>
where
    Store: Debug,
    Bug: Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LoadedSlice")
            .field("op", &self.op)
            .field("blocks", &self.blocks)
            .field("buf", &self.buf)
            .finish()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SliceBound {
    SliceStart,
    SliceEnd,
    BlockOffset {
        index: BlockIndex,
        offset: BlockOffset,
    },
}

impl SliceBound {
    pub fn block_index(&self, blocks: &[IndexedBlock]) -> BlockIndex {
        match self {
            SliceBound::SliceStart => 0,
            SliceBound::SliceEnd => blocks
                .iter()
                .next_back()
                .map(|block| block.index)
                .unwrap_or_default(),
            SliceBound::BlockOffset {
                index: block_index, ..
            } => *block_index,
        }
    }

    pub fn absolute_offset(&self, blocks: &[IndexedBlock]) -> BlockOffset {
        match self {
            SliceBound::SliceStart => 0,
            SliceBound::SliceEnd => blocks.iter().map(|b| b.block.len()).sum(),
            SliceBound::BlockOffset {
                index: block_index,
                offset: offset_in_block,
            } => blocks.iter().fold(0, |total, block| {
                if block.index < *block_index {
                    total + block.block.len()
                } else if block.index == *block_index {
                    total + *offset_in_block
                } else {
                    total
                }
            }),
        }
    }

    pub fn offset_in_block(&self, blocks: &[IndexedBlock]) -> BlockOffset {
        let last_block = blocks.iter().next_back();

        match self {
            SliceBound::BlockOffset { offset, .. } => *offset,
            SliceBound::SliceStart => 0,
            SliceBound::SliceEnd => match last_block {
                Some(block) => block.block.len(),
                None => 0,
            },
        }
    }
}

impl<'a, Store, Buf> LoadedSlice<'a, Store, Buf>
where
    Store: BlockStore,
    Buf: EjectBuf,
{
    pub fn new(
        op: WriteOperation<'a, Store>,
        chunker: &'a mut ChunkerGuard,
        blocks: Vec<IndexedBlock>,
        buf: Buf,
    ) -> Self {
        Self {
            op,
            chunker,
            blocks,
            buf,
        }
    }

    /// Return the buffer back to the write operation and return it.
    fn into_operation(self) -> WriteOperation<'a, Store> {
        self.op.with_buf(self.buf.eject())
    }
}

impl<'a, Store> LoadedSlice<'a, Store, FullSliceBuf>
where
    Store: BlockStore,
{
    pub fn into_slice_buf(self) -> LoadedSlice<'a, Store, SliceBuf> {
        LoadedSlice {
            op: self.op,
            chunker: self.chunker,
            blocks: self.blocks,
            buf: SliceBuf::Full(self.buf),
        }
    }

    // Splice the new data/hole into this slice and re-chunk the entire slice.
    fn splice(
        mut self,
        start: SliceBound,
        end: SliceBound,
        new: BlockStoreInput,
    ) -> crate::Result<WriteOperation<'a, Store>> {
        let start_offset = start.absolute_offset(&self.blocks);
        let end_offset = end.absolute_offset(&self.blocks);

        let start_block_index = start.block_index(&self.blocks);
        let end_block_index = end.block_index(&self.blocks);

        match new {
            BlockStoreInput::Data { bytes, .. } => {
                self.chunker.with_rewind(|chunker| {
                    self.buf
                        .splice(start_offset..end_offset, bytes.iter().copied());

                    chunker.push_bytes(self.buf.as_ref());

                    let mut i = 0usize;

                    while let Some(chunk) = chunker.next_chunk() {
                        let input = chunk.as_ref().into();

                        if i == 0 {
                            self.op
                                .replace_blocks(start_block_index..=end_block_index, input)?;
                        } else {
                            self.op.insert_block(start_block_index + i, input)?;
                        }

                        i += 1;
                    }

                    let input = chunker.remaining().into();

                    if i == 0 {
                        self.op
                            .replace_blocks(start_block_index..=end_block_index, input)?;
                    } else {
                        self.op.insert_block(start_block_index + i, input)?;
                    }

                    crate::Result::Ok(())
                })?;
            }
            BlockStoreInput::Hole { len } => {
                let before_split = self
                    .buf
                    .as_ref()
                    .get(0..start_offset)
                    .expect("Unexpected out of bounds when splicing a slice.");

                let after_split = self
                    .buf
                    .as_ref()
                    .get(end_offset..)
                    .expect("Unexpected out of bounds when splicing a slice.");

                // If splicing in this hole involved splitting blocks, those block fragments remain
                // as-is and are not re-chunked or joined with other blocks.
                self.op
                    .replace_blocks(start_block_index..=end_block_index, before_split.into())?;
                self.op
                    .insert_block(start_block_index + 1, BlockStoreInput::Hole { len })?;
                self.op
                    .insert_block(start_block_index + 2, after_split.into())?;
            }
        }

        Ok(self.into_operation())
    }
}

impl<'a, Store> LoadedSlice<'a, Store, PartialSliceBuf>
where
    Store: BlockStore,
{
    pub fn into_slice_buf(self) -> LoadedSlice<'a, Store, SliceBuf> {
        LoadedSlice {
            op: self.op,
            chunker: self.chunker,
            blocks: self.blocks,
            buf: SliceBuf::Partial(self.buf),
        }
    }

    // Rather than re-chunking the entire slice, we insert the new data/hole as a standalone chunk.
    // This splits the first and last blocks as necessary, leaving them as fragments and replacing
    // any blocks in between. This is an optimization for the case where we're only writing a small
    // number of bytes and want to avoid re-chunking and re-hashing the entire slice. We refer to
    // this as the "small write optimization".
    //
    // Hole blocks at the first/last positions are split logically (without materializing their
    // implicit zeros), so writing into a sparse region preserves sparsity on either side of the
    // write.
    fn splice(
        mut self,
        start: SliceBound,
        end: SliceBound,
        new: BlockStoreInput,
    ) -> crate::Result<WriteOperation<'a, Store>> {
        let start_block_index = start.block_index(&self.blocks);
        let end_block_index = end.block_index(&self.blocks);

        let start_block_offset = start.offset_in_block(&self.blocks);
        let end_block_offset = end.offset_in_block(&self.blocks);

        let first_signature = self
            .buf
            .first_block_signature()
            .expect("Partial slice buffer unexpectedly empty.")
            .clone();
        let last_signature = self
            .buf
            .last_block_signature()
            .expect("Partial slice buffer unexpectedly empty.")
            .clone();

        if start_block_index == end_block_index {
            // The write is contained entirely within a single block.
            //
            // We need to split it into: [fragment_before, new_data, fragment_after] but only
            // create non-empty fragments.
            let (fragment_before, _) = split_block(
                &first_signature,
                self.buf.first_block_data(),
                start_block_offset,
            );
            let (_, fragment_after) = split_block(
                &last_signature,
                self.buf.last_block_data(),
                end_block_offset,
            );

            let has_fragment_before = !fragment_before.is_empty();
            let has_fragment_after = !fragment_after.is_empty();

            match (has_fragment_before, has_fragment_after) {
                (true, true) => {
                    // Both fragments exist: [fragment_before, new_data, fragment_after].
                    self.op.replace_blocks(
                        start_block_index..=start_block_index,
                        fragment_before.as_input(),
                    )?;
                    self.op.insert_block(start_block_index + 1, new)?;
                    self.op
                        .insert_block(start_block_index + 2, fragment_after.as_input())?;
                }
                (true, false) => {
                    // Only the fragment before exists: [fragment_before, new_data].
                    self.op.replace_blocks(
                        start_block_index..=start_block_index,
                        fragment_before.as_input(),
                    )?;
                    self.op.insert_block(start_block_index + 1, new)?;
                }
                (false, true) => {
                    // Only the fragment after exists: [new_data, fragment_after].
                    self.op
                        .replace_blocks(start_block_index..=start_block_index, new)?;
                    self.op
                        .insert_block(start_block_index + 1, fragment_after.as_input())?;
                }
                (false, false) => {
                    // No fragments (entire block replaced): [new_data]
                    self.op
                        .replace_blocks(start_block_index..=start_block_index, new)?;
                }
            }
        } else {
            // The write spans multiple blocks.
            let (start_fragment, _) = split_block(
                &first_signature,
                self.buf.first_block_data(),
                start_block_offset,
            );
            let (_, end_fragment) = split_block(
                &last_signature,
                self.buf.last_block_data(),
                end_block_offset,
            );

            let has_start_fragment = !start_fragment.is_empty();
            let has_end_fragment = !end_fragment.is_empty();

            // Create fragments only if they're non-empty.
            if has_start_fragment {
                self.op.replace_blocks(
                    start_block_index..=start_block_index,
                    start_fragment.as_input(),
                )?;
            }

            if has_end_fragment {
                self.op
                    .replace_blocks(end_block_index..=end_block_index, end_fragment.as_input())?;
            }

            // Determine the range of blocks to replace with new data.
            let new_data_start = if has_start_fragment {
                start_block_index + 1
            } else {
                start_block_index
            };

            let new_data_end = if has_end_fragment {
                end_block_index
            } else {
                end_block_index + 1
            };

            // Replace the blocks in the determined range with new data.
            self.op.replace_blocks(new_data_start..new_data_end, new)?;
        }

        Ok(self.into_operation())
    }
}

impl<'a, Store> LoadedSlice<'a, Store, SliceBuf>
where
    Store: BlockStore,
{
    /// Replace a range of bytes in this slice with new data or a hole.
    pub fn splice(
        self,
        start: SliceBound,
        end: SliceBound,
        new: BlockStoreInput,
    ) -> crate::Result<WriteOperation<'a, Store>> {
        if (start == SliceBound::SliceEnd && end != SliceBound::SliceEnd)
            || (end == SliceBound::SliceStart && start != SliceBound::SliceStart)
        {
            panic!("Invalid splice bounds: start is after end.");
        }

        match self.buf {
            SliceBuf::Full(buf) => {
                let loaded = LoadedSlice {
                    op: self.op,
                    chunker: self.chunker,
                    blocks: self.blocks,
                    buf,
                };

                loaded.splice(start, end, new)
            }
            SliceBuf::Partial(buf) => {
                let loaded = LoadedSlice {
                    op: self.op,
                    chunker: self.chunker,
                    blocks: self.blocks,
                    buf,
                };

                loaded.splice(start, end, new)
            }
        }
    }
}

/// An abstraction representing mutations that can happen with a block loaded in memory.
#[derive(Debug)]
pub struct LoadedBlock<'a, Store> {
    operation: WriteOperation<'a, Store>,
    signature: BlockSignature,
    buf: Vec<u8>,
}

impl<'a, Store> LoadedBlock<'a, Store>
where
    Store: BlockStore,
{
    pub fn new(
        operation: WriteOperation<'a, Store>,
        signature: BlockSignature,
        buf: Vec<u8>,
    ) -> Self {
        Self {
            operation,
            signature,
            buf,
        }
    }

    /// Return the buffer back to the writer operation and return it.
    fn into_operation(self) -> WriteOperation<'a, Store> {
        self.operation.with_buf(self.buf)
    }

    /// Truncate the loaded block at the given offset, replacing the block at the given index with
    /// the truncated data and removing any subsequent blocks.
    pub fn truncate(
        mut self,
        index: usize,
        offset: BlockOffset,
    ) -> crate::Result<WriteOperation<'a, Store>> {
        match self.signature {
            BlockSignature::Data { .. } => {
                let bytes_before_split = self
                    .buf
                    .get(0..offset)
                    .expect("Unexpected out of bounds when truncating a block.");

                self.operation
                    .replace_blocks(index..=index, bytes_before_split.into())?;
            }
            BlockSignature::Hole { .. } => {
                self.operation.replace_blocks(
                    index..=index,
                    BlockStoreInput::Hole {
                        len: offset as HoleLen,
                    },
                )?;
            }
        };

        self.operation
            .remove_blocks((index + 1)..self.operation.len())?;

        Ok(self.into_operation())
    }
}