1use std::{borrow::Borrow, collections::HashMap, hash::Hash};
18
19use eyeball_im::{ObservableVector, Vector, VectorDiff};
20use futures_util::Stream;
21
22#[derive(Debug)]
44pub(crate) struct ObservableMap<K, V>
45where
46 V: Clone + 'static,
47{
48 mapping: HashMap<K, usize>,
50
51 values: ObservableVector<V>,
54}
55
56impl<K, V> ObservableMap<K, V>
57where
58 K: Hash + Eq,
59 V: Clone + 'static,
60{
61 pub(crate) fn new() -> Self {
63 Self { mapping: HashMap::new(), values: ObservableVector::new() }
64 }
65
66 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 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 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 pub(crate) fn iter(&self) -> impl Iterator<Item = &V> {
124 self.values.iter()
125 }
126
127 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 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 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 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 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 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 map.insert('b', 'f');
196
197 assert_eq!(map.get_or_create(&'a', || 'E'), &'E');
199 assert_eq!(map.get_or_create(&'b', || 'F'), &'f'); 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 assert_eq!(map.remove(&'b'), Some('f'));
208
209 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 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 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 assert_eq!(map.remove(&'c'), None);
240
241 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 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 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 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 map.insert('b', 'F');
285 assert_next_eq!(stream, vec![VectorDiff::Set { index: 0, value: 'F' }]);
286
287 assert_pending!(stream);
288
289 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}