liteboxfs 0.1.0

use crate::block::MerkleNodeStore;

use super::{
    store::{BlockStore, BlockStoreInput, ReadBlock},
    types::{BlockEncoding, BlockId, BlockIndex, BlockList, FileId},
};

pub trait Encode {
    const CODE: BlockEncoding;

    fn encode(&self, data: &[u8], out: &mut Vec<u8>) -> crate::Result<()>;
    fn decode(&self, data: &[u8], out: &mut Vec<u8>) -> crate::Result<()>;
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EncodingMode {
    /// Encode writes and decode reads.
    EncodeDecode,

    /// Pass writes through unchanged and decode reads based on the per-block encoding flag.
    ///
    /// This is used when compression is disabled, but existing blocks may have been written with
    /// compression enabled.
    DecodeOnly,
}

/// A block store which wraps another block store and compresses data blocks.
#[derive(Debug)]
pub struct EncodingBlockStore<Store, Encoder> {
    inner: Store,
    encoder: Encoder,
    mode: EncodingMode,
    buf: Vec<u8>,
}

impl<Store, Encoder> EncodingBlockStore<Store, Encoder> {
    pub fn new(inner: Store, encoder: Encoder, mode: EncodingMode) -> Self {
        Self {
            inner,
            encoder,
            mode,
            buf: Vec::new(),
        }
    }
}

impl<Store, Encoder> BlockStore for EncodingBlockStore<Store, Encoder>
where
    Store: BlockStore,
    Encoder: Encode,
{
    fn list_blocks(&self, file: FileId) -> crate::Result<BlockList> {
        self.inner.list_blocks(file)
    }

    fn read_block(&mut self, block: super::BlockId, buf: &mut Vec<u8>) -> crate::Result<ReadBlock> {
        self.buf.clear();
        let read_block = self.inner.read_block(block, &mut self.buf)?;

        if read_block.encoding.contains(Encoder::CODE) {
            self.encoder.decode(&self.buf, buf)?;
        } else {
            buf.extend_from_slice(&self.buf);
        }

        Ok(read_block)
    }

    fn replace_block(
        &mut self,
        file: super::FileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        if self.mode == EncodingMode::DecodeOnly {
            return self.inner.replace_block(file, index, input);
        }

        let encoded = match input {
            BlockStoreInput::Data { bytes, hash, .. } => {
                let logical_len = bytes.len();
                self.buf.clear();
                self.encoder.encode(bytes, &mut self.buf)?;

                BlockStoreInput::Data {
                    bytes: &self.buf,
                    hash,
                    encoding: Encoder::CODE,
                    len: Some(logical_len),
                }
            }

            hole @ BlockStoreInput::Hole { .. } => hole,
        };

        self.inner.replace_block(file, index, encoded)
    }

    fn insert_block(
        &mut self,
        file: FileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        if self.mode == EncodingMode::DecodeOnly {
            return self.inner.insert_block(file, index, input);
        }

        let encoded = match input {
            BlockStoreInput::Data { bytes, hash, .. } => {
                let logical_len = bytes.len();
                self.buf.clear();
                self.encoder.encode(bytes, &mut self.buf)?;

                BlockStoreInput::Data {
                    bytes: &self.buf,
                    hash,
                    encoding: Encoder::CODE,
                    len: Some(logical_len),
                }
            }

            hole @ BlockStoreInput::Hole { .. } => hole,
        };

        self.inner.insert_block(file, index, encoded)
    }

    fn append_block(
        &mut self,
        file: FileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        if self.mode == EncodingMode::DecodeOnly {
            return self.inner.append_block(file, index, input);
        }

        let encoded = match input {
            BlockStoreInput::Data { bytes, hash, .. } => {
                let logical_len = bytes.len();
                self.buf.clear();
                self.encoder.encode(bytes, &mut self.buf)?;

                BlockStoreInput::Data {
                    bytes: &self.buf,
                    hash,
                    encoding: Encoder::CODE,
                    len: Some(logical_len),
                }
            }

            hole @ BlockStoreInput::Hole { .. } => hole,
        };

        self.inner.append_block(file, index, encoded)
    }

    fn remove_blocks<R: std::ops::RangeBounds<usize>>(
        &mut self,
        file: FileId,
        blocks: R,
    ) -> crate::Result<Vec<BlockId>> {
        self.inner.remove_blocks(file, blocks)
    }
}

#[cfg_attr(coverage_nightly, coverage(off))]
impl<Store, Encoder> MerkleNodeStore for EncodingBlockStore<Store, Encoder>
where
    Store: MerkleNodeStore,
{
    fn get_root(&mut self, file: FileId) -> crate::Result<Option<super::MerkleNode>> {
        self.inner.get_root(file)
    }

    fn node_indices(
        &mut self,
        file: FileId,
        level: u32,
        filter: super::DirtyFilter,
    ) -> crate::Result<Vec<usize>> {
        self.inner.node_indices(file, level, filter)
    }

    fn list_nodes<R>(
        &mut self,
        file: FileId,
        level: u32,
        indices: R,
    ) -> crate::Result<Vec<super::MerkleNode>>
    where
        R: std::ops::RangeBounds<usize>,
    {
        self.inner.list_nodes(file, level, indices)
    }

    fn put_node(
        &mut self,
        pos: super::MerkleNodePosition,
        hash: crate::hash::MerkleHash,
    ) -> crate::Result<()> {
        self.inner.put_node(pos, hash)
    }

    fn delete_node(&mut self, file: FileId, index: usize) -> crate::Result<()> {
        self.inner.delete_node(file, index)
    }

    fn mark_dirty<R>(&mut self, file: FileId, indices: R) -> crate::Result<()>
    where
        R: std::ops::RangeBounds<usize>,
    {
        self.inner.mark_dirty(file, indices)
    }
}