use alloc::vec::Vec;
use core::mem;
use hashbrown::HashMap;
#[derive(Debug, Clone)]
pub struct OrderedMap<K, V> {
values: Vec<V>,
map: HashMap<K, usize>,
}
impl<K, V> Default for OrderedMap<K, V> {
fn default() -> Self {
Self {
values: Default::default(),
map: Default::default(),
}
}
}
impl<K, V> OrderedMap<K, V>
where
K: PartialEq + Eq + core::hash::Hash,
{
pub fn new() -> Self {
Self::default()
}
#[allow(dead_code)]
pub fn len(&self) -> usize {
self.values.len()
}
#[allow(dead_code)]
pub fn is_empty(&self) -> bool {
self.values.is_empty()
}
pub fn retain<F>(&mut self, mut f: F)
where
F: FnMut(&V) -> bool,
{
let values = mem::take(&mut self.values);
let map = mem::take(&mut self.map);
let mut new_values = Vec::new();
let mut old_pos_to_new_pos = HashMap::new();
for (pos, value) in values.into_iter().enumerate().filter(|(_, v)| f(v)) {
old_pos_to_new_pos.insert(pos, new_values.len());
new_values.push(value);
}
self.values = new_values;
self.map = map
.into_iter()
.filter_map(|(k, v)| old_pos_to_new_pos.get(&v).map(|v2| (k, *v2)))
.collect();
}
pub fn push_insert(&mut self, key: K, value: V) {
let idx = self.values.len();
self.values.push(value);
self.map.insert(key, idx);
}
pub fn get_by_key<Q>(&self, key: &Q) -> Option<&V>
where
K: alloc::borrow::Borrow<Q>,
Q: core::hash::Hash + Eq + ?Sized,
{
self.map.get(key).and_then(|&v| self.values.get(v))
}
pub fn get_by_index(&self, i: usize) -> Option<&V> {
self.values.get(i)
}
pub fn values(&self) -> &[V] {
&self.values
}
pub fn values_mut(&mut self) -> &mut [V] {
&mut self.values
}
pub fn into_values(self) -> Vec<V> {
self.values
}
}
impl<K, V> FromIterator<(K, V)> for OrderedMap<K, V>
where
K: PartialEq + Eq + core::hash::Hash,
{
fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
let mut map = OrderedMap::new();
for (k, v) in iter {
map.push_insert(k, v)
}
map
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn retain() {
let mut m = OrderedMap::from_iter([(1, 'a'), (2, 'b'), (3, 'c')]);
m.retain(|v| *v != 'b');
assert_eq!(m.get_by_key(&1), Some(&'a'));
assert_eq!(m.get_by_key(&2), None);
assert_eq!(m.get_by_key(&3), Some(&'c'));
assert_eq!(m.values(), &['a', 'c'])
}
}