infinitree/fields/
map.rs

1use super::{Collection, Intent, Key, Load, LocalField, SparseField, Store, Strategy, Value};
2use crate::{
3    index::{FieldWriter, Transaction},
4    object::{self, serializer::SizedPointer, ObjectError},
5};
6use scc::HashMap;
7use std::{borrow::Borrow, hash::Hash, sync::Arc};
8
9/// Multithreaded hash map that is always saved atomically
10///
11/// Calling [`clone()`](Map::clone) will create a reference to the
12/// same instance, and can be easily shared between threads.
13///
14/// For tracking incremental changes to your data structure, look at
15/// [`VersionedMap`](crate::fields::VersionedMap)
16#[derive(Clone, Default)]
17pub struct Map<K: 'static + Key, V: 'static + Value>(Arc<HashMap<K, Arc<V>>>);
18
19impl<K, V> Map<K, V>
20where
21    K: Key,
22    V: Value,
23{
24    /// Insert a new value for the given key in the map.
25    #[inline(always)]
26    pub fn insert(&self, key: K, value: impl Into<Arc<V>>) -> Arc<V> {
27        match self.get(&key) {
28            Some(existing) => existing,
29            None => {
30                let new: Arc<_> = value.into();
31                let _ = self.0.insert(key, new.clone());
32                new
33            }
34        }
35    }
36
37    /// Update a value with the given key, and return the new value if
38    /// the key exists.
39    #[inline(always)]
40    pub fn update_with<Q, R>(&self, key: &Q, fun: impl FnOnce(&K, &mut Arc<V>) -> R) -> Option<R>
41    where
42        K: Borrow<Q>,
43        Q: Hash + Eq + ?Sized,
44    {
45        self.0.update(key, fun)
46    }
47
48    /// Call the given function to insert a value if it doesn't exist.
49    /// Return with the current value to the key.
50    #[inline(always)]
51    pub fn insert_with(&self, key: K, mut fun: impl FnMut() -> V) -> Arc<V> {
52        match self.get(&key) {
53            Some(existing) => existing,
54            None => {
55                let new: Arc<_> = (fun)().into();
56                self.insert(key, new.clone());
57                new
58            }
59        }
60    }
61
62    /// Returns the stored value for a key, or `None`.
63    #[inline(always)]
64    pub fn get<Q>(&self, key: &Q) -> Option<Arc<V>>
65    where
66        K: Borrow<Q>,
67        Q: Hash + Eq + ?Sized,
68    {
69        self.0.read(key, |_, v| v.clone())
70    }
71
72    /// Sets the key as removed in the map.
73    #[inline(always)]
74    pub fn remove<Q>(&self, key: &Q)
75    where
76        K: Borrow<Q>,
77        Q: Hash + Eq + ?Sized,
78    {
79        self.0.remove(key);
80    }
81
82    /// Returns if there's an addition for the specified key.
83    #[inline(always)]
84    pub fn contains<Q>(&self, key: &Q) -> bool
85    where
86        K: Borrow<Q>,
87        Q: Hash + Eq + ?Sized,
88    {
89        self.0.contains(key)
90    }
91
92    /// Call the function for all additive keys.
93    #[inline(always)]
94    pub fn for_each(&self, mut callback: impl FnMut(&K, &Arc<V>)) {
95        let mut current = self.0.first_entry();
96        while let Some(entry) = current {
97            (callback)(entry.key(), entry.get());
98            current = entry.next();
99        }
100    }
101
102    /// Returns the number of keys.
103    #[inline(always)]
104    pub fn len(&self) -> usize {
105        self.0.len()
106    }
107
108    /// Return the size of all allocated items.
109    #[inline(always)]
110    pub fn capacity(&self) -> usize {
111        self.0.capacity()
112    }
113
114    /// True if the map doesn't contain any items.
115    #[inline(always)]
116    pub fn is_empty(&self) -> bool {
117        self.0.len() == 0
118    }
119}
120
121impl<K, V> Store for LocalField<Map<K, V>>
122where
123    K: Key,
124    V: Value,
125{
126    fn store(&mut self, mut transaction: &mut dyn Transaction, _object: &mut dyn object::Writer) {
127        self.field.for_each(|k, v| {
128            transaction.write_next((k, v));
129        })
130    }
131}
132
133impl<K, V> Collection for LocalField<Map<K, V>>
134where
135    K: Key,
136    V: Value,
137{
138    type Depth = super::depth::Snapshot;
139    type Key = K;
140    type Serialized = (K, V);
141    type Item = (K, V);
142
143    fn key(from: &Self::Serialized) -> &Self::Key {
144        &from.0
145    }
146
147    fn load(from: Self::Serialized, _object: &mut dyn object::Reader) -> Self::Item {
148        from
149    }
150
151    fn insert(&mut self, record: Self::Item) {
152        self.field.insert(record.0, record.1);
153    }
154}
155
156impl<K, V> Store for SparseField<Map<K, V>>
157where
158    K: Key,
159    V: Value,
160{
161    fn store(&mut self, mut transaction: &mut dyn Transaction, writer: &mut dyn object::Writer) {
162        self.field.for_each(|key, value| {
163            let ptr = object::serializer::write(
164                writer,
165                |x| {
166                    crate::serialize_to_vec(x).map_err(|e| ObjectError::Serialize {
167                        source: Box::new(e),
168                    })
169                },
170                value,
171            )
172            .unwrap();
173            transaction.write_next((key, ptr));
174        })
175    }
176}
177
178impl<K, V> Collection for SparseField<Map<K, V>>
179where
180    K: Key,
181    V: Value,
182{
183    type Depth = super::depth::Snapshot;
184    type Key = K;
185    type Serialized = (K, SizedPointer);
186    type Item = (K, V);
187
188    fn key(from: &Self::Serialized) -> &Self::Key {
189        &from.0
190    }
191
192    fn load(from: Self::Serialized, object: &mut dyn object::Reader) -> Self::Item {
193        let (key, ptr) = from;
194
195        let value = object::serializer::read(
196            object,
197            |x| {
198                crate::deserialize_from_slice(x).map_err(|e| ObjectError::Deserialize {
199                    source: Box::new(e),
200                })
201            },
202            ptr,
203        )
204        .unwrap();
205
206        (key, value)
207    }
208
209    fn insert(&mut self, record: Self::Item) {
210        self.field.insert(record.0, record.1);
211    }
212}
213
214impl<K, V> crate::Index for Map<K, V>
215where
216    K: Key + Clone,
217    V: Value + Clone,
218{
219    fn store_all(&self) -> anyhow::Result<Vec<Intent<Box<dyn Store>>>> {
220        Ok(vec![Intent::new(
221            "root",
222            Box::new(LocalField::for_field(self)),
223        )])
224    }
225
226    fn load_all(&self) -> anyhow::Result<Vec<Intent<Box<dyn Load>>>> {
227        Ok(vec![Intent::new(
228            "root",
229            Box::new(LocalField::for_field(self)),
230        )])
231    }
232}
233
234#[cfg(test)]
235mod test {
236    use super::Map;
237    use crate::{
238        fields::{LocalField, SparseField, Strategy},
239        index::test::store_then_load,
240    };
241
242    #[test]
243    fn duplicate_insert_is_noop() {
244        let m = Map::<usize, String>::default();
245        assert_eq!(m.insert(1, "first".to_owned()), "first".to_owned().into());
246        assert_eq!(m.insert(1, "second".to_owned()), "first".to_owned().into());
247    }
248
249    #[test]
250    fn updating_empty_is_noop() {
251        let m = Map::<usize, String>::default();
252        assert_eq!(
253            m.update_with(&1, |_, v| *v = "first".to_owned().into()),
254            None
255        );
256    }
257
258    #[test]
259    fn store_then_confirm_then_remove() {
260        let m = Map::<usize, String>::default();
261        let first = "first".to_owned();
262        let updated = "updated".to_owned();
263        let second = "second".to_owned();
264
265        // insert
266        assert_eq!(m.insert_with(1, || first.clone()), first.clone().into());
267        assert_eq!(m.insert_with(2, || second.clone()), second.clone().into());
268
269        // get first
270        assert_eq!(m.get(&1), Some(first.into()));
271
272        // contains
273        assert!(m.contains(&1));
274        assert!(m.contains(&2));
275
276        // update
277        assert_eq!(
278            m.update_with(&1, |_, v| {
279                *v = updated.clone().into();
280                v.clone()
281            }),
282            Some(updated.clone().into())
283        );
284        assert_eq!(m.get(&1), Some(updated.into()));
285
286        // removed
287        m.remove(&1);
288        assert_eq!(m.get(&1), None);
289
290        // second still fine
291        assert_eq!(m.get(&2), Some(second.into()));
292    }
293
294    type TestMap = Map<usize, String>;
295    fn init_map(store: &TestMap) {
296        store.insert(1, "one".to_owned());
297        store.insert(2, "two".to_owned());
298    }
299    crate::len_check_test!(TestMap, LocalField, init_map, |m: TestMap| m.len());
300    crate::len_check_test!(TestMap, SparseField, init_map, |m: TestMap| m.len());
301}