cellular_raza_core/storage/
memory_storage.rs

1use super::concepts::StorageError;
2use super::concepts::{StorageInterfaceLoad, StorageInterfaceOpen, StorageInterfaceStore};
3use std::sync::{Arc, Mutex};
4
5use serde::{Deserialize, Serialize};
6
7use std::collections::{BTreeMap, HashMap};
8
9/// Use the [sled] database to save results to an embedded database.
10// TODO use custom field for config [](https://docs.rs/sled/latest/sled/struct.Config.html) to let the user control these parameters
11#[derive(Clone, Debug)]
12pub struct MemoryStorageInterface<Id, Element>
13where
14    Id: Sized,
15    Element: Sized,
16{
17    map: Arc<Mutex<BTreeMap<u64, HashMap<Id, Element>>>>,
18}
19
20impl<Id, Element> StorageInterfaceOpen for MemoryStorageInterface<Id, Element> {
21    fn open_or_create(
22        _location: &std::path::Path,
23        _storage_instance: u64,
24    ) -> Result<Self, StorageError> {
25        Ok(Self {
26            map: Arc::new(Mutex::new(BTreeMap::new())),
27        })
28    }
29
30    fn clone_to_new_instance(&self, _storage_instance: u64) -> Self {
31        Self {
32            map: Arc::clone(&self.map),
33        }
34    }
35}
36
37impl<Id, Element> StorageInterfaceStore<Id, Element> for MemoryStorageInterface<Id, Element>
38where
39    Id: Clone + std::hash::Hash + std::cmp::Eq,
40    Element: Clone,
41{
42    fn store_single_element(
43        &mut self,
44        iteration: u64,
45        identifier: &Id,
46        element: &Element,
47    ) -> Result<(), StorageError>
48    where
49        Id: Serialize,
50        Element: Serialize,
51    {
52        self.map
53            .lock()?
54            .entry(iteration)
55            .and_modify(|e| {
56                e.entry(identifier.clone()).or_insert(element.clone());
57            })
58            .or_insert(HashMap::from([(identifier.clone(), element.clone())]));
59        Ok(())
60    }
61
62    fn store_batch_elements<'a, I>(
63        &'a mut self,
64        iteration: u64,
65        identifiers_elements: I,
66    ) -> Result<(), StorageError>
67    where
68        Id: 'a + Serialize,
69        Element: 'a + Serialize,
70        I: Clone + IntoIterator<Item = (&'a Id, &'a Element)>,
71    {
72        let identifiers_elements = identifiers_elements
73            .clone()
74            .into_iter()
75            .map(|(id, el)| (id.clone(), el.clone()))
76            .collect::<Vec<_>>();
77        self.map
78            .lock()?
79            .entry(iteration)
80            .or_insert(HashMap::new())
81            .extend(identifiers_elements.into_iter().collect::<HashMap<_, _>>());
82        Ok(())
83    }
84}
85
86impl<Id, Element> StorageInterfaceLoad<Id, Element> for MemoryStorageInterface<Id, Element>
87where
88    Id: core::hash::Hash + core::cmp::Eq + Clone,
89    Element: Clone,
90{
91    fn load_single_element(
92        &self,
93        iteration: u64,
94        identifier: &Id,
95    ) -> Result<Option<Element>, StorageError>
96    where
97        Id: Serialize + for<'a> Deserialize<'a>,
98        Element: for<'a> Deserialize<'a>,
99    {
100        Ok(self
101            .map
102            .lock()?
103            .get(&iteration)
104            .and_then(|elements| elements.get(identifier).and_then(|x| Some(x.clone()))))
105    }
106
107    fn load_all_elements_at_iteration(
108        &self,
109        iteration: u64,
110    ) -> Result<HashMap<Id, Element>, StorageError>
111    where
112        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
113        Element: for<'a> Deserialize<'a>,
114    {
115        match self.map.lock()?.get(&iteration) {
116            Some(x) => Ok(x.clone()),
117            None => Ok(HashMap::new()),
118        }
119    }
120
121    fn load_all_elements(&self) -> Result<BTreeMap<u64, HashMap<Id, Element>>, StorageError>
122    where
123        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
124        Element: for<'a> Deserialize<'a>,
125    {
126        Ok(self.map.lock()?.clone())
127    }
128
129    fn get_all_iterations(&self) -> Result<Vec<u64>, StorageError> {
130        Ok(self.map.lock()?.keys().into_iter().map(|&k| k).collect())
131    }
132}