matrix_sdk_base/store/
observable_map.rs

1// Copyright 2024 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! An [`ObservableMap`] implementation.
16
17use std::{borrow::Borrow, collections::HashMap, hash::Hash};
18
19use eyeball_im::{ObservableVector, Vector, VectorDiff};
20use futures_util::Stream;
21
22/// An observable map.
23///
24/// This is an “observable map” naive implementation. Just like regular
25/// hashmap, we have a redirection from a key to a position, and from a
26/// position to a value. The (key, position) tuples are stored in an
27/// [`HashMap`]. The (position, value) tuples are stored in an
28/// [`ObservableVector`]. The (key, position) tuple is only provided for
29/// fast _reading_ implementations, like `Self::get` and
30/// `Self::get_or_create`. The (position, value) tuples are observable,
31/// this is what interests us the most here.
32///
33/// Why not implementing a new `ObservableMap` type in `eyeball-im` instead
34/// of this custom implementation? Because we want to continue providing
35/// `VectorDiff` when observing the changes, so that the rest of the API in
36/// the Matrix Rust SDK aren't broken. Indeed, an `ObservableMap` must
37/// produce `MapDiff`, which would be quite different.
38/// Plus, we would like to re-use all our existing code, test, stream
39/// adapters and so on.
40///
41/// This is a trade-off. This implementation is simple enough for the
42/// moment, and basically does the job.
43#[derive(Debug)]
44pub(crate) struct ObservableMap<K, V>
45where
46    V: Clone + 'static,
47{
48    /// The (key, position) tuples.
49    mapping: HashMap<K, usize>,
50
51    /// The values where the indices are the `position` part of
52    /// `Self::mapping`.
53    values: ObservableVector<V>,
54}
55
56impl<K, V> ObservableMap<K, V>
57where
58    K: Hash + Eq,
59    V: Clone + 'static,
60{
61    /// Create a new `Self`.
62    pub(crate) fn new() -> Self {
63        Self { mapping: HashMap::new(), values: ObservableVector::new() }
64    }
65
66    /// Insert a new `V` in the collection.
67    ///
68    /// If the `V` value already exists, it will be updated to the new one.
69    pub(crate) fn insert(&mut self, key: K, value: V) -> usize {
70        match self.mapping.get(&key) {
71            Some(position) => {
72                self.values.set(*position, value);
73
74                *position
75            }
76            None => {
77                let position = self.values.len();
78
79                self.values.push_back(value);
80                self.mapping.insert(key, position);
81
82                position
83            }
84        }
85    }
86
87    /// Reading one `V` value based on their ID, if it exists.
88    pub(crate) fn get<L>(&self, key: &L) -> Option<&V>
89    where
90        K: Borrow<L>,
91        L: Hash + Eq + ?Sized,
92    {
93        self.mapping.get(key).and_then(|position| self.values.get(*position))
94    }
95
96    /// Reading one `V` value based on their ID, or create a new one (by
97    /// using `default`).
98    pub(crate) fn get_or_create<L, F>(&mut self, key: &L, default: F) -> &V
99    where
100        K: Borrow<L>,
101        L: Hash + Eq + ?Sized + ToOwned<Owned = K>,
102        F: FnOnce() -> V,
103    {
104        let position = match self.mapping.get(key) {
105            Some(position) => *position,
106            None => {
107                let value = default();
108                let position = self.values.len();
109
110                self.values.push_back(value);
111                self.mapping.insert(key.to_owned(), position);
112
113                position
114            }
115        };
116
117        self.values
118            .get(position)
119            .expect("Value should be present or has just been inserted, but it's missing")
120    }
121
122    /// Return an iterator over the existing values.
123    pub(crate) fn iter(&self) -> impl Iterator<Item = &V> {
124        self.values.iter()
125    }
126
127    /// Get a [`Stream`] of the values.
128    pub(crate) fn stream(&self) -> (Vector<V>, impl Stream<Item = Vec<VectorDiff<V>>>) {
129        self.values.subscribe().into_values_and_batched_stream()
130    }
131
132    /// Remove a `V` value based on their ID, if it exists.
133    ///
134    /// Returns the removed value.
135    pub(crate) fn remove<L>(&mut self, key: &L) -> Option<V>
136    where
137        K: Borrow<L>,
138        L: Hash + Eq + ?Sized,
139    {
140        let position = self.mapping.remove(key)?;
141
142        // Reindex every mapped entry that is after the position we're looking to
143        // remove.
144        for mapped_pos in self.mapping.values_mut().filter(|pos| **pos > position) {
145            *mapped_pos = mapped_pos.saturating_sub(1);
146        }
147
148        Some(self.values.remove(position))
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use eyeball_im::VectorDiff;
155    use stream_assert::{assert_closed, assert_next_eq, assert_pending};
156
157    use super::ObservableMap;
158
159    #[test]
160    fn test_insert_and_get() {
161        let mut map = ObservableMap::<char, char>::new();
162
163        assert!(map.get(&'a').is_none());
164        assert!(map.get(&'b').is_none());
165        assert!(map.get(&'c').is_none());
166
167        // new items
168        map.insert('a', 'e');
169        map.insert('b', 'f');
170
171        assert_eq!(map.get(&'a'), Some(&'e'));
172        assert_eq!(map.get(&'b'), Some(&'f'));
173        assert!(map.get(&'c').is_none());
174
175        // one new item
176        map.insert('c', 'g');
177
178        assert_eq!(map.get(&'a'), Some(&'e'));
179        assert_eq!(map.get(&'b'), Some(&'f'));
180        assert_eq!(map.get(&'c'), Some(&'g'));
181
182        // update one item
183        map.insert('b', 'F');
184
185        assert_eq!(map.get(&'a'), Some(&'e'));
186        assert_eq!(map.get(&'b'), Some(&'F'));
187        assert_eq!(map.get(&'c'), Some(&'g'));
188    }
189
190    #[test]
191    fn test_get_or_create() {
192        let mut map = ObservableMap::<char, char>::new();
193
194        // insert one item
195        map.insert('b', 'f');
196
197        // get or create many items
198        assert_eq!(map.get_or_create(&'a', || 'E'), &'E');
199        assert_eq!(map.get_or_create(&'b', || 'F'), &'f'); // this one already exists
200        assert_eq!(map.get_or_create(&'c', || 'G'), &'G');
201
202        assert_eq!(map.get(&'a'), Some(&'E'));
203        assert_eq!(map.get(&'b'), Some(&'f'));
204        assert_eq!(map.get(&'c'), Some(&'G'));
205
206        // remove non-last item
207        assert_eq!(map.remove(&'b'), Some('f'));
208
209        // get_or_create item after the removed one
210        assert_eq!(map.get_or_create(&'c', || 'G'), &'G');
211    }
212
213    #[test]
214    fn test_remove() {
215        let mut map = ObservableMap::<char, char>::new();
216
217        assert!(map.get(&'a').is_none());
218        assert!(map.get(&'b').is_none());
219        assert!(map.get(&'c').is_none());
220
221        // new items
222        map.insert('a', 'e');
223        map.insert('b', 'f');
224        map.insert('c', 'g');
225
226        assert_eq!(map.get(&'a'), Some(&'e'));
227        assert_eq!(map.get(&'b'), Some(&'f'));
228        assert_eq!(map.get(&'c'), Some(&'g'));
229        assert!(map.get(&'d').is_none());
230
231        // remove last item
232        assert_eq!(map.remove(&'c'), Some('g'));
233
234        assert_eq!(map.get(&'a'), Some(&'e'));
235        assert_eq!(map.get(&'b'), Some(&'f'));
236        assert_eq!(map.get(&'c'), None);
237
238        // remove a non-existent item
239        assert_eq!(map.remove(&'c'), None);
240
241        // remove a non-last item
242        assert_eq!(map.remove(&'a'), Some('e'));
243        assert_eq!(map.get(&'b'), Some(&'f'));
244    }
245
246    #[test]
247    fn test_iter() {
248        let mut map = ObservableMap::<char, char>::new();
249
250        // new items
251        map.insert('a', 'e');
252        map.insert('b', 'f');
253        map.insert('c', 'g');
254
255        assert_eq!(
256            map.iter().map(|c| c.to_ascii_uppercase()).collect::<Vec<_>>(),
257            &['E', 'F', 'G']
258        );
259    }
260
261    #[test]
262    fn test_stream() {
263        let mut map = ObservableMap::<char, char>::new();
264
265        // insert one item
266        map.insert('b', 'f');
267
268        let (initial_values, mut stream) = map.stream();
269        assert_eq!(initial_values.iter().copied().collect::<Vec<_>>(), &['f']);
270
271        assert_pending!(stream);
272
273        // insert two items
274        map.insert('c', 'g');
275        map.insert('a', 'e');
276        assert_next_eq!(
277            stream,
278            vec![VectorDiff::PushBack { value: 'g' }, VectorDiff::PushBack { value: 'e' }]
279        );
280
281        assert_pending!(stream);
282
283        // update one item
284        map.insert('b', 'F');
285        assert_next_eq!(stream, vec![VectorDiff::Set { index: 0, value: 'F' }]);
286
287        assert_pending!(stream);
288
289        // remove one item
290        map.remove(&'b');
291        assert_next_eq!(stream, vec![VectorDiff::Remove { index: 0 }]);
292
293        assert_pending!(stream);
294
295        drop(map);
296        assert_closed!(stream);
297    }
298}