btree_ondisk/loader/
memory.rs

1use std::io::{Error, ErrorKind, Result};
2use std::collections::HashMap;
3#[cfg(feature = "rc")]
4use std::rc::Rc;
5#[cfg(feature = "rc")]
6use std::cell::RefCell;
7#[cfg(feature = "arc")]
8use atomic_refcell::AtomicRefCell;
9#[cfg(feature = "arc")]
10use std::sync::Arc;
11use crate::BlockLoader;
12
13// memory block loader for test purpose
14#[derive(Clone)]
15pub struct MemoryBlockLoader<V> {
16    #[cfg(feature = "rc")]
17    inner: Rc<RefCell<HashMap<V, Vec<u8>>>>,
18    #[cfg(feature = "arc")]
19    inner: Arc<AtomicRefCell<HashMap<V, Vec<u8>>>>,
20    meta_node_size: usize,
21}
22
23// force impl Sync for life easy
24unsafe impl<V> Sync for MemoryBlockLoader<V> {}
25
26impl<V: Send + Sync + Eq + std::hash::Hash + std::fmt::Display> BlockLoader<V> for MemoryBlockLoader<V>
27    where
28        u64: From<V>,
29{
30    async fn read(&self, v: V, buf: &mut [u8], _user_data: u32) -> Result<Vec<(V, Vec<u8>)>> {
31        assert!(buf.len() == self.meta_node_size);
32        if let Some(data) =  self.inner.borrow().get(&v) {
33            buf.copy_from_slice(data);
34            return Ok(Vec::new());
35        }
36        let msg = format!("requested key {v} not exists");
37        return Err(Error::new(ErrorKind::NotFound, msg));
38    }
39
40    fn from_new_path(self, _: &str) -> Self {
41        todo!()
42    }
43}
44
45impl<V: Eq + Clone + std::hash::Hash + std::fmt::Display> MemoryBlockLoader<V> {
46    pub fn new(meta_node_size: usize) -> Self {
47        Self {
48            #[cfg(feature = "rc")]
49            inner: Rc::new(RefCell::new(HashMap::new())),
50            #[cfg(feature = "arc")]
51            inner: Arc::new(AtomicRefCell::new(HashMap::new())),
52            meta_node_size,
53        }
54    }
55
56    pub fn write(&self, v: V, buf: &[u8]) {
57        assert!(buf.len() == self.meta_node_size);
58        let data = buf.to_vec();
59        let key = v.clone();
60        if let Some(_) = self.inner.borrow_mut().insert(v, data) {
61            panic!("key {key} already exists!");
62        }
63    }
64}