dynomite/util/rbtree.rs
1//! Ordered key/value map with a `lower_bound` operation.
2//!
3//! Used for token-ring lookups, where callers want the smallest key
4//! greater than or equal to a probe value. This wraps
5//! [`std::collections::BTreeMap`], which supports both ordered
6//! iteration and `range` queries efficiently.
7
8use std::collections::BTreeMap;
9
10/// Ordered key/value map backed by [`BTreeMap`].
11///
12/// # Examples
13///
14/// ```
15/// use dynomite::util::rbtree::OrderedMap;
16///
17/// let mut t: OrderedMap<u64, &'static str> = OrderedMap::new();
18/// t.insert(10, "ten");
19/// t.insert(30, "thirty");
20/// t.insert(20, "twenty");
21/// assert_eq!(t.lower_bound(&15), Some((&20, &"twenty")));
22/// assert_eq!(t.min(), Some((&10, &"ten")));
23/// ```
24#[derive(Debug, Default, Clone)]
25pub struct OrderedMap<K: Ord, V> {
26 inner: BTreeMap<K, V>,
27}
28
29impl<K: Ord, V> OrderedMap<K, V> {
30 /// Construct an empty map.
31 ///
32 /// # Examples
33 ///
34 /// ```
35 /// use dynomite::util::rbtree::OrderedMap;
36 /// let m: OrderedMap<u32, u32> = OrderedMap::new();
37 /// assert!(m.is_empty());
38 /// ```
39 pub fn new() -> Self {
40 Self {
41 inner: BTreeMap::new(),
42 }
43 }
44
45 /// Insert `value` at `key`, returning the previous value (if any).
46 ///
47 /// # Examples
48 ///
49 /// ```
50 /// use dynomite::util::rbtree::OrderedMap;
51 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
52 /// assert_eq!(m.insert(1, 10), None);
53 /// assert_eq!(m.insert(1, 20), Some(10));
54 /// ```
55 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
56 self.inner.insert(key, value)
57 }
58
59 /// Remove the entry at `key`.
60 ///
61 /// # Examples
62 ///
63 /// ```
64 /// use dynomite::util::rbtree::OrderedMap;
65 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
66 /// m.insert(1, 10);
67 /// assert_eq!(m.remove(&1), Some(10));
68 /// ```
69 pub fn remove(&mut self, key: &K) -> Option<V> {
70 self.inner.remove(key)
71 }
72
73 /// Look up the value at `key`.
74 ///
75 /// # Examples
76 ///
77 /// ```
78 /// use dynomite::util::rbtree::OrderedMap;
79 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
80 /// m.insert(1, 10);
81 /// assert_eq!(m.get(&1), Some(&10));
82 /// ```
83 pub fn get(&self, key: &K) -> Option<&V> {
84 self.inner.get(key)
85 }
86
87 /// Number of entries.
88 ///
89 /// # Examples
90 ///
91 /// ```
92 /// use dynomite::util::rbtree::OrderedMap;
93 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
94 /// m.insert(1, 1);
95 /// assert_eq!(m.len(), 1);
96 /// ```
97 pub fn len(&self) -> usize {
98 self.inner.len()
99 }
100
101 /// Whether the map is empty.
102 ///
103 /// # Examples
104 ///
105 /// ```
106 /// use dynomite::util::rbtree::OrderedMap;
107 /// let m: OrderedMap<u32, u32> = OrderedMap::new();
108 /// assert!(m.is_empty());
109 /// ```
110 pub fn is_empty(&self) -> bool {
111 self.inner.is_empty()
112 }
113
114 /// Drop every entry.
115 ///
116 /// # Examples
117 ///
118 /// ```
119 /// use dynomite::util::rbtree::OrderedMap;
120 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
121 /// m.insert(1, 1);
122 /// m.clear();
123 /// assert!(m.is_empty());
124 /// ```
125 pub fn clear(&mut self) {
126 self.inner.clear();
127 }
128
129 /// Smallest entry, or [`None`] if the map is empty.
130 ///
131 /// # Examples
132 ///
133 /// ```
134 /// use dynomite::util::rbtree::OrderedMap;
135 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
136 /// m.insert(3, 30);
137 /// m.insert(1, 10);
138 /// assert_eq!(m.min(), Some((&1, &10)));
139 /// ```
140 pub fn min(&self) -> Option<(&K, &V)> {
141 self.inner.iter().next()
142 }
143
144 /// Largest entry, or [`None`] if the map is empty.
145 ///
146 /// # Examples
147 ///
148 /// ```
149 /// use dynomite::util::rbtree::OrderedMap;
150 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
151 /// m.insert(3, 30);
152 /// m.insert(1, 10);
153 /// assert_eq!(m.max(), Some((&3, &30)));
154 /// ```
155 pub fn max(&self) -> Option<(&K, &V)> {
156 self.inner.iter().next_back()
157 }
158
159 /// Smallest entry whose key is `>= probe`. Used by the token ring
160 /// for "find the owner of this key" queries.
161 ///
162 /// # Examples
163 ///
164 /// ```
165 /// use dynomite::util::rbtree::OrderedMap;
166 /// let mut t: OrderedMap<u64, &str> = OrderedMap::new();
167 /// t.insert(10, "a");
168 /// t.insert(20, "b");
169 /// assert_eq!(t.lower_bound(&15), Some((&20, &"b")));
170 /// assert_eq!(t.lower_bound(&25), None);
171 /// assert_eq!(t.lower_bound(&20), Some((&20, &"b")));
172 /// ```
173 pub fn lower_bound(&self, probe: &K) -> Option<(&K, &V)> {
174 self.inner.range(probe..).next()
175 }
176
177 /// In-order iterator over the entries.
178 ///
179 /// # Examples
180 ///
181 /// ```
182 /// use dynomite::util::rbtree::OrderedMap;
183 /// let mut m: OrderedMap<u32, u32> = OrderedMap::new();
184 /// m.insert(2, 0);
185 /// m.insert(1, 0);
186 /// let keys: Vec<u32> = m.iter().map(|(k, _)| *k).collect();
187 /// assert_eq!(keys, vec![1, 2]);
188 /// ```
189 pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
190 self.inner.iter()
191 }
192}
193
194#[cfg(test)]
195mod tests {
196 use super::*;
197
198 #[test]
199 fn ordered_iteration() {
200 let mut t: OrderedMap<u32, u32> = OrderedMap::new();
201 for k in [3, 1, 5, 2, 4] {
202 t.insert(k, k * 10);
203 }
204 let keys: Vec<u32> = t.iter().map(|(k, _)| *k).collect();
205 assert_eq!(keys, vec![1, 2, 3, 4, 5]);
206 }
207
208 #[test]
209 fn min_and_max() {
210 let mut t: OrderedMap<u32, u32> = OrderedMap::new();
211 assert!(t.min().is_none());
212 t.insert(7, 0);
213 t.insert(3, 0);
214 t.insert(11, 0);
215 assert_eq!(t.min().map(|(k, _)| *k), Some(3));
216 assert_eq!(t.max().map(|(k, _)| *k), Some(11));
217 }
218
219 #[test]
220 fn lower_bound_wraps_at_end() {
221 let mut t: OrderedMap<u32, u32> = OrderedMap::new();
222 t.insert(10, 1);
223 t.insert(20, 2);
224 t.insert(30, 3);
225 assert_eq!(t.lower_bound(&0).map(|(k, _)| *k), Some(10));
226 assert_eq!(t.lower_bound(&15).map(|(k, _)| *k), Some(20));
227 assert_eq!(t.lower_bound(&30).map(|(k, _)| *k), Some(30));
228 assert!(t.lower_bound(&31).is_none());
229 }
230}