db_rs/
lookup_set.rs

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