Skip to main content

nectar_primitives/store/
memory.rs

1//! In-memory chunk storage.
2
3use std::collections::HashMap;
4
5use parking_lot::RwLock;
6
7use crate::bmt::DEFAULT_BODY_SIZE;
8use crate::chunk::{AnyChunk, ChunkAddress};
9
10use super::ChunkStoreError;
11use super::typed::{SyncChunkGet, SyncChunkHas, SyncChunkPut};
12
13/// In-memory chunk storage using a `RwLock<HashMap>`.
14///
15/// Uses interior mutability so `SyncChunkPut::put(&self)` works without
16/// external synchronization.
17#[derive(Debug)]
18pub struct MemoryStore<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
19    chunks: RwLock<HashMap<ChunkAddress, AnyChunk<BODY_SIZE>>>,
20}
21
22impl<const BODY_SIZE: usize> Clone for MemoryStore<BODY_SIZE> {
23    fn clone(&self) -> Self {
24        Self {
25            chunks: RwLock::new(self.chunks.read().clone()),
26        }
27    }
28}
29
30impl<const BODY_SIZE: usize> Default for MemoryStore<BODY_SIZE> {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl<const BODY_SIZE: usize> MemoryStore<BODY_SIZE> {
37    /// Create an empty memory store.
38    pub fn new() -> Self {
39        Self {
40            chunks: RwLock::new(HashMap::new()),
41        }
42    }
43
44    /// Get a cloned chunk by address.
45    pub fn get(&self, address: &ChunkAddress) -> Option<AnyChunk<BODY_SIZE>> {
46        self.chunks.read().get(address).cloned()
47    }
48
49    /// Number of stored chunks.
50    pub fn len(&self) -> usize {
51        self.chunks.read().len()
52    }
53
54    /// Whether the store is empty.
55    pub fn is_empty(&self) -> bool {
56        self.chunks.read().is_empty()
57    }
58
59    /// Consume the store and return all chunks.
60    pub fn into_chunks(self) -> HashMap<ChunkAddress, AnyChunk<BODY_SIZE>> {
61        self.chunks.into_inner()
62    }
63}
64
65impl<const BODY_SIZE: usize> SyncChunkPut<BODY_SIZE> for MemoryStore<BODY_SIZE> {
66    type Error = std::convert::Infallible;
67
68    fn put(&self, chunk: AnyChunk<BODY_SIZE>) -> Result<(), Self::Error> {
69        self.chunks.write().insert(*chunk.address(), chunk);
70        Ok(())
71    }
72}
73
74impl<const BODY_SIZE: usize> SyncChunkGet<BODY_SIZE> for MemoryStore<BODY_SIZE> {
75    type Error = ChunkStoreError;
76
77    fn get(&self, address: &ChunkAddress) -> Result<AnyChunk<BODY_SIZE>, Self::Error> {
78        self.chunks
79            .read()
80            .get(address)
81            .cloned()
82            .ok_or_else(|| ChunkStoreError::not_found(address))
83    }
84}
85
86impl<const BODY_SIZE: usize> SyncChunkHas<BODY_SIZE> for MemoryStore<BODY_SIZE> {
87    fn has(&self, address: &ChunkAddress) -> bool {
88        self.chunks.read().contains_key(address)
89    }
90}
91
92impl<const BODY_SIZE: usize> SyncChunkGet<BODY_SIZE>
93    for HashMap<ChunkAddress, AnyChunk<BODY_SIZE>>
94{
95    type Error = ChunkStoreError;
96
97    fn get(&self, address: &ChunkAddress) -> Result<AnyChunk<BODY_SIZE>, Self::Error> {
98        self.get(address)
99            .cloned()
100            .ok_or_else(|| ChunkStoreError::not_found(address))
101    }
102}
103
104impl<const BODY_SIZE: usize> SyncChunkHas<BODY_SIZE>
105    for HashMap<ChunkAddress, AnyChunk<BODY_SIZE>>
106{
107    fn has(&self, address: &ChunkAddress) -> bool {
108        self.contains_key(address)
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::chunk::{Chunk, ContentChunk};
116
117    #[test]
118    fn test_memory_store() {
119        let store = MemoryStore::<DEFAULT_BODY_SIZE>::new();
120        assert!(store.is_empty());
121
122        let chunk = ContentChunk::new(b"hello".as_slice()).unwrap();
123        let addr = *chunk.address();
124        let any: AnyChunk = chunk.into();
125
126        SyncChunkPut::put(&store, any.clone()).unwrap();
127        assert_eq!(store.len(), 1);
128        assert!(SyncChunkHas::has(&store, &addr));
129        assert_eq!(store.get(&addr), Some(any));
130    }
131}