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
86/// An iterator over the potential value in a [`UniqueMap`] for a given key.
87pub struct UniqueMapPointIter<'a, V> {
88    /// The iterator seeking for matching keys in the range.
89    iter: IntoIter<&'a V>,
90}
91
92impl<'a, V> Iterator for UniqueMapPointIter<'a, V> {
93    type Item = &'a V;
94
95    fn next(&mut self) -> Option<Self::Item> {
96        self.iter.next()
97    }
98}
99
100/// An iterator over values in a [`UniqueMap`] where the keys are in a certain range.
101pub struct UniqueMapRangeIter<'a, K, V> {
102    /// The iterator seeking for matching keys in the range.
103    iter: Range<'a, K, V>,
104}
105
106impl<'a, K, V> Iterator for UniqueMapRangeIter<'a, K, V> {
107    type Item = &'a V;
108
109    fn next(&mut self) -> Option<Self::Item> {
110        self.iter.next().map(|(_, v)| v)
111    }
112}