db_rs/
lookup_map.rs

1use crate::table::Table;
2use crate::{DbResult, Logger, TableId};
3use serde::de::DeserializeOwned;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::hash::Hash;
7
8/// A special case of [crate::lookup::LookupTable] where the value of the [HashMap] is another `HashMap<K2, V>`.
9#[derive(Debug)]
10#[cfg_attr(feature = "clone", derive(Clone))]
11pub struct LookupMap<K1, K2, V>
12where
13    K1: Hash + Eq + Serialize,
14    K2: Hash + Eq + Serialize,
15    V: Serialize + DeserializeOwned,
16{
17    table_id: TableId,
18    inner: HashMap<K1, HashMap<K2, V>>,
19    pub logger: Logger,
20}
21
22#[derive(Serialize, Deserialize)]
23pub enum LogEntry<K1, K2, V> {
24    Insert(K1, K2, V),
25    Remove(K1, K2),
26    CreateKey(K1),
27    ClearKey(K1),
28    Clear,
29}
30
31impl<K1, K2, V> Table for LookupMap<K1, K2, V>
32where
33    K1: Hash + Eq + Serialize + DeserializeOwned,
34    K2: Hash + Eq + Serialize + DeserializeOwned,
35    V: Serialize + DeserializeOwned,
36{
37    fn init(table_id: TableId, logger: Logger) -> Self {
38        Self { table_id, inner: HashMap::default(), logger }
39    }
40
41    fn handle_event(&mut self, bytes: &[u8]) -> DbResult<()> {
42        match bincode::deserialize::<LogEntry<K1, K2, V>>(bytes)? {
43            LogEntry::Insert(k1, k2, v) => {
44                self.insert_inner(k1, k2, v);
45            }
46            LogEntry::Remove(k1, k2) => {
47                if let Some(map) = self.inner.get_mut(&k1) {
48                    map.remove(&k2);
49                };
50            }
51            LogEntry::CreateKey(k1) => {
52                self.inner.insert(k1, HashMap::new());
53            }
54            LogEntry::ClearKey(k1) => {
55                self.inner.remove(&k1);
56            }
57            LogEntry::Clear => {
58                self.inner.clear();
59            }
60        };
61
62        Ok(())
63    }
64
65    fn compact_repr(&self) -> DbResult<Vec<u8>> {
66        let mut repr = vec![];
67        for (k1, values) in &self.inner {
68            if values.is_empty() {
69                let data = bincode::serialize(&LogEntry::<&K1, &K2, &V>::CreateKey(k1))?;
70                let mut data = Logger::log_entry(self.table_id, data);
71                repr.append(&mut data);
72                continue;
73            }
74            for (k2, v) in values {
75                let data = bincode::serialize(&LogEntry::Insert(k1, k2, v))?;
76                let mut data = Logger::log_entry(self.table_id, data);
77                repr.append(&mut data);
78            }
79        }
80
81        Ok(repr)
82    }
83}
84
85impl<K1, K2, V> LookupMap<K1, K2, V>
86where
87    K1: Hash + Eq + Serialize + DeserializeOwned,
88    K2: Hash + Eq + Serialize + DeserializeOwned,
89    V: Serialize + DeserializeOwned,
90{
91    pub(crate) fn insert_inner(&mut self, k1: K1, k2: K2, v: V) -> Option<V> {
92        if let Some(map) = self.inner.get_mut(&k1) {
93            map.insert(k2, v)
94        } else {
95            let mut map = HashMap::new();
96            map.insert(k2, v);
97            self.inner.insert(k1, map);
98            None
99        }
100    }
101
102    pub fn insert(&mut self, k1: K1, k2: K2, v: V) -> DbResult<Option<V>> {
103        let log_entry = LogEntry::Insert(&k1, &k2, &v);
104        let data = bincode::serialize(&log_entry)?;
105        let ret = self.insert_inner(k1, k2, v);
106
107        self.logger.write(self.table_id, data)?;
108        Ok(ret)
109    }
110
111    pub fn create_key(&mut self, key: K1) -> DbResult<Option<HashMap<K2, V>>> {
112        let log_entry = LogEntry::<&K1, K2, &V>::CreateKey(&key);
113        let data = bincode::serialize(&log_entry)?;
114
115        let ret = self.inner.insert(key, HashMap::new());
116
117        self.logger.write(self.table_id, data)?;
118        Ok(ret)
119    }
120
121    pub fn remove(&mut self, k1: &K1, k2: &K2) -> DbResult<Option<V>> {
122        if let Some(map) = self.inner.get_mut(k1) {
123            let log_entry = LogEntry::Remove::<&K1, &K2, V>(k1, k2);
124            let data = bincode::serialize(&log_entry)?;
125            self.logger.write(self.table_id, data)?;
126            Ok(map.remove(k2))
127        } else {
128            Ok(None)
129        }
130    }
131
132    pub fn clear_key(&mut self, k1: &K1) -> DbResult<Option<HashMap<K2, V>>> {
133        let log_entry = LogEntry::ClearKey::<&K1, K2, V>(k1);
134        let data = bincode::serialize(&log_entry)?;
135
136        let ret = self.inner.remove(k1);
137
138        self.logger.write(self.table_id, data)?;
139
140        Ok(ret)
141    }
142
143    pub fn get(&self) -> &HashMap<K1, HashMap<K2, V>> {
144        &self.inner
145    }
146
147    pub fn clear(&mut self) -> DbResult<()> {
148        self.inner.clear();
149        let log_entry = LogEntry::<K1, K2, V>::Clear;
150        let data = bincode::serialize(&log_entry)?;
151        self.logger.write(self.table_id, data)?;
152
153        Ok(())
154    }
155}
db_rs/lookup_map.rs

db_rs/
lookup_map.rs
