la_arena/
map.rs

1use std::iter::Enumerate;
2use std::marker::PhantomData;
3
4use crate::Idx;
5
6/// A map from arena indexes to some other type.
7/// Space requirement is O(highest index).
8#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub struct ArenaMap<IDX, V> {
10    v: Vec<Option<V>>,
11    _ty: PhantomData<IDX>,
12}
13
14impl<T, V> ArenaMap<Idx<T>, V> {
15    /// Creates a new empty map.
16    pub const fn new() -> Self {
17        Self { v: Vec::new(), _ty: PhantomData }
18    }
19
20    /// Create a new empty map with specific capacity.
21    pub fn with_capacity(capacity: usize) -> Self {
22        Self { v: Vec::with_capacity(capacity), _ty: PhantomData }
23    }
24
25    /// Reserves capacity for at least additional more elements to be inserted in the map.
26    pub fn reserve(&mut self, additional: usize) {
27        self.v.reserve(additional);
28    }
29
30    /// Clears the map, removing all elements.
31    pub fn clear(&mut self) {
32        self.v.clear();
33    }
34
35    /// Shrinks the capacity of the map as much as possible.
36    pub fn shrink_to_fit(&mut self) {
37        let min_len = self.v.iter().rposition(|slot| slot.is_some()).map_or(0, |i| i + 1);
38        self.v.truncate(min_len);
39        self.v.shrink_to_fit();
40    }
41
42    /// Returns whether the map contains a value for the specified index.
43    pub fn contains_idx(&self, idx: Idx<T>) -> bool {
44        matches!(self.v.get(Self::to_idx(idx)), Some(Some(_)))
45    }
46
47    /// Removes an index from the map, returning the value at the index if the index was previously in the map.
48    pub fn remove(&mut self, idx: Idx<T>) -> Option<V> {
49        self.v.get_mut(Self::to_idx(idx))?.take()
50    }
51
52    /// Inserts a value associated with a given arena index into the map.
53    ///
54    /// If the map did not have this index present, None is returned.
55    /// Otherwise, the value is updated, and the old value is returned.
56    pub fn insert(&mut self, idx: Idx<T>, t: V) -> Option<V> {
57        let idx = Self::to_idx(idx);
58
59        self.v.resize_with((idx + 1).max(self.v.len()), || None);
60        self.v[idx].replace(t)
61    }
62
63    /// Returns a reference to the value associated with the provided index
64    /// if it is present.
65    pub fn get(&self, idx: Idx<T>) -> Option<&V> {
66        self.v.get(Self::to_idx(idx)).and_then(|it| it.as_ref())
67    }
68
69    /// Returns a mutable reference to the value associated with the provided index
70    /// if it is present.
71    pub fn get_mut(&mut self, idx: Idx<T>) -> Option<&mut V> {
72        self.v.get_mut(Self::to_idx(idx)).and_then(|it| it.as_mut())
73    }
74
75    /// Returns an iterator over the values in the map.
76    pub fn values(&self) -> impl DoubleEndedIterator<Item = &V> {
77        self.v.iter().filter_map(|o| o.as_ref())
78    }
79
80    /// Returns an iterator over mutable references to the values in the map.
81    pub fn values_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut V> {
82        self.v.iter_mut().filter_map(|o| o.as_mut())
83    }
84
85    /// Returns an iterator over the arena indexes and values in the map.
86    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Idx<T>, &V)> {
87        self.v.iter().enumerate().filter_map(|(idx, o)| Some((Self::from_idx(idx), o.as_ref()?)))
88    }
89
90    /// Returns an iterator over the arena indexes and values in the map.
91    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Idx<T>, &mut V)> {
92        self.v
93            .iter_mut()
94            .enumerate()
95            .filter_map(|(idx, o)| Some((Self::from_idx(idx), o.as_mut()?)))
96    }
97
98    /// Gets the given key's corresponding entry in the map for in-place manipulation.
99    pub fn entry(&mut self, idx: Idx<T>) -> Entry<'_, Idx<T>, V> {
100        let idx = Self::to_idx(idx);
101        self.v.resize_with((idx + 1).max(self.v.len()), || None);
102        match &mut self.v[idx] {
103            slot @ Some(_) => Entry::Occupied(OccupiedEntry { slot, _ty: PhantomData }),
104            slot @ None => Entry::Vacant(VacantEntry { slot, _ty: PhantomData }),
105        }
106    }
107
108    fn to_idx(idx: Idx<T>) -> usize {
109        u32::from(idx.into_raw()) as usize
110    }
111
112    fn from_idx(idx: usize) -> Idx<T> {
113        Idx::from_raw((idx as u32).into())
114    }
115}
116
117impl<T, V> std::ops::Index<Idx<V>> for ArenaMap<Idx<V>, T> {
118    type Output = T;
119    fn index(&self, idx: Idx<V>) -> &T {
120        self.v[Self::to_idx(idx)].as_ref().unwrap()
121    }
122}
123
124impl<T, V> std::ops::IndexMut<Idx<V>> for ArenaMap<Idx<V>, T> {
125    fn index_mut(&mut self, idx: Idx<V>) -> &mut T {
126        self.v[Self::to_idx(idx)].as_mut().unwrap()
127    }
128}
129
130impl<T, V> Default for ArenaMap<Idx<V>, T> {
131    fn default() -> Self {
132        Self::new()
133    }
134}
135
136impl<T, V> Extend<(Idx<V>, T)> for ArenaMap<Idx<V>, T> {
137    fn extend<I: IntoIterator<Item = (Idx<V>, T)>>(&mut self, iter: I) {
138        iter.into_iter().for_each(move |(k, v)| {
139            self.insert(k, v);
140        });
141    }
142}
143
144impl<T, V> FromIterator<(Idx<V>, T)> for ArenaMap<Idx<V>, T> {
145    fn from_iter<I: IntoIterator<Item = (Idx<V>, T)>>(iter: I) -> Self {
146        let mut this = Self::new();
147        this.extend(iter);
148        this
149    }
150}
151
152pub struct ArenaMapIter<IDX, V> {
153    iter: Enumerate<std::vec::IntoIter<Option<V>>>,
154    _ty: PhantomData<IDX>,
155}
156
157impl<T, V> IntoIterator for ArenaMap<Idx<T>, V> {
158    type Item = (Idx<T>, V);
159
160    type IntoIter = ArenaMapIter<Idx<T>, V>;
161
162    fn into_iter(self) -> Self::IntoIter {
163        let iter = self.v.into_iter().enumerate();
164        Self::IntoIter { iter, _ty: PhantomData }
165    }
166}
167
168impl<T, V> ArenaMapIter<Idx<T>, V> {
169    fn mapper((idx, o): (usize, Option<V>)) -> Option<(Idx<T>, V)> {
170        Some((ArenaMap::<Idx<T>, V>::from_idx(idx), o?))
171    }
172}
173
174impl<T, V> Iterator for ArenaMapIter<Idx<T>, V> {
175    type Item = (Idx<T>, V);
176
177    #[inline]
178    fn next(&mut self) -> Option<Self::Item> {
179        for next in self.iter.by_ref() {
180            match Self::mapper(next) {
181                Some(r) => return Some(r),
182                None => continue,
183            }
184        }
185
186        None
187    }
188
189    #[inline]
190    fn size_hint(&self) -> (usize, Option<usize>) {
191        self.iter.size_hint()
192    }
193}
194
195impl<T, V> DoubleEndedIterator for ArenaMapIter<Idx<T>, V> {
196    #[inline]
197    fn next_back(&mut self) -> Option<Self::Item> {
198        while let Some(next_back) = self.iter.next_back() {
199            match Self::mapper(next_back) {
200                Some(r) => return Some(r),
201                None => continue,
202            }
203        }
204
205        None
206    }
207}
208
209/// A view into a single entry in a map, which may either be vacant or occupied.
210///
211/// This `enum` is constructed from the [`entry`] method on [`ArenaMap`].
212///
213/// [`entry`]: ArenaMap::entry
214pub enum Entry<'a, IDX, V> {
215    /// A vacant entry.
216    Vacant(VacantEntry<'a, IDX, V>),
217    /// An occupied entry.
218    Occupied(OccupiedEntry<'a, IDX, V>),
219}
220
221impl<'a, IDX, V> Entry<'a, IDX, V> {
222    /// Ensures a value is in the entry by inserting the default if empty, and returns a mutable reference to
223    /// the value in the entry.
224    pub fn or_insert(self, default: V) -> &'a mut V {
225        match self {
226            Self::Vacant(ent) => ent.insert(default),
227            Self::Occupied(ent) => ent.into_mut(),
228        }
229    }
230
231    /// Ensures a value is in the entry by inserting the result of the default function if empty, and returns
232    /// a mutable reference to the value in the entry.
233    pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
234        match self {
235            Self::Vacant(ent) => ent.insert(default()),
236            Self::Occupied(ent) => ent.into_mut(),
237        }
238    }
239
240    /// Provides in-place mutable access to an occupied entry before any potential inserts into the map.
241    pub fn and_modify<F: FnOnce(&mut V)>(mut self, f: F) -> Self {
242        if let Self::Occupied(ent) = &mut self {
243            f(ent.get_mut());
244        }
245        self
246    }
247}
248
249impl<'a, IDX, V> Entry<'a, IDX, V>
250where
251    V: Default,
252{
253    /// Ensures a value is in the entry by inserting the default value if empty, and returns a mutable reference
254    /// to the value in the entry.
255    // BUG this clippy lint is a false positive
256    #[allow(clippy::unwrap_or_default)]
257    pub fn or_default(self) -> &'a mut V {
258        self.or_insert_with(Default::default)
259    }
260}
261
262/// A view into an vacant entry in a [`ArenaMap`]. It is part of the [`Entry`] enum.
263pub struct VacantEntry<'a, IDX, V> {
264    slot: &'a mut Option<V>,
265    _ty: PhantomData<IDX>,
266}
267
268impl<'a, IDX, V> VacantEntry<'a, IDX, V> {
269    /// Sets the value of the entry with the `VacantEntry`’s key, and returns a mutable reference to it.
270    pub fn insert(self, value: V) -> &'a mut V {
271        self.slot.insert(value)
272    }
273}
274
275/// A view into an occupied entry in a [`ArenaMap`]. It is part of the [`Entry`] enum.
276pub struct OccupiedEntry<'a, IDX, V> {
277    slot: &'a mut Option<V>,
278    _ty: PhantomData<IDX>,
279}
280
281impl<'a, IDX, V> OccupiedEntry<'a, IDX, V> {
282    /// Gets a reference to the value in the entry.
283    pub fn get(&self) -> &V {
284        self.slot.as_ref().expect("Occupied")
285    }
286
287    /// Gets a mutable reference to the value in the entry.
288    pub fn get_mut(&mut self) -> &mut V {
289        self.slot.as_mut().expect("Occupied")
290    }
291
292    /// Converts the entry into a mutable reference to its value.
293    pub fn into_mut(self) -> &'a mut V {
294        self.slot.as_mut().expect("Occupied")
295    }
296
297    /// Sets the value of the entry with the `OccupiedEntry`’s key, and returns the entry’s old value.
298    pub fn insert(&mut self, value: V) -> V {
299        self.slot.replace(value).expect("Occupied")
300    }
301
302    /// Takes the value of the entry out of the map, and returns it.
303    pub fn remove(self) -> V {
304        self.slot.take().expect("Occupied")
305    }
306}
la_arena/map.rs

la_arena/
map.rs
