composable_indexes/core/
collection.rs

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