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>`](Handle) 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. (When examining the serialized form of a `Module`, note that the first
9element of an `Arena` has an index of 1, not 0.)
10
11A [`UniqueArena`] is just like an `Arena`, except that it stores only a single
12instance of each value. The value type must implement `Eq` and `Hash`. Like an
13`Arena`, inserting a value into a `UniqueArena` returns a `Handle` which can be
14used to efficiently access the value, without a hash lookup. Inserting a value
15multiple times returns the same `Handle`.
16
17If the `span` feature is enabled, both `Arena` and `UniqueArena` can associate a
18source code span with each element.
19
20[`Handle<T>`]: Handle
21*/
22
23mod handle;
24mod handle_set;
25mod handlevec;
26mod range;
27mod unique_arena;
28
29pub use handle::{BadHandle, Handle};
30pub(crate) use handle_set::HandleSet;
31pub(crate) use handlevec::HandleVec;
32pub use range::{BadRangeError, Range};
33pub use unique_arena::UniqueArena;
34
35use crate::Span;
36
37use handle::Index;
38
39use std::{fmt, ops};
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    #[allow(clippy::missing_const_for_fn)] // ignore due to requirement of #![feature(const_precise_live_drops)]
82    pub fn into_inner(self) -> Vec<T> {
83        self.data
84    }
85
86    /// Returns the current number of items stored in this arena.
87    pub fn len(&self) -> usize {
88        self.data.len()
89    }
90
91    /// Returns `true` if the arena contains no elements.
92    pub fn is_empty(&self) -> bool {
93        self.data.is_empty()
94    }
95
96    /// Returns an iterator over the items stored in this arena, returning both
97    /// the item's handle and a reference to it.
98    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> {
99        self.data
100            .iter()
101            .enumerate()
102            .map(|(i, v)| unsafe { (Handle::from_usize_unchecked(i), v) })
103    }
104
105    /// Drains the arena, returning an iterator over the items stored.
106    pub fn drain(&mut self) -> impl DoubleEndedIterator<Item = (Handle<T>, T, Span)> {
107        let arena = std::mem::take(self);
108        arena
109            .data
110            .into_iter()
111            .zip(arena.span_info)
112            .enumerate()
113            .map(|(i, (v, span))| unsafe { (Handle::from_usize_unchecked(i), v, span) })
114    }
115
116    /// Returns a iterator over the items stored in this arena,
117    /// returning both the item's handle and a mutable reference to it.
118    pub fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = (Handle<T>, &mut T)> {
119        self.data
120            .iter_mut()
121            .enumerate()
122            .map(|(i, v)| unsafe { (Handle::from_usize_unchecked(i), v) })
123    }
124
125    /// Adds a new value to the arena, returning a typed handle.
126    pub fn append(&mut self, value: T, span: Span) -> Handle<T> {
127        let index = self.data.len();
128        self.data.push(value);
129        self.span_info.push(span);
130        Handle::from_usize(index)
131    }
132
133    /// Fetch a handle to an existing type.
134    pub fn fetch_if<F: Fn(&T) -> bool>(&self, fun: F) -> Option<Handle<T>> {
135        self.data
136            .iter()
137            .position(fun)
138            .map(|index| unsafe { Handle::from_usize_unchecked(index) })
139    }
140
141    /// Adds a value with a custom check for uniqueness:
142    /// returns a handle pointing to
143    /// an existing element if the check succeeds, or adds a new
144    /// element otherwise.
145    pub fn fetch_if_or_append<F: Fn(&T, &T) -> bool>(
146        &mut self,
147        value: T,
148        span: Span,
149        fun: F,
150    ) -> Handle<T> {
151        if let Some(index) = self.data.iter().position(|d| fun(d, &value)) {
152            unsafe { Handle::from_usize_unchecked(index) }
153        } else {
154            self.append(value, span)
155        }
156    }
157
158    /// Adds a value with a check for uniqueness, where the check is plain comparison.
159    pub fn fetch_or_append(&mut self, value: T, span: Span) -> Handle<T>
160    where
161        T: PartialEq,
162    {
163        self.fetch_if_or_append(value, span, T::eq)
164    }
165
166    pub fn try_get(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
167        self.data
168            .get(handle.index())
169            .ok_or_else(|| BadHandle::new(handle))
170    }
171
172    /// Get a mutable reference to an element in the arena.
173    pub fn get_mut(&mut self, handle: Handle<T>) -> &mut T {
174        self.data.get_mut(handle.index()).unwrap()
175    }
176
177    /// Get the range of handles from a particular number of elements to the end.
178    pub fn range_from(&self, old_length: usize) -> Range<T> {
179        let range = old_length as u32..self.data.len() as u32;
180        Range::from_index_range(range, self)
181    }
182
183    /// Clears the arena keeping all allocations
184    pub fn clear(&mut self) {
185        self.data.clear()
186    }
187
188    pub fn get_span(&self, handle: Handle<T>) -> Span {
189        *self
190            .span_info
191            .get(handle.index())
192            .unwrap_or(&Span::default())
193    }
194
195    /// Assert that `handle` is valid for this arena.
196    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
197        if handle.index() < self.data.len() {
198            Ok(())
199        } else {
200            Err(BadHandle::new(handle))
201        }
202    }
203
204    /// Assert that `range` is valid for this arena.
205    pub fn check_contains_range(&self, range: &Range<T>) -> Result<(), BadRangeError> {
206        // Since `range.inner` is a `Range<u32>`, we only need to check that the
207        // start precedes the end, and that the end is in range.
208        if range.inner.start > range.inner.end {
209            return Err(BadRangeError::new(range.clone()));
210        }
211
212        // Empty ranges are tolerated: they can be produced by compaction.
213        if range.inner.start == range.inner.end {
214            return Ok(());
215        }
216
217        let last_handle = Handle::new(Index::new(range.inner.end - 1).unwrap());
218        if self.check_contains_handle(last_handle).is_err() {
219            return Err(BadRangeError::new(range.clone()));
220        }
221
222        Ok(())
223    }
224
225    #[cfg(feature = "compact")]
226    pub(crate) fn retain_mut<P>(&mut self, mut predicate: P)
227    where
228        P: FnMut(Handle<T>, &mut T) -> bool,
229    {
230        let mut index = 0;
231        let mut retained = 0;
232        self.data.retain_mut(|elt| {
233            let handle = Handle::from_usize(index);
234            let keep = predicate(handle, elt);
235
236            // Since `predicate` needs mutable access to each element,
237            // we can't feasibly call it twice, so we have to compact
238            // spans by hand in parallel as part of this iteration.
239            if keep {
240                self.span_info[retained] = self.span_info[index];
241                retained += 1;
242            }
243
244            index += 1;
245            keep
246        });
247
248        self.span_info.truncate(retained);
249    }
250}
251
252#[cfg(feature = "deserialize")]
253impl<'de, T> serde::Deserialize<'de> for Arena<T>
254where
255    T: serde::Deserialize<'de>,
256{
257    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
258    where
259        D: serde::Deserializer<'de>,
260    {
261        let data = Vec::deserialize(deserializer)?;
262        let span_info = std::iter::repeat(Span::default())
263            .take(data.len())
264            .collect();
265
266        Ok(Self { data, span_info })
267    }
268}
269
270impl<T> ops::Index<Handle<T>> for Arena<T> {
271    type Output = T;
272    fn index(&self, handle: Handle<T>) -> &T {
273        &self.data[handle.index()]
274    }
275}
276
277impl<T> ops::IndexMut<Handle<T>> for Arena<T> {
278    fn index_mut(&mut self, handle: Handle<T>) -> &mut T {
279        &mut self.data[handle.index()]
280    }
281}
282
283impl<T> ops::Index<Range<T>> for Arena<T> {
284    type Output = [T];
285    fn index(&self, range: Range<T>) -> &[T] {
286        &self.data[range.inner.start as usize..range.inner.end as usize]
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    #[test]
295    fn append_non_unique() {
296        let mut arena: Arena<u8> = Arena::new();
297        let t1 = arena.append(0, Default::default());
298        let t2 = arena.append(0, Default::default());
299        assert!(t1 != t2);
300        assert!(arena[t1] == arena[t2]);
301    }
302
303    #[test]
304    fn append_unique() {
305        let mut arena: Arena<u8> = Arena::new();
306        let t1 = arena.append(0, Default::default());
307        let t2 = arena.append(1, Default::default());
308        assert!(t1 != t2);
309        assert!(arena[t1] != arena[t2]);
310    }
311
312    #[test]
313    fn fetch_or_append_non_unique() {
314        let mut arena: Arena<u8> = Arena::new();
315        let t1 = arena.fetch_or_append(0, Default::default());
316        let t2 = arena.fetch_or_append(0, Default::default());
317        assert!(t1 == t2);
318        assert!(arena[t1] == arena[t2])
319    }
320
321    #[test]
322    fn fetch_or_append_unique() {
323        let mut arena: Arena<u8> = Arena::new();
324        let t1 = arena.fetch_or_append(0, Default::default());
325        let t2 = arena.fetch_or_append(1, Default::default());
326        assert!(t1 != t2);
327        assert!(arena[t1] != arena[t2]);
328    }
329}
naga/arena/mod.rs

naga/arena/
mod.rs
