dfir_rs/util/
slot_vec.rs

Help
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    /// Inserts a value into the `SlotVec` and returns the key.
80    pub fn insert(&mut self, value: Val) -> Key<Tag> {
81        let key = Key::from_raw(self.slots.len());
82        self.slots.push(value);
83        key
84    }
85
86    /// Use the provided function to generate a value given the key and insert it into the `SlotVec`.
87    pub fn insert_with_key<F>(&mut self, func: F) -> Key<Tag>
88    where
89        F: FnOnce(Key<Tag>) -> Val,
90    {
91        let key = Key::from_raw(self.slots.len());
92        self.slots.push((func)(key));
93        key
94    }
95
96    /// Returns a reference to the value associated with the key.
97    pub fn get(&self, key: Key<Tag>) -> Option<&Val> {
98        self.slots.get(key.index)
99    }
100
101    /// Returns a mutable reference to the value associated with the key.
102    pub fn get_mut(&mut self, key: Key<Tag>) -> Option<&mut Val> {
103        self.slots.get_mut(key.index)
104    }
105
106    /// Returns the number of elements in the `SlotVec`.
107    pub fn len(&self) -> usize {
108        self.slots.len()
109    }
110
111    /// Returns true if the `SlotVec` is empty.
112    pub fn is_empty(&self) -> bool {
113        self.slots.is_empty()
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    /// Inserts a value into the `SecondarySlotVec` and returns the previous value associated with the key.
191    pub fn insert(&mut self, key: Key<Tag>, value: Val) -> Option<Val> {
192        if key.index >= self.slots.len() {
193            self.slots.resize_with(key.index + 1, || None);
194        }
195        self.slots[key.index].replace(value)
196    }
197
198    /// Removes a value associated with the key from the `SecondarySlotVec` and returns it.
199    pub fn remove(&mut self, key: Key<Tag>) -> Option<Val> {
200        // TODO(mingwei): Shrink the vector?
201        self.slots[key.index].take()
202    }
203
204    /// Returns a reference to the value associated with the key.
205    pub fn get(&self, key: Key<Tag>) -> Option<&Val> {
206        self.slots.get(key.index).and_then(|v| v.as_ref())
207    }
208
209    /// Returns a mutable reference to the value associated with the key.
210    pub fn get_mut(&mut self, key: Key<Tag>) -> Option<&mut Val> {
211        self.slots.get_mut(key.index).and_then(|v| v.as_mut())
212    }
213
214    /// Returns a mutable reference to the value associated with the key, inserting a default value
215    /// if it doesn't yet exist.
216    pub fn get_or_insert_with(&mut self, key: Key<Tag>, default: impl FnOnce() -> Val) -> &mut Val {
217        if key.index >= self.slots.len() {
218            self.slots.resize_with(key.index + 1, || None);
219        }
220        self.slots[key.index].get_or_insert_with(default)
221    }
222
223    /// Iterate the key-value pairs, where the value is a shared reference.
224    pub fn iter(
225        &self,
226    ) -> impl DoubleEndedIterator<Item = (Key<Tag>, &'_ Val)> + FusedIterator + Clone {
227        self.slots
228            .iter()
229            .enumerate()
230            .filter_map(|(idx, opt_val)| Some((Key::from_raw(idx), opt_val.as_ref()?)))
231    }
232
233    /// Iterate the key-value pairs, where the value is a exclusive reference.
234    pub fn iter_mut(
235        &mut self,
236    ) -> impl DoubleEndedIterator<Item = (Key<Tag>, &'_ mut Val)> + FusedIterator {
237        self.slots
238            .iter_mut()
239            .enumerate()
240            .filter_map(|(idx, opt_val)| Some((Key::from_raw(idx), opt_val.as_mut()?)))
241    }
242
243    /// Iterate over the values by shared reference.
244    pub fn values(&self) -> impl DoubleEndedIterator<Item = &'_ Val> + FusedIterator + Clone {
245        self.slots.iter().filter_map(Option::as_ref)
246    }
247
248    /// Iterate over the values by exclusive reference.
249    pub fn values_mut(&mut self) -> impl DoubleEndedIterator<Item = &'_ mut Val> + FusedIterator {
250        self.slots.iter_mut().filter_map(Option::as_mut)
251    }
252
253    /// Iterate over the keys.
254    pub fn keys(&self) -> impl '_ + DoubleEndedIterator<Item = Key<Tag>> + FusedIterator + Clone {
255        self.iter().map(|(key, _)| key)
256    }
257}
258impl<Tag: ?Sized, Val> Default for SecondarySlotVec<Tag, Val> {
259    fn default() -> Self {
260        Self::new()
261    }
262}
263impl<Tag: ?Sized, Val> Index<Key<Tag>> for SecondarySlotVec<Tag, Val> {
264    type Output = Val;
265
266    fn index(&self, key: Key<Tag>) -> &Self::Output {
267        self.get(key).unwrap()
268    }
269}
270impl<Tag: ?Sized, Val> IndexMut<Key<Tag>> for SecondarySlotVec<Tag, Val> {
271    fn index_mut(&mut self, key: Key<Tag>) -> &mut Self::Output {
272        self.get_mut(key).unwrap()
273    }
274}
dfir_rs/util/slot_vec.rs

dfir_rs/util/
slot_vec.rs
