nectar-primitives 0.4.0

//! In-memory chunk storage.

use std::collections::HashMap;

use parking_lot::RwLock;

use crate::bmt::DEFAULT_BODY_SIZE;
use crate::chunk::{AnyChunk, ChunkAddress};

use super::ChunkStoreError;
use super::typed::{ChunkGet, ChunkHas, ChunkPut};

/// In-memory chunk storage using a `RwLock<HashMap>`.
///
/// Uses interior mutability so `ChunkPut::put(&self)` works without
/// external synchronization.
#[derive(Debug)]
pub struct MemoryStore<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
    chunks: RwLock<HashMap<ChunkAddress, AnyChunk<BODY_SIZE>>>,
}

impl<const BODY_SIZE: usize> Clone for MemoryStore<BODY_SIZE> {
    fn clone(&self) -> Self {
        Self {
            chunks: RwLock::new(self.chunks.read().clone()),
        }
    }
}

impl<const BODY_SIZE: usize> Default for MemoryStore<BODY_SIZE> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const BODY_SIZE: usize> MemoryStore<BODY_SIZE> {
    /// Create an empty memory store.
    pub fn new() -> Self {
        Self {
            chunks: RwLock::new(HashMap::new()),
        }
    }

    /// Build a store from a collection of chunks, keyed by address.
    pub fn from_chunks(chunks: impl IntoIterator<Item = AnyChunk<BODY_SIZE>>) -> Self {
        Self {
            chunks: RwLock::new(chunks.into_iter().map(|c| (*c.address(), c)).collect()),
        }
    }

    /// Get a cloned chunk by address.
    pub fn get(&self, address: &ChunkAddress) -> Option<AnyChunk<BODY_SIZE>> {
        self.chunks.read().get(address).cloned()
    }

    /// Number of stored chunks.
    pub fn len(&self) -> usize {
        self.chunks.read().len()
    }

    /// Whether the store is empty.
    pub fn is_empty(&self) -> bool {
        self.chunks.read().is_empty()
    }

    /// Consume the store and return all chunks.
    pub fn into_chunks(self) -> HashMap<ChunkAddress, AnyChunk<BODY_SIZE>> {
        self.chunks.into_inner()
    }
}

impl<const BODY_SIZE: usize> ChunkPut<BODY_SIZE> for MemoryStore<BODY_SIZE> {
    type Error = std::convert::Infallible;

    async fn put(&self, chunk: AnyChunk<BODY_SIZE>) -> Result<(), Self::Error> {
        self.chunks.write().insert(*chunk.address(), chunk);
        Ok(())
    }
}

impl<const BODY_SIZE: usize> ChunkGet<BODY_SIZE> for MemoryStore<BODY_SIZE> {
    type Error = ChunkStoreError;

    async fn get(&self, address: &ChunkAddress) -> Result<AnyChunk<BODY_SIZE>, Self::Error> {
        self.chunks
            .read()
            .get(address)
            .cloned()
            .ok_or_else(|| ChunkStoreError::not_found(address))
    }
}

impl<const BODY_SIZE: usize> ChunkHas<BODY_SIZE> for MemoryStore<BODY_SIZE> {
    async fn has(&self, address: &ChunkAddress) -> bool {
        self.chunks.read().contains_key(address)
    }
}

impl<const BODY_SIZE: usize> ChunkGet<BODY_SIZE> for HashMap<ChunkAddress, AnyChunk<BODY_SIZE>> {
    type Error = ChunkStoreError;

    async fn get(&self, address: &ChunkAddress) -> Result<AnyChunk<BODY_SIZE>, Self::Error> {
        self.get(address)
            .cloned()
            .ok_or_else(|| ChunkStoreError::not_found(address))
    }
}

impl<const BODY_SIZE: usize> ChunkHas<BODY_SIZE> for HashMap<ChunkAddress, AnyChunk<BODY_SIZE>> {
    async fn has(&self, address: &ChunkAddress) -> bool {
        self.contains_key(address)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::chunk::{Chunk, ContentChunk};
    use futures::executor::block_on;

    #[test]
    fn test_memory_store() {
        let store = MemoryStore::<DEFAULT_BODY_SIZE>::new();
        assert!(store.is_empty());

        let chunk = ContentChunk::new(b"hello".as_slice()).unwrap();
        let addr = *chunk.address();
        let any: AnyChunk = chunk.into();

        block_on(ChunkPut::put(&store, any.clone())).unwrap();
        assert_eq!(store.len(), 1);
        assert!(block_on(ChunkHas::has(&store, &addr)));
        assert_eq!(store.get(&addr), Some(any));
    }
}