taple_core/database/
memory.rs

1use super::Error;
2use crate::{test_database_manager_trait, DatabaseCollection, DatabaseManager};
3use std::{
4    collections::{btree_map::Iter, BTreeMap, HashMap},
5    iter::Rev,
6    sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
7};
8
9pub struct DataStore {
10    data: RwLock<BTreeMap<String, Vec<u8>>>,
11}
12
13impl DataStore {
14    fn new() -> Self {
15        Self {
16            data: RwLock::new(BTreeMap::new()),
17        }
18    }
19
20    fn _get_inner_read_lock<'a>(&'a self) -> RwLockReadGuard<'a, BTreeMap<String, Vec<u8>>> {
21        self.data.read().unwrap()
22    }
23
24    fn _get_inner_write_lock<'a>(&'a self) -> RwLockWriteGuard<'a, BTreeMap<String, Vec<u8>>> {
25        self.data.write().unwrap()
26    }
27}
28
29impl DataStore {
30    fn iter(&self, prefix: String) -> MemoryIterator {
31        MemoryIterator::new(&self, prefix)
32    }
33
34    fn rev_iter(&self, prefix: String) -> RevMemoryIterator {
35        RevMemoryIterator::new(&self, prefix)
36    }
37}
38
39/// In-memory database implementation for TAPLE.
40pub struct MemoryManager {
41    data: RwLock<HashMap<String, Arc<DataStore>>>,
42}
43
44impl MemoryManager {
45    pub fn new() -> Self {
46        Self {
47            data: RwLock::new(HashMap::new()),
48        }
49    }
50}
51
52impl DatabaseManager<MemoryCollection> for MemoryManager {
53    fn default() -> Self {
54        Self {
55            data: RwLock::new(HashMap::new()),
56        }
57    }
58
59    fn create_collection(&self, _identifier: &str) -> MemoryCollection {
60        let mut lock = self.data.write().unwrap();
61        let db: Arc<DataStore> = match lock.get("") {
62            Some(map) => map.clone(),
63            None => {
64                let db: Arc<DataStore> = Arc::new(DataStore::new());
65                lock.insert("".to_string(), db.clone());
66                db
67            }
68        };
69        MemoryCollection { data: db }
70    }
71}
72
73/// Collection for in-memory database implementation. It must be created through [MemoryManager].
74pub struct MemoryCollection {
75    data: Arc<DataStore>,
76}
77
78impl DatabaseCollection for MemoryCollection {
79    fn get(&self, key: &str) -> Result<Vec<u8>, Error> {
80        let lock = self.data._get_inner_read_lock();
81        let Some(data) = lock.get(key) else {
82            return Err(Error::EntryNotFound);
83        };
84        Ok(data.clone())
85    }
86
87    fn put(&self, key: &str, data: Vec<u8>) -> Result<(), Error> {
88        let mut lock = self.data._get_inner_write_lock();
89        lock.insert(key.to_string(), data);
90        Ok(())
91    }
92
93    fn del(&self, key: &str) -> Result<(), Error> {
94        let mut lock = self.data._get_inner_write_lock();
95        lock.remove(key);
96        Ok(())
97    }
98
99    fn iter<'a>(
100        &'a self,
101        reverse: bool,
102        prefix: String,
103    ) -> Box<dyn Iterator<Item = (String, Vec<u8>)> + 'a> {
104        if reverse {
105            Box::new(self.data.rev_iter(prefix))
106        } else {
107            Box::new(self.data.iter(prefix))
108        }
109    }
110}
111
112type GuardIter<'a, K, V> = (Arc<RwLockReadGuard<'a, BTreeMap<K, V>>>, Iter<'a, K, V>);
113
114pub struct MemoryIterator<'a> {
115    map: &'a DataStore,
116    current: Option<GuardIter<'a, String, Vec<u8>>>,
117    table_name: String,
118}
119
120impl<'a> MemoryIterator<'a> {
121    fn new(map: &'a DataStore, table_name: String) -> Self {
122        Self {
123            map,
124            current: None,
125            table_name,
126        }
127    }
128}
129
130impl<'a> Iterator for MemoryIterator<'a> {
131    type Item = (String, Vec<u8>);
132    fn next(&mut self) -> Option<Self::Item> {
133        let iter = if let Some((_, iter)) = self.current.as_mut() {
134            iter
135        } else {
136            let guard = self.map._get_inner_read_lock();
137            let sref: &BTreeMap<String, Vec<u8>> = unsafe { change_lifetime_const(&*guard) };
138            let iter = sref.iter();
139            self.current = Some((Arc::new(guard), iter));
140            &mut self.current.as_mut().unwrap().1
141        };
142
143        let Some(item) = iter.next() else {
144                return None;
145            };
146        let key = {
147            let value = item.0.clone();
148            if !value.starts_with(&self.table_name) {
149                return None;
150            }
151            value.replace(&self.table_name, "")
152        };
153        return Some((key, item.1.clone()));
154    }
155}
156
157type GuardRevIter<'a> = (
158    Arc<RwLockReadGuard<'a, BTreeMap<String, Vec<u8>>>>,
159    Rev<Iter<'a, String, Vec<u8>>>,
160);
161
162pub struct RevMemoryIterator<'a> {
163    map: &'a DataStore,
164    current: Option<GuardRevIter<'a>>,
165    table_name: String,
166}
167
168impl<'a> RevMemoryIterator<'a> {
169    fn new(map: &'a DataStore, table_name: String) -> Self {
170        Self {
171            map,
172            current: None,
173            table_name,
174        }
175    }
176}
177
178impl<'a> Iterator for RevMemoryIterator<'a> {
179    type Item = (String, Vec<u8>);
180    fn next(&mut self) -> Option<Self::Item> {
181        let iter = if let Some((_, iter)) = self.current.as_mut() {
182            iter
183        } else {
184            let guard = self.map._get_inner_read_lock();
185            let sref: &BTreeMap<String, Vec<u8>> = unsafe { change_lifetime_const(&*guard) };
186            let iter = sref.iter().rev();
187            self.current = Some((Arc::new(guard), iter));
188            &mut self.current.as_mut().unwrap().1
189        };
190        let Some(item) = iter.next() else {
191                return None;
192            };
193        let key = {
194            let value = item.0.clone();
195            if !value.starts_with(&self.table_name) {
196                return None;
197            }
198            value.replace(&self.table_name, "")
199        };
200        return Some((key, item.1.clone()));
201    }
202}
203
204unsafe fn change_lifetime_const<'a, 'b, T>(x: &'a T) -> &'b T {
205    &*(x as *const T)
206}
207
208test_database_manager_trait! {
209    unit_test_memory_manager:crate::MemoryManager:MemoryCollection
210}
taple_core/database/memory.rs

taple_core/database/
memory.rs
