composable_indexes/core/
collection.rs

1use crate::{
2    ShallowClone,
3    core::{DefaultStore, SEAL, store::Store},
4};
5
6use super::{
7    QueryResult, QueryResultDistinct,
8    index::{Index, Insert, Remove, Update},
9};
10
11/// Unique identifier for an item in a collection.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
13pub struct Key {
14    id: u64,
15}
16
17impl Key {
18    pub fn unsafe_from_u64(id: u64) -> Self {
19        Key { id }
20    }
21
22    pub fn as_u64(&self) -> u64 {
23        self.id
24    }
25}
26
27/// A collection of items, with an index that is automatically kept up-to-date.
28#[derive(Clone)]
29pub struct Collection<In, Ix, S = DefaultStore<In>> {
30    index: Ix,
31    data: S,
32    next_key_id: u64,
33    _marker: core::marker::PhantomData<fn() -> In>,
34}
35
36impl<In, Ix> Collection<In, Ix> {
37    /// Create an empty collection.
38    pub fn new(ix: Ix) -> Self
39    where
40        In: 'static,
41        Ix: Index<In>,
42        DefaultStore<In>: Store<In>,
43    {
44        Collection::new_with_empty_store(ix)
45    }
46}
47
48impl<In, Ix, S> ShallowClone for Collection<In, Ix, S>
49where
50    In: Clone,
51    Ix: ShallowClone,
52    S: ShallowClone,
53{
54}
55
56impl<In, Ix, S: Default> Collection<In, Ix, S> {
57    /// Create an empty collection.
58    pub fn new_with_empty_store(ix: Ix) -> Self {
59        Collection {
60            data: S::default(),
61            next_key_id: 0,
62            index: ix,
63            _marker: core::marker::PhantomData,
64        }
65    }
66}
67
68impl<In, Ix, S> Collection<In, Ix, S>
69where
70    S: Store<In>,
71    In: 'static,
72    Ix: Index<In>,
73{
74    /// Create a collection with a custom store.
75    ///
76    /// If the store is non-empty, the indexes will be populated from the store's contents. So
77    /// this may be an expensive operation depending on the size of the store.
78    pub fn new_with_store(store: S, mut ix: Ix) -> Self {
79        // Get the indexes up-to-date with the initial store contents
80        let mut max_key_id = 0;
81        for (k, v) in store.iter() {
82            max_key_id = max_key_id.max(k.id);
83            ix.insert(SEAL, &Insert { key: k, new: v });
84        }
85
86        // Then create the collection
87        Collection {
88            data: store,
89            next_key_id: max_key_id + 1,
90            index: ix,
91            _marker: core::marker::PhantomData,
92        }
93    }
94}
95
96impl<In, Ix, S> Collection<In, Ix, S> {
97    /// Access the underlying store.
98    pub fn store(&self) -> &S {
99        &self.data
100    }
101
102    /// Extract the underlying store.
103    pub fn extract_store(self) -> S {
104        self.data
105    }
106
107    /// Mutably access the underlying store.
108    /// Warning: Mutating the elements inside the store may lead to inconsistencies with the index
109    /// and cause undefined behavior. Use with caution. Ideally - only use it in a way that doesn't
110    /// modify the elements.
111    pub fn unsafe_mut_store(&mut self) -> &mut S {
112        &mut self.data
113    }
114}
115
116impl<In, Ix, S> Collection<In, Ix, S>
117where
118    In: 'static,
119    Ix: Index<In>,
120    S: Store<In>,
121{
122    /// Lookup an item in the collection by its key.
123    pub fn get_by_key(&self, key: Key) -> Option<&In> {
124        self.data.get(key)
125    }
126
127    /// Insert a new item into the collection.
128    pub fn insert(&mut self, value: In) -> Key {
129        let key = self.mk_key();
130        let existing = self.data.insert(key, value);
131
132        // There shouldn't be an existing key, as we just generated it
133        debug_assert!(existing.is_none());
134
135        self.index.insert(
136            SEAL,
137            &Insert {
138                key,
139                new: self.data.get_unwrapped(key),
140            },
141        );
142
143        key
144    }
145
146    /// Insert all items from an iterator into the collection.
147    pub fn insert_all<I>(&mut self, iter: I)
148    where
149        I: Iterator<Item = In>,
150    {
151        for item in iter {
152            self.insert(item);
153        }
154    }
155
156    /// Iterate over all items in the collection.
157    pub fn iter(&self) -> impl Iterator<Item = (Key, &In)> {
158        self.data.iter()
159    }
160
161    /// Mutate (or alter the presence of) the item in the collection.
162    pub fn update_by_key_mut<F>(&mut self, key: Key, f: F)
163    where
164        F: FnOnce(&mut Option<In>),
165    {
166        let mut existing = self.delete_by_key(key);
167        f(&mut existing);
168
169        if let Some(existing) = existing {
170            self.data.insert(key, existing);
171            self.index.insert(
172                SEAL,
173                &Insert {
174                    key,
175                    new: self.data.get_unwrapped(key),
176                },
177            );
178        }
179    }
180
181    /// Update the item in the collection.
182    pub fn update_by_key<F>(&mut self, key: Key, f: F)
183    where
184        F: FnOnce(Option<&In>) -> In,
185    {
186        let existing = self.data.get(key);
187        let new = f(existing);
188
189        match existing {
190            Some(existing) => {
191                self.index.update(
192                    SEAL,
193                    &Update {
194                        key,
195                        new: &new,
196                        existing,
197                    },
198                );
199                self.data.insert(key, new);
200            }
201            None => {
202                self.index.insert(SEAL, &Insert { key, new: &new });
203                self.data.insert(key, new);
204            }
205        };
206    }
207
208    /// Mutate the item in the collection, if it exists.
209    pub fn adjust_by_key_mut<F>(&mut self, key: Key, f: F)
210    where
211        F: FnOnce(&mut In),
212    {
213        if let Some(mut existing) = self.delete_by_key(key) {
214            f(&mut existing);
215            self.data.insert(key, existing);
216            self.index.insert(
217                SEAL,
218                &Insert {
219                    key,
220                    new: self.data.get_unwrapped(key),
221                },
222            );
223        }
224    }
225
226    /// Adjust the item in the collection, if it exists.
227    pub fn adjust_by_key<F>(&mut self, key: Key, f: F)
228    where
229        F: FnOnce(&In) -> In,
230    {
231        if let Some(existing) = self.data.get(key) {
232            let new = f(existing);
233            self.index.update(
234                SEAL,
235                &Update {
236                    key,
237                    new: &new,
238                    existing,
239                },
240            );
241            self.data.insert(key, new);
242        }
243    }
244
245    /// Remove an item from the collection, returning it if it exists.
246    pub fn delete_by_key(&mut self, key: Key) -> Option<In> {
247        let existing = self.data.remove(key);
248
249        if let Some(ref existing) = existing {
250            self.index.remove(SEAL, &Remove { key, existing });
251        }
252
253        existing
254    }
255
256    /// Perform a query using the collection's indexes, returning references to the items.
257    pub fn query<Res>(&self, f: impl FnOnce(&Ix) -> Res) -> Res::Resolved<&In>
258    where
259        Res: QueryResult,
260    {
261        let res = f(&self.index);
262        res.map(|k| self.data.get_unwrapped(k))
263    }
264
265    /// Perform a query using the collection's indexes, returning the keys of the items.
266    pub fn query_keys<Res>(&self, f: impl FnOnce(&Ix) -> Res) -> Res::Resolved<Key>
267    where
268        Res: QueryResult,
269    {
270        let res = f(&self.index);
271        res.map(|k| k)
272    }
273
274    /// Perform a query using the collection's indexes, returning both the keys and the items.
275    pub fn query_with_keys<Res>(&self, f: impl FnOnce(&Ix) -> Res) -> Res::Resolved<(Key, &In)>
276    where
277        Res: QueryResult,
278    {
279        let res = f(&self.index);
280        res.map(|k| (k, self.data.get_unwrapped(k)))
281    }
282
283    /// Delete all items matching the query.
284    pub fn delete<Res>(&mut self, f: impl FnOnce(&Ix) -> Res) -> usize
285    where
286        Res: QueryResult,
287    {
288        let mut affected_count = 0;
289        let res = f(&self.index);
290        res.map(|key| {
291            self.delete_by_key(key);
292            affected_count += 1;
293        });
294        affected_count
295    }
296
297    /// Update all items matching the query with the provided function.
298    pub fn update<Res, F>(
299        &mut self,
300        f: impl FnOnce(&Ix) -> Res,
301        update_fn: impl Fn(&In) -> In,
302    ) -> Res::Resolved<()>
303    where
304        Res: QueryResultDistinct,
305    {
306        let res = f(&self.index);
307        res.map(|key| {
308            self.data.update(key, |existing| {
309                let new = update_fn(existing);
310                self.index.update(
311                    SEAL,
312                    &Update {
313                        key,
314                        new: &new,
315                        existing,
316                    },
317                );
318                new
319            });
320        })
321    }
322
323    /// Remove and return all items matching the query.
324    pub fn take<Res>(&mut self, f: impl FnOnce(&Ix) -> Res) -> Res::Resolved<In>
325    where
326        Res: QueryResultDistinct,
327    {
328        let res = f(&self.index);
329        res.map(|k| self.delete_by_key(k).unwrap())
330    }
331
332    /// Number of items in the collection.
333    pub fn len(&self) -> usize {
334        self.data.len()
335    }
336
337    pub fn is_empty(&self) -> bool {
338        self.data.is_empty()
339    }
340
341    fn mk_key(&mut self) -> Key {
342        let k = Key {
343            id: self.next_key_id,
344        };
345        self.next_key_id += 1;
346        k
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    #[test]
355    fn test_len() {
356        let mut collection = Collection::new(());
357        assert_eq!(collection.len(), 0);
358
359        collection.insert(1);
360        assert_eq!(collection.len(), 1);
361
362        collection.insert(2);
363        assert_eq!(collection.len(), 2);
364
365        let key = collection.insert(3);
366        assert_eq!(collection.len(), 3);
367
368        collection.delete_by_key(key);
369        assert_eq!(collection.len(), 2);
370    }
371
372    // See composable-indexes/src/lib.rs for more tests
373}
composable_indexes/core/collection.rs

composable_indexes/core/
collection.rs
