indexed_arena/
lib.rs

1#![doc = include_str!("../README.md")]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3#![no_std]
4
5extern crate alloc;
6
7use alloc::vec::Vec;
8use core::{
9    fmt,
10    hash::{Hash, Hasher},
11    iter,
12    marker::PhantomData,
13    num::NonZero,
14    ops::{Index, IndexMut, Range},
15    slice,
16};
17
18/// A trait for index types used in arenas.
19///
20/// An [`Id`] represents both the internal index in an arena and a type-level distinction
21/// (for example, when using multiple arenas with the same underlying numeric index type).
22pub trait Id: Copy + Ord {
23    /// The maximum value (as a usize) this id type can represent.
24    const MAX: usize;
25
26    /// Converts a `usize` value to this id type.
27    ///
28    /// The input `idx` (should / is guaranteed to) be less than `Self::MAX`.
29    fn from_usize(idx: usize) -> Self;
30
31    /// Converts this id type into a `usize`.
32    ///
33    /// The returned value (should / is guaranteed to) be less than `Self::MAX`.
34    fn into_usize(self) -> usize;
35}
36
37macro_rules! impl_id_for_nums {
38    ($($ty:ty),*) => {$(
39        impl Id for $ty {
40            const MAX: usize = <$ty>::MAX as usize;
41            #[inline]
42            fn from_usize(idx: usize) -> Self {
43                idx as $ty
44            }
45            #[inline]
46            fn into_usize(self) -> usize {
47                self as usize
48            }
49        }
50        impl Id for NonZero<$ty> {
51            const MAX: usize = (<$ty>::MAX - 1) as usize;
52            #[inline]
53            fn from_usize(idx: usize) -> Self {
54                unsafe { NonZero::new_unchecked((idx + 1) as $ty) }
55            }
56            #[inline]
57            fn into_usize(self) -> usize {
58                (self.get() - 1) as usize
59            }
60        }
61    )*};
62}
63impl_id_for_nums!(u8, u16, u32, u64, usize);
64
65/// A typed index for referencing elements in an [`Arena`].
66///
67/// The [`Idx<T, I>`] type wraps an underlying id of type `I` and carries a phantom type `T`
68/// to ensure type safety when indexing into an arena.
69///
70/// # Examples
71///
72/// ```
73/// use indexed_arena::{Arena, Idx};
74///
75/// let mut arena: Arena<&str, u32> = Arena::new();
76/// let idx: Idx<&str, u32> = arena.alloc("hello");
77/// assert_eq!(arena[idx], "hello");
78/// ```
79pub struct Idx<T, I: Id> {
80    raw: I,
81    phantom: PhantomData<fn() -> T>,
82}
83
84impl<T, I: Id> Idx<T, I> {
85    /// Consumes the index and returns its underlying raw value.
86    ///
87    /// # Examples
88    ///
89    /// ```
90    /// use indexed_arena::{Arena, Idx};
91    ///
92    /// let mut arena: Arena<i32, u32> = Arena::new();
93    /// let idx = arena.alloc(10);
94    /// let raw = idx.into_raw();
95    /// // raw is a u32 representing the index inside the arena.
96    /// assert_eq!(raw, 0u32);
97    /// ```
98    #[inline]
99    pub const fn into_raw(self) -> I {
100        self.raw
101    }
102}
103
104impl<T, I: Id> Clone for Idx<T, I> {
105    #[inline]
106    fn clone(&self) -> Self {
107        *self
108    }
109}
110impl<T, I: Id> Copy for Idx<T, I> {}
111
112impl<T, I: Id + fmt::Debug> fmt::Debug for Idx<T, I> {
113    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
114        let mut type_name = core::any::type_name::<T>();
115        if let Some(idx) = type_name.rfind(':') {
116            type_name = &type_name[idx + 1..]
117        }
118        write!(fmt, "Idx::<{}>({:?})", type_name, self.raw)
119    }
120}
121
122impl<T, I: Id> PartialEq for Idx<T, I> {
123    #[inline]
124    fn eq(&self, other: &Self) -> bool {
125        self.raw == other.raw
126    }
127}
128impl<T, I: Id> Eq for Idx<T, I> {}
129
130impl<T, I: Id> Ord for Idx<T, I> {
131    #[inline]
132    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
133        self.raw.cmp(&other.raw)
134    }
135}
136
137impl<T, I: Id> PartialOrd for Idx<T, I> {
138    #[inline]
139    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
140        Some(self.cmp(other))
141    }
142}
143
144impl<T, I: Id + Hash> Hash for Idx<T, I> {
145    #[inline]
146    fn hash<H: Hasher>(&self, state: &mut H) {
147        self.raw.hash(state)
148    }
149}
150
151/// A range of indices within an `Arena`.
152///
153/// This type represents a contiguous range of allocated indices in an arena.
154///
155/// # Examples
156///
157/// ```
158/// use indexed_arena::{Arena, IdxRange};
159///
160/// let mut arena: Arena<i32, u32> = Arena::new();
161/// let range: IdxRange<i32, u32> = arena.alloc_many(vec![1, 2, 3, 4]);
162/// assert_eq!(range.len(), 4);
163/// assert!(!range.is_empty());
164/// ```
165pub struct IdxRange<T, I: Id> {
166    start: I,
167    end: I,
168    phantom: PhantomData<fn() -> T>,
169}
170
171impl<T, I: Id> IdxRange<T, I> {
172    /// Creates a new [`IdxRange`] from the given range of raw indices.
173    #[inline]
174    pub const fn new(range: Range<I>) -> Self {
175        Self { start: range.start, end: range.end, phantom: PhantomData }
176    }
177
178    /// Returns the starting raw index.
179    #[inline]
180    pub const fn start(&self) -> I {
181        self.start
182    }
183
184    /// Returns the ending raw index.
185    #[inline]
186    pub const fn end(&self) -> I {
187        self.end
188    }
189
190    /// Returns the number of indices in the range.
191    #[inline]
192    pub fn len(&self) -> usize {
193        self.end.into_usize() - self.start.into_usize()
194    }
195
196    /// Returns true if the range is empty.
197    #[inline]
198    pub fn is_empty(&self) -> bool {
199        self.start == self.end
200    }
201}
202
203impl<T, I: Id> Clone for IdxRange<T, I> {
204    #[inline]
205    fn clone(&self) -> Self {
206        *self
207    }
208}
209impl<T, I: Id> Copy for IdxRange<T, I> {}
210
211impl<T, I: Id + fmt::Debug> fmt::Debug for IdxRange<T, I> {
212    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
213        let mut type_name = core::any::type_name::<T>();
214        if let Some(idx) = type_name.rfind(':') {
215            type_name = &type_name[idx + 1..]
216        }
217        write!(fmt, "IdxRange::<{}>({:?}..{:?})", type_name, self.start, self.end)
218    }
219}
220
221impl<T, I: Id> PartialEq for IdxRange<T, I> {
222    #[inline]
223    fn eq(&self, other: &Self) -> bool {
224        self.start == other.start && self.end == other.end
225    }
226}
227impl<T, I: Id> Eq for IdxRange<T, I> {}
228
229impl<T, I: Id> Ord for IdxRange<T, I> {
230    #[inline]
231    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
232        self.start.cmp(&other.start).then_with(|| self.end.cmp(&other.end))
233    }
234}
235
236impl<T, I: Id> PartialOrd for IdxRange<T, I> {
237    #[inline]
238    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
239        Some(self.cmp(other))
240    }
241}
242
243impl<T, I: Id + Hash> Hash for IdxRange<T, I> {
244    #[inline]
245    fn hash<H: Hasher>(&self, state: &mut H) {
246        self.start.hash(state);
247        self.end.hash(state);
248    }
249}
250
251/// A index-based arena.
252///
253/// [`Arena`] provides a mechanism to allocate objects and refer to them by a
254/// strongly-typed index ([`Idx<T, I>`]). The index not only represents the position
255/// in the underlying vector but also leverages the type system to prevent accidental misuse
256/// across different arenas.
257pub struct Arena<T, I: Id> {
258    data: Vec<T>,
259    phantom: PhantomData<(I, T)>,
260}
261
262impl<T, I: Id> Arena<T, I> {
263    /// Creates a new empty arena.
264    ///
265    /// # Examples
266    ///
267    /// ```
268    /// # use indexed_arena::Arena;
269    /// let arena: Arena<i32, u32> = Arena::new();
270    /// assert!(arena.is_empty());
271    /// ```
272    #[inline]
273    pub const fn new() -> Self {
274        Self { data: Vec::new(), phantom: PhantomData }
275    }
276
277    /// Creates a new arena with the specified capacity.
278    ///
279    /// # Examples
280    ///
281    /// ```
282    /// # use indexed_arena::Arena;
283    /// let arena: Arena<i32, u32> = Arena::with_capacity(10);
284    /// assert!(arena.is_empty());
285    /// assert!(arena.capacity() >= 10);
286    /// ```
287    #[inline]
288    pub fn with_capacity(capacity: usize) -> Self {
289        Self { data: Vec::with_capacity(capacity), phantom: PhantomData }
290    }
291
292    /// Returns the number of elements stored in the arena.
293    ///
294    /// # Examples
295    ///
296    /// ```
297    /// # use indexed_arena::Arena;
298    /// let mut arena = Arena::<_, u32>::new();
299    /// assert_eq!(arena.len(), 0);
300    ///
301    /// arena.alloc("foo");
302    /// assert_eq!(arena.len(), 1);
303    ///
304    /// arena.alloc("bar");
305    /// assert_eq!(arena.len(), 2);
306    ///
307    /// arena.alloc("baz");
308    /// assert_eq!(arena.len(), 3);
309    /// ```
310    #[inline]
311    pub fn len(&self) -> usize {
312        self.data.len()
313    }
314
315    /// Returns the capacity of the arena.
316    ///
317    /// # Examples
318    ///
319    /// ```
320    /// use indexed_arena::Arena;
321    ///
322    /// let arena: Arena<String, u32> = Arena::with_capacity(10);
323    /// assert!(arena.capacity() >= 10);
324    /// ```
325    #[inline]
326    pub fn capacity(&self) -> usize {
327        self.data.capacity()
328    }
329
330    /// Returns `true` if the arena contains no elements.
331    ///
332    /// # Examples
333    ///
334    /// ```
335    /// # use indexed_arena::Arena;
336    /// let mut arena = Arena::<_, u32>::new();
337    /// assert!(arena.is_empty());
338    ///
339    /// arena.alloc(0.9);
340    /// assert!(!arena.is_empty());
341    /// ```
342    #[inline]
343    pub fn is_empty(&self) -> bool {
344        self.data.is_empty()
345    }
346
347    /// Allocates an element in the arena and returns its index.
348    ///
349    /// # Panics
350    ///
351    /// Panics if the arena is full (i.e. if the number of elements exceeds `I::MAX`).
352    /// If you hnadle this case, use [`Arena::try_alloc`] instead.
353    ///
354    /// # Examples
355    ///
356    /// ```
357    /// use indexed_arena::{Arena, Idx};
358    ///
359    /// let mut arena: Arena<&str, u32> = Arena::new();
360    /// let idx: Idx<&str, u32> = arena.alloc("hello");
361    /// assert_eq!(arena[idx], "hello");
362    /// ```
363    #[inline]
364    pub fn alloc(&mut self, value: T) -> Idx<T, I> {
365        self.try_alloc(value).expect("arena is full")
366    }
367
368    /// Fallible version of [`Arena::alloc`].
369    ///
370    /// This method returns `None` if the arena is full.
371    #[inline]
372    pub fn try_alloc(&mut self, value: T) -> Option<Idx<T, I>> {
373        if self.data.len() > I::MAX {
374            None
375        } else {
376            let id = I::from_usize(self.data.len());
377            self.data.push(value);
378            Some(Idx { raw: id, phantom: PhantomData })
379        }
380    }
381
382    /// Allocates multiple elements in the arena and returns the index range covering them.
383    ///
384    /// # Panics
385    ///
386    /// Panics if the arena cannot allocate all elements (i.e. if the arena becomes full).
387    ///
388    /// # Examples
389    ///
390    /// ```
391    /// use indexed_arena::{Arena, IdxRange};
392    ///
393    /// let mut arena: Arena<i32, u32> = Arena::new();
394    /// let range: IdxRange<i32, u32> = arena.alloc_many(vec![10, 20, 30]);
395    /// assert_eq!(&arena[range], &[10, 20, 30]);
396    /// ```
397    #[inline]
398    pub fn alloc_many(&mut self, values: impl IntoIterator<Item = T>) -> IdxRange<T, I> {
399        self.try_alloc_many(values).expect("arena is full")
400    }
401
402    /// Fallible version of [`Arena::alloc_many`].
403    ///
404    /// This method returns `None` if the arena becomes full.
405    #[inline]
406    pub fn try_alloc_many(
407        &mut self,
408        values: impl IntoIterator<Item = T>,
409    ) -> Option<IdxRange<T, I>> {
410        let start = I::from_usize(self.data.len());
411        let mut len = 0;
412        for value in values {
413            if len > I::MAX {
414                return None;
415            }
416            self.data.push(value);
417            len += 1;
418        }
419        let end = I::from_usize(len);
420        Some(IdxRange::new(start..end))
421    }
422
423    /// Returns a iterator over the elements and their indices in the arena.
424    ///
425    /// # Examples
426    ///
427    /// ```
428    /// # use indexed_arena::Arena;
429    /// let mut arena = Arena::<_, u32>::new();
430    ///
431    /// let idx1 = arena.alloc(20);
432    /// let idx2 = arena.alloc(40);
433    /// let idx3 = arena.alloc(60);
434    ///
435    /// let mut iter = arena.iter();
436    /// assert_eq!(iter.next(), Some((idx1, &20)));
437    /// assert_eq!(iter.next(), Some((idx2, &40)));
438    /// assert_eq!(iter.next(), Some((idx3, &60)));
439    /// assert_eq!(iter.next(), None);
440    /// ```
441    #[inline]
442    pub fn iter(&self) -> Iter<'_, T, I> {
443        Iter { iter: self.data.iter().enumerate(), phantom: PhantomData }
444    }
445
446    /// Returns a mutable iterator over the elements and their indices in the arena.
447    ///
448    /// # Examples
449    ///
450    /// ```
451    /// # use indexed_arena::Arena;
452    /// let mut arena = Arena::<_, u32>::new();
453    /// let idx1 = arena.alloc(20);
454    ///
455    /// assert_eq!(arena[idx1], 20);
456    ///
457    /// let mut iterator = arena.iter_mut();
458    /// *iterator.next().unwrap().1 = 10;
459    /// drop(iterator);
460    ///
461    /// assert_eq!(arena[idx1], 10);
462    /// ```
463    #[inline]
464    pub fn iter_mut(&mut self) -> IterMut<'_, T, I> {
465        IterMut { iter: self.data.iter_mut().enumerate(), phantom: PhantomData }
466    }
467
468    /// Shrinks the capacity of the arena to fit the number of elements.
469    ///
470    /// # Examples
471    ///
472    /// ```
473    /// # use indexed_arena::Arena;
474    /// let mut arena = Arena::<_, u32>::with_capacity(10);
475    /// arena.alloc_many(&[1, 2, 3]);
476    /// assert!(arena.capacity() >= 10);
477    ///
478    /// arena.shrink_to_fit();
479    /// assert!(arena.capacity() >= 3);
480    /// ```
481    #[inline]
482    pub fn shrink_to_fit(&mut self) {
483        self.data.shrink_to_fit();
484    }
485}
486
487impl<T, I: Id> Default for Arena<T, I> {
488    #[inline]
489    fn default() -> Self {
490        Self::new()
491    }
492}
493
494impl<T, I: Id> Index<Idx<T, I>> for Arena<T, I> {
495    type Output = T;
496
497    #[inline]
498    fn index(&self, idx: Idx<T, I>) -> &Self::Output {
499        &self.data[idx.raw.into_usize()]
500    }
501}
502
503impl<T, I: Id> Index<IdxRange<T, I>> for Arena<T, I> {
504    type Output = [T];
505
506    #[inline]
507    fn index(&self, range: IdxRange<T, I>) -> &Self::Output {
508        &self.data[range.start.into_usize()..range.end.into_usize()]
509    }
510}
511
512impl<T, I: Id> IndexMut<Idx<T, I>> for Arena<T, I> {
513    #[inline]
514    fn index_mut(&mut self, index: Idx<T, I>) -> &mut Self::Output {
515        &mut self.data[index.raw.into_usize()]
516    }
517}
518
519impl<T: Clone, I: Id> Clone for Arena<T, I> {
520    #[inline]
521    fn clone(&self) -> Self {
522        Self { data: self.data.clone(), phantom: PhantomData }
523    }
524}
525
526impl<T: fmt::Debug, I: Id> fmt::Debug for Arena<T, I> {
527    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
528        f.debug_struct("Arena").field("len", &self.len()).field("data", &self.data).finish()
529    }
530}
531
532impl<T: PartialEq, I: Id> PartialEq for Arena<T, I> {
533    #[inline]
534    fn eq(&self, other: &Self) -> bool {
535        self.data == other.data
536    }
537}
538
539impl<T: Eq, I: Id> Eq for Arena<T, I> {}
540
541impl<T: Hash, I: Id> Hash for Arena<T, I> {
542    #[inline]
543    fn hash<H: Hasher>(&self, state: &mut H) {
544        self.data.hash(state)
545    }
546}
547
548impl<'a, T, I: Id> IntoIterator for &'a Arena<T, I> {
549    type Item = (Idx<T, I>, &'a T);
550    type IntoIter = Iter<'a, T, I>;
551
552    #[inline]
553    fn into_iter(self) -> Self::IntoIter {
554        self.iter()
555    }
556}
557
558impl<T, I: Id> IntoIterator for Arena<T, I> {
559    type Item = (Idx<T, I>, T);
560    type IntoIter = IntoIter<T, I>;
561
562    #[inline]
563    fn into_iter(self) -> Self::IntoIter {
564        IntoIter { iter: self.data.into_iter().enumerate(), phantom: PhantomData }
565    }
566}
567
568macro_rules! iterator_impls {
569    ($ty:ty, type Item = $item_ty:ty;) => {
570        impl<'a, T, I: Id> Iterator for $ty {
571            type Item = $item_ty;
572            #[inline]
573            fn next(&mut self) -> Option<Self::Item> {
574                let (id, value) = self.iter.next().map(|(i, v)| (I::from_usize(i), v))?;
575                Some((Idx { raw: id, phantom: PhantomData }, value))
576            }
577            #[inline]
578            fn size_hint(&self) -> (usize, Option<usize>) {
579                self.iter.size_hint()
580            }
581            #[inline]
582            fn count(self) -> usize {
583                self.iter.count()
584            }
585            #[inline]
586            fn nth(&mut self, n: usize) -> Option<Self::Item> {
587                let (id, value) = self.iter.nth(n).map(|(i, v)| (I::from_usize(i), v))?;
588                Some((Idx { raw: id, phantom: PhantomData }, value))
589            }
590        }
591        impl<'a, T, I: Id> DoubleEndedIterator for $ty {
592            #[inline]
593            fn next_back(&mut self) -> Option<Self::Item> {
594                let (id, value) = self.iter.next_back().map(|(i, v)| (I::from_usize(i), v))?;
595                Some((Idx { raw: id, phantom: PhantomData }, value))
596            }
597        }
598        impl<'a, T, I: Id> ExactSizeIterator for $ty {
599            #[inline]
600            fn len(&self) -> usize {
601                self.iter.len()
602            }
603        }
604        impl<'a, T, I: Id> iter::FusedIterator for $ty {}
605    };
606}
607
608pub struct Iter<'a, T, I: Id> {
609    iter: iter::Enumerate<slice::Iter<'a, T>>,
610    phantom: PhantomData<I>,
611}
612
613iterator_impls! {
614    Iter<'a, T, I>,
615    type Item = (Idx<T, I>, &'a T);
616}
617
618impl<T: Clone, I: Id> Clone for Iter<'_, T, I> {
619    #[inline]
620    fn clone(&self) -> Self {
621        Self { iter: self.iter.clone(), phantom: PhantomData }
622    }
623}
624
625pub struct IterMut<'a, T, I: Id> {
626    iter: iter::Enumerate<slice::IterMut<'a, T>>,
627    phantom: PhantomData<I>,
628}
629
630iterator_impls! {
631    IterMut<'a, T, I>,
632    type Item = (Idx<T, I>, &'a mut T);
633}
634
635pub struct IntoIter<T, I: Id> {
636    iter: iter::Enumerate<alloc::vec::IntoIter<T>>,
637    phantom: PhantomData<I>,
638}
639
640iterator_impls! {
641    IntoIter<T, I>,
642    type Item = (Idx<T, I>, T);
643}
644
645impl<T: Clone, I: Id> Clone for IntoIter<T, I> {
646    #[inline]
647    fn clone(&self) -> Self {
648        Self { iter: self.iter.clone(), phantom: PhantomData }
649    }
650}