prefix_tree/map.rs
1use crate::tree::Tree;
2use std::hash::{Hash, Hasher};
3use std::iter::{FromIterator, FusedIterator};
4use std::ops::Index;
5
6/// A map implemented with prefix tree.
7#[derive(Debug, Clone, Default)]
8pub struct PrefixMap<K, V> {
9 root: Tree<K, V>,
10 length: usize,
11}
12
13impl<K: Eq + Clone, V> PrefixMap<K, V> {
14 /// Creates an empty `PrefixMap`.
15 ///
16 /// # Examples
17 ///
18 /// ```
19 /// use prefix_tree::PrefixMap;
20 ///
21 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
22 /// ```
23 pub fn new() -> PrefixMap<K, V> {
24 PrefixMap {
25 root: Tree::empty(),
26 length: 0,
27 }
28 }
29
30 /// Returns `true` if the map contains a value for the specified key.
31 ///
32 /// # Examples
33 ///
34 /// ```
35 /// use prefix_tree::PrefixMap;
36 ///
37 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
38 /// map.insert("foo", 1);
39 /// assert_eq!(map.contains_key("foo"), true);
40 /// assert_eq!(map.contains_key("bar"), false);
41 /// ```
42 pub fn contains_key<Q>(&self, key: Q) -> bool
43 where
44 Q: AsRef<[K]>,
45 {
46 self.get(key).is_some()
47 }
48
49 /// Clears the map, removing all key-value pairs.
50 ///
51 /// # Examples
52 ///
53 /// ```
54 /// use prefix_tree::PrefixMap;
55 ///
56 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
57 /// map.insert("foo", 1);
58 /// map.clear();
59 /// assert!(map.is_empty());
60 /// ```
61 pub fn clear(&mut self) {
62 *self = PrefixMap::new();
63 }
64
65 /// Returns a reference to the value corresponding to the key.
66 ///
67 /// # Examples
68 ///
69 /// ```
70 /// use prefix_tree::PrefixMap;
71 ///
72 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
73 /// map.insert("foo", 1);
74 /// assert_eq!(map.get("foo"), Some(&1));
75 /// assert_eq!(map.get("bar"), None);
76 /// ```
77 pub fn get<Q>(&self, key: Q) -> Option<&V>
78 where
79 Q: AsRef<[K]>,
80 {
81 self.root.find(key.as_ref()).and_then(|x| x.value())
82 }
83
84 /// Returns a mutable reference to the value corresponding to the key.
85 ///
86 /// # Examples
87 ///
88 /// ```
89 /// use prefix_tree::PrefixMap;
90 ///
91 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
92 /// map.insert("foo", 1);
93 /// if let Some(x) = map.get_mut("foo") {
94 /// *x = 2;
95 /// }
96 /// assert_eq!(map.get("foo"), Some(&2));
97 /// ```
98 pub fn get_mut<Q>(&mut self, key: Q) -> Option<&mut V>
99 where
100 Q: AsRef<[K]>,
101 {
102 self.root.find_mut(key.as_ref()).and_then(|x| x.value_mut())
103 }
104
105 /// Inserts a key-value pair into the map.
106 ///
107 /// # Examples
108 ///
109 /// ```
110 /// use prefix_tree::PrefixMap;
111 ///
112 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
113 /// assert_eq!(map.insert("a", 42), None);
114 /// assert_eq!(map.is_empty(), false);
115 /// assert_eq!(map.insert("a", 5), Some(42));
116 /// assert_eq!(map.get("a"), Some(&5));
117 /// ```
118 pub fn insert<Q>(&mut self, key: Q, value: V) -> Option<V>
119 where
120 Q: AsRef<[K]>,
121 {
122 let old = self.root.insert(key.as_ref(), value);
123 if old.is_none() {
124 self.length += 1;
125 }
126 old
127 }
128
129 /// Removes a key from the map, returning the value at the key
130 /// if the key was previously in the map.
131 ///
132 /// # Examples
133 ///
134 ///
135 /// ```
136 /// use prefix_tree::PrefixMap;
137 ///
138 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
139 /// assert_eq!(map.insert("a", 42), None);
140 /// assert_eq!(map.remove("a"), Some(42));
141 /// assert_eq!(map.get("a"), None);
142 /// ```
143 pub fn remove<Q>(&mut self, key: Q) -> Option<V>
144 where
145 Q: AsRef<[K]>,
146 {
147 self.root.remove(key.as_ref())
148 }
149
150 /// Returns `true` if the map contains no elements.
151 ///
152 /// # Examples
153 ///
154 /// ```
155 /// use prefix_tree::PrefixMap;
156 ///
157 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
158 /// assert_eq!(map.is_empty(), true);
159 /// map.insert("foo", 1);
160 /// assert_eq!(map.is_empty(), false);
161 /// ```
162 pub fn is_empty(&self) -> bool {
163 self.len() == 0
164 }
165
166 /// Returns the number of elements in the map.
167 ///
168 /// # Examples
169 ///
170 /// ```
171 /// use prefix_tree::PrefixMap;
172 ///
173 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
174 /// assert_eq!(map.len(), 0);
175 /// map.insert("foo", 1);
176 /// assert_eq!(map.len(), 1);
177 /// ```
178 pub fn len(&self) -> usize {
179 self.length
180 }
181
182 /// Gets an iterator over the entries of the map, in arbitrary order.
183 ///
184 /// # Examples
185 ///
186 /// ```
187 /// use prefix_tree::PrefixMap;
188 ///
189 /// let mut map: PrefixMap<u8, i32> = PrefixMap::new();
190 /// map.insert("1", 9);
191 /// map.insert("2", 8);
192 /// map.insert("3", 7);
193 ///
194 /// for (key, value) in map.iter() {
195 /// println!("{:?}: {:?}", key, value);
196 /// }
197 /// ```
198 pub fn iter(&self) -> Iter<K, V> {
199 Iter {
200 root: &self.root,
201 stack: vec![IterStackItem {
202 iter: self.root.children().iter(),
203 key_fragment: &self.root.key(),
204 }],
205 length: self.length,
206 }
207 }
208
209 /// Gets an iterator over the keys of the map, in arbitrary order.
210 ///
211 /// # Examples
212 ///
213 /// ```
214 /// use prefix_tree::PrefixMap;
215 ///
216 /// let mut map: PrefixMap<i32, i32> = PrefixMap::new();
217 /// map.insert([1], 2);
218 /// map.insert([2], 3);
219 ///
220 /// assert_eq!(map.keys().collect::<Vec<_>>(), vec![vec![1], vec![2]]);
221 /// ```
222 pub fn keys(&self) -> Keys<K, V> {
223 Keys { inner: self.iter() }
224 }
225
226 /// Gets an iterator over the values of the map, in arbitrary order.
227 ///
228 /// # Examples
229 ///
230 /// ```
231 /// use prefix_tree::PrefixMap;
232 ///
233 /// let mut map: PrefixMap<i32, i32> = PrefixMap::new();
234 /// map.insert([1], 2);
235 /// map.insert([2], 3);
236 ///
237 /// assert_eq!(map.values().cloned().collect::<Vec<_>>(), vec![2, 3]);
238 /// ```
239 pub fn values(&self) -> Values<K, V> {
240 Values { inner: self.iter() }
241 }
242}
243
244impl<'a, K: 'a + Eq + Clone, V: 'a> FromIterator<(&'a [K], V)> for PrefixMap<K, V> {
245 fn from_iter<I>(iter: I) -> PrefixMap<K, V>
246 where
247 I: IntoIterator<Item = (&'a [K], V)>,
248 {
249 let mut map = PrefixMap::new();
250 iter.into_iter().for_each(|(k, v)| {
251 map.insert(k, v);
252 });
253 map
254 }
255}
256
257impl<'a, K: 'a + Eq + Clone, V: 'a> IntoIterator for &'a PrefixMap<K, V> {
258 type Item = (Vec<K>, &'a V);
259
260 type IntoIter = Iter<'a, K, V>;
261
262 fn into_iter(self) -> Self::IntoIter {
263 self.iter()
264 }
265}
266
267struct IterStackItem<'a, K: 'a, V: 'a> {
268 iter: std::slice::Iter<'a, Tree<K, V>>,
269 key_fragment: &'a [K],
270}
271
272pub struct Iter<'a, K: 'a, V: 'a> {
273 root: &'a Tree<K, V>,
274 stack: Vec<IterStackItem<'a, K, V>>,
275 length: usize,
276}
277
278impl<'a, K: 'a + Eq + Clone, V: 'a> Iterator for Iter<'a, K, V> {
279 type Item = (Vec<K>, &'a V);
280
281 fn next(&mut self) -> Option<Self::Item> {
282 if self.length == 1 && self.root.value().is_some() {
283 self.length = 0;
284 return self.root.value().map(|x| (vec![], x));
285 }
286 while let Some(IterStackItem { iter, .. }) = self.stack.last_mut() {
287 if let Some(tree) = iter.next() {
288 self.stack.push(IterStackItem {
289 iter: tree.children().iter(),
290 key_fragment: tree.key(),
291 });
292 if tree.value().is_some() {
293 self.length -= 1;
294 return Some((
295 self.stack
296 .iter()
297 .map(|x| x.key_fragment)
298 .flatten()
299 .cloned()
300 .collect(),
301 tree.value().unwrap(),
302 ));
303 }
304 } else {
305 self.stack.pop();
306 }
307 }
308 None
309 }
310
311 fn size_hint(&self) -> (usize, Option<usize>) {
312 (self.length, Some(self.length))
313 }
314}
315
316impl<'a, K: 'a + Eq + Clone, V: 'a> ExactSizeIterator for Iter<'a, K, V> {
317 fn len(&self) -> usize {
318 self.length
319 }
320}
321
322impl<'a, K: 'a + Eq + Clone, V: 'a> FusedIterator for Iter<'a, K, V> {}
323
324pub struct Keys<'a, K: 'a, V: 'a> {
325 inner: Iter<'a, K, V>,
326}
327
328impl<'a, K: 'a + Eq + Clone, V: 'a> Iterator for Keys<'a, K, V> {
329 type Item = Vec<K>;
330
331 fn next(&mut self) -> Option<Self::Item> {
332 self.inner.next().map(|(k, _)| k)
333 }
334}
335
336impl<'a, K: 'a + Eq + Clone, V: 'a> ExactSizeIterator for Keys<'a, K, V> {
337 fn len(&self) -> usize {
338 self.inner.length
339 }
340}
341
342impl<'a, K: 'a + Eq + Clone, V: 'a> FusedIterator for Keys<'a, K, V> {}
343
344pub struct Values<'a, K: 'a, V: 'a> {
345 inner: Iter<'a, K, V>,
346}
347
348impl<'a, K: 'a + Eq + Clone, V: 'a> Iterator for Values<'a, K, V> {
349 type Item = &'a V;
350
351 fn next(&mut self) -> Option<Self::Item> {
352 self.inner.next().map(|(_, v)| v)
353 }
354}
355
356impl<'a, K: 'a + Eq + Clone, V: 'a> ExactSizeIterator for Values<'a, K, V> {
357 fn len(&self) -> usize {
358 self.inner.length
359 }
360}
361
362impl<'a, K: 'a + Eq + Clone, V: 'a> FusedIterator for Values<'a, K, V> {}
363
364impl<K: Eq + Clone, V, Q: AsRef<[K]>> Index<Q> for PrefixMap<K, V> {
365 type Output = V;
366
367 fn index(&self, index: Q) -> &Self::Output {
368 self.get(index).expect("no entry found for key")
369 }
370}
371
372impl<K: Eq + Clone, V: Eq> PartialEq<PrefixMap<K, V>> for PrefixMap<K, V> {
373 fn eq(&self, other: &PrefixMap<K, V>) -> bool {
374 self.len() == other.len() && self.iter().zip(other).all(|(a, b)| a == b)
375 }
376}
377
378impl<K: Eq + Clone, V: Eq> Eq for PrefixMap<K, V> {}
379
380impl<K: Eq + Clone + Hash, V: Hash> Hash for PrefixMap<K, V> {
381 fn hash<H: Hasher>(&self, state: &mut H) {
382 self.iter().for_each(|x| x.hash(state))
383 }
384}