liteboxfs 0.2.0

use std::{collections::HashMap, ops::RangeBounds};

use super::id_pool::IdPool;
use crate::block::{
    Block, BlockData, BlockEncoding, BlockId, BlockIndex, BlockList, BlockSignature, BlockStore,
    BlockStoreInput, FileId, ReadBlock,
};
use crate::errors::InternalError;

#[derive(Debug)]
struct MemoryBlockStoreBlock {
    data: BlockData,
    signature: BlockSignature,
    encoding: BlockEncoding,
}

/// An in-memory block store implementation.
///
/// In practice, the block store is backed by a SQLite database. This in-memory implementation is
/// primarily intended for testing.
#[derive(Debug)]
pub struct MemoryBlockStore {
    files: HashMap<FileId, Vec<BlockId>>,
    blocks: HashMap<BlockId, MemoryBlockStoreBlock>,
    block_id_pool: IdPool,
}

impl MemoryBlockStore {
    pub fn new() -> Self {
        Self {
            files: HashMap::new(),
            blocks: HashMap::new(),
            block_id_pool: IdPool::new(),
        }
    }

    fn input_to_block(input: BlockStoreInput) -> MemoryBlockStoreBlock {
        let encoding = input.encoding();

        MemoryBlockStoreBlock {
            data: match input {
                BlockStoreInput::Hole { len } => BlockData::Hole { len },
                BlockStoreInput::Data { bytes, .. } => BlockData::Data {
                    bytes: bytes.to_vec(),
                },
            },
            signature: match input {
                BlockStoreInput::Data {
                    bytes,
                    hash,
                    len: logical_len,
                    ..
                } => BlockSignature::Data {
                    hash,
                    len: logical_len.unwrap_or(bytes.len()),
                },
                BlockStoreInput::Hole { len } => BlockSignature::Hole { len },
            },
            encoding,
        }
    }
}

impl BlockStore for MemoryBlockStore {
    fn list_blocks(&self, file: FileId) -> crate::Result<BlockList> {
        self
            .files
            .get(&file)
            .ok_or(InternalError::FileNotFound { id: file })?
            .iter()
            .map(|block_id| {
                match self.blocks.get(block_id) {
                    Some(block) => Ok(Block {
                        id: *block_id,
                        signature: block.signature.clone(),
                    }),
                    None => panic!("Inconsistent state in in-memory data store: block with ID {:?} listed for file with ID {:?} but not found in the blocks map.", block_id, file),
                }
            })
            .collect::<crate::Result<Vec<Block>>>()
            .map(BlockList::new)
    }

    fn read_block(&mut self, block: BlockId, buf: &mut Vec<u8>) -> crate::Result<ReadBlock> {
        match self.blocks.get(&block) {
            Some(store_block) => {
                match &store_block.data {
                    BlockData::Data { bytes } => {
                        buf.reserve(bytes.len());
                        buf.extend_from_slice(bytes);
                    }
                    BlockData::Hole { len: _ } => {
                        // For holes, we don't fill the buffer.
                    }
                }
                Ok(ReadBlock {
                    signature: store_block.signature.clone(),
                    encoding: store_block.encoding,
                })
            }
            None => Err(InternalError::BlockNotFound { id: block }.into()),
        }
    }

    fn replace_block(
        &mut self,
        file: FileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        let block_id = BlockId::from(self.block_id_pool.next() as i64);
        let block = Self::input_to_block(input);

        self.blocks.insert(block_id, block);

        match self.files.entry(file).or_default().get_mut(index) {
            Some(existing_block_id) => {
                *existing_block_id = block_id;
            }
            None => {
                panic!(
                    "Tried to replace a block at index {} for file ID {:?}, but no block exists at that index.",
                    index, file
                );
            }
        }

        Ok(block_id)
    }

    fn insert_block(
        &mut self,
        file: FileId,
        index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        let block_id = BlockId::from(self.block_id_pool.next() as i64);
        let block = Self::input_to_block(input);

        self.blocks.insert(block_id, block);

        self.files.entry(file).or_default().insert(index, block_id);

        Ok(block_id)
    }

    fn append_block(
        &mut self,
        file: FileId,
        _index: BlockIndex,
        input: BlockStoreInput,
    ) -> crate::Result<BlockId> {
        let block_id = BlockId::from(self.block_id_pool.next() as i64);
        let block = Self::input_to_block(input);

        self.blocks.insert(block_id, block);

        self.files.entry(file).or_default().push(block_id);

        Ok(block_id)
    }

    fn remove_blocks<R: RangeBounds<usize>>(
        &mut self,
        file: FileId,
        blocks: R,
    ) -> crate::Result<Vec<BlockId>> {
        let file_blocks = match self.files.get_mut(&file) {
            Some(blocks) => blocks,
            None => return Err(InternalError::FileNotFound { id: file }.into()),
        };

        let removed_block_ids = file_blocks.drain(blocks).collect::<Vec<_>>();

        for block_id in &removed_block_ids {
            self.blocks.remove(block_id);
            self.block_id_pool.recycle(i64::from(*block_id) as u64);
        }

        Ok(removed_block_ids)
    }
}