kwap_common/
map.rs

1use core::borrow::Borrow;
2use std::collections::{btree_map, hash_map, BTreeMap, HashMap};
3use std::hash::Hash;
4use std::{iter, slice};
5
6use crate::result::ResultExt;
7use crate::{GetSize, Reserve};
8
9/// Things that can go unhappily when trying to insert into a map
10#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Ord, Eq)]
11pub enum InsertError<V> {
12  /// The new value was inserted successfully but there was already a value in the map for that key.
13  Exists(V),
14  /// The map is at capacity and cannot fit any more pairs.
15  CapacityExhausted,
16}
17
18/// An collection of key-value pairs
19///
20/// # Provided implementations
21/// - [`HashMap`]`<K, V>`
22/// - [`tinyvec::ArrayVec`]`<(K, V)>`
23/// - [`Vec`]`<(K, V)>`
24///
25/// # Requirements
26/// - [`Default`] for creating the map
27/// - [`Extend`] for adding new entries to the map
28/// - [`Reserve`] for reserving space ahead of time
29/// - [`GetSize`] for bound checks, empty checks, and accessing the length
30/// - [`FromIterator`] for [`collect`](core::iter::Iterator#method.collect)ing into the map
31/// - [`IntoIterator`] for iterating and destroying the map
32pub trait Map<K: Ord + Eq + Hash, V>:
33  Default + GetSize + Reserve + Extend<(K, V)> + FromIterator<(K, V)> + IntoIterator<Item = (K, V)>
34{
35  /// See [`HashMap.insert`]
36  fn insert(&mut self, key: K, val: V) -> Result<(), InsertError<V>>;
37
38  /// See [`HashMap.remove`]
39  fn remove<Q: Hash + Eq + Ord>(&mut self, key: &Q) -> Option<V>
40    where K: Borrow<Q>;
41
42  /// See [`HashMap.get`]
43  fn get<'a, Q: Hash + Eq + Ord>(&'a self, key: &Q) -> Option<&'a V>
44    where K: Borrow<Q> + 'a;
45
46  /// See [`HashMap.get_mut`]
47  fn get_mut<'a, Q: Hash + Eq + Ord>(&'a mut self, key: &Q) -> Option<&'a mut V>
48    where K: Borrow<Q> + 'a;
49
50  /// See [`HashMap.contains_key`]
51  fn has<Q: Hash + Eq + Ord>(&self, key: &Q) -> bool
52    where K: Borrow<Q>
53  {
54    self.get(key).is_some()
55  }
56
57  /// See [`HashMap.iter`]
58  fn iter(&self) -> Iter<'_, K, V>;
59
60  /// See [`HashMap.iter_mut`]
61  fn iter_mut(&mut self) -> IterMut<'_, K, V>;
62}
63
64impl<K: Eq + Hash + Ord, V> Map<K, V> for BTreeMap<K, V> {
65  fn insert(&mut self, key: K, val: V) -> Result<(), InsertError<V>> {
66    self.insert(key, val)
67        .map(InsertError::Exists)
68        .ok_or(())
69        .swap()
70  }
71
72  fn remove<Q: Ord>(&mut self, key: &Q) -> Option<V>
73    where K: Borrow<Q>
74  {
75    self.remove(key)
76  }
77
78  fn get<'a, Q: Hash + Eq + Ord>(&'a self, key: &Q) -> Option<&'a V>
79    where K: Borrow<Q> + 'a
80  {
81    self.get(key)
82  }
83
84  fn get_mut<'a, Q: Hash + Eq + Ord>(&'a mut self, key: &Q) -> Option<&'a mut V>
85    where K: Borrow<Q> + 'a
86  {
87    self.get_mut(key)
88  }
89
90  fn iter(&self) -> Iter<'_, K, V> {
91    Iter { array_iter: None,
92           hashmap_iter: None,
93           btreemap_iter: Some(self.iter()) }
94  }
95
96  fn iter_mut(&mut self) -> IterMut<'_, K, V> {
97    IterMut { array_iter: None,
98              hashmap_iter: None,
99              btreemap_iter: Some(self.iter_mut()) }
100  }
101}
102
103impl<K: Eq + Hash + Ord, V> Map<K, V> for HashMap<K, V> {
104  fn iter(&self) -> Iter<'_, K, V> {
105    Iter { array_iter: None,
106           btreemap_iter: None,
107           hashmap_iter: Some(self.iter()) }
108  }
109
110  fn iter_mut(&mut self) -> IterMut<'_, K, V> {
111    IterMut { array_iter: None,
112              btreemap_iter: None,
113              hashmap_iter: Some(self.iter_mut()) }
114  }
115
116  fn get<'a, Q: Hash + Eq + Ord>(&'a self, key: &Q) -> Option<&'a V>
117    where K: Borrow<Q> + 'a
118  {
119    self.get(key)
120  }
121
122  fn get_mut<'a, Q: Hash + Eq + Ord>(&'a mut self, key: &Q) -> Option<&'a mut V>
123    where K: Borrow<Q> + 'a
124  {
125    self.get_mut(key)
126  }
127
128  fn insert(&mut self, key: K, val: V) -> Result<(), InsertError<V>> {
129    self.insert(key, val)
130        .map(InsertError::Exists)
131        .ok_or(())
132        .swap()
133  }
134
135  fn remove<Q: Hash + Eq + Ord>(&mut self, key: &Q) -> Option<V>
136    where K: Borrow<Q>
137  {
138    self.remove(key)
139  }
140}
141
142impl<T: crate::Array<Item = (K, V)>, K: Eq + Hash + Ord, V> Map<K, V> for T {
143  fn insert(&mut self, key: K, mut val: V) -> Result<(), InsertError<V>> {
144    match self.iter_mut().find(|(k, _)| k == &&key) {
145      | Some((_, exist)) => {
146        std::mem::swap(exist, &mut val);
147        Err(InsertError::Exists(val))
148      },
149      | None => match self.is_full() {
150        | true => Err(InsertError::CapacityExhausted),
151        | false => {
152          self.push((key, val));
153          Ok(())
154        },
155      },
156    }
157  }
158
159  fn remove<Q: Hash + Eq + Ord>(&mut self, key: &Q) -> Option<V>
160    where K: Borrow<Q>
161  {
162    match self.iter()
163              .enumerate()
164              .find(|(_, (k, _))| Borrow::<Q>::borrow(*k) == key)
165    {
166      | Some((ix, _)) => self.remove(ix).map(|(_, v)| v),
167      | None => None,
168    }
169  }
170
171  fn get<'a, Q: Hash + Eq + Ord>(&'a self, key: &Q) -> Option<&'a V>
172    where K: Borrow<Q> + 'a
173  {
174    match self.iter().find(|(k, _)| Borrow::<Q>::borrow(*k) == key) {
175      | Some((_, v)) => Some(v),
176      | None => None,
177    }
178  }
179
180  fn get_mut<'a, Q: Hash + Eq + Ord>(&'a mut self, key: &Q) -> Option<&'a mut V>
181    where K: Borrow<Q> + 'a
182  {
183    match self.iter_mut()
184              .find(|(k, _)| Borrow::<Q>::borrow(*k) == key)
185    {
186      | Some((_, v)) => Some(v),
187      | None => None,
188    }
189  }
190
191  fn iter(&self) -> Iter<'_, K, V> {
192    Iter { array_iter: Some(self.deref().iter().map(Iter::coerce_array_iter)),
193           btreemap_iter: None,
194           hashmap_iter: None }
195  }
196
197  fn iter_mut(&mut self) -> IterMut<'_, K, V> {
198    IterMut { array_iter: Some(self.deref_mut().iter_mut().map(IterMut::coerce_array_iter)),
199              btreemap_iter: None,
200              hashmap_iter: None }
201  }
202}
203
204impl<K: Eq + Hash, V> GetSize for HashMap<K, V> {
205  fn get_size(&self) -> usize {
206    self.len()
207  }
208
209  fn max_size(&self) -> Option<usize> {
210    None
211  }
212}
213
214impl<K, V> Reserve for BTreeMap<K, V> {}
215impl<K, V> GetSize for BTreeMap<K, V> {
216  fn get_size(&self) -> usize {
217    self.len()
218  }
219
220  fn max_size(&self) -> Option<usize> {
221    None
222  }
223}
224
225impl<K: Eq + Hash, V> Reserve for HashMap<K, V> {
226  fn reserve(n: usize) -> Self {
227    Self::with_capacity(n)
228  }
229}
230
231type ArrayIterCoercer<'a, K, V> = fn(&'a (K, V)) -> (&'a K, &'a V);
232type ArrayIterMapped<'a, K, V> = iter::Map<slice::Iter<'a, (K, V)>, ArrayIterCoercer<'a, K, V>>;
233
234type ArrayIterMutCoercer<'a, K, V> = fn(&'a mut (K, V)) -> (&'a K, &'a mut V);
235type ArrayIterMutMapped<'a, K, V> =
236  iter::Map<slice::IterMut<'a, (K, V)>, ArrayIterMutCoercer<'a, K, V>>;
237
238/// An iterator over the entries of a `Map`.
239///
240/// This `struct` is created by the [`iter`] method on [`Map`].
241/// See its documentation for more.
242///
243/// [`iter`]: Map::iter
244///
245/// # Example
246///
247/// ```
248/// use std::collections::HashMap;
249///
250/// use kwap_common::Map;
251///
252/// let mut map = HashMap::from([("a", 1)]);
253///
254/// fn do_stuff(map: &impl Map<&'static str, usize>) {
255///   let iter = map.iter();
256/// }
257/// ```
258#[derive(Debug)]
259pub struct Iter<'a, K: Eq + Hash, V> {
260  // TODO: #[cfg(not(no_std))]?
261  hashmap_iter: Option<hash_map::Iter<'a, K, V>>,
262  // TODO: #[cfg(alloc)]?
263  btreemap_iter: Option<btree_map::Iter<'a, K, V>>,
264  array_iter: Option<ArrayIterMapped<'a, K, V>>,
265}
266
267impl<'a, K: Eq + Hash, V> Iter<'a, K, V> {
268  #[inline(always)]
269  fn coerce_array_iter((k, v): &'a (K, V)) -> (&'a K, &'a V) {
270    (k, v)
271  }
272
273  fn get_iter(&mut self) -> &mut dyn Iterator<Item = (&'a K, &'a V)> {
274    let (a, b, c) = (self.hashmap_iter.as_mut().map(|a| a as &mut _),
275                     self.array_iter.as_mut().map(|a| a as &mut _),
276                     self.btreemap_iter.as_mut().map(|a| a as &mut _));
277    a.or(b).or(c).unwrap()
278  }
279}
280
281impl<'a, K: Eq + Hash, V> Iterator for Iter<'a, K, V> {
282  type Item = (&'a K, &'a V);
283
284  fn next(&mut self) -> Option<Self::Item> {
285    self.get_iter().next()
286  }
287}
288
289/// A mutable iterator over the entries of a `Map`.
290///
291/// This `struct` is created by the [`iter_mut`] method on [`Map`]. See its
292/// documentation for more.
293///
294/// [`iter_mut`]: Map::iter_mut
295///
296/// # Example
297///
298/// ```
299/// use std::collections::HashMap;
300///
301/// use kwap_common::Map;
302///
303/// let mut map = HashMap::from([("a", 1)]);
304///
305/// fn do_stuff(map: &mut impl Map<&'static str, usize>) {
306///   let iter = map.iter_mut();
307/// }
308/// ```
309#[derive(Debug)]
310pub struct IterMut<'a, K: Eq + Hash, V> {
311  // TODO: #[cfg(not(no_std))]?
312  hashmap_iter: Option<hash_map::IterMut<'a, K, V>>,
313  // TODO: #[cfg(alloc)]?
314  btreemap_iter: Option<btree_map::IterMut<'a, K, V>>,
315  array_iter: Option<ArrayIterMutMapped<'a, K, V>>,
316}
317
318impl<'a, K: Eq + Hash, V> IterMut<'a, K, V> {
319  #[inline(always)]
320  fn coerce_array_iter((k, v): &'a mut (K, V)) -> (&'a K, &'a mut V) {
321    (k, v)
322  }
323
324  fn get_iter(&mut self) -> &mut dyn Iterator<Item = (&'a K, &'a mut V)> {
325    let (a, b, c) = (self.hashmap_iter.as_mut().map(|a| a as &mut _),
326                     self.array_iter.as_mut().map(|a| a as &mut _),
327                     self.btreemap_iter.as_mut().map(|a| a as &mut _));
328    a.or(b).or(c).unwrap()
329  }
330}
331
332impl<'a, K: Eq + Hash, V> Iterator for IterMut<'a, K, V> {
333  type Item = (&'a K, &'a mut V);
334
335  fn next(&mut self) -> Option<Self::Item> {
336    self.get_iter().next()
337  }
338}
339
340#[cfg(test)]
341mod tests {
342  use super::*;
343
344  fn impls(
345    )
346      -> (impl Map<String, String>,
347          impl Map<String, String>,
348          impl Map<String, String>,
349          impl Map<String, String>)
350  {
351    (HashMap::<String, String>::from([("foo".into(), "bar".into())]),
352     BTreeMap::<String, String>::from([("foo".into(), "bar".into())]),
353     tinyvec::array_vec!([(String, String); 16] => ("foo".into(), "bar".into())),
354     vec![("foo".to_string(), "bar".to_string())])
355  }
356
357  macro_rules! each_impl {
358    ($work:expr) => {{
359      let (hm, bt, av, vc) = impls();
360      println!("hashmap");
361      $work(hm);
362      println!("btreemap");
363      $work(bt);
364      println!("arrayvec");
365      $work(av);
366      println!("vec");
367      $work(vc);
368    }};
369  }
370
371  #[test]
372  fn get() {
373    fn test_get<M: Map<String, String>>(map: M) {
374      assert_eq!(map.get(&"foo".to_string()), Some(&"bar".into()));
375      assert_eq!(map.get(&"foot".to_string()), None);
376    }
377
378    each_impl!(test_get);
379  }
380
381  #[test]
382  fn get_mut() {
383    fn test_get_mut<M: Map<String, String>>(mut map: M) {
384      let old = map.get_mut(&"foo".to_string()).unwrap();
385      *old = format!("{}f", old);
386
387      assert_eq!(map.get(&"foo".to_string()), Some(&"barf".into()));
388    }
389
390    each_impl!(test_get_mut);
391  }
392
393  #[test]
394  fn insert() {
395    fn test_insert<M: Map<String, String>>(mut map: M) {
396      let old = map.insert("foot".to_string(), "butt".to_string());
397
398      assert_eq!(old, Ok(()));
399      assert_eq!(map.get(&"foo".to_string()).unwrap().as_str(), "bar");
400      assert_eq!(map.get(&"foot".to_string()).unwrap().as_str(), "butt");
401
402      let old = map.insert("foot".to_string(), "squat".to_string());
403      assert_eq!(old, Err(InsertError::Exists("butt".to_string())));
404      assert_eq!(map.get(&"foot".to_string()).unwrap().as_str(), "squat");
405    }
406
407    each_impl!(test_insert);
408  }
409
410  #[test]
411  fn remove() {
412    fn test_remove<M: Map<String, String>>(mut map: M) {
413      let old = map.remove(&"foo".to_string());
414      assert_eq!(old, Some("bar".to_string()));
415
416      let old = map.remove(&"foo".to_string());
417      assert_eq!(old, None);
418    }
419
420    each_impl!(test_remove);
421  }
422
423  #[test]
424  fn has() {
425    fn test_has<M: Map<String, String>>(map: M) {
426      assert!(map.has(&"foo".to_string()));
427      assert!(!map.has(&"foot".to_string()));
428    }
429
430    each_impl!(test_has);
431  }
432
433  #[test]
434  fn into_iter() {
435    fn test_into_iter<M: Map<String, String>>(mut map: M) {
436      map.insert("a".into(), "a".into()).unwrap();
437      map.insert("b".into(), "b".into()).unwrap();
438      map.insert("c".into(), "c".into()).unwrap();
439
440      let mut kvs = map.into_iter().collect::<Vec<_>>();
441      kvs.sort();
442
443      assert_eq!(kvs,
444                 vec![("a".into(), "a".into()),
445                      ("b".into(), "b".into()),
446                      ("c".into(), "c".into()),
447                      ("foo".into(), "bar".into()),]);
448    }
449
450    each_impl!(test_into_iter);
451  }
452
453  #[test]
454  fn iter() {
455    fn test_iter<M: Map<String, String>>(mut map: M) {
456      map.insert("a".into(), "a".into()).unwrap();
457      map.insert("b".into(), "b".into()).unwrap();
458      map.insert("c".into(), "c".into()).unwrap();
459
460      let mut kvs = map.iter().collect::<Vec<_>>();
461      kvs.sort();
462
463      assert_eq!(kvs,
464                 vec![(&"a".into(), &"a".into()),
465                      (&"b".into(), &"b".into()),
466                      (&"c".into(), &"c".into()),
467                      (&"foo".into(), &"bar".into()),]);
468    }
469
470    each_impl!(test_iter);
471  }
472
473  #[test]
474  fn iter_mut() {
475    fn test_iter_mut<M: Map<String, String>>(mut map: M) {
476      map.insert("a".into(), "a".into()).unwrap();
477      map.insert("b".into(), "b".into()).unwrap();
478      map.insert("c".into(), "c".into()).unwrap();
479
480      let mut kvs = map.iter_mut().collect::<Vec<_>>();
481      kvs.sort();
482
483      assert_eq!(kvs,
484                 vec![(&"a".into(), &mut "a".into()),
485                      (&"b".into(), &mut "b".into()),
486                      (&"c".into(), &mut "c".into()),
487                      (&"foo".into(), &mut "bar".into()),]);
488    }
489
490    each_impl!(test_iter_mut);
491  }
492}