anathema_store/
smallmap.rs

1use std::borrow::Borrow;
2
3use crate::slab::{Slab, SlabIndex};
4
5type NumType = u16;
6
7#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Hash)]
8pub struct SmallIndex(NumType);
9
10impl SmallIndex {
11    pub const MAX: Self = Self(NumType::MAX);
12    pub const ONE: Self = Self(1);
13    pub const ZERO: Self = Self(0);
14}
15
16impl From<SmallIndex> for NumType {
17    fn from(value: SmallIndex) -> Self {
18        value.0
19    }
20}
21
22impl SlabIndex for SmallIndex {
23    const MAX: usize = NumType::MAX as usize;
24
25    fn as_usize(&self) -> usize {
26        self.0 as usize
27    }
28
29    fn from_usize(index: usize) -> Self
30    where
31        Self: Sized,
32    {
33        Self(index as NumType)
34    }
35}
36
37// impl From<u8> for SmallIndex {
38//     fn from(value: u8) -> Self {
39//         Self(value)
40//     }
41// }
42// impl From<usize> for SmallIndex {
43//     fn from(value: usize) -> Self {
44//         Self(value as u8)
45//     }
46// }
47
48// impl From<SmallIndex> for usize {
49//     fn from(value: SmallIndex) -> Self {
50//         value.0 as usize
51//     }
52// }
53
54/// A small map used to store a small amount of values.
55///
56/// The `SmallMap` can store up to 256 values.
57/// ```
58/// # use anathema_store::smallmap::*;
59///
60/// let mut map = SmallMap::empty();
61/// map.set("a", 1);
62/// map.set("b", 2);
63/// let Some(1) = map.set("a", 5) else { unreachable!("we know there is a one there") };
64///
65/// let value = map.get("b").unwrap();
66/// assert_eq!(2, *value);
67///
68/// let value = map.get("a").unwrap();
69/// assert_eq!(5, *value);
70/// ```
71#[derive(Debug, Clone, PartialEq)]
72pub struct SmallMap<K, V>(Slab<SmallIndex, (K, V)>);
73
74impl<K, V> SmallMap<K, V>
75where
76    K: PartialEq,
77{
78    /// Create a en empty map
79    pub fn empty() -> Self {
80        Self(Slab::empty())
81    }
82
83    /// Set a value in the map.
84    /// If there is already a value with the same key,
85    /// the old value will be returned
86    pub fn set(&mut self, key: K, mut value: V) -> Option<V> {
87        let old_value = self.get_mut(&key);
88        match old_value {
89            Some(old) => {
90                std::mem::swap(old, &mut value);
91                Some(value)
92            }
93            None => {
94                self.0.insert((key, value));
95                None
96            }
97        }
98    }
99
100    pub fn insert_with<F>(&mut self, key: K, f: F) -> SmallIndex
101    where
102        F: FnOnce(SmallIndex) -> V,
103    {
104        let id = self.0.next_id();
105        let value = f(id);
106        assert_eq!(self.0.insert((key, value)), id);
107        id
108    }
109
110    /// Get a reference to a value in the map
111    pub fn get<Q>(&self, key: &Q) -> Option<&V>
112    where
113        K: Borrow<Q>,
114        Q: PartialEq + ?Sized,
115    {
116        self.0.iter().find_map(|(_, (k, v))| match k.borrow() == key {
117            true => Some(v),
118            false => None,
119        })
120    }
121
122    /// Get a mutable reference to a value in the map
123    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
124    where
125        K: Borrow<Q>,
126        Q: PartialEq + ?Sized,
127    {
128        self.0.iter_mut().find_map(|(_, (k, v))| match (*k).borrow() == key {
129            true => Some(v),
130            false => None,
131        })
132    }
133
134    pub fn get_index<Q>(&self, key: &Q) -> Option<SmallIndex>
135    where
136        K: Borrow<Q>,
137        Q: PartialEq + ?Sized,
138    {
139        self.0.iter().find_map(|(i, (k, _))| match k.borrow() == key {
140            true => Some(i),
141            false => None,
142        })
143    }
144
145    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
146    where
147        K: Borrow<Q>,
148        Q: PartialEq + ?Sized,
149    {
150        let idx = self.get_index(key)?;
151        self.0.try_remove(idx).map(|(_, v)| v)
152    }
153
154    /// Iterate over the key-value pairs of the map.
155    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> + '_ {
156        self.0.iter().map(|(_, (k, v))| (k, v))
157    }
158
159    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&mut K, &mut V)> + '_ {
160        self.0.iter_mut().map(|(_, (k, v))| (k, v))
161    }
162
163    /// Get a value ref by the value index instead of the key
164    pub fn get_with_index(&self, idx: SmallIndex) -> Option<&V> {
165        self.0.get(idx).map(|(_, v)| v)
166    }
167
168    /// Get a mutable value ref by the value index instead of the key
169    pub fn get_mut_with_index(&mut self, idx: SmallIndex) -> Option<&mut V> {
170        self.0.get_mut(idx).map(|(_, v)| v)
171    }
172}
173
174#[cfg(test)]
175mod test {
176    use super::*;
177
178    #[test]
179    fn insert_twice() {
180        let mut map = SmallMap::<&str, u8>::empty();
181        map.set("a", 1);
182        map.set("b", 2);
183
184        assert_eq!(1, *map.get("a").unwrap());
185    }
186
187    #[test]
188    fn get_and_get_mut() {
189        let mut map = SmallMap::<&str, u8>::empty();
190        map.set("a", 1);
191        map.set("b", 2);
192
193        *map.get_mut("a").unwrap() += 1;
194
195        assert_eq!(2, *map.get("b").unwrap());
196        assert_eq!(2, *map.get("a").unwrap());
197    }
198
199    #[test]
200    fn double_set() {
201        let mut map = SmallMap::<&str, u8>::empty();
202        assert_eq!(None, map.set("a", 1));
203        assert_eq!(Some(1), map.set("a", 2));
204    }
205
206    #[test]
207    fn get_by_index() {
208        let mut map = SmallMap::<&str, u8>::empty();
209        map.set("a", 1);
210        let idx = map.get_index("a").unwrap();
211        assert_eq!(1, *map.get_with_index(idx).unwrap());
212    }
213
214    #[test]
215    fn get_by_index_mut() {
216        let mut map = SmallMap::<&str, u8>::empty();
217        map.set("a", 1);
218        let idx = map.get_index("a").unwrap();
219        assert_eq!(1, *map.get_mut_with_index(idx).unwrap());
220    }
221}
anathema_store/smallmap.rs

anathema_store/
smallmap.rs
