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    /// Return the span associated with `handle`.
65    ///
66    /// If a value has been inserted multiple times, the span returned is the
67    /// one provided with the first insertion.
68    pub fn get_span(&self, handle: Handle<T>) -> Span {
69        *self
70            .span_info
71            .get(handle.index())
72            .unwrap_or(&Span::default())
73    }
74
75    pub(crate) fn drain_all(&mut self) -> UniqueArenaDrain<'_, T> {
76        UniqueArenaDrain {
77            inner_elts: self.set.drain(..),
78            inner_spans: self.span_info.drain(..),
79            index: Index::new(0).unwrap(),
80        }
81    }
82}
83
84pub struct UniqueArenaDrain<'a, T> {
85    inner_elts: indexmap::set::Drain<'a, T>,
86    inner_spans: alloc::vec::Drain<'a, Span>,
87    index: Index,
88}
89
90impl<T> Iterator for UniqueArenaDrain<'_, T> {
91    type Item = (Handle<T>, T, Span);
92
93    fn next(&mut self) -> Option<Self::Item> {
94        match self.inner_elts.next() {
95            Some(elt) => {
96                let handle = Handle::new(self.index);
97                self.index = self.index.checked_add(1).unwrap();
98                let span = self.inner_spans.next().unwrap();
99                Some((handle, elt, span))
100            }
101            None => None,
102        }
103    }
104}
105
106impl<T: Eq + hash::Hash> UniqueArena<T> {
107    /// Returns an iterator over the items stored in this arena, returning both
108    /// the item's handle and a reference to it.
109    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> + ExactSizeIterator {
110        self.set.iter().enumerate().map(|(i, v)| {
111            let index = unsafe { Index::new_unchecked(i as u32) };
112            (Handle::new(index), v)
113        })
114    }
115
116    /// Insert a new value into the arena.
117    ///
118    /// Return a [`Handle<T>`], which can be used to index this arena to get a
119    /// shared reference to the element.
120    ///
121    /// If this arena already contains an element that is `Eq` to `value`,
122    /// return a `Handle` to the existing element, and drop `value`.
123    ///
124    /// If `value` is inserted into the arena, associate `span` with
125    /// it. An element's span can be retrieved with the [`get_span`]
126    /// method.
127    ///
128    /// [`Handle<T>`]: Handle
129    /// [`get_span`]: UniqueArena::get_span
130    pub fn insert(&mut self, value: T, span: Span) -> Handle<T> {
131        let (index, added) = self.set.insert_full(value);
132
133        if added {
134            debug_assert!(index == self.span_info.len());
135            self.span_info.push(span);
136        }
137
138        debug_assert!(self.set.len() == self.span_info.len());
139
140        Handle::from_usize(index)
141    }
142
143    /// Replace an old value with a new value.
144    ///
145    /// # Panics
146    ///
147    /// - if the old value is not in the arena
148    /// - if the new value already exists in the arena
149    pub fn replace(&mut self, old: Handle<T>, new: T) {
150        let (index, added) = self.set.insert_full(new);
151        assert!(added && index == self.set.len() - 1);
152
153        self.set.swap_remove_index(old.index()).unwrap();
154    }
155
156    /// Return this arena's handle for `value`, if present.
157    ///
158    /// If this arena already contains an element equal to `value`,
159    /// return its handle. Otherwise, return `None`.
160    pub fn get(&self, value: &T) -> Option<Handle<T>> {
161        self.set
162            .get_index_of(value)
163            .map(|index| unsafe { Handle::from_usize_unchecked(index) })
164    }
165
166    /// Return this arena's value at `handle`, if that is a valid handle.
167    pub fn get_handle(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
168        self.set
169            .get_index(handle.index())
170            .ok_or_else(|| BadHandle::new(handle))
171    }
172
173    /// Assert that `handle` is valid for this arena.
174    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
175        if handle.index() < self.set.len() {
176            Ok(())
177        } else {
178            Err(BadHandle::new(handle))
179        }
180    }
181}
182
183impl<T> Default for UniqueArena<T> {
184    fn default() -> Self {
185        Self::new()
186    }
187}
188
189impl<T: fmt::Debug + Eq + hash::Hash> fmt::Debug for UniqueArena<T> {
190    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
191        f.debug_map().entries(self.iter()).finish()
192    }
193}
194
195impl<T> ops::Index<Handle<T>> for UniqueArena<T> {
196    type Output = T;
197    fn index(&self, handle: Handle<T>) -> &T {
198        &self.set[handle.index()]
199    }
200}
201
202#[cfg(feature = "serialize")]
203impl<T> serde::Serialize for UniqueArena<T>
204where
205    T: Eq + hash::Hash + serde::Serialize,
206{
207    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
208    where
209        S: serde::Serializer,
210    {
211        self.set.serialize(serializer)
212    }
213}
214
215#[cfg(feature = "deserialize")]
216impl<'de, T> serde::Deserialize<'de> for UniqueArena<T>
217where
218    T: Eq + hash::Hash + serde::Deserialize<'de>,
219{
220    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
221    where
222        D: serde::Deserializer<'de>,
223    {
224        let set = FastIndexSet::deserialize(deserializer)?;
225        let span_info = core::iter::repeat_n(Span::default(), set.len()).collect();
226
227        Ok(Self { set, span_info })
228    }
229}
230
231//Note: largely borrowed from `HashSet` implementation
232#[cfg(feature = "arbitrary")]
233impl<'a, T> arbitrary::Arbitrary<'a> for UniqueArena<T>
234where
235    T: Eq + hash::Hash + arbitrary::Arbitrary<'a>,
236{
237    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
238        let mut arena = Self::default();
239        for elem in u.arbitrary_iter()? {
240            arena.set.insert(elem?);
241            arena.span_info.push(Span::UNDEFINED);
242        }
243        Ok(arena)
244    }
245
246    fn arbitrary_take_rest(u: arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
247        let mut arena = Self::default();
248        for elem in u.arbitrary_take_rest_iter()? {
249            arena.set.insert(elem?);
250            arena.span_info.push(Span::UNDEFINED);
251        }
252        Ok(arena)
253    }
254
255    #[inline]
256    fn size_hint(depth: usize) -> (usize, Option<usize>) {
257        let depth_hint = <usize as arbitrary::Arbitrary>::size_hint(depth);
258        arbitrary::size_hint::and(depth_hint, (0, None))
259    }
260}
naga/arena/unique_arena.rs

naga/arena/
unique_arena.rs
