dfir_rs/util/
slot_vec.rs

1//! A Vec-based SlotMap-esque datastructure and corresponding Key type.
2
3use std::fmt::{Debug, Display, Formatter};
4use std::hash::{Hash, Hasher};
5use std::iter::FusedIterator;
6use std::marker::PhantomData;
7use std::ops::{Index, IndexMut};
8
9/// A key into a SlotVec.
10#[repr(transparent)]
11pub struct Key<Tag: ?Sized> {
12    index: usize,
13    _phantom: PhantomData<Tag>,
14}
15impl<Tag: ?Sized> Key<Tag> {
16    /// Creates a Key from a raw index. Avoid using this function directly.
17    pub const fn from_raw(index: usize) -> Self {
18        Key {
19            index,
20            _phantom: PhantomData,
21        }
22    }
23}
24impl<Tag: ?Sized> Clone for Key<Tag> {
25    fn clone(&self) -> Self {
26        *self
27    }
28}
29impl<Tag: ?Sized> Copy for Key<Tag> {}
30impl<Tag: ?Sized> PartialOrd for Key<Tag> {
31    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
32        Some(self.cmp(other))
33    }
34}
35impl<Tag: ?Sized> Ord for Key<Tag> {
36    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
37        self.index.cmp(&other.index)
38    }
39}
40impl<Tag: ?Sized> PartialEq for Key<Tag> {
41    fn eq(&self, other: &Self) -> bool {
42        self.index == other.index
43    }
44}
45impl<Tag: ?Sized> Eq for Key<Tag> {}
46impl<Tag: ?Sized> Hash for Key<Tag> {
47    fn hash<H: Hasher>(&self, state: &mut H) {
48        self.index.hash(state);
49    }
50}
51impl<Tag: ?Sized> Debug for Key<Tag> {
52    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
53        write!(f, "Key({})", self.index)
54    }
55}
56impl<Tag: ?Sized> Display for Key<Tag> {
57    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
58        write!(f, "{}", self.index)
59    }
60}
61
62/// A Vec-based SlotMap-esque datastructure without removes.
63///
64/// Analogous to [`slotmap::SlotMap`], but avoids the overhead of tracking removed keys.
65#[repr(transparent)]
66pub struct SlotVec<Tag: ?Sized, Val> {
67    slots: Vec<Val>,
68    _phantom: PhantomData<Tag>,
69}
70impl<Tag: ?Sized, Val> SlotVec<Tag, Val> {
71    /// Creates a new `SlotVec`.
72    pub const fn new() -> Self {
73        Self {
74            slots: Vec::new(),
75            _phantom: PhantomData,
76        }
77    }
78
79    /// Returns the number of elements in the `SlotVec`.
80    pub fn len(&self) -> usize {
81        self.slots.len()
82    }
83
84    /// Returns true if the `SlotVec` is empty.
85    pub fn is_empty(&self) -> bool {
86        self.slots.is_empty()
87    }
88
89    /// Inserts a value into the `SlotVec` and returns the key.
90    pub fn insert(&mut self, value: Val) -> Key<Tag> {
91        let key = Key::from_raw(self.slots.len());
92        self.slots.push(value);
93        key
94    }
95
96    /// Use the provided function to generate a value given the key and insert it into the `SlotVec`.
97    pub fn insert_with_key<F>(&mut self, func: F) -> Key<Tag>
98    where
99        F: FnOnce(Key<Tag>) -> Val,
100    {
101        let key = Key::from_raw(self.slots.len());
102        self.slots.push((func)(key));
103        key
104    }
105
106    /// Returns a reference to the value associated with the key.
107    pub fn get(&self, key: Key<Tag>) -> Option<&Val> {
108        self.slots.get(key.index)
109    }
110
111    /// Returns a mutable reference to the value associated with the key.
112    pub fn get_mut(&mut self, key: Key<Tag>) -> Option<&mut Val> {
113        self.slots.get_mut(key.index)
114    }
115
116    /// Iterate the key-value pairs, where the value is a shared reference.
117    pub fn iter(
118        &self,
119    ) -> impl DoubleEndedIterator<Item = (Key<Tag>, &'_ Val)> + ExactSizeIterator + FusedIterator + Clone
120    {
121        self.slots
122            .iter()
123            .enumerate()
124            .map(|(idx, val)| (Key::from_raw(idx), val))
125    }
126
127    /// Iterate the key-value pairs, where the value is a exclusive reference.
128    pub fn iter_mut(
129        &mut self,
130    ) -> impl DoubleEndedIterator<Item = (Key<Tag>, &'_ mut Val)> + ExactSizeIterator + FusedIterator
131    {
132        self.slots
133            .iter_mut()
134            .enumerate()
135            .map(|(idx, val)| (Key::from_raw(idx), val))
136    }
137
138    /// Iterate over the values by shared reference.
139    pub fn values(&self) -> std::slice::Iter<'_, Val> {
140        self.slots.iter()
141    }
142
143    /// Iterate over the values by exclusive reference.
144    pub fn values_mut(&mut self) -> std::slice::IterMut<'_, Val> {
145        self.slots.iter_mut()
146    }
147
148    /// Iterate over the keys.
149    pub fn keys(
150        &self,
151    ) -> impl '_ + DoubleEndedIterator<Item = Key<Tag>> + ExactSizeIterator + FusedIterator + Clone
152    {
153        self.iter().map(|(key, _)| key)
154    }
155}
156impl<Tag: ?Sized, Val> Index<Key<Tag>> for SlotVec<Tag, Val> {
157    type Output = Val;
158
159    fn index(&self, key: Key<Tag>) -> &Self::Output {
160        self.get(key).unwrap()
161    }
162}
163impl<Tag: ?Sized, Val> IndexMut<Key<Tag>> for SlotVec<Tag, Val> {
164    fn index_mut(&mut self, key: Key<Tag>) -> &mut Self::Output {
165        self.get_mut(key).unwrap()
166    }
167}
168impl<Key: ?Sized, Val> Default for SlotVec<Key, Val> {
169    fn default() -> Self {
170        Self::new()
171    }
172}
173
174/// A secondary map used to associated data with keys from elements in an existing [`SlotVec`].
175///
176/// Analogous to [`slotmap::SecondaryMap`].
177pub struct SecondarySlotVec<Tag: ?Sized, Val> {
178    slots: Vec<Option<Val>>,
179    _phantom: PhantomData<Tag>,
180}
181impl<Tag: ?Sized, Val> SecondarySlotVec<Tag, Val> {
182    /// Creates a new `SecondarySlotVec`.
183    pub const fn new() -> Self {
184        Self {
185            slots: Vec::new(),
186            _phantom: PhantomData,
187        }
188    }
189
190    /// Returns the number of elements in the secondary map.
191    pub fn len(&self) -> usize {
192        self.iter().count()
193    }
194
195    /// Returns if the secondary map is empty.
196    pub fn is_empty(&self) -> bool {
197        self.iter().next().is_none()
198    }
199
200    /// Inserts a value into the `SecondarySlotVec` and returns the previous value associated with the key.
201    pub fn insert(&mut self, key: Key<Tag>, value: Val) -> Option<Val> {
202        if key.index >= self.slots.len() {
203            self.slots.resize_with(key.index + 1, || None);
204        }
205        self.slots[key.index].replace(value)
206    }
207
208    /// Removes a value associated with the key from the `SecondarySlotVec` and returns it.
209    pub fn remove(&mut self, key: Key<Tag>) -> Option<Val> {
210        // TODO(mingwei): Shrink the vector?
211        self.slots[key.index].take()
212    }
213
214    /// Returns a reference to the value associated with the key.
215    pub fn get(&self, key: Key<Tag>) -> Option<&Val> {
216        self.slots.get(key.index).and_then(|v| v.as_ref())
217    }
218
219    /// Returns a mutable reference to the value associated with the key.
220    pub fn get_mut(&mut self, key: Key<Tag>) -> Option<&mut Val> {
221        self.slots.get_mut(key.index).and_then(|v| v.as_mut())
222    }
223
224    /// Returns a mutable reference to the value associated with the key, inserting a default value
225    /// if it doesn't yet exist.
226    pub fn get_or_insert_with(&mut self, key: Key<Tag>, default: impl FnOnce() -> Val) -> &mut Val {
227        if key.index >= self.slots.len() {
228            self.slots.resize_with(key.index + 1, || None);
229        }
230        self.slots[key.index].get_or_insert_with(default)
231    }
232
233    /// Iterate the key-value pairs, where the value is a shared reference.
234    pub fn iter(
235        &self,
236    ) -> impl DoubleEndedIterator<Item = (Key<Tag>, &'_ Val)> + FusedIterator + Clone {
237        self.slots
238            .iter()
239            .enumerate()
240            .filter_map(|(idx, opt_val)| Some((Key::from_raw(idx), opt_val.as_ref()?)))
241    }
242
243    /// Iterate the key-value pairs, where the value is a exclusive reference.
244    pub fn iter_mut(
245        &mut self,
246    ) -> impl DoubleEndedIterator<Item = (Key<Tag>, &'_ mut Val)> + FusedIterator {
247        self.slots
248            .iter_mut()
249            .enumerate()
250            .filter_map(|(idx, opt_val)| Some((Key::from_raw(idx), opt_val.as_mut()?)))
251    }
252
253    /// Iterate over the values by shared reference.
254    pub fn values(&self) -> impl DoubleEndedIterator<Item = &'_ Val> + FusedIterator + Clone {
255        self.slots.iter().filter_map(Option::as_ref)
256    }
257
258    /// Iterate over the values by exclusive reference.
259    pub fn values_mut(&mut self) -> impl DoubleEndedIterator<Item = &'_ mut Val> + FusedIterator {
260        self.slots.iter_mut().filter_map(Option::as_mut)
261    }
262
263    /// Iterate over the keys.
264    pub fn keys(&self) -> impl '_ + DoubleEndedIterator<Item = Key<Tag>> + FusedIterator + Clone {
265        self.iter().map(|(key, _)| key)
266    }
267}
268impl<Tag: ?Sized, Val> Default for SecondarySlotVec<Tag, Val> {
269    fn default() -> Self {
270        Self::new()
271    }
272}
273impl<Tag: ?Sized, Val> Clone for SecondarySlotVec<Tag, Val>
274where
275    Val: Clone,
276{
277    fn clone(&self) -> Self {
278        Self {
279            slots: self.slots.clone(),
280            _phantom: PhantomData,
281        }
282    }
283}
284impl<Tag: ?Sized, Val> Index<Key<Tag>> for SecondarySlotVec<Tag, Val> {
285    type Output = Val;
286
287    fn index(&self, key: Key<Tag>) -> &Self::Output {
288        self.get(key).unwrap()
289    }
290}
291impl<Tag: ?Sized, Val> IndexMut<Key<Tag>> for SecondarySlotVec<Tag, Val> {
292    fn index_mut(&mut self, key: Key<Tag>) -> &mut Self::Output {
293        self.get_mut(key).unwrap()
294    }
295}
296impl<Tag: ?Sized, Val> FromIterator<(Key<Tag>, Val)> for SecondarySlotVec<Tag, Val> {
297    fn from_iter<T: IntoIterator<Item = (Key<Tag>, Val)>>(iter: T) -> Self {
298        let mut map = SecondarySlotVec::new();
299        for (key, val) in iter {
300            map.insert(key, val);
301        }
302        map
303    }
304}
dfir_rs/util/slot_vec.rs

dfir_rs/util/
slot_vec.rs
