naga/arena/
mod.rs

1/*! The [`Arena`], [`UniqueArena`], and [`Handle`] types.
2
3To improve translator performance and reduce memory usage, most structures are
4stored in an [`Arena`]. An `Arena<T>` stores a series of `T` values, indexed by
5[`Handle<T>`] values, which are just wrappers around integer indexes.
6For example, a `Function`'s expressions are stored in an `Arena<Expression>`,
7and compound expressions refer to their sub-expressions via `Handle<Expression>`
8values.
9
10A [`UniqueArena`] is just like an `Arena`, except that it stores only a single
11instance of each value. The value type must implement `Eq` and `Hash`. Like an
12`Arena`, inserting a value into a `UniqueArena` returns a `Handle` which can be
13used to efficiently access the value, without a hash lookup. Inserting a value
14multiple times returns the same `Handle`.
15
16If the `span` feature is enabled, both `Arena` and `UniqueArena` can associate a
17source code span with each element.
18
19[`Handle<T>`]: Handle
20*/
21
22mod handle;
23mod handle_set;
24mod handlevec;
25mod range;
26mod unique_arena;
27
28pub use handle::{BadHandle, Handle};
29pub(crate) use handle_set::HandleSet;
30pub(crate) use handlevec::HandleVec;
31pub use range::{BadRangeError, Range};
32pub use unique_arena::UniqueArena;
33
34use alloc::vec::Vec;
35use core::{fmt, ops};
36
37use crate::Span;
38
39use handle::Index;
40
41/// An arena holding some kind of component (e.g., type, constant,
42/// instruction, etc.) that can be referenced.
43///
44/// Adding new items to the arena produces a strongly-typed [`Handle`].
45/// The arena can be indexed using the given handle to obtain
46/// a reference to the stored item.
47#[derive(Clone)]
48#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
49#[cfg_attr(feature = "serialize", serde(transparent))]
50#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
51#[cfg_attr(test, derive(PartialEq))]
52pub struct Arena<T> {
53    /// Values of this arena.
54    data: Vec<T>,
55    #[cfg_attr(feature = "serialize", serde(skip))]
56    span_info: Vec<Span>,
57}
58
59impl<T> Default for Arena<T> {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65impl<T: fmt::Debug> fmt::Debug for Arena<T> {
66    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
67        f.debug_map().entries(self.iter()).finish()
68    }
69}
70
71impl<T> Arena<T> {
72    /// Create a new arena with no initial capacity allocated.
73    pub const fn new() -> Self {
74        Arena {
75            data: Vec::new(),
76            span_info: Vec::new(),
77        }
78    }
79
80    /// Extracts the inner vector.
81    pub fn into_inner(self) -> Vec<T> {
82        self.data
83    }
84
85    /// Returns the current number of items stored in this arena.
86    pub const fn len(&self) -> usize {
87        self.data.len()
88    }
89
90    /// Returns `true` if the arena contains no elements.
91    pub const fn is_empty(&self) -> bool {
92        self.data.is_empty()
93    }
94
95    /// Returns an iterator over the items stored in this arena, returning both
96    /// the item's handle and a reference to it.
97    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> + ExactSizeIterator {
98        self.data
99            .iter()
100            .enumerate()
101            .map(|(i, v)| (Handle::from_usize(i), v))
102    }
103
104    /// Returns an iterator over the items stored in this arena, returning both
105    /// the item's handle and a reference to it.
106    pub fn iter_mut_span(
107        &mut self,
108    ) -> impl DoubleEndedIterator<Item = (Handle<T>, &mut T, &Span)> + ExactSizeIterator {
109        self.data
110            .iter_mut()
111            .zip(self.span_info.iter())
112            .enumerate()
113            .map(|(i, (v, span))| (Handle::from_usize(i), v, span))
114    }
115
116    /// Drains the arena, returning an iterator over the items stored.
117    pub fn drain(&mut self) -> impl DoubleEndedIterator<Item = (Handle<T>, T, Span)> {
118        let arena = core::mem::take(self);
119        arena
120            .data
121            .into_iter()
122            .zip(arena.span_info)
123            .enumerate()
124            .map(|(i, (v, span))| (Handle::from_usize(i), v, span))
125    }
126
127    /// Returns a iterator over the items stored in this arena,
128    /// returning both the item's handle and a mutable reference to it.
129    pub fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = (Handle<T>, &mut T)> {
130        self.data
131            .iter_mut()
132            .enumerate()
133            .map(|(i, v)| (Handle::from_usize(i), v))
134    }
135
136    /// Adds a new value to the arena, returning a typed handle.
137    pub fn append(&mut self, value: T, span: Span) -> Handle<T> {
138        let index = self.data.len();
139        self.data.push(value);
140        self.span_info.push(span);
141        Handle::from_usize(index)
142    }
143
144    /// Fetch a handle to an existing type.
145    pub fn fetch_if<F: Fn(&T) -> bool>(&self, fun: F) -> Option<Handle<T>> {
146        self.data
147            .iter()
148            .position(fun)
149            .map(|index| Handle::from_usize(index))
150    }
151
152    /// Adds a value with a custom check for uniqueness:
153    /// returns a handle pointing to
154    /// an existing element if the check succeeds, or adds a new
155    /// element otherwise.
156    pub fn fetch_if_or_append<F: Fn(&T, &T) -> bool>(
157        &mut self,
158        value: T,
159        span: Span,
160        fun: F,
161    ) -> Handle<T> {
162        if let Some(index) = self.data.iter().position(|d| fun(d, &value)) {
163            Handle::from_usize(index)
164        } else {
165            self.append(value, span)
166        }
167    }
168
169    /// Adds a value with a check for uniqueness, where the check is plain comparison.
170    pub fn fetch_or_append(&mut self, value: T, span: Span) -> Handle<T>
171    where
172        T: PartialEq,
173    {
174        self.fetch_if_or_append(value, span, T::eq)
175    }
176
177    pub fn try_get(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
178        self.data
179            .get(handle.index())
180            .ok_or_else(|| BadHandle::new(handle))
181    }
182
183    /// Get a mutable reference to an element in the arena.
184    pub fn get_mut(&mut self, handle: Handle<T>) -> &mut T {
185        self.data.get_mut(handle.index()).unwrap()
186    }
187
188    /// Get the range of handles from a particular number of elements to the end.
189    pub fn range_from(&self, old_length: usize) -> Range<T> {
190        let range = old_length as u32..self.data.len() as u32;
191        Range::from_index_range(range, self)
192    }
193
194    /// Clears the arena keeping all allocations
195    pub fn clear(&mut self) {
196        self.data.clear()
197    }
198
199    pub fn get_span(&self, handle: Handle<T>) -> Span {
200        *self
201            .span_info
202            .get(handle.index())
203            .unwrap_or(&Span::default())
204    }
205
206    /// Assert that `handle` is valid for this arena.
207    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
208        if handle.index() < self.data.len() {
209            Ok(())
210        } else {
211            Err(BadHandle::new(handle))
212        }
213    }
214
215    /// Assert that `range` is valid for this arena.
216    pub fn check_contains_range(&self, range: &Range<T>) -> Result<(), BadRangeError> {
217        // Since `range.inner` is a `Range<u32>`, we only need to check that the
218        // start precedes the end, and that the end is in range.
219        if range.inner.start > range.inner.end {
220            return Err(BadRangeError::new(range.clone()));
221        }
222
223        // Empty ranges are tolerated: they can be produced by compaction.
224        if range.inner.start == range.inner.end {
225            return Ok(());
226        }
227
228        let last_handle = Handle::new(Index::new(range.inner.end - 1).unwrap());
229        if self.check_contains_handle(last_handle).is_err() {
230            return Err(BadRangeError::new(range.clone()));
231        }
232
233        Ok(())
234    }
235
236    pub(crate) fn retain_mut<P>(&mut self, mut predicate: P)
237    where
238        P: FnMut(Handle<T>, &mut T) -> bool,
239    {
240        let mut index = 0;
241        let mut retained = 0;
242        self.data.retain_mut(|elt| {
243            let handle = Handle::from_usize(index);
244            let keep = predicate(handle, elt);
245
246            // Since `predicate` needs mutable access to each element,
247            // we can't feasibly call it twice, so we have to compact
248            // spans by hand in parallel as part of this iteration.
249            if keep {
250                self.span_info[retained] = self.span_info[index];
251                retained += 1;
252            }
253
254            index += 1;
255            keep
256        });
257
258        self.span_info.truncate(retained);
259    }
260}
261
262#[cfg(feature = "deserialize")]
263impl<'de, T> serde::Deserialize<'de> for Arena<T>
264where
265    T: serde::Deserialize<'de>,
266{
267    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
268    where
269        D: serde::Deserializer<'de>,
270    {
271        let data = Vec::deserialize(deserializer)?;
272        let span_info = core::iter::repeat_n(Span::default(), data.len()).collect();
273
274        Ok(Self { data, span_info })
275    }
276}
277
278impl<T> ops::Index<Handle<T>> for Arena<T> {
279    type Output = T;
280    fn index(&self, handle: Handle<T>) -> &T {
281        &self.data[handle.index()]
282    }
283}
284
285impl<T> ops::IndexMut<Handle<T>> for Arena<T> {
286    fn index_mut(&mut self, handle: Handle<T>) -> &mut T {
287        &mut self.data[handle.index()]
288    }
289}
290
291impl<T> ops::Index<Range<T>> for Arena<T> {
292    type Output = [T];
293    fn index(&self, range: Range<T>) -> &[T] {
294        &self.data[range.inner.start as usize..range.inner.end as usize]
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    #[test]
303    fn append_non_unique() {
304        let mut arena: Arena<u8> = Arena::new();
305        let t1 = arena.append(0, Default::default());
306        let t2 = arena.append(0, Default::default());
307        assert!(t1 != t2);
308        assert!(arena[t1] == arena[t2]);
309    }
310
311    #[test]
312    fn append_unique() {
313        let mut arena: Arena<u8> = Arena::new();
314        let t1 = arena.append(0, Default::default());
315        let t2 = arena.append(1, Default::default());
316        assert!(t1 != t2);
317        assert!(arena[t1] != arena[t2]);
318    }
319
320    #[test]
321    fn fetch_or_append_non_unique() {
322        let mut arena: Arena<u8> = Arena::new();
323        let t1 = arena.fetch_or_append(0, Default::default());
324        let t2 = arena.fetch_or_append(0, Default::default());
325        assert!(t1 == t2);
326        assert!(arena[t1] == arena[t2])
327    }
328
329    #[test]
330    fn fetch_or_append_unique() {
331        let mut arena: Arena<u8> = Arena::new();
332        let t1 = arena.fetch_or_append(0, Default::default());
333        let t2 = arena.fetch_or_append(1, Default::default());
334        assert!(t1 != t2);
335        assert!(arena[t1] != arena[t2]);
336    }
337}
naga/arena/mod.rs

naga/arena/
mod.rs
