irox_tools/
map.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2025 IROX Contributors
3//
4
5extern crate alloc;
6
7use alloc::collections::vec_deque::Iter;
8use alloc::collections::VecDeque;
9use core::borrow::Borrow;
10use core::hash::Hash;
11use std::collections::HashMap;
12
13///
14/// Tracks the insertion order of key/value pairs.  Backed by a [`HashMap`] for storage
15/// and a [`VecDeque`] for key insertion order.
16#[derive(Default)]
17pub struct OrderedHashMap<K, V> {
18    map: HashMap<K, V>,
19    key_order: VecDeque<K>,
20}
21
22impl<K, V> OrderedHashMap<K, V> {
23    pub fn new() -> OrderedHashMap<K, V> {
24        Self {
25            map: HashMap::new(),
26            key_order: VecDeque::new(),
27        }
28    }
29    pub fn with_capacity(capacity: usize) -> OrderedHashMap<K, V> {
30        Self {
31            map: HashMap::with_capacity(capacity),
32            key_order: VecDeque::with_capacity(capacity),
33        }
34    }
35    pub fn capacity(&self) -> usize {
36        self.map.capacity()
37    }
38    pub fn keys(&self) -> impl Iterator<Item = &K> {
39        self.key_order.iter()
40    }
41    pub fn clear(&mut self) {
42        self.map.clear();
43        self.key_order.clear();
44    }
45    pub fn len(&self) -> usize {
46        self.map.len()
47    }
48    pub fn is_empty(&self) -> bool {
49        self.map.is_empty()
50    }
51}
52impl<K: Eq + Hash + Clone, V> OrderedHashMap<K, V> {
53    pub fn get<Q: ?Sized + Hash + Eq>(&self, k: &Q) -> Option<&V>
54    where
55        K: Borrow<Q>,
56    {
57        self.map.get(k)
58    }
59
60    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
61        let old = self.map.remove(&k);
62        if old.is_some() {
63            self.key_order.push_back(k.clone());
64        }
65        self.map.insert(k, v);
66        old
67    }
68}
69impl<'a, K: Eq + Hash, V> OrderedHashMap<K, V> {
70    pub fn iter(&'a self) -> OrderedMapIter<'a, K, V> {
71        OrderedMapIter {
72            key_iter: self.key_order.iter(),
73            values: &self.map,
74        }
75    }
76    pub fn drain(self) -> OrderedDrain<K, V> {
77        OrderedDrain {
78            key_iter: self.key_order,
79            values: self.map,
80        }
81    }
82}
83
84pub struct OrderedMapIter<'a, K, V> {
85    key_iter: Iter<'a, K>,
86    values: &'a HashMap<K, V>,
87}
88impl<'a, K: Eq + Hash, V> Iterator for OrderedMapIter<'a, K, V> {
89    type Item = (&'a K, &'a V);
90    fn next(&mut self) -> Option<Self::Item> {
91        let key = self.key_iter.next()?;
92        let v = self.values.get(key)?;
93        Some((key, v))
94    }
95}
96pub struct OrderedMapValues<'a, K, V> {
97    iter: OrderedMapIter<'a, K, V>,
98}
99impl<'a, K: Eq + Hash, V> Iterator for OrderedMapValues<'a, K, V> {
100    type Item = &'a V;
101    fn next(&mut self) -> Option<Self::Item> {
102        self.iter.next().map(|(_k, v)| v)
103    }
104}
105
106pub struct OrderedDrain<K, V> {
107    key_iter: VecDeque<K>,
108    values: HashMap<K, V>,
109}
110impl<K: Eq + Hash, V> Iterator for OrderedDrain<K, V> {
111    type Item = (K, V);
112    fn next(&mut self) -> Option<Self::Item> {
113        let key = self.key_iter.pop_front()?;
114        let v = self.values.remove(&key)?;
115        Some((key, v))
116    }
117}