spacetimedb_table/table_index/
uniquemap.rs

1use crate::MemoryUsage;
2use core::{ops::RangeBounds, option::IntoIter};
3use std::collections::btree_map::{BTreeMap, Entry, Range};
4
5/// A "unique map" that relates a `K` to a `V`.
6///
7/// (This is just a `BTreeMap<K, V>`) with a slightly modified interface.
8#[derive(Debug, PartialEq, Eq, Clone)]
9pub struct UniqueMap<K, V> {
10    /// The map is backed by a `BTreeMap` for relating a key to a value.
11    map: BTreeMap<K, V>,
12}
13
14impl<K, V> Default for UniqueMap<K, V> {
15    fn default() -> Self {
16        Self { map: BTreeMap::new() }
17    }
18}
19
20impl<K: MemoryUsage, V: MemoryUsage> MemoryUsage for UniqueMap<K, V> {
21    fn heap_usage(&self) -> usize {
22        let Self { map } = self;
23        map.heap_usage()
24    }
25}
26
27impl<K: Ord, V: Ord> UniqueMap<K, V> {
28    /// Inserts the relation `key -> val` to this map.
29    ///
30    /// If `key` was already present in the map, does not add an association with `val`.
31    /// Returns the existing associated value instead.
32    pub fn insert(&mut self, key: K, val: V) -> Result<(), &V> {
33        match self.map.entry(key) {
34            Entry::Vacant(e) => {
35                e.insert(val);
36                Ok(())
37            }
38            Entry::Occupied(e) => Err(e.into_mut()),
39        }
40    }
41
42    /// Deletes `key` from this map.
43    ///
44    /// Returns whether `key` was present.
45    pub fn delete(&mut self, key: &K) -> bool {
46        self.map.remove(key).is_some()
47    }
48
49    /// Returns an iterator over the map that yields all the `V`s
50    /// of the `K`s that fall within the specified `range`.
51    pub fn values_in_range(&self, range: &impl RangeBounds<K>) -> UniqueMapRangeIter<'_, K, V> {
52        UniqueMapRangeIter {
53            iter: self.map.range((range.start_bound(), range.end_bound())),
54        }
55    }
56
57    /// Returns an iterator over the map that yields the potential `V` of the `key: &K`.
58    pub fn values_in_point(&self, key: &K) -> UniqueMapPointIter<'_, V> {
59        let iter = self.map.get(key).into_iter();
60        UniqueMapPointIter { iter }
61    }
62
63    /// Returns the number of unique keys in the map.
64    pub fn num_keys(&self) -> usize {
65        self.len()
66    }
67
68    /// Returns the total number of entries in the map.s
69    pub fn len(&self) -> usize {
70        self.map.len()
71    }
72
73    /// Returns whether there are any entries in the map.
74    #[allow(unused)] // No use for this currently.
75    pub fn is_empty(&self) -> bool {
76        self.len() == 0
77    }
78
79    /// Deletes all entries from the map, leaving it empty.
80    /// This will not deallocate the outer map.
81    pub fn clear(&mut self) {
82        self.map.clear();
83    }
84
85    /// Returns whether `other` can be merged into `self`
86    /// with an error containing the element in `self` that caused the violation.
87    ///
88    /// The closure `ignore` indicates whether a row in `self` should be ignored.
89    pub(crate) fn can_merge(&self, other: &Self, ignore: impl Fn(&V) -> bool) -> Result<(), &V> {
90        let Some(found) = other
91            .map
92            .keys()
93            .find_map(|key| self.map.get(key).filter(|val| !ignore(val)))
94        else {
95            return Ok(());
96        };
97        Err(found)
98    }
99}
100
101/// An iterator over the potential value in a [`UniqueMap`] for a given key.
102pub struct UniqueMapPointIter<'a, V> {
103    /// The iterator seeking for matching keys in the range.
104    iter: IntoIter<&'a V>,
105}
106
107impl<'a, V> Iterator for UniqueMapPointIter<'a, V> {
108    type Item = &'a V;
109
110    fn next(&mut self) -> Option<Self::Item> {
111        self.iter.next()
112    }
113}
114
115/// An iterator over values in a [`UniqueMap`] where the keys are in a certain range.
116pub struct UniqueMapRangeIter<'a, K, V> {
117    /// The iterator seeking for matching keys in the range.
118    iter: Range<'a, K, V>,
119}
120
121impl<'a, K, V> Iterator for UniqueMapRangeIter<'a, K, V> {
122    type Item = &'a V;
123
124    fn next(&mut self) -> Option<Self::Item> {
125        self.iter.next().map(|(_, v)| v)
126    }
127}