atrium_repo/blockstore/
memory.rs

1use std::collections::HashMap;
2
3use ipld_core::cid::{multihash::Multihash, Cid};
4use sha2::Digest;
5
6use super::{AsyncBlockStoreRead, AsyncBlockStoreWrite, Error, SHA2_256};
7
8/// Basic in-memory blockstore. This is primarily used for testing.
9pub struct MemoryBlockStore {
10    blocks: HashMap<Cid, Vec<u8>>,
11}
12
13impl Default for MemoryBlockStore {
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl MemoryBlockStore {
20    pub fn new() -> Self {
21        Self { blocks: HashMap::new() }
22    }
23
24    pub fn contains(&self, block: Cid) -> bool {
25        self.blocks.contains_key(&block)
26    }
27}
28
29impl AsyncBlockStoreRead for MemoryBlockStore {
30    async fn read_block_into(&mut self, cid: Cid, contents: &mut Vec<u8>) -> Result<(), Error> {
31        contents.clear();
32        contents.extend_from_slice(self.blocks.get(&cid).ok_or(Error::CidNotFound)?);
33        Ok(())
34    }
35}
36
37impl AsyncBlockStoreWrite for MemoryBlockStore {
38    async fn write_block(&mut self, codec: u64, hash: u64, contents: &[u8]) -> Result<Cid, Error> {
39        let digest = match hash {
40            SHA2_256 => sha2::Sha256::digest(contents),
41            _ => return Err(Error::UnsupportedHash(hash)),
42        };
43        let hash =
44            Multihash::wrap(hash, digest.as_slice()).expect("internal error encoding multihash");
45        let cid = Cid::new_v1(codec, hash);
46
47        // Insert the block. We're explicitly ignoring the case where it's already present inside the hashmap.
48        self.blocks.insert(cid, contents.to_vec());
49        Ok(cid)
50    }
51}