p3_util/
linear_map.rs

1use alloc::vec::Vec;
2use core::mem;
3
4use crate::VecExt;
5
6/// O(n) Vec-backed map for keys that only implement Eq.
7/// Only use this for a very small number of keys, operations
8/// on it can easily become O(n^2).
9#[derive(Debug)]
10pub struct LinearMap<K, V>(Vec<(K, V)>);
11
12impl<K, V> Default for LinearMap<K, V> {
13    fn default() -> Self {
14        Self(Default::default())
15    }
16}
17
18impl<K: Eq, V> LinearMap<K, V> {
19    pub fn new() -> Self {
20        Default::default()
21    }
22    pub fn get(&self, k: &K) -> Option<&V> {
23        self.0.iter().find(|(kk, _)| kk == k).map(|(_, v)| v)
24    }
25    pub fn get_mut(&mut self, k: &K) -> Option<&mut V> {
26        self.0.iter_mut().find(|(kk, _)| kk == k).map(|(_, v)| v)
27    }
28    /// This is O(n), because we check for an existing entry.
29    pub fn insert(&mut self, k: K, mut v: V) -> Option<V> {
30        if let Some(vv) = self.get_mut(&k) {
31            mem::swap(&mut v, vv);
32            Some(v)
33        } else {
34            self.0.push((k, v));
35            None
36        }
37    }
38    pub fn get_or_insert_with(&mut self, k: K, f: impl FnOnce() -> V) -> &mut V {
39        let existing = self.0.iter().position(|(kk, _)| kk == &k);
40        if let Some(idx) = existing {
41            &mut self.0[idx].1
42        } else {
43            let slot = self.0.pushed_mut((k, f()));
44            &mut slot.1
45        }
46    }
47    pub fn values(&self) -> impl Iterator<Item = &V> {
48        self.0.iter().map(|(_, v)| v)
49    }
50}
51
52impl<K: Eq, V> FromIterator<(K, V)> for LinearMap<K, V> {
53    /// This calls `insert` in a loop, so is O(n^2)!!
54    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
55        let mut me = LinearMap::default();
56        for (k, v) in iter {
57            me.insert(k, v);
58        }
59        me
60    }
61}
62
63impl<K, V> IntoIterator for LinearMap<K, V> {
64    type Item = (K, V);
65    type IntoIter = <Vec<(K, V)> as IntoIterator>::IntoIter;
66    fn into_iter(self) -> Self::IntoIter {
67        self.0.into_iter()
68    }
69}