naga/
arena.rs

1use std::{cmp::Ordering, fmt, hash, marker::PhantomData, num::NonZeroU32, ops};
2
3/// An unique index in the arena array that a handle points to.
4/// The "non-zero" part ensures that an `Option<Handle<T>>` has
5/// the same size and representation as `Handle<T>`.
6type Index = NonZeroU32;
7
8use crate::Span;
9use indexmap::set::IndexSet;
10
11#[derive(Clone, Copy, Debug, thiserror::Error, PartialEq)]
12#[error("Handle {index} of {kind} is either not present, or inaccessible yet")]
13pub struct BadHandle {
14    pub kind: &'static str,
15    pub index: usize,
16}
17
18impl BadHandle {
19    fn new<T>(handle: Handle<T>) -> Self {
20        Self {
21            kind: std::any::type_name::<T>(),
22            index: handle.index(),
23        }
24    }
25}
26
27/// A strongly typed reference to an arena item.
28///
29/// A `Handle` value can be used as an index into an [`Arena`] or [`UniqueArena`].
30#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
31#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
32#[cfg_attr(
33    any(feature = "serialize", feature = "deserialize"),
34    serde(transparent)
35)]
36#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
37pub struct Handle<T> {
38    index: Index,
39    #[cfg_attr(any(feature = "serialize", feature = "deserialize"), serde(skip))]
40    marker: PhantomData<T>,
41}
42
43impl<T> Clone for Handle<T> {
44    fn clone(&self) -> Self {
45        Handle {
46            index: self.index,
47            marker: self.marker,
48        }
49    }
50}
51
52impl<T> Copy for Handle<T> {}
53
54impl<T> PartialEq for Handle<T> {
55    fn eq(&self, other: &Self) -> bool {
56        self.index == other.index
57    }
58}
59
60impl<T> Eq for Handle<T> {}
61
62impl<T> PartialOrd for Handle<T> {
63    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
64        self.index.partial_cmp(&other.index)
65    }
66}
67
68impl<T> Ord for Handle<T> {
69    fn cmp(&self, other: &Self) -> Ordering {
70        self.index.cmp(&other.index)
71    }
72}
73
74impl<T> fmt::Debug for Handle<T> {
75    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
76        write!(formatter, "[{}]", self.index)
77    }
78}
79
80impl<T> hash::Hash for Handle<T> {
81    fn hash<H: hash::Hasher>(&self, hasher: &mut H) {
82        self.index.hash(hasher)
83    }
84}
85
86impl<T> Handle<T> {
87    #[cfg(test)]
88    pub const DUMMY: Self = Handle {
89        index: unsafe { NonZeroU32::new_unchecked(u32::MAX) },
90        marker: PhantomData,
91    };
92
93    pub(crate) const fn new(index: Index) -> Self {
94        Handle {
95            index,
96            marker: PhantomData,
97        }
98    }
99
100    /// Returns the zero-based index of this handle.
101    pub const fn index(self) -> usize {
102        let index = self.index.get() - 1;
103        index as usize
104    }
105
106    /// Convert a `usize` index into a `Handle<T>`.
107    fn from_usize(index: usize) -> Self {
108        let handle_index = u32::try_from(index + 1)
109            .ok()
110            .and_then(Index::new)
111            .expect("Failed to insert into arena. Handle overflows");
112        Handle::new(handle_index)
113    }
114
115    /// Convert a `usize` index into a `Handle<T>`, without range checks.
116    const unsafe fn from_usize_unchecked(index: usize) -> Self {
117        Handle::new(Index::new_unchecked((index + 1) as u32))
118    }
119}
120
121/// A strongly typed range of handles.
122#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
123#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
124#[cfg_attr(
125    any(feature = "serialize", feature = "deserialize"),
126    serde(transparent)
127)]
128#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
129pub struct Range<T> {
130    inner: ops::Range<u32>,
131    #[cfg_attr(any(feature = "serialize", feature = "deserialize"), serde(skip))]
132    marker: PhantomData<T>,
133}
134
135impl<T> Range<T> {
136    pub(crate) const fn erase_type(self) -> Range<()> {
137        let Self { inner, marker: _ } = self;
138        Range {
139            inner,
140            marker: PhantomData,
141        }
142    }
143}
144
145// NOTE: Keep this diagnostic in sync with that of [`BadHandle`].
146#[derive(Clone, Debug, thiserror::Error)]
147#[error("Handle range {range:?} of {kind} is either not present, or inaccessible yet")]
148pub struct BadRangeError {
149    // This error is used for many `Handle` types, but there's no point in making this generic, so
150    // we just flatten them all to `Handle<()>` here.
151    kind: &'static str,
152    range: Range<()>,
153}
154
155impl BadRangeError {
156    pub fn new<T>(range: Range<T>) -> Self {
157        Self {
158            kind: std::any::type_name::<T>(),
159            range: range.erase_type(),
160        }
161    }
162}
163
164impl<T> Clone for Range<T> {
165    fn clone(&self) -> Self {
166        Range {
167            inner: self.inner.clone(),
168            marker: self.marker,
169        }
170    }
171}
172
173impl<T> fmt::Debug for Range<T> {
174    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
175        write!(formatter, "[{}..{}]", self.inner.start + 1, self.inner.end)
176    }
177}
178
179impl<T> Iterator for Range<T> {
180    type Item = Handle<T>;
181    fn next(&mut self) -> Option<Self::Item> {
182        if self.inner.start < self.inner.end {
183            self.inner.start += 1;
184            Some(Handle {
185                index: NonZeroU32::new(self.inner.start).unwrap(),
186                marker: self.marker,
187            })
188        } else {
189            None
190        }
191    }
192}
193
194impl<T> Range<T> {
195    pub fn new_from_bounds(first: Handle<T>, last: Handle<T>) -> Self {
196        Self {
197            inner: (first.index() as u32)..(last.index() as u32 + 1),
198            marker: Default::default(),
199        }
200    }
201}
202
203/// An arena holding some kind of component (e.g., type, constant,
204/// instruction, etc.) that can be referenced.
205///
206/// Adding new items to the arena produces a strongly-typed [`Handle`].
207/// The arena can be indexed using the given handle to obtain
208/// a reference to the stored item.
209#[cfg_attr(feature = "clone", derive(Clone))]
210#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
211#[cfg_attr(feature = "serialize", serde(transparent))]
212#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
213#[cfg_attr(test, derive(PartialEq))]
214pub struct Arena<T> {
215    /// Values of this arena.
216    data: Vec<T>,
217    #[cfg(feature = "span")]
218    #[cfg_attr(feature = "serialize", serde(skip))]
219    span_info: Vec<Span>,
220}
221
222impl<T> Default for Arena<T> {
223    fn default() -> Self {
224        Self::new()
225    }
226}
227
228impl<T: fmt::Debug> fmt::Debug for Arena<T> {
229    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
230        f.debug_map().entries(self.iter()).finish()
231    }
232}
233
234impl<T> Arena<T> {
235    /// Create a new arena with no initial capacity allocated.
236    pub const fn new() -> Self {
237        Arena {
238            data: Vec::new(),
239            #[cfg(feature = "span")]
240            span_info: Vec::new(),
241        }
242    }
243
244    /// Extracts the inner vector.
245    #[allow(clippy::missing_const_for_fn)] // ignore due to requirement of #![feature(const_precise_live_drops)]
246    pub fn into_inner(self) -> Vec<T> {
247        self.data
248    }
249
250    /// Returns the current number of items stored in this arena.
251    pub fn len(&self) -> usize {
252        self.data.len()
253    }
254
255    /// Returns `true` if the arena contains no elements.
256    pub fn is_empty(&self) -> bool {
257        self.data.is_empty()
258    }
259
260    /// Returns an iterator over the items stored in this arena, returning both
261    /// the item's handle and a reference to it.
262    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> {
263        self.data
264            .iter()
265            .enumerate()
266            .map(|(i, v)| unsafe { (Handle::from_usize_unchecked(i), v) })
267    }
268
269    /// Returns a iterator over the items stored in this arena,
270    /// returning both the item's handle and a mutable reference to it.
271    pub fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = (Handle<T>, &mut T)> {
272        self.data
273            .iter_mut()
274            .enumerate()
275            .map(|(i, v)| unsafe { (Handle::from_usize_unchecked(i), v) })
276    }
277
278    /// Adds a new value to the arena, returning a typed handle.
279    pub fn append(&mut self, value: T, span: Span) -> Handle<T> {
280        #[cfg(not(feature = "span"))]
281        let _ = span;
282        let index = self.data.len();
283        self.data.push(value);
284        #[cfg(feature = "span")]
285        self.span_info.push(span);
286        Handle::from_usize(index)
287    }
288
289    /// Fetch a handle to an existing type.
290    pub fn fetch_if<F: Fn(&T) -> bool>(&self, fun: F) -> Option<Handle<T>> {
291        self.data
292            .iter()
293            .position(fun)
294            .map(|index| unsafe { Handle::from_usize_unchecked(index) })
295    }
296
297    /// Adds a value with a custom check for uniqueness:
298    /// returns a handle pointing to
299    /// an existing element if the check succeeds, or adds a new
300    /// element otherwise.
301    pub fn fetch_if_or_append<F: Fn(&T, &T) -> bool>(
302        &mut self,
303        value: T,
304        span: Span,
305        fun: F,
306    ) -> Handle<T> {
307        if let Some(index) = self.data.iter().position(|d| fun(d, &value)) {
308            unsafe { Handle::from_usize_unchecked(index) }
309        } else {
310            self.append(value, span)
311        }
312    }
313
314    /// Adds a value with a check for uniqueness, where the check is plain comparison.
315    pub fn fetch_or_append(&mut self, value: T, span: Span) -> Handle<T>
316    where
317        T: PartialEq,
318    {
319        self.fetch_if_or_append(value, span, T::eq)
320    }
321
322    pub fn try_get(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
323        self.data
324            .get(handle.index())
325            .ok_or_else(|| BadHandle::new(handle))
326    }
327
328    /// Get a mutable reference to an element in the arena.
329    pub fn get_mut(&mut self, handle: Handle<T>) -> &mut T {
330        self.data.get_mut(handle.index()).unwrap()
331    }
332
333    /// Get the range of handles from a particular number of elements to the end.
334    pub fn range_from(&self, old_length: usize) -> Range<T> {
335        Range {
336            inner: old_length as u32..self.data.len() as u32,
337            marker: PhantomData,
338        }
339    }
340
341    /// Clears the arena keeping all allocations
342    pub fn clear(&mut self) {
343        self.data.clear()
344    }
345
346    pub fn get_span(&self, handle: Handle<T>) -> Span {
347        #[cfg(feature = "span")]
348        {
349            *self
350                .span_info
351                .get(handle.index())
352                .unwrap_or(&Span::default())
353        }
354        #[cfg(not(feature = "span"))]
355        {
356            let _ = handle;
357            Span::default()
358        }
359    }
360
361    /// Assert that `handle` is valid for this arena.
362    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
363        if handle.index() < self.data.len() {
364            Ok(())
365        } else {
366            Err(BadHandle::new(handle))
367        }
368    }
369
370    /// Assert that `range` is valid for this arena.
371    pub fn check_contains_range(&self, range: &Range<T>) -> Result<(), BadRangeError> {
372        // Since `range.inner` is a `Range<u32>`, we only need to
373        // check that the start precedes the end, and that the end is
374        // in range.
375        if range.inner.start > range.inner.end
376            || self
377                .check_contains_handle(Handle::new(range.inner.end.try_into().unwrap()))
378                .is_err()
379        {
380            Err(BadRangeError::new(range.clone()))
381        } else {
382            Ok(())
383        }
384    }
385}
386
387#[cfg(feature = "deserialize")]
388impl<'de, T> serde::Deserialize<'de> for Arena<T>
389where
390    T: serde::Deserialize<'de>,
391{
392    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
393    where
394        D: serde::Deserializer<'de>,
395    {
396        let data = Vec::deserialize(deserializer)?;
397        #[cfg(feature = "span")]
398        let span_info = std::iter::repeat(Span::default())
399            .take(data.len())
400            .collect();
401
402        Ok(Self {
403            data,
404            #[cfg(feature = "span")]
405            span_info,
406        })
407    }
408}
409
410impl<T> ops::Index<Handle<T>> for Arena<T> {
411    type Output = T;
412    fn index(&self, handle: Handle<T>) -> &T {
413        &self.data[handle.index()]
414    }
415}
416
417impl<T> ops::IndexMut<Handle<T>> for Arena<T> {
418    fn index_mut(&mut self, handle: Handle<T>) -> &mut T {
419        &mut self.data[handle.index()]
420    }
421}
422
423impl<T> ops::Index<Range<T>> for Arena<T> {
424    type Output = [T];
425    fn index(&self, range: Range<T>) -> &[T] {
426        &self.data[range.inner.start as usize..range.inner.end as usize]
427    }
428}
429
430#[cfg(test)]
431mod tests {
432    use super::*;
433
434    #[test]
435    fn append_non_unique() {
436        let mut arena: Arena<u8> = Arena::new();
437        let t1 = arena.append(0, Default::default());
438        let t2 = arena.append(0, Default::default());
439        assert!(t1 != t2);
440        assert!(arena[t1] == arena[t2]);
441    }
442
443    #[test]
444    fn append_unique() {
445        let mut arena: Arena<u8> = Arena::new();
446        let t1 = arena.append(0, Default::default());
447        let t2 = arena.append(1, Default::default());
448        assert!(t1 != t2);
449        assert!(arena[t1] != arena[t2]);
450    }
451
452    #[test]
453    fn fetch_or_append_non_unique() {
454        let mut arena: Arena<u8> = Arena::new();
455        let t1 = arena.fetch_or_append(0, Default::default());
456        let t2 = arena.fetch_or_append(0, Default::default());
457        assert!(t1 == t2);
458        assert!(arena[t1] == arena[t2])
459    }
460
461    #[test]
462    fn fetch_or_append_unique() {
463        let mut arena: Arena<u8> = Arena::new();
464        let t1 = arena.fetch_or_append(0, Default::default());
465        let t2 = arena.fetch_or_append(1, Default::default());
466        assert!(t1 != t2);
467        assert!(arena[t1] != arena[t2]);
468    }
469}
470
471/// An arena whose elements are guaranteed to be unique.
472///
473/// A `UniqueArena` holds a set of unique values of type `T`, each with an
474/// associated [`Span`]. Inserting a value returns a `Handle<T>`, which can be
475/// used to index the `UniqueArena` and obtain shared access to the `T` element.
476/// Access via a `Handle` is an array lookup - no hash lookup is necessary.
477///
478/// The element type must implement `Eq` and `Hash`. Insertions of equivalent
479/// elements, according to `Eq`, all return the same `Handle`.
480///
481/// Once inserted, elements may not be mutated.
482///
483/// `UniqueArena` is similar to [`Arena`]: If `Arena` is vector-like,
484/// `UniqueArena` is `HashSet`-like.
485#[cfg_attr(feature = "clone", derive(Clone))]
486pub struct UniqueArena<T> {
487    set: IndexSet<T>,
488
489    /// Spans for the elements, indexed by handle.
490    ///
491    /// The length of this vector is always equal to `set.len()`. `IndexSet`
492    /// promises that its elements "are indexed in a compact range, without
493    /// holes in the range 0..set.len()", so we can always use the indices
494    /// returned by insertion as indices into this vector.
495    #[cfg(feature = "span")]
496    span_info: Vec<Span>,
497}
498
499impl<T> UniqueArena<T> {
500    /// Create a new arena with no initial capacity allocated.
501    pub fn new() -> Self {
502        UniqueArena {
503            set: IndexSet::new(),
504            #[cfg(feature = "span")]
505            span_info: Vec::new(),
506        }
507    }
508
509    /// Return the current number of items stored in this arena.
510    pub fn len(&self) -> usize {
511        self.set.len()
512    }
513
514    /// Return `true` if the arena contains no elements.
515    pub fn is_empty(&self) -> bool {
516        self.set.is_empty()
517    }
518
519    /// Clears the arena, keeping all allocations.
520    pub fn clear(&mut self) {
521        self.set.clear();
522        #[cfg(feature = "span")]
523        self.span_info.clear();
524    }
525
526    /// Return the span associated with `handle`.
527    ///
528    /// If a value has been inserted multiple times, the span returned is the
529    /// one provided with the first insertion.
530    ///
531    /// If the `span` feature is not enabled, always return `Span::default`.
532    /// This can be detected with [`Span::is_defined`].
533    pub fn get_span(&self, handle: Handle<T>) -> Span {
534        #[cfg(feature = "span")]
535        {
536            *self
537                .span_info
538                .get(handle.index())
539                .unwrap_or(&Span::default())
540        }
541        #[cfg(not(feature = "span"))]
542        {
543            let _ = handle;
544            Span::default()
545        }
546    }
547}
548
549impl<T: Eq + hash::Hash> UniqueArena<T> {
550    /// Returns an iterator over the items stored in this arena, returning both
551    /// the item's handle and a reference to it.
552    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> {
553        self.set.iter().enumerate().map(|(i, v)| {
554            let position = i + 1;
555            let index = unsafe { Index::new_unchecked(position as u32) };
556            (Handle::new(index), v)
557        })
558    }
559
560    /// Insert a new value into the arena.
561    ///
562    /// Return a [`Handle<T>`], which can be used to index this arena to get a
563    /// shared reference to the element.
564    ///
565    /// If this arena already contains an element that is `Eq` to `value`,
566    /// return a `Handle` to the existing element, and drop `value`.
567    ///
568    /// When the `span` feature is enabled, if `value` is inserted into the
569    /// arena, associate `span` with it. An element's span can be retrieved with
570    /// the [`get_span`] method.
571    ///
572    /// [`Handle<T>`]: Handle
573    /// [`get_span`]: UniqueArena::get_span
574    pub fn insert(&mut self, value: T, span: Span) -> Handle<T> {
575        let (index, added) = self.set.insert_full(value);
576
577        #[cfg(feature = "span")]
578        {
579            if added {
580                debug_assert!(index == self.span_info.len());
581                self.span_info.push(span);
582            }
583
584            debug_assert!(self.set.len() == self.span_info.len());
585        }
586
587        #[cfg(not(feature = "span"))]
588        let _ = (span, added);
589
590        Handle::from_usize(index)
591    }
592
593    /// Replace an old value with a new value.
594    ///
595    /// # Panics
596    ///
597    /// - if the old value is not in the arena
598    /// - if the new value already exists in the arena
599    pub fn replace(&mut self, old: Handle<T>, new: T) {
600        let (index, added) = self.set.insert_full(new);
601        assert!(added && index == self.set.len() - 1);
602
603        self.set.swap_remove_index(old.index()).unwrap();
604    }
605
606    /// Return this arena's handle for `value`, if present.
607    ///
608    /// If this arena already contains an element equal to `value`,
609    /// return its handle. Otherwise, return `None`.
610    pub fn get(&self, value: &T) -> Option<Handle<T>> {
611        self.set
612            .get_index_of(value)
613            .map(|index| unsafe { Handle::from_usize_unchecked(index) })
614    }
615
616    /// Return this arena's value at `handle`, if that is a valid handle.
617    pub fn get_handle(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
618        self.set
619            .get_index(handle.index())
620            .ok_or_else(|| BadHandle::new(handle))
621    }
622
623    /// Assert that `handle` is valid for this arena.
624    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
625        if handle.index() < self.set.len() {
626            Ok(())
627        } else {
628            Err(BadHandle::new(handle))
629        }
630    }
631}
632
633impl<T> Default for UniqueArena<T> {
634    fn default() -> Self {
635        Self::new()
636    }
637}
638
639impl<T: fmt::Debug + Eq + hash::Hash> fmt::Debug for UniqueArena<T> {
640    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
641        f.debug_map().entries(self.iter()).finish()
642    }
643}
644
645impl<T> ops::Index<Handle<T>> for UniqueArena<T> {
646    type Output = T;
647    fn index(&self, handle: Handle<T>) -> &T {
648        &self.set[handle.index()]
649    }
650}
651
652#[cfg(feature = "serialize")]
653impl<T> serde::Serialize for UniqueArena<T>
654where
655    T: Eq + hash::Hash + serde::Serialize,
656{
657    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
658    where
659        S: serde::Serializer,
660    {
661        self.set.serialize(serializer)
662    }
663}
664
665#[cfg(feature = "deserialize")]
666impl<'de, T> serde::Deserialize<'de> for UniqueArena<T>
667where
668    T: Eq + hash::Hash + serde::Deserialize<'de>,
669{
670    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
671    where
672        D: serde::Deserializer<'de>,
673    {
674        let set = IndexSet::deserialize(deserializer)?;
675        #[cfg(feature = "span")]
676        let span_info = std::iter::repeat(Span::default()).take(set.len()).collect();
677
678        Ok(Self {
679            set,
680            #[cfg(feature = "span")]
681            span_info,
682        })
683    }
684}
685
686//Note: largely borrowed from `HashSet` implementation
687#[cfg(feature = "arbitrary")]
688impl<'a, T> arbitrary::Arbitrary<'a> for UniqueArena<T>
689where
690    T: Eq + hash::Hash + arbitrary::Arbitrary<'a>,
691{
692    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
693        let mut arena = Self::default();
694        for elem in u.arbitrary_iter()? {
695            arena.set.insert(elem?);
696            #[cfg(feature = "span")]
697            arena.span_info.push(Span::UNDEFINED);
698        }
699        Ok(arena)
700    }
701
702    fn arbitrary_take_rest(u: arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
703        let mut arena = Self::default();
704        for elem in u.arbitrary_take_rest_iter()? {
705            arena.set.insert(elem?);
706            #[cfg(feature = "span")]
707            arena.span_info.push(Span::UNDEFINED);
708        }
709        Ok(arena)
710    }
711
712    #[inline]
713    fn size_hint(depth: usize) -> (usize, Option<usize>) {
714        let depth_hint = <usize as arbitrary::Arbitrary>::size_hint(depth);
715        arbitrary::size_hint::and(depth_hint, (0, None))
716    }
717}
naga/arena.rs

naga/
arena.rs
