Skip to main content

light_magic/
table.rs

1use serde::de::{Error as DeError, MapAccess, SeqAccess, Visitor};
2use serde::ser::{SerializeMap, SerializeSeq};
3use serde::{Deserialize, Serialize};
4use std::collections::BTreeMap;
5use std::fmt::{Debug, Display, Formatter};
6use std::marker::PhantomData;
7use std::str::FromStr;
8use std::{
9    clone::Clone,
10    collections::btree_map::{Values, ValuesMut},
11};
12
13/// Trait for getting the value of the primary key
14pub trait PrimaryKey {
15    type PrimaryKeyType;
16    fn primary_key(&self) -> &Self::PrimaryKeyType;
17}
18
19/// Trait alias-like bound for primary key types supported by [`Table`].
20pub trait TableKey: Ord + FromStr + Display + Debug + Clone {
21    fn parse_key(value: &str) -> Result<Self, String>
22    where
23        Self: Sized;
24}
25
26impl<T> TableKey for T
27where
28    T: Ord + FromStr + Display + Debug + Clone,
29    T::Err: Display,
30{
31    fn parse_key(value: &str) -> Result<Self, String> {
32        T::from_str(value).map_err(|error| error.to_string())
33    }
34}
35
36/// Represents a database table utilizing a `BTreeMap` for underlying data storage.
37/// Needs the `PrimaryKey` trait to be implemented for the value type. Offers
38/// enhanced methods for manipulating records, including `add`, `edit`, `delete`, `get`, and `search`.
39/// ```
40/// use light_magic::{
41///     serde::{Deserialize, Serialize},
42///     table::{PrimaryKey, Table},
43/// };
44///
45/// #[derive(Clone, Debug, Serialize, Deserialize)]
46/// struct User {
47///     id: usize,
48///     name: String,
49///     age: usize,
50/// }
51///
52/// impl PrimaryKey for User {
53///     type PrimaryKeyType = usize;
54///
55///     fn primary_key(&self) -> &Self::PrimaryKeyType {
56///         &self.id
57///     }
58/// }
59/// ```
60pub struct Table<V>
61where
62    V: PrimaryKey,
63{
64    inner: BTreeMap<<V as PrimaryKey>::PrimaryKeyType, V>,
65}
66
67impl<V> Debug for Table<V>
68where
69    V: PrimaryKey + Debug,
70    V::PrimaryKeyType: Debug,
71{
72    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
73        f.debug_struct("Table").field("inner", &self.inner).finish()
74    }
75}
76
77impl<V> Clone for Table<V>
78where
79    V: PrimaryKey + Clone,
80    V::PrimaryKeyType: Clone,
81{
82    fn clone(&self) -> Self {
83        Table {
84            inner: self.inner.clone(),
85        }
86    }
87}
88
89impl<V> Default for Table<V>
90where
91    V: PrimaryKey,
92{
93    fn default() -> Self {
94        Table {
95            inner: BTreeMap::new(),
96        }
97    }
98}
99
100impl<V> Serialize for Table<V>
101where
102    V: PrimaryKey + Serialize + for<'a> Deserialize<'a>,
103    V::PrimaryKeyType: TableKey,
104{
105    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
106    where
107        S: serde::Serializer,
108    {
109        if serializer.is_human_readable() {
110            // Human-readable: emit as a map<String, V>
111            let mut map = serializer.serialize_map(Some(self.inner.len()))?;
112            for (k, v) in &self.inner {
113                map.serialize_entry(&k.to_string(), v)?;
114            }
115            map.end()
116        } else {
117            // Binary (e.g., bincode): emit as a sequence of V with known length
118            let mut seq = serializer.serialize_seq(Some(self.inner.len()))?;
119            for v in self.inner.values() {
120                seq.serialize_element(v)?;
121            }
122            seq.end()
123        }
124    }
125}
126
127impl<'de, V> Deserialize<'de> for Table<V>
128where
129    V: PrimaryKey + Serialize + Deserialize<'de>,
130    V::PrimaryKeyType: TableKey,
131{
132    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
133    where
134        D: serde::Deserializer<'de>,
135    {
136        if deserializer.is_human_readable() {
137            // Human-readable: expect a map<String, V>
138            struct MapVisitor<V>(PhantomData<V>);
139
140            impl<'de, V> Visitor<'de> for MapVisitor<V>
141            where
142                V: PrimaryKey + Serialize + Deserialize<'de>,
143                V::PrimaryKeyType: TableKey,
144            {
145                type Value = Table<V>;
146
147                fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
148                    f.write_str("a map of stringified primary keys to rows")
149                }
150
151                fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
152                where
153                    A: MapAccess<'de>,
154                {
155                    let mut inner = BTreeMap::new();
156                    while let Some((k_str, v)) = map.next_entry::<String, V>()? {
157                        let k = V::PrimaryKeyType::parse_key(&k_str).map_err(|e| {
158                            A::Error::custom(format!(
159                                "failed to parse primary key '{}': {}",
160                                k_str, e
161                            ))
162                        })?;
163                        // Optional: sanity check that v.primary_key() matches k
164                        inner.insert(k, v);
165                    }
166                    Ok(Table { inner })
167                }
168            }
169
170            deserializer.deserialize_map(MapVisitor::<V>(PhantomData))
171        } else {
172            // Binary: expect a sequence of V; rebuild keys from PrimaryKey
173            struct SeqVisitor<V>(PhantomData<V>);
174
175            impl<'de, V> Visitor<'de> for SeqVisitor<V>
176            where
177                V: PrimaryKey + Serialize + Deserialize<'de>,
178                V::PrimaryKeyType: TableKey,
179            {
180                type Value = Table<V>;
181
182                fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
183                    f.write_str("a sequence of table rows")
184                }
185
186                fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
187                where
188                    A: SeqAccess<'de>,
189                {
190                    let mut inner = BTreeMap::new();
191                    while let Some(v) = seq.next_element::<V>()? {
192                        let k = v.primary_key().clone();
193                        inner.insert(k, v);
194                    }
195                    Ok(Table { inner })
196                }
197            }
198
199            deserializer.deserialize_seq(SeqVisitor::<V>(PhantomData))
200        }
201    }
202}
203
204impl<V> Table<V>
205where
206    V: PrimaryKey + Serialize + for<'a> Deserialize<'a>,
207    V::PrimaryKeyType: TableKey,
208{
209    /// Adds an entry to the table, returns the `value` or `None` if the `key` already exists in that table.
210    pub fn add(&mut self, value: V) -> Option<V>
211    where
212        V: Clone,
213        V::PrimaryKeyType: Clone,
214    {
215        let key = value.primary_key();
216        if !self.inner.contains_key(key) {
217            self.inner.insert(key.clone(), value.clone());
218            return Some(value);
219        }
220        None
221    }
222
223    /// Gets an entry from the table, returns the `value` or `None` if it couldn't find the `value`.
224    pub fn get(&self, key: &V::PrimaryKeyType) -> Option<&V> {
225        self.inner.get(key)
226    }
227
228    /// Gets a mutable entry from the table, returns the `value` or `None` if it couldn't find the `value`.
229    pub fn get_mut(&mut self, key: &V::PrimaryKeyType) -> Option<&mut V> {
230        self.inner.get_mut(key)
231    }
232
233    /// Edits an entry in the table, returns the `new_value` or `None` if the entry couldn't be found.
234    pub fn edit(&mut self, key: &V::PrimaryKeyType, new_value: V) -> Option<V>
235    where
236        V: Clone,
237        V::PrimaryKeyType: Clone,
238    {
239        let new_key = new_value.primary_key();
240        if (key == new_key || !self.inner.contains_key(new_key)) && self.inner.remove(key).is_some()
241        {
242            self.inner.insert(new_key.clone(), new_value.clone());
243            return Some(new_value);
244        }
245        None
246    }
247
248    /// Deletes an entry from the table, returns the `value` or `None` if the `key` wasn't found.
249    pub fn delete(&mut self, key: &V::PrimaryKeyType) -> Option<V> {
250        self.inner.remove(key)
251    }
252
253    /// Searches the table by a predicate function.
254    pub fn search<F>(&self, predicate: F) -> Vec<&V>
255    where
256        F: Fn(&V) -> bool,
257    {
258        self.inner.values().filter(|&val| predicate(val)).collect()
259    }
260
261    /// Searches the table by a predicate function and a custom ordering with a comparator function.
262    pub fn search_ordered<F, O>(&self, predicate: F, comparator: O) -> Vec<&V>
263    where
264        F: Fn(&V) -> bool,
265        O: Fn(&&V, &&V) -> std::cmp::Ordering,
266    {
267        let mut result = self.search(predicate);
268        result.sort_by(comparator);
269        result
270    }
271
272    /// Gets an iterator over the values of the map, in order by key.
273    pub fn values(&self) -> Values<'_, V::PrimaryKeyType, V> {
274        self.inner.values()
275    }
276
277    /// Gets a mutable iterator over the values of the map, in order by key.
278    pub fn values_mut(&mut self) -> ValuesMut<'_, V::PrimaryKeyType, V> {
279        self.inner.values_mut()
280    }
281}
282
283#[cfg(test)]
284mod test {
285    use super::{PrimaryKey, Table};
286    use serde::{Deserialize, Serialize};
287
288    #[derive(Clone, Debug, Serialize, Deserialize)]
289    struct User {
290        id: usize,
291        name: String,
292        age: usize,
293    }
294
295    impl PrimaryKey for User {
296        type PrimaryKeyType = usize;
297        fn primary_key(&self) -> &Self::PrimaryKeyType {
298            &self.id
299        }
300    }
301
302    #[test]
303    fn json_roundtrip_as_map() {
304        let mut table = Table::default();
305        table.add(User {
306            id: 0,
307            name: "".into(),
308            age: 0,
309        });
310        let s = serde_json::to_string(&table).unwrap();
311        assert_eq!(s, r#"{"0":{"id":0,"name":"","age":0}}"#);
312        let back: Table<User> = serde_json::from_str(&s).unwrap();
313        assert!(back.get(&0).is_some());
314    }
315
316    #[test]
317    #[cfg(feature = "encrypted")]
318    fn bincode_roundtrip_as_seq() {
319        use crate::encrypted::bincode_cfg;
320
321        let mut table = Table::default();
322        for i in 0..3 {
323            table.add(User {
324                id: i,
325                name: format!("u{i}"),
326                age: i,
327            });
328        }
329        let bytes = bincode::serde::encode_to_vec(&table, bincode_cfg()).unwrap();
330        let (back, _): (Table<User>, usize) =
331            bincode::serde::decode_from_slice(&bytes, bincode_cfg()).unwrap();
332        assert_eq!(table.values().count(), back.values().count());
333        for i in 0..3 {
334            assert_eq!(table.get(&i).unwrap().name, back.get(&i).unwrap().name);
335        }
336    }
337}