zed_sum_tree/
tree_map.rs

1use std::{cmp::Ordering, fmt::Debug};
2
3use crate::{Bias, ContextLessSummary, Dimension, Edit, Item, KeyedItem, SeekTarget, SumTree};
4
5/// A cheaply-cloneable ordered map based on a [SumTree](crate::SumTree).
6#[derive(Clone, PartialEq, Eq)]
7pub struct TreeMap<K, V>(SumTree<MapEntry<K, V>>)
8where
9    K: Clone + Ord,
10    V: Clone;
11
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct MapEntry<K, V> {
14    key: K,
15    value: V,
16}
17
18#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
19pub struct MapKey<K>(Option<K>);
20
21impl<K> Default for MapKey<K> {
22    fn default() -> Self {
23        Self(None)
24    }
25}
26
27#[derive(Clone, Debug)]
28pub struct MapKeyRef<'a, K>(Option<&'a K>);
29
30impl<K> Default for MapKeyRef<'_, K> {
31    fn default() -> Self {
32        Self(None)
33    }
34}
35
36#[derive(Clone, Debug, PartialEq, Eq)]
37pub struct TreeSet<K>(TreeMap<K, ()>)
38where
39    K: Clone + Ord;
40
41impl<K: Clone + Ord, V: Clone> TreeMap<K, V> {
42    pub fn from_ordered_entries(entries: impl IntoIterator<Item = (K, V)>) -> Self {
43        let tree = SumTree::from_iter(
44            entries
45                .into_iter()
46                .map(|(key, value)| MapEntry { key, value }),
47            (),
48        );
49        Self(tree)
50    }
51
52    pub fn is_empty(&self) -> bool {
53        self.0.is_empty()
54    }
55
56    pub fn get(&self, key: &K) -> Option<&V> {
57        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(());
58        cursor.seek(&MapKeyRef(Some(key)), Bias::Left);
59        if let Some(item) = cursor.item() {
60            if Some(key) == item.key().0.as_ref() {
61                Some(&item.value)
62            } else {
63                None
64            }
65        } else {
66            None
67        }
68    }
69
70    pub fn insert(&mut self, key: K, value: V) {
71        self.0.insert_or_replace(MapEntry { key, value }, ());
72    }
73
74    pub fn extend(&mut self, iter: impl IntoIterator<Item = (K, V)>) {
75        let edits: Vec<_> = iter
76            .into_iter()
77            .map(|(key, value)| Edit::Insert(MapEntry { key, value }))
78            .collect();
79        self.0.edit(edits, ());
80    }
81
82    pub fn clear(&mut self) {
83        self.0 = SumTree::default();
84    }
85
86    pub fn remove(&mut self, key: &K) -> Option<V> {
87        let mut removed = None;
88        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(());
89        let key = MapKeyRef(Some(key));
90        let mut new_tree = cursor.slice(&key, Bias::Left);
91        if key.cmp(&cursor.end(), ()) == Ordering::Equal {
92            removed = Some(cursor.item().unwrap().value.clone());
93            cursor.next();
94        }
95        new_tree.append(cursor.suffix(), ());
96        drop(cursor);
97        self.0 = new_tree;
98        removed
99    }
100
101    pub fn remove_range(&mut self, start: &impl MapSeekTarget<K>, end: &impl MapSeekTarget<K>) {
102        let start = MapSeekTargetAdaptor(start);
103        let end = MapSeekTargetAdaptor(end);
104        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(());
105        let mut new_tree = cursor.slice(&start, Bias::Left);
106        cursor.seek(&end, Bias::Left);
107        new_tree.append(cursor.suffix(), ());
108        drop(cursor);
109        self.0 = new_tree;
110    }
111
112    /// Returns the key-value pair with the greatest key less than or equal to the given key.
113    pub fn closest(&self, key: &K) -> Option<(&K, &V)> {
114        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(());
115        let key = MapKeyRef(Some(key));
116        cursor.seek(&key, Bias::Right);
117        cursor.prev();
118        cursor.item().map(|item| (&item.key, &item.value))
119    }
120
121    pub fn iter_from<'a>(&'a self, from: &K) -> impl Iterator<Item = (&'a K, &'a V)> + 'a {
122        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(());
123        let from_key = MapKeyRef(Some(from));
124        cursor.seek(&from_key, Bias::Left);
125
126        cursor.map(|map_entry| (&map_entry.key, &map_entry.value))
127    }
128
129    pub fn update<F, T>(&mut self, key: &K, f: F) -> Option<T>
130    where
131        F: FnOnce(&mut V) -> T,
132    {
133        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(());
134        let key = MapKeyRef(Some(key));
135        let mut new_tree = cursor.slice(&key, Bias::Left);
136        let mut result = None;
137        if key.cmp(&cursor.end(), ()) == Ordering::Equal {
138            let mut updated = cursor.item().unwrap().clone();
139            result = Some(f(&mut updated.value));
140            new_tree.push(updated, ());
141            cursor.next();
142        }
143        new_tree.append(cursor.suffix(), ());
144        drop(cursor);
145        self.0 = new_tree;
146        result
147    }
148
149    pub fn retain<F: FnMut(&K, &V) -> bool>(&mut self, mut predicate: F) {
150        let mut new_map = SumTree::<MapEntry<K, V>>::default();
151
152        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>(());
153        cursor.next();
154        while let Some(item) = cursor.item() {
155            if predicate(&item.key, &item.value) {
156                new_map.push(item.clone(), ());
157            }
158            cursor.next();
159        }
160        drop(cursor);
161
162        self.0 = new_map;
163    }
164
165    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> + '_ {
166        self.0.iter().map(|entry| (&entry.key, &entry.value))
167    }
168
169    pub fn values(&self) -> impl Iterator<Item = &V> + '_ {
170        self.0.iter().map(|entry| &entry.value)
171    }
172
173    pub fn first(&self) -> Option<(&K, &V)> {
174        self.0.first().map(|entry| (&entry.key, &entry.value))
175    }
176
177    pub fn last(&self) -> Option<(&K, &V)> {
178        self.0.last().map(|entry| (&entry.key, &entry.value))
179    }
180
181    pub fn insert_tree(&mut self, other: TreeMap<K, V>) {
182        let edits = other
183            .iter()
184            .map(|(key, value)| {
185                Edit::Insert(MapEntry {
186                    key: key.to_owned(),
187                    value: value.to_owned(),
188                })
189            })
190            .collect();
191
192        self.0.edit(edits, ());
193    }
194}
195
196impl<K, V> Debug for TreeMap<K, V>
197where
198    K: Clone + Debug + Ord,
199    V: Clone + Debug,
200{
201    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202        f.debug_map().entries(self.iter()).finish()
203    }
204}
205
206#[derive(Debug)]
207struct MapSeekTargetAdaptor<'a, T>(&'a T);
208
209impl<'a, K: Clone + Ord, T: MapSeekTarget<K>> SeekTarget<'a, MapKey<K>, MapKeyRef<'a, K>>
210    for MapSeekTargetAdaptor<'_, T>
211{
212    fn cmp(&self, cursor_location: &MapKeyRef<K>, _: ()) -> Ordering {
213        if let Some(key) = &cursor_location.0 {
214            MapSeekTarget::cmp_cursor(self.0, key)
215        } else {
216            Ordering::Greater
217        }
218    }
219}
220
221pub trait MapSeekTarget<K> {
222    fn cmp_cursor(&self, cursor_location: &K) -> Ordering;
223}
224
225impl<K: Ord> MapSeekTarget<K> for K {
226    fn cmp_cursor(&self, cursor_location: &K) -> Ordering {
227        self.cmp(cursor_location)
228    }
229}
230
231impl<K, V> Default for TreeMap<K, V>
232where
233    K: Clone + Ord,
234    V: Clone,
235{
236    fn default() -> Self {
237        Self(Default::default())
238    }
239}
240
241impl<K, V> Item for MapEntry<K, V>
242where
243    K: Clone + Ord,
244    V: Clone,
245{
246    type Summary = MapKey<K>;
247
248    fn summary(&self, _cx: ()) -> Self::Summary {
249        self.key()
250    }
251}
252
253impl<K, V> KeyedItem for MapEntry<K, V>
254where
255    K: Clone + Ord,
256    V: Clone,
257{
258    type Key = MapKey<K>;
259
260    fn key(&self) -> Self::Key {
261        MapKey(Some(self.key.clone()))
262    }
263}
264
265impl<K> ContextLessSummary for MapKey<K>
266where
267    K: Clone,
268{
269    fn zero() -> Self {
270        Default::default()
271    }
272
273    fn add_summary(&mut self, summary: &Self) {
274        *self = summary.clone()
275    }
276}
277
278impl<'a, K> Dimension<'a, MapKey<K>> for MapKeyRef<'a, K>
279where
280    K: Clone + Ord,
281{
282    fn zero(_cx: ()) -> Self {
283        Default::default()
284    }
285
286    fn add_summary(&mut self, summary: &'a MapKey<K>, _: ()) {
287        self.0 = summary.0.as_ref();
288    }
289}
290
291impl<'a, K> SeekTarget<'a, MapKey<K>, MapKeyRef<'a, K>> for MapKeyRef<'_, K>
292where
293    K: Clone + Ord,
294{
295    fn cmp(&self, cursor_location: &MapKeyRef<K>, _: ()) -> Ordering {
296        Ord::cmp(&self.0, &cursor_location.0)
297    }
298}
299
300impl<K> Default for TreeSet<K>
301where
302    K: Clone + Ord,
303{
304    fn default() -> Self {
305        Self(Default::default())
306    }
307}
308
309impl<K> TreeSet<K>
310where
311    K: Clone + Ord,
312{
313    pub fn from_ordered_entries(entries: impl IntoIterator<Item = K>) -> Self {
314        Self(TreeMap::from_ordered_entries(
315            entries.into_iter().map(|key| (key, ())),
316        ))
317    }
318
319    pub fn is_empty(&self) -> bool {
320        self.0.is_empty()
321    }
322
323    pub fn insert(&mut self, key: K) {
324        self.0.insert(key, ());
325    }
326
327    pub fn remove(&mut self, key: &K) -> bool {
328        self.0.remove(key).is_some()
329    }
330
331    pub fn extend(&mut self, iter: impl IntoIterator<Item = K>) {
332        self.0.extend(iter.into_iter().map(|key| (key, ())));
333    }
334
335    pub fn contains(&self, key: &K) -> bool {
336        self.0.get(key).is_some()
337    }
338
339    pub fn iter(&self) -> impl Iterator<Item = &K> + '_ {
340        self.0.iter().map(|(k, _)| k)
341    }
342
343    pub fn iter_from<'a>(&'a self, key: &K) -> impl Iterator<Item = &'a K> + 'a {
344        self.0.iter_from(key).map(move |(k, _)| k)
345    }
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351
352    #[test]
353    fn test_basic() {
354        let mut map = TreeMap::default();
355        assert_eq!(map.iter().collect::<Vec<_>>(), vec![]);
356
357        map.insert(3, "c");
358        assert_eq!(map.get(&3), Some(&"c"));
359        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&3, &"c")]);
360
361        map.insert(1, "a");
362        assert_eq!(map.get(&1), Some(&"a"));
363        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&1, &"a"), (&3, &"c")]);
364
365        map.insert(2, "b");
366        assert_eq!(map.get(&2), Some(&"b"));
367        assert_eq!(map.get(&1), Some(&"a"));
368        assert_eq!(map.get(&3), Some(&"c"));
369        assert_eq!(
370            map.iter().collect::<Vec<_>>(),
371            vec![(&1, &"a"), (&2, &"b"), (&3, &"c")]
372        );
373
374        assert_eq!(map.closest(&0), None);
375        assert_eq!(map.closest(&1), Some((&1, &"a")));
376        assert_eq!(map.closest(&10), Some((&3, &"c")));
377
378        map.remove(&2);
379        assert_eq!(map.get(&2), None);
380        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&1, &"a"), (&3, &"c")]);
381
382        assert_eq!(map.closest(&2), Some((&1, &"a")));
383
384        map.remove(&3);
385        assert_eq!(map.get(&3), None);
386        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&1, &"a")]);
387
388        map.remove(&1);
389        assert_eq!(map.get(&1), None);
390        assert_eq!(map.iter().collect::<Vec<_>>(), vec![]);
391
392        map.insert(4, "d");
393        map.insert(5, "e");
394        map.insert(6, "f");
395        map.retain(|key, _| *key % 2 == 0);
396        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&4, &"d"), (&6, &"f")]);
397    }
398
399    #[test]
400    fn test_iter_from() {
401        let mut map = TreeMap::default();
402
403        map.insert("a", 1);
404        map.insert("b", 2);
405        map.insert("baa", 3);
406        map.insert("baaab", 4);
407        map.insert("c", 5);
408
409        let result = map
410            .iter_from(&"ba")
411            .take_while(|(key, _)| key.starts_with("ba"))
412            .collect::<Vec<_>>();
413
414        assert_eq!(result.len(), 2);
415        assert!(result.iter().any(|(k, _)| k == &&"baa"));
416        assert!(result.iter().any(|(k, _)| k == &&"baaab"));
417
418        let result = map
419            .iter_from(&"c")
420            .take_while(|(key, _)| key.starts_with("c"))
421            .collect::<Vec<_>>();
422
423        assert_eq!(result.len(), 1);
424        assert!(result.iter().any(|(k, _)| k == &&"c"));
425    }
426
427    #[test]
428    fn test_insert_tree() {
429        let mut map = TreeMap::default();
430        map.insert("a", 1);
431        map.insert("b", 2);
432        map.insert("c", 3);
433
434        let mut other = TreeMap::default();
435        other.insert("a", 2);
436        other.insert("b", 2);
437        other.insert("d", 4);
438
439        map.insert_tree(other);
440
441        assert_eq!(map.iter().count(), 4);
442        assert_eq!(map.get(&"a"), Some(&2));
443        assert_eq!(map.get(&"b"), Some(&2));
444        assert_eq!(map.get(&"c"), Some(&3));
445        assert_eq!(map.get(&"d"), Some(&4));
446    }
447
448    #[test]
449    fn test_extend() {
450        let mut map = TreeMap::default();
451        map.insert("a", 1);
452        map.insert("b", 2);
453        map.insert("c", 3);
454        map.extend([("a", 2), ("b", 2), ("d", 4)]);
455        assert_eq!(map.iter().count(), 4);
456        assert_eq!(map.get(&"a"), Some(&2));
457        assert_eq!(map.get(&"b"), Some(&2));
458        assert_eq!(map.get(&"c"), Some(&3));
459        assert_eq!(map.get(&"d"), Some(&4));
460    }
461
462    #[test]
463    fn test_remove_between_and_path_successor() {
464        use std::path::{Path, PathBuf};
465
466        #[derive(Debug)]
467        pub struct PathDescendants<'a>(&'a Path);
468
469        impl MapSeekTarget<PathBuf> for PathDescendants<'_> {
470            fn cmp_cursor(&self, key: &PathBuf) -> Ordering {
471                if key.starts_with(self.0) {
472                    Ordering::Greater
473                } else {
474                    self.0.cmp(key)
475                }
476            }
477        }
478
479        let mut map = TreeMap::default();
480
481        map.insert(PathBuf::from("a"), 1);
482        map.insert(PathBuf::from("a/a"), 1);
483        map.insert(PathBuf::from("b"), 2);
484        map.insert(PathBuf::from("b/a/a"), 3);
485        map.insert(PathBuf::from("b/a/a/a/b"), 4);
486        map.insert(PathBuf::from("c"), 5);
487        map.insert(PathBuf::from("c/a"), 6);
488
489        map.remove_range(
490            &PathBuf::from("b/a"),
491            &PathDescendants(&PathBuf::from("b/a")),
492        );
493
494        assert_eq!(map.get(&PathBuf::from("a")), Some(&1));
495        assert_eq!(map.get(&PathBuf::from("a/a")), Some(&1));
496        assert_eq!(map.get(&PathBuf::from("b")), Some(&2));
497        assert_eq!(map.get(&PathBuf::from("b/a/a")), None);
498        assert_eq!(map.get(&PathBuf::from("b/a/a/a/b")), None);
499        assert_eq!(map.get(&PathBuf::from("c")), Some(&5));
500        assert_eq!(map.get(&PathBuf::from("c/a")), Some(&6));
501
502        map.remove_range(&PathBuf::from("c"), &PathDescendants(&PathBuf::from("c")));
503
504        assert_eq!(map.get(&PathBuf::from("a")), Some(&1));
505        assert_eq!(map.get(&PathBuf::from("a/a")), Some(&1));
506        assert_eq!(map.get(&PathBuf::from("b")), Some(&2));
507        assert_eq!(map.get(&PathBuf::from("c")), None);
508        assert_eq!(map.get(&PathBuf::from("c/a")), None);
509
510        map.remove_range(&PathBuf::from("a"), &PathDescendants(&PathBuf::from("a")));
511
512        assert_eq!(map.get(&PathBuf::from("a")), None);
513        assert_eq!(map.get(&PathBuf::from("a/a")), None);
514        assert_eq!(map.get(&PathBuf::from("b")), Some(&2));
515
516        map.remove_range(&PathBuf::from("b"), &PathDescendants(&PathBuf::from("b")));
517
518        assert_eq!(map.get(&PathBuf::from("b")), None);
519    }
520}
zed_sum_tree/tree_map.rs

zed_sum_tree/
tree_map.rs
