naga/arena/
unique_arena.rs

1//! The [`UniqueArena`] type and supporting definitions.
2
3use alloc::vec::Vec;
4use core::{fmt, hash, ops};
5
6use super::handle::{BadHandle, Handle, Index};
7use crate::{FastIndexSet, Span};
8
9/// An arena whose elements are guaranteed to be unique.
10///
11/// A `UniqueArena` holds a set of unique values of type `T`, each with an
12/// associated [`Span`]. Inserting a value returns a `Handle<T>`, which can be
13/// used to index the `UniqueArena` and obtain shared access to the `T` element.
14/// Access via a `Handle` is an array lookup - no hash lookup is necessary.
15///
16/// The element type must implement `Eq` and `Hash`. Insertions of equivalent
17/// elements, according to `Eq`, all return the same `Handle`.
18///
19/// Once inserted, elements generally may not be mutated, although a `replace`
20/// method exists to support rare cases.
21///
22/// `UniqueArena` is similar to [`Arena`]: If `Arena` is vector-like,
23/// `UniqueArena` is `HashSet`-like.
24///
25/// [`Arena`]: super::Arena
26#[derive(Clone)]
27pub struct UniqueArena<T> {
28    set: FastIndexSet<T>,
29
30    /// Spans for the elements, indexed by handle.
31    ///
32    /// The length of this vector is always equal to `set.len()`. `FastIndexSet`
33    /// promises that its elements "are indexed in a compact range, without
34    /// holes in the range 0..set.len()", so we can always use the indices
35    /// returned by insertion as indices into this vector.
36    span_info: Vec<Span>,
37}
38
39impl<T> UniqueArena<T> {
40    /// Create a new arena with no initial capacity allocated.
41    pub fn new() -> Self {
42        UniqueArena {
43            set: FastIndexSet::default(),
44            span_info: Vec::new(),
45        }
46    }
47
48    /// Return the current number of items stored in this arena.
49    pub fn len(&self) -> usize {
50        self.set.len()
51    }
52
53    /// Return `true` if the arena contains no elements.
54    pub fn is_empty(&self) -> bool {
55        self.set.is_empty()
56    }
57
58    /// Clears the arena, keeping all allocations.
59    pub fn clear(&mut self) {
60        self.set.clear();
61        self.span_info.clear();
62    }
63
64    /// Replaces this arena with an empty one, returning the old data.
65    pub fn take(&mut self) -> Self {
66        core::mem::take(self)
67    }
68
69    /// Return the span associated with `handle`.
70    ///
71    /// If a value has been inserted multiple times, the span returned is the
72    /// one provided with the first insertion.
73    pub fn get_span(&self, handle: Handle<T>) -> Span {
74        *self
75            .span_info
76            .get(handle.index())
77            .unwrap_or(&Span::default())
78    }
79
80    pub(crate) fn drain_all(&mut self) -> UniqueArenaDrain<'_, T> {
81        UniqueArenaDrain {
82            inner_elts: self.set.drain(..),
83            inner_spans: self.span_info.drain(..),
84            index: Index::new(0).unwrap(),
85        }
86    }
87}
88
89pub struct UniqueArenaDrain<'a, T> {
90    inner_elts: indexmap::set::Drain<'a, T>,
91    inner_spans: alloc::vec::Drain<'a, Span>,
92    index: Index,
93}
94
95impl<T> Iterator for UniqueArenaDrain<'_, T> {
96    type Item = (Handle<T>, T, Span);
97
98    fn next(&mut self) -> Option<Self::Item> {
99        match self.inner_elts.next() {
100            Some(elt) => {
101                let handle = Handle::new(self.index);
102                self.index = self.index.checked_add(1).unwrap();
103                let span = self.inner_spans.next().unwrap();
104                Some((handle, elt, span))
105            }
106            None => None,
107        }
108    }
109}
110
111impl<T: Eq + hash::Hash> UniqueArena<T> {
112    /// Returns an iterator over the items stored in this arena, returning both
113    /// the item's handle and a reference to it.
114    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> + ExactSizeIterator {
115        self.set.iter().enumerate().map(|(i, v)| {
116            let index = Index::new(i as u32).unwrap();
117            (Handle::new(index), v)
118        })
119    }
120
121    /// Insert a new value into the arena.
122    ///
123    /// Return a [`Handle<T>`], which can be used to index this arena to get a
124    /// shared reference to the element.
125    ///
126    /// If this arena already contains an element that is `Eq` to `value`,
127    /// return a `Handle` to the existing element, and drop `value`.
128    ///
129    /// If `value` is inserted into the arena, associate `span` with
130    /// it. An element's span can be retrieved with the [`get_span`]
131    /// method.
132    ///
133    /// [`Handle<T>`]: Handle
134    /// [`get_span`]: UniqueArena::get_span
135    pub fn insert(&mut self, value: T, span: Span) -> Handle<T> {
136        let (index, added) = self.set.insert_full(value);
137
138        if added {
139            debug_assert!(index == self.span_info.len());
140            self.span_info.push(span);
141        }
142
143        debug_assert!(self.set.len() == self.span_info.len());
144
145        Handle::from_usize(index)
146    }
147
148    /// Replace an old value with a new value.
149    ///
150    /// # Panics
151    ///
152    /// - if the old value is not in the arena
153    /// - if the new value already exists in the arena
154    pub fn replace(&mut self, old: Handle<T>, new: T) {
155        let (index, added) = self.set.insert_full(new);
156        assert!(added && index == self.set.len() - 1);
157
158        self.set.swap_remove_index(old.index()).unwrap();
159    }
160
161    /// Return this arena's handle for `value`, if present.
162    ///
163    /// If this arena already contains an element equal to `value`,
164    /// return its handle. Otherwise, return `None`.
165    pub fn get(&self, value: &T) -> Option<Handle<T>> {
166        self.set
167            .get_index_of(value)
168            .map(|index| Handle::from_usize(index))
169    }
170
171    /// Return this arena's value at `handle`, if that is a valid handle.
172    pub fn get_handle(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
173        self.set
174            .get_index(handle.index())
175            .ok_or_else(|| BadHandle::new(handle))
176    }
177
178    /// Assert that `handle` is valid for this arena.
179    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
180        if handle.index() < self.set.len() {
181            Ok(())
182        } else {
183            Err(BadHandle::new(handle))
184        }
185    }
186}
187
188impl<T> Default for UniqueArena<T> {
189    fn default() -> Self {
190        Self::new()
191    }
192}
193
194impl<T: fmt::Debug + Eq + hash::Hash> fmt::Debug for UniqueArena<T> {
195    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
196        f.debug_map().entries(self.iter()).finish()
197    }
198}
199
200impl<T> ops::Index<Handle<T>> for UniqueArena<T> {
201    type Output = T;
202    fn index(&self, handle: Handle<T>) -> &T {
203        &self.set[handle.index()]
204    }
205}
206
207#[cfg(feature = "serialize")]
208impl<T> serde::Serialize for UniqueArena<T>
209where
210    T: Eq + hash::Hash + serde::Serialize,
211{
212    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
213    where
214        S: serde::Serializer,
215    {
216        self.set.serialize(serializer)
217    }
218}
219
220#[cfg(feature = "deserialize")]
221impl<'de, T> serde::Deserialize<'de> for UniqueArena<T>
222where
223    T: Eq + hash::Hash + serde::Deserialize<'de>,
224{
225    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
226    where
227        D: serde::Deserializer<'de>,
228    {
229        let set = FastIndexSet::deserialize(deserializer)?;
230        let span_info = core::iter::repeat_n(Span::default(), set.len()).collect();
231
232        Ok(Self { set, span_info })
233    }
234}
235
236//Note: largely borrowed from `HashSet` implementation
237#[cfg(feature = "arbitrary")]
238impl<'a, T> arbitrary::Arbitrary<'a> for UniqueArena<T>
239where
240    T: Eq + hash::Hash + arbitrary::Arbitrary<'a>,
241{
242    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
243        let mut arena = Self::default();
244        for elem in u.arbitrary_iter()? {
245            arena.set.insert(elem?);
246            arena.span_info.push(Span::UNDEFINED);
247        }
248        Ok(arena)
249    }
250
251    fn arbitrary_take_rest(u: arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
252        let mut arena = Self::default();
253        for elem in u.arbitrary_take_rest_iter()? {
254            arena.set.insert(elem?);
255            arena.span_info.push(Span::UNDEFINED);
256        }
257        Ok(arena)
258    }
259
260    #[inline]
261    fn size_hint(depth: usize) -> (usize, Option<usize>) {
262        let depth_hint = <usize as arbitrary::Arbitrary>::size_hint(depth);
263        arbitrary::size_hint::and(depth_hint, (0, None))
264    }
265}
naga/arena/unique_arena.rs

naga/arena/
unique_arena.rs
