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::marker::PhantomData;
6use std::ops::{Index, IndexMut};
7
8/// A key into a SlotVec.
9#[repr(transparent)]
10pub struct Key<Tag: ?Sized> {
11    index: usize,
12    _phantom: PhantomData<Tag>,
13}
14impl<Tag: ?Sized> Key<Tag> {
15    /// Creates a Key from a raw index. Avoid using this function directly.
16    pub fn from_raw(index: usize) -> Self {
17        Key {
18            index,
19            _phantom: PhantomData,
20        }
21    }
22}
23impl<Tag: ?Sized> Clone for Key<Tag> {
24    fn clone(&self) -> Self {
25        *self
26    }
27}
28impl<Tag: ?Sized> Copy for Key<Tag> {}
29impl<Tag: ?Sized> PartialOrd for Key<Tag> {
30    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
31        Some(self.cmp(other))
32    }
33}
34impl<Tag: ?Sized> Ord for Key<Tag> {
35    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
36        self.index.cmp(&other.index)
37    }
38}
39impl<Tag: ?Sized> PartialEq for Key<Tag> {
40    fn eq(&self, other: &Self) -> bool {
41        self.index == other.index
42    }
43}
44impl<Tag: ?Sized> Eq for Key<Tag> {}
45impl<Tag: ?Sized> Hash for Key<Tag> {
46    fn hash<H: Hasher>(&self, state: &mut H) {
47        self.index.hash(state);
48    }
49}
50impl<Tag: ?Sized> Debug for Key<Tag> {
51    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
52        write!(f, "Key({})", self.index)
53    }
54}
55impl<Tag: ?Sized> Display for Key<Tag> {
56    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
57        write!(f, "{}", self.index)
58    }
59}
60
61/// A Vec-based SlotMap-esque datastructure without removes.
62///
63/// Analogous to [`slotmap::SlotMap`], but avoids the overhead of tracking removed keys.
64#[repr(transparent)]
65pub struct SlotVec<Tag: ?Sized, Val> {
66    slots: Vec<Val>,
67    _phantom: PhantomData<Tag>,
68}
69impl<Tag: ?Sized, Val> SlotVec<Tag, Val> {
70    /// Creates a new `SlotVec`.
71    pub fn new() -> Self {
72        Self {
73            slots: Vec::default(),
74            _phantom: PhantomData,
75        }
76    }
77
78    /// Inserts a value into the `SlotVec` and returns the key.
79    pub fn insert(&mut self, value: Val) -> Key<Tag> {
80        let key = Key::from_raw(self.slots.len());
81        self.slots.push(value);
82        key
83    }
84
85    /// Use the provided function to generate a value given the key and insert it into the `SlotVec`.
86    pub fn insert_with_key<F>(&mut self, func: F) -> Key<Tag>
87    where
88        F: FnOnce(Key<Tag>) -> Val,
89    {
90        let key = Key::from_raw(self.slots.len());
91        self.slots.push((func)(key));
92        key
93    }
94
95    /// Returns a reference to the value associated with the key.
96    pub fn get(&self, key: Key<Tag>) -> Option<&Val> {
97        self.slots.get(key.index)
98    }
99
100    /// Returns a mutable reference to the value associated with the key.
101    pub fn get_mut(&mut self, key: Key<Tag>) -> Option<&mut Val> {
102        self.slots.get_mut(key.index)
103    }
104
105    /// Returns the number of elements in the `SlotVec`.
106    pub fn len(&self) -> usize {
107        self.slots.len()
108    }
109
110    /// Returns true if the `SlotVec` is empty.
111    pub fn is_empty(&self) -> bool {
112        self.slots.is_empty()
113    }
114}
115impl<Tag: ?Sized, Val> Index<Key<Tag>> for SlotVec<Tag, Val> {
116    type Output = Val;
117
118    fn index(&self, key: Key<Tag>) -> &Self::Output {
119        self.get(key).unwrap()
120    }
121}
122impl<Tag: ?Sized, Val> IndexMut<Key<Tag>> for SlotVec<Tag, Val> {
123    fn index_mut(&mut self, key: Key<Tag>) -> &mut Self::Output {
124        self.get_mut(key).unwrap()
125    }
126}
127impl<Key: ?Sized, Val> Default for SlotVec<Key, Val> {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133/// A secondary map used to associated data with keys from elements in an existing [`SlotVec`].
134///
135/// Analogous to [`slotmap::SecondaryMap`].
136pub struct SecondarySlotVec<Tag: ?Sized, Val> {
137    slots: Vec<Option<Val>>,
138    _phantom: PhantomData<Tag>,
139}
140impl<Tag: ?Sized, Val> SecondarySlotVec<Tag, Val> {
141    /// Creates a new `SecondarySlotVec`.
142    pub fn new() -> Self {
143        Self {
144            slots: Vec::default(),
145            _phantom: PhantomData,
146        }
147    }
148
149    /// Inserts a value into the `SecondarySlotVec` and returns the previous value associated with the key.
150    pub fn insert(&mut self, key: Key<Tag>, value: Val) -> Option<Val> {
151        if key.index >= self.slots.len() {
152            self.slots.resize_with(key.index + 1, || None);
153        }
154        self.slots[key.index].replace(value)
155    }
156
157    /// Removes a value associated with the key from the `SecondarySlotVec` and returns it.
158    pub fn remove(&mut self, key: Key<Tag>) -> Option<Val> {
159        // TODO(mingwei): Shrink the vector?
160        self.slots[key.index].take()
161    }
162
163    /// Returns a reference to the value associated with the key.
164    pub fn get(&self, key: Key<Tag>) -> Option<&Val> {
165        self.slots.get(key.index).and_then(|v| v.as_ref())
166    }
167
168    /// Returns a mutable reference to the value associated with the key.
169    pub fn get_mut(&mut self, key: Key<Tag>) -> Option<&mut Val> {
170        self.slots.get_mut(key.index).and_then(|v| v.as_mut())
171    }
172}
173impl<Tag: ?Sized, Val> Default for SecondarySlotVec<Tag, Val> {
174    fn default() -> Self {
175        Self::new()
176    }
177}
178impl<Tag: ?Sized, Val> Index<Key<Tag>> for SecondarySlotVec<Tag, Val> {
179    type Output = Val;
180
181    fn index(&self, key: Key<Tag>) -> &Self::Output {
182        self.get(key).unwrap()
183    }
184}
185impl<Tag: ?Sized, Val> IndexMut<Key<Tag>> for SecondarySlotVec<Tag, Val> {
186    fn index_mut(&mut self, key: Key<Tag>) -> &mut Self::Output {
187        self.get_mut(key).unwrap()
188    }
189}
dfir_rs/util/slot_vec.rs

dfir_rs/util/
slot_vec.rs
