use std::{cmp::Ordering, fmt, hash, marker::PhantomData, num::NonZeroU32, ops}; /// An unique index in the arena array that a handle points to. /// The "non-zero" part ensures that an `Option>` has /// the same size and representation as `Handle`. type Index = NonZeroU32; use crate::{FastIndexSet, Span}; #[derive(Clone, Copy, Debug, thiserror::Error, PartialEq)] #[error("Handle {index} of {kind} is either not present, or inaccessible yet")] pub struct BadHandle { pub kind: &'static str, pub index: usize, } impl BadHandle { fn new(handle: Handle) -> Self { Self { kind: std::any::type_name::(), index: handle.index(), } } } /// A strongly typed reference to an arena item. /// /// A `Handle` value can be used as an index into an [`Arena`] or [`UniqueArena`]. #[cfg_attr(feature = "serialize", derive(serde::Serialize))] #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))] #[cfg_attr( any(feature = "serialize", feature = "deserialize"), serde(transparent) )] #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))] pub struct Handle { index: Index, #[cfg_attr(any(feature = "serialize", feature = "deserialize"), serde(skip))] marker: PhantomData, } impl Clone for Handle { fn clone(&self) -> Self { *self } } impl Copy for Handle {} impl PartialEq for Handle { fn eq(&self, other: &Self) -> bool { self.index == other.index } } impl Eq for Handle {} impl PartialOrd for Handle { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) } } impl Ord for Handle { fn cmp(&self, other: &Self) -> Ordering { self.index.cmp(&other.index) } } impl fmt::Debug for Handle { fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { write!(formatter, "[{}]", self.index) } } impl hash::Hash for Handle { fn hash(&self, hasher: &mut H) { self.index.hash(hasher) } } impl Handle { #[cfg(test)] pub const DUMMY: Self = Handle { index: unsafe { NonZeroU32::new_unchecked(u32::MAX) }, marker: PhantomData, }; pub(crate) const fn new(index: Index) -> Self { Handle { index, marker: PhantomData, } } /// Returns the zero-based index of this handle. pub const fn index(self) -> usize { let index = self.index.get() - 1; index as usize } /// Convert a `usize` index into a `Handle`. fn from_usize(index: usize) -> Self { let handle_index = u32::try_from(index + 1) .ok() .and_then(Index::new) .expect("Failed to insert into arena. Handle overflows"); Handle::new(handle_index) } /// Convert a `usize` index into a `Handle`, without range checks. const unsafe fn from_usize_unchecked(index: usize) -> Self { Handle::new(Index::new_unchecked((index + 1) as u32)) } } /// A strongly typed range of handles. #[cfg_attr(feature = "serialize", derive(serde::Serialize))] #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))] #[cfg_attr( any(feature = "serialize", feature = "deserialize"), serde(transparent) )] #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))] #[cfg_attr(test, derive(PartialEq))] pub struct Range { inner: ops::Range, #[cfg_attr(any(feature = "serialize", feature = "deserialize"), serde(skip))] marker: PhantomData, } impl Range { pub(crate) const fn erase_type(self) -> Range<()> { let Self { inner, marker: _ } = self; Range { inner, marker: PhantomData, } } } // NOTE: Keep this diagnostic in sync with that of [`BadHandle`]. #[derive(Clone, Debug, thiserror::Error)] #[cfg_attr(test, derive(PartialEq))] #[error("Handle range {range:?} of {kind} is either not present, or inaccessible yet")] pub struct BadRangeError { // This error is used for many `Handle` types, but there's no point in making this generic, so // we just flatten them all to `Handle<()>` here. kind: &'static str, range: Range<()>, } impl BadRangeError { pub fn new(range: Range) -> Self { Self { kind: std::any::type_name::(), range: range.erase_type(), } } } impl Clone for Range { fn clone(&self) -> Self { Range { inner: self.inner.clone(), marker: self.marker, } } } impl fmt::Debug for Range { fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { write!(formatter, "[{}..{}]", self.inner.start + 1, self.inner.end) } } impl Iterator for Range { type Item = Handle; fn next(&mut self) -> Option { if self.inner.start < self.inner.end { self.inner.start += 1; Some(Handle { index: NonZeroU32::new(self.inner.start).unwrap(), marker: self.marker, }) } else { None } } } impl Range { /// Return a range enclosing handles `first` through `last`, inclusive. pub fn new_from_bounds(first: Handle, last: Handle) -> Self { Self { inner: (first.index() as u32)..(last.index() as u32 + 1), marker: Default::default(), } } /// return the first and last handles included in `self`. /// /// If `self` is an empty range, there are no handles included, so /// return `None`. pub fn first_and_last(&self) -> Option<(Handle, Handle)> { if self.inner.start < self.inner.end { Some(( // `Range::new_from_bounds` expects a 1-based, start- and // end-inclusive range, but `self.inner` is a zero-based, // end-exclusive range. Handle::new(Index::new(self.inner.start + 1).unwrap()), Handle::new(Index::new(self.inner.end).unwrap()), )) } else { None } } /// Return the zero-based index range covered by `self`. pub fn zero_based_index_range(&self) -> ops::Range { self.inner.clone() } /// Construct a `Range` that covers the zero-based indices in `inner`. pub fn from_zero_based_index_range(inner: ops::Range, arena: &Arena) -> Self { // Since `inner` is a `Range`, we only need to check that // the start and end are well-ordered, and that the end fits // within `arena`. assert!(inner.start <= inner.end); assert!(inner.end as usize <= arena.len()); Self { inner, marker: Default::default(), } } } /// An arena holding some kind of component (e.g., type, constant, /// instruction, etc.) that can be referenced. /// /// Adding new items to the arena produces a strongly-typed [`Handle`]. /// The arena can be indexed using the given handle to obtain /// a reference to the stored item. #[derive(Clone)] #[cfg_attr(feature = "serialize", derive(serde::Serialize))] #[cfg_attr(feature = "serialize", serde(transparent))] #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))] #[cfg_attr(test, derive(PartialEq))] pub struct Arena { /// Values of this arena. data: Vec, #[cfg_attr(feature = "serialize", serde(skip))] span_info: Vec, } impl Default for Arena { fn default() -> Self { Self::new() } } impl fmt::Debug for Arena { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_map().entries(self.iter()).finish() } } impl Arena { /// Create a new arena with no initial capacity allocated. pub const fn new() -> Self { Arena { data: Vec::new(), span_info: Vec::new(), } } /// Extracts the inner vector. #[allow(clippy::missing_const_for_fn)] // ignore due to requirement of #![feature(const_precise_live_drops)] pub fn into_inner(self) -> Vec { self.data } /// Returns the current number of items stored in this arena. pub fn len(&self) -> usize { self.data.len() } /// Returns `true` if the arena contains no elements. pub fn is_empty(&self) -> bool { self.data.is_empty() } /// Returns an iterator over the items stored in this arena, returning both /// the item's handle and a reference to it. pub fn iter(&self) -> impl DoubleEndedIterator, &T)> { self.data .iter() .enumerate() .map(|(i, v)| unsafe { (Handle::from_usize_unchecked(i), v) }) } /// Drains the arena, returning an iterator over the items stored. pub fn drain(&mut self) -> impl DoubleEndedIterator, T, Span)> { let arena = std::mem::take(self); arena .data .into_iter() .zip(arena.span_info) .enumerate() .map(|(i, (v, span))| unsafe { (Handle::from_usize_unchecked(i), v, span) }) } /// Returns a iterator over the items stored in this arena, /// returning both the item's handle and a mutable reference to it. pub fn iter_mut(&mut self) -> impl DoubleEndedIterator, &mut T)> { self.data .iter_mut() .enumerate() .map(|(i, v)| unsafe { (Handle::from_usize_unchecked(i), v) }) } /// Adds a new value to the arena, returning a typed handle. pub fn append(&mut self, value: T, span: Span) -> Handle { let index = self.data.len(); self.data.push(value); self.span_info.push(span); Handle::from_usize(index) } /// Fetch a handle to an existing type. pub fn fetch_if bool>(&self, fun: F) -> Option> { self.data .iter() .position(fun) .map(|index| unsafe { Handle::from_usize_unchecked(index) }) } /// Adds a value with a custom check for uniqueness: /// returns a handle pointing to /// an existing element if the check succeeds, or adds a new /// element otherwise. pub fn fetch_if_or_append bool>( &mut self, value: T, span: Span, fun: F, ) -> Handle { if let Some(index) = self.data.iter().position(|d| fun(d, &value)) { unsafe { Handle::from_usize_unchecked(index) } } else { self.append(value, span) } } /// Adds a value with a check for uniqueness, where the check is plain comparison. pub fn fetch_or_append(&mut self, value: T, span: Span) -> Handle where T: PartialEq, { self.fetch_if_or_append(value, span, T::eq) } pub fn try_get(&self, handle: Handle) -> Result<&T, BadHandle> { self.data .get(handle.index()) .ok_or_else(|| BadHandle::new(handle)) } /// Get a mutable reference to an element in the arena. pub fn get_mut(&mut self, handle: Handle) -> &mut T { self.data.get_mut(handle.index()).unwrap() } /// Get the range of handles from a particular number of elements to the end. pub fn range_from(&self, old_length: usize) -> Range { Range { inner: old_length as u32..self.data.len() as u32, marker: PhantomData, } } /// Clears the arena keeping all allocations pub fn clear(&mut self) { self.data.clear() } pub fn get_span(&self, handle: Handle) -> Span { *self .span_info .get(handle.index()) .unwrap_or(&Span::default()) } /// Assert that `handle` is valid for this arena. pub fn check_contains_handle(&self, handle: Handle) -> Result<(), BadHandle> { if handle.index() < self.data.len() { Ok(()) } else { Err(BadHandle::new(handle)) } } /// Assert that `range` is valid for this arena. pub fn check_contains_range(&self, range: &Range) -> Result<(), BadRangeError> { // Since `range.inner` is a `Range`, we only need to check that the // start precedes the end, and that the end is in range. if range.inner.start > range.inner.end { return Err(BadRangeError::new(range.clone())); } // Empty ranges are tolerated: they can be produced by compaction. if range.inner.start == range.inner.end { return Ok(()); } // `range.inner` is zero-based, but end-exclusive, so `range.inner.end` // is actually the right one-based index for the last handle within the // range. let last_handle = Handle::new(range.inner.end.try_into().unwrap()); if self.check_contains_handle(last_handle).is_err() { return Err(BadRangeError::new(range.clone())); } Ok(()) } #[cfg(feature = "compact")] pub(crate) fn retain_mut

(&mut self, mut predicate: P) where P: FnMut(Handle, &mut T) -> bool, { let mut index = 0; let mut retained = 0; self.data.retain_mut(|elt| { let handle = Handle::new(Index::new(index as u32 + 1).unwrap()); let keep = predicate(handle, elt); // Since `predicate` needs mutable access to each element, // we can't feasibly call it twice, so we have to compact // spans by hand in parallel as part of this iteration. if keep { self.span_info[retained] = self.span_info[index]; retained += 1; } index += 1; keep }); self.span_info.truncate(retained); } } #[cfg(feature = "deserialize")] impl<'de, T> serde::Deserialize<'de> for Arena where T: serde::Deserialize<'de>, { fn deserialize(deserializer: D) -> Result where D: serde::Deserializer<'de>, { let data = Vec::deserialize(deserializer)?; let span_info = std::iter::repeat(Span::default()) .take(data.len()) .collect(); Ok(Self { data, span_info }) } } impl ops::Index> for Arena { type Output = T; fn index(&self, handle: Handle) -> &T { &self.data[handle.index()] } } impl ops::IndexMut> for Arena { fn index_mut(&mut self, handle: Handle) -> &mut T { &mut self.data[handle.index()] } } impl ops::Index> for Arena { type Output = [T]; fn index(&self, range: Range) -> &[T] { &self.data[range.inner.start as usize..range.inner.end as usize] } } #[cfg(test)] mod tests { use super::*; #[test] fn append_non_unique() { let mut arena: Arena = Arena::new(); let t1 = arena.append(0, Default::default()); let t2 = arena.append(0, Default::default()); assert!(t1 != t2); assert!(arena[t1] == arena[t2]); } #[test] fn append_unique() { let mut arena: Arena = Arena::new(); let t1 = arena.append(0, Default::default()); let t2 = arena.append(1, Default::default()); assert!(t1 != t2); assert!(arena[t1] != arena[t2]); } #[test] fn fetch_or_append_non_unique() { let mut arena: Arena = Arena::new(); let t1 = arena.fetch_or_append(0, Default::default()); let t2 = arena.fetch_or_append(0, Default::default()); assert!(t1 == t2); assert!(arena[t1] == arena[t2]) } #[test] fn fetch_or_append_unique() { let mut arena: Arena = Arena::new(); let t1 = arena.fetch_or_append(0, Default::default()); let t2 = arena.fetch_or_append(1, Default::default()); assert!(t1 != t2); assert!(arena[t1] != arena[t2]); } } /// An arena whose elements are guaranteed to be unique. /// /// A `UniqueArena` holds a set of unique values of type `T`, each with an /// associated [`Span`]. Inserting a value returns a `Handle`, which can be /// used to index the `UniqueArena` and obtain shared access to the `T` element. /// Access via a `Handle` is an array lookup - no hash lookup is necessary. /// /// The element type must implement `Eq` and `Hash`. Insertions of equivalent /// elements, according to `Eq`, all return the same `Handle`. /// /// Once inserted, elements may not be mutated. /// /// `UniqueArena` is similar to [`Arena`]: If `Arena` is vector-like, /// `UniqueArena` is `HashSet`-like. #[derive(Clone)] pub struct UniqueArena { set: FastIndexSet, /// Spans for the elements, indexed by handle. /// /// The length of this vector is always equal to `set.len()`. `FastIndexSet` /// promises that its elements "are indexed in a compact range, without /// holes in the range 0..set.len()", so we can always use the indices /// returned by insertion as indices into this vector. span_info: Vec, } impl UniqueArena { /// Create a new arena with no initial capacity allocated. pub fn new() -> Self { UniqueArena { set: FastIndexSet::default(), span_info: Vec::new(), } } /// Return the current number of items stored in this arena. pub fn len(&self) -> usize { self.set.len() } /// Return `true` if the arena contains no elements. pub fn is_empty(&self) -> bool { self.set.is_empty() } /// Clears the arena, keeping all allocations. pub fn clear(&mut self) { self.set.clear(); self.span_info.clear(); } /// Return the span associated with `handle`. /// /// If a value has been inserted multiple times, the span returned is the /// one provided with the first insertion. pub fn get_span(&self, handle: Handle) -> Span { *self .span_info .get(handle.index()) .unwrap_or(&Span::default()) } #[cfg(feature = "compact")] pub(crate) fn drain_all(&mut self) -> UniqueArenaDrain { UniqueArenaDrain { inner_elts: self.set.drain(..), inner_spans: self.span_info.drain(..), index: Index::new(1).unwrap(), } } } #[cfg(feature = "compact")] pub(crate) struct UniqueArenaDrain<'a, T> { inner_elts: indexmap::set::Drain<'a, T>, inner_spans: std::vec::Drain<'a, Span>, index: Index, } #[cfg(feature = "compact")] impl<'a, T> Iterator for UniqueArenaDrain<'a, T> { type Item = (Handle, T, Span); fn next(&mut self) -> Option { match self.inner_elts.next() { Some(elt) => { let handle = Handle::new(self.index); self.index = self.index.checked_add(1).unwrap(); let span = self.inner_spans.next().unwrap(); Some((handle, elt, span)) } None => None, } } } impl UniqueArena { /// Returns an iterator over the items stored in this arena, returning both /// the item's handle and a reference to it. pub fn iter(&self) -> impl DoubleEndedIterator, &T)> { self.set.iter().enumerate().map(|(i, v)| { let position = i + 1; let index = unsafe { Index::new_unchecked(position as u32) }; (Handle::new(index), v) }) } /// Insert a new value into the arena. /// /// Return a [`Handle`], which can be used to index this arena to get a /// shared reference to the element. /// /// If this arena already contains an element that is `Eq` to `value`, /// return a `Handle` to the existing element, and drop `value`. /// /// If `value` is inserted into the arena, associate `span` with /// it. An element's span can be retrieved with the [`get_span`] /// method. /// /// [`Handle`]: Handle /// [`get_span`]: UniqueArena::get_span pub fn insert(&mut self, value: T, span: Span) -> Handle { let (index, added) = self.set.insert_full(value); if added { debug_assert!(index == self.span_info.len()); self.span_info.push(span); } debug_assert!(self.set.len() == self.span_info.len()); Handle::from_usize(index) } /// Replace an old value with a new value. /// /// # Panics /// /// - if the old value is not in the arena /// - if the new value already exists in the arena pub fn replace(&mut self, old: Handle, new: T) { let (index, added) = self.set.insert_full(new); assert!(added && index == self.set.len() - 1); self.set.swap_remove_index(old.index()).unwrap(); } /// Return this arena's handle for `value`, if present. /// /// If this arena already contains an element equal to `value`, /// return its handle. Otherwise, return `None`. pub fn get(&self, value: &T) -> Option> { self.set .get_index_of(value) .map(|index| unsafe { Handle::from_usize_unchecked(index) }) } /// Return this arena's value at `handle`, if that is a valid handle. pub fn get_handle(&self, handle: Handle) -> Result<&T, BadHandle> { self.set .get_index(handle.index()) .ok_or_else(|| BadHandle::new(handle)) } /// Assert that `handle` is valid for this arena. pub fn check_contains_handle(&self, handle: Handle) -> Result<(), BadHandle> { if handle.index() < self.set.len() { Ok(()) } else { Err(BadHandle::new(handle)) } } } impl Default for UniqueArena { fn default() -> Self { Self::new() } } impl fmt::Debug for UniqueArena { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_map().entries(self.iter()).finish() } } impl ops::Index> for UniqueArena { type Output = T; fn index(&self, handle: Handle) -> &T { &self.set[handle.index()] } } #[cfg(feature = "serialize")] impl serde::Serialize for UniqueArena where T: Eq + hash::Hash + serde::Serialize, { fn serialize(&self, serializer: S) -> Result where S: serde::Serializer, { self.set.serialize(serializer) } } #[cfg(feature = "deserialize")] impl<'de, T> serde::Deserialize<'de> for UniqueArena where T: Eq + hash::Hash + serde::Deserialize<'de>, { fn deserialize(deserializer: D) -> Result where D: serde::Deserializer<'de>, { let set = FastIndexSet::deserialize(deserializer)?; let span_info = std::iter::repeat(Span::default()).take(set.len()).collect(); Ok(Self { set, span_info }) } } //Note: largely borrowed from `HashSet` implementation #[cfg(feature = "arbitrary")] impl<'a, T> arbitrary::Arbitrary<'a> for UniqueArena where T: Eq + hash::Hash + arbitrary::Arbitrary<'a>, { fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result { let mut arena = Self::default(); for elem in u.arbitrary_iter()? { arena.set.insert(elem?); arena.span_info.push(Span::UNDEFINED); } Ok(arena) } fn arbitrary_take_rest(u: arbitrary::Unstructured<'a>) -> arbitrary::Result { let mut arena = Self::default(); for elem in u.arbitrary_take_rest_iter()? { arena.set.insert(elem?); arena.span_info.push(Span::UNDEFINED); } Ok(arena) } #[inline] fn size_hint(depth: usize) -> (usize, Option) { let depth_hint = ::size_hint(depth); arbitrary::size_hint::and(depth_hint, (0, None)) } }