composable_indexes_core/
collection.rs

1use std::collections::HashMap;
2
3use crate::{index::Index, QueryResult, QueryResultDistinct};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
6pub struct Key {
7    pub id: u64,
8}
9
10/// A collection of items, with an index that is automatically kept up-to-date.
11pub struct Collection<In, Ix, S = std::hash::RandomState> {
12    index: Ix,
13    data: HashMap<Key, In, S>,
14    next_key_id: u64,
15}
16
17impl<In, Ix> Collection<In, Ix> {
18    /// Create an empty collection.
19    pub fn new(ix: Ix) -> Self {
20        Collection {
21            data: HashMap::new(),
22            next_key_id: 0,
23            index: ix,
24        }
25    }
26}
27
28impl<In, Ix, S> Collection<In, Ix, S>
29where
30    S: std::hash::BuildHasher,
31{
32    /// Create an empty collection with a custom hasher.
33    pub fn new_with_hasher(hasher: S, ix: Ix) -> Self {
34        Collection {
35            data: HashMap::with_hasher(hasher),
36            next_key_id: 0,
37            index: ix,
38        }
39    }
40}
41
42impl<In, Ix> Collection<In, Ix>
43where
44    Ix: Index<In>,
45{
46    /// Lookup an item in the collection by its key.
47    pub fn get_by_key(&self, key: Key) -> Option<&In> {
48        self.data.get(&key)
49    }
50
51    /// Insert a new item into the collection.
52    pub fn insert(&mut self, value: In) -> Key {
53        let key = self.mk_key();
54        let existing = self.data.insert(key, value);
55
56        // There shouldn't be an existing key, as we just generated it
57        debug_assert!(existing.is_none());
58
59        self.index.insert(&Insert {
60            key,
61            new: &self.data[&key],
62        });
63
64        key
65    }
66
67    /// Iterate over all items in the collection.
68    pub fn iter(&self) -> impl Iterator<Item = (&Key, &In)> {
69        self.data.iter()
70    }
71
72    /// Mutate (or alter the presence of) the item in the collection.
73    pub fn update_by_key_mut<F>(&mut self, key: Key, f: F)
74    where
75        F: FnOnce(&mut Option<In>),
76    {
77        let mut existing = self.delete_by_key(&key);
78        f(&mut existing);
79
80        if let Some(existing) = existing {
81            self.data.insert(key, existing);
82            self.index.insert(&Insert {
83                key,
84                new: &self.data[&key],
85            });
86        }
87    }
88
89    /// Update the item in the collection.
90    pub fn update_by_key<F>(&mut self, key: Key, f: F)
91    where
92        F: FnOnce(Option<&In>) -> In,
93    {
94        let existing = self.data.get(&key);
95        let new = f(existing);
96
97        match existing {
98            Some(existing) => {
99                self.index.update(&Update {
100                    key,
101                    new: &new,
102                    existing,
103                });
104                self.data.insert(key, new);
105            }
106            None => {
107                self.index.insert(&Insert { key, new: &new });
108                self.data.insert(key, new);
109            }
110        };
111    }
112
113    /// Mutate the item in the collection, if it exists.
114    pub fn adjust_by_key_mut<F>(&mut self, key: Key, f: F)
115    where
116        F: FnOnce(&mut In),
117    {
118        if let Some(mut existing) = self.delete_by_key(&key) {
119            f(&mut existing);
120            self.data.insert(key, existing);
121            self.index.insert(&Insert {
122                key,
123                new: &self.data[&key],
124            });
125        }
126    }
127
128    /// Adjust the item in the collection, if it exists.
129    pub fn adjust_by_key<F>(&mut self, key: Key, f: F)
130    where
131        F: FnOnce(&In) -> In,
132    {
133        if let Some(existing) = self.data.get(&key) {
134            let new = f(existing);
135            self.index.update(&Update {
136                key,
137                new: &new,
138                existing,
139            });
140            self.data.insert(key, new);
141        }
142    }
143
144    /// Remove an item from the collection, returning it if it exists.
145    pub fn delete_by_key(&mut self, key: &Key) -> Option<In> {
146        let existing = self.data.remove_entry(key);
147        if let Some((key, ref existing)) = existing {
148            self.index.remove(&Remove { key, existing });
149        }
150        existing.map(|(_, v)| v)
151    }
152
153    /// Query the collection using its index(es).
154    pub fn query<Res>(&self, f: impl FnOnce(&Ix) -> Res) -> Res::Resolved<&In>
155    where
156        Res: QueryResult,
157    {
158        let res = f(&self.index);
159        res.map(|k| &self.data[&k])
160    }
161
162    pub fn delete<Res>(&mut self, f: impl FnOnce(&Ix) -> Res) -> usize
163    where
164        Res: QueryResult,
165    {
166        let mut affected_count = 0;
167        let res = f(&self.index);
168        res.map(|key| {
169            self.delete_by_key(&key);
170            affected_count += 1;
171        });
172        affected_count
173    }
174
175    pub fn update<Res, F>(
176        &mut self,
177        f: impl FnOnce(&Ix) -> Res,
178        update_fn: impl Fn(&In) -> In,
179    ) -> Res::Resolved<()>
180    where
181        Res: QueryResultDistinct,
182    {
183        let res = f(&self.index);
184        res.map(|key| {
185            self.data.entry(key).and_modify(|existing| {
186                let new = update_fn(existing);
187                self.index.update(&Update {
188                    key,
189                    new: &new,
190                    existing,
191                });
192                *existing = new;
193            });
194        })
195    }
196
197    pub fn take<Res>(&mut self, f: impl FnOnce(&Ix) -> Res) -> Res::Resolved<In>
198    where
199        Res: QueryResultDistinct,
200    {
201        let res = f(&self.index);
202        res.map(|k| self.delete_by_key(&k).unwrap())
203    }
204
205    /// Number of items in the collection.
206    pub fn len(&self) -> usize {
207        self.data.len()
208    }
209
210    pub fn is_empty(&self) -> bool {
211        self.data.is_empty()
212    }
213
214    fn mk_key(&mut self) -> Key {
215        let k = Key {
216            id: self.next_key_id,
217        };
218        self.next_key_id += 1;
219        k
220    }
221}
222
223#[derive(Clone)]
224pub struct Insert<'t, In> {
225    pub key: Key,
226    pub new: &'t In,
227}
228
229#[derive(Clone)]
230pub struct Update<'t, In> {
231    pub key: Key,
232    pub new: &'t In,
233    pub existing: &'t In,
234}
235
236#[derive(Clone)]
237pub struct Remove<'t, In> {
238    pub key: Key,
239    pub existing: &'t In,
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    struct TrivialIndex;
247    impl<In> Index<In> for TrivialIndex {
248        fn insert(&mut self, _op: &Insert<In>) {}
249        fn remove(&mut self, _op: &Remove<In>) {}
250    }
251
252    #[test]
253    fn test_len() {
254        let mut collection = Collection::new(TrivialIndex);
255        assert_eq!(collection.len(), 0);
256
257        collection.insert(1);
258        assert_eq!(collection.len(), 1);
259
260        collection.insert(2);
261        assert_eq!(collection.len(), 2);
262
263        let key = collection.insert(3);
264        assert_eq!(collection.len(), 3);
265
266        collection.delete_by_key(&key);
267        assert_eq!(collection.len(), 2);
268    }
269
270    // See composable-indexes/src/lib.rs for more tests
271}
composable_indexes_core/collection.rs

composable_indexes_core/
collection.rs
