toad_common/
map.rs

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