naga/arena/
unique_arena.rs

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

naga/arena/
unique_arena.rs
