1extern crate alloc as alloc_crate;
2
3use alloc_crate::alloc::Allocator;
4use alloc_crate::collections::{TryReserveError, TryReserveErrorKind};
5use alloc_crate::vec::{self, Vec};
6use composable_allocators::Global as Global;
7use const_default::ConstDefault;
8use core::fmt::Debug;
9use core::hint::unreachable_unchecked;
10use core::iter::{self, FusedIterator};
11use core::mem::{align_of, replace, size_of};
12use core::ops::{Index, IndexMut};
13use core::slice::{self};
14use either::{Either, Left, Right};
15
16pub use crate::index::*;
17
18type ArenaItem<I, T> = Either<<I as ArenaIndex>::Optional, T>;
19
20#[derive(Debug, Clone)]
24pub struct ArenaItems<I: ArenaIndex, T, A: Allocator = Global> {
25 vec: Vec<ArenaItem<I, T>, A>,
26 vacancy: I::Optional,
27}
28
29impl<I: ArenaIndex, T, A: Allocator> ArenaItems<I, T, A> {
30 pub const fn item_size() -> usize {
34 size_of::<ArenaItem<I, T>>()
35 }
36
37 pub const fn item_align() -> usize {
41 align_of::<ArenaItem<I, T>>()
42 }
43
44 const fn new_in(alloc: A) -> Self {
45 ArenaItems {
46 vec: Vec::new_in(alloc),
47 vacancy: I::NONE
48 }
49 }
50
51 fn with_capacity_in(capacity: usize, alloc: A) -> Self {
52 ArenaItems {
53 vec: Vec::with_capacity_in(capacity, alloc),
54 vacancy: I::NONE
55 }
56 }
57
58 pub fn allocator(&self) -> &A { self.vec.allocator() }
60
61 pub fn capacity(&self) -> usize { self.vec.capacity() }
63
64 pub fn len(&self) -> usize {
68 let mut vacancies = 0;
69 let mut vacancy = self.vacancy;
70 while let Some(i) = I::to_option(vacancy) {
71 vacancies += 1;
72 vacancy = *self.vec[I::try_to_usize(i).unwrap()].as_ref().left().unwrap();
73 }
74 self.vec.len() - vacancies
75 }
76
77 pub fn len_equals_to_min_capacity(&self) -> bool {
85 I::to_option(self.vacancy).is_none()
86 }
87
88 pub fn is_empty(&self) -> bool { self.vec.iter().all(|x| x.is_left()) }
92
93 pub fn min_capacity(&self) -> usize { self.vec.len() }
121
122 pub fn reserve(&mut self, additional: usize) { self.vec.reserve(additional) }
131
132 pub fn reserve_exact(&mut self, additional: usize) { self.vec.reserve_exact(additional) }
144
145 pub fn shrink_to(&mut self, min_capacity: usize) { self.vec.shrink_to(min_capacity) }
150
151 pub fn shrink_to_fit(&mut self) { self.vec.shrink_to_fit() }
156
157 pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
166 self.vec.try_reserve(additional)
167 }
168
169 pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> {
182 self.vec.try_reserve_exact(additional)
183 }
184
185 pub fn get_value(&self, index: usize) -> Option<&T> {
191 self.vec[index].as_ref().right()
192 }
193
194 pub fn get_value_mut(&mut self, index: usize) -> Option<&mut T> {
200 self.vec[index].as_mut().right()
201 }
202
203 pub fn indices(&self) -> ArenaItemsIndices<'_, I, T> {
205 ArenaItemsIndices(self.vec.iter().enumerate())
206 }
207
208 pub fn values(&self) -> ArenaItemsValues<'_, I, T> {
210 ArenaItemsValues(self.vec.iter())
211 }
212
213 pub fn values_mut(&mut self) -> ArenaItemsValuesMut<'_, I, T> {
215 ArenaItemsValuesMut(self.vec.iter_mut())
216 }
217
218 pub fn iter(&self) -> ArenaItemsIter<'_, I, T> {
220 ArenaItemsIter(self.vec.iter().enumerate())
221 }
222
223 pub fn iter_mut(&mut self) -> ArenaItemsIterMut<'_, I, T> {
225 ArenaItemsIterMut(self.vec.iter_mut().enumerate())
226 }
227
228 pub fn into_indices(self) -> ArenaItemsIntoIndices<I, T, A> {
230 ArenaItemsIntoIndices(self.vec.into_iter().enumerate())
231 }
232
233 pub fn into_values(self) -> ArenaItemsIntoValues<I, T, A> {
235 ArenaItemsIntoValues(self.vec.into_iter())
236 }
237}
238
239#[derive(Debug, Clone)]
243pub struct ArenaItemsIter<'a, I: ArenaIndex, T>(
244 iter::Enumerate<slice::Iter<'a, Either<I::Optional, T>>>
245);
246
247impl<'a, I: ArenaIndex, T> Iterator for ArenaItemsIter<'a, I, T> {
248 type Item = (I, &'a T);
249
250 fn next(&mut self) -> Option<Self::Item> {
251 loop {
252 match self.0.next() {
253 None => return None,
254 Some((_, Left(_))) => { },
255 Some((index, Right(item))) =>
256 return Some((I::try_from_usize(index).unwrap(), item)),
257 }
258 }
259 }
260
261 fn size_hint(&self) -> (usize, Option<usize>) {
262 (0, self.0.size_hint().1)
263 }
264}
265
266impl<'a, I: ArenaIndex, T> DoubleEndedIterator for ArenaItemsIter<'a, I, T> {
267 fn next_back(&mut self) -> Option<Self::Item> {
268 loop {
269 match self.0.next_back() {
270 None => return None,
271 Some((_, Left(_))) => { },
272 Some((index, Right(item))) =>
273 return Some((I::try_from_usize(index).unwrap(), item)),
274 }
275 }
276 }
277}
278
279impl<'a, I: ArenaIndex, T> FusedIterator for ArenaItemsIter<'a, I, T> { }
280
281#[derive(Debug)]
285pub struct ArenaItemsIterMut<'a, I: ArenaIndex, T>(
286 iter::Enumerate<slice::IterMut<'a, Either<I::Optional, T>>>
287);
288
289impl<'a, I: ArenaIndex, T> Iterator for ArenaItemsIterMut<'a, I, T> {
290 type Item = (I, &'a mut T);
291
292 fn next(&mut self) -> Option<Self::Item> {
293 loop {
294 match self.0.next() {
295 None => return None,
296 Some((_, Left(_))) => { },
297 Some((index, Right(item))) =>
298 return Some((I::try_from_usize(index).unwrap(), item)),
299 }
300 }
301 }
302
303 fn size_hint(&self) -> (usize, Option<usize>) {
304 (0, self.0.size_hint().1)
305 }
306}
307
308impl<'a, I: ArenaIndex, T> DoubleEndedIterator for ArenaItemsIterMut<'a, I, T> {
309 fn next_back(&mut self) -> Option<Self::Item> {
310 loop {
311 match self.0.next_back() {
312 None => return None,
313 Some((_, Left(_))) => { },
314 Some((index, Right(item))) =>
315 return Some((I::try_from_usize(index).unwrap(), item)),
316 }
317 }
318 }
319}
320
321impl<'a, I: ArenaIndex, T> FusedIterator for ArenaItemsIterMut<'a, I, T> { }
322
323#[derive(Debug, Clone)]
327pub struct ArenaItemsIndices<'a, I: ArenaIndex, T>(
328 iter::Enumerate<slice::Iter<'a, Either<I::Optional, T>>>
329);
330
331impl<'a, I: ArenaIndex, T> Iterator for ArenaItemsIndices<'a, I, T> {
332 type Item = I;
333
334 fn next(&mut self) -> Option<Self::Item> {
335 loop {
336 match self.0.next() {
337 None => return None,
338 Some((_, Left(_))) => { },
339 Some((index, Right(_))) => return Some(I::try_from_usize(index).unwrap()),
340 }
341 }
342 }
343
344 fn size_hint(&self) -> (usize, Option<usize>) {
345 (0, self.0.size_hint().1)
346 }
347}
348
349impl<'a, I: ArenaIndex, T> DoubleEndedIterator for ArenaItemsIndices<'a, I, T> {
350 fn next_back(&mut self) -> Option<Self::Item> {
351 loop {
352 match self.0.next_back() {
353 None => return None,
354 Some((_, Left(_))) => { },
355 Some((index, Right(_))) => return Some(I::try_from_usize(index).unwrap()),
356 }
357 }
358 }
359}
360
361impl<'a, I: ArenaIndex, T> FusedIterator for ArenaItemsIndices<'a, I, T> { }
362
363#[derive(Debug, Clone)]
367pub struct ArenaItemsValues<'a, I: ArenaIndex, T>(
368 slice::Iter<'a, Either<I::Optional, T>>
369);
370
371impl<'a, I: ArenaIndex, T> Iterator for ArenaItemsValues<'a, I, T> {
372 type Item = &'a T;
373
374 fn next(&mut self) -> Option<Self::Item> {
375 loop {
376 match self.0.next() {
377 None => return None,
378 Some(Left(_)) => { },
379 Some(Right(item)) => return Some(item),
380 }
381 }
382 }
383
384 fn size_hint(&self) -> (usize, Option<usize>) {
385 (0, self.0.size_hint().1)
386 }
387}
388
389impl<'a, I: ArenaIndex, T> DoubleEndedIterator for ArenaItemsValues<'a, I, T> {
390 fn next_back(&mut self) -> Option<Self::Item> {
391 loop {
392 match self.0.next_back() {
393 None => return None,
394 Some(Left(_)) => { },
395 Some(Right(item)) => return Some(item),
396 }
397 }
398 }
399}
400
401impl<'a, I: ArenaIndex, T> FusedIterator for ArenaItemsValues<'a, I, T> { }
402
403#[derive(Debug)]
407pub struct ArenaItemsValuesMut<'a, I: ArenaIndex, T>(
408 slice::IterMut<'a, Either<I::Optional, T>>
409);
410
411impl<'a, I: ArenaIndex, T> Iterator for ArenaItemsValuesMut<'a, I, T> {
412 type Item = &'a mut T;
413
414 fn next(&mut self) -> Option<Self::Item> {
415 loop {
416 match self.0.next() {
417 None => return None,
418 Some(Left(_)) => { },
419 Some(Right(item)) => return Some(item),
420 }
421 }
422 }
423
424 fn size_hint(&self) -> (usize, Option<usize>) {
425 (0, self.0.size_hint().1)
426 }
427}
428
429impl<'a, I: ArenaIndex, T> DoubleEndedIterator for ArenaItemsValuesMut<'a, I, T> {
430 fn next_back(&mut self) -> Option<Self::Item> {
431 loop {
432 match self.0.next_back() {
433 None => return None,
434 Some(Left(_)) => { },
435 Some(Right(item)) => return Some(item),
436 }
437 }
438 }
439}
440
441impl<'a, I: ArenaIndex, T> FusedIterator for ArenaItemsValuesMut<'a, I, T> { }
442
443#[derive(Debug)]
447pub struct ArenaItemsIntoIndices<I: ArenaIndex, T, A: Allocator = Global>(
448 iter::Enumerate<vec::IntoIter<Either<I::Optional, T>, A>>,
449);
450
451impl<I: ArenaIndex, T, A: Allocator> Iterator for ArenaItemsIntoIndices<I, T, A> {
452 type Item = I;
453
454 fn next(&mut self) -> Option<Self::Item> {
455 loop {
456 match self.0.next() {
457 None => return None,
458 Some((_, Left(_))) => { },
459 Some((index, Right(_))) => return Some(I::try_from_usize(index).unwrap()),
460 }
461 }
462 }
463
464 fn size_hint(&self) -> (usize, Option<usize>) {
465 (0, self.0.size_hint().1)
466 }
467}
468
469impl<I: ArenaIndex, T, A: Allocator> DoubleEndedIterator for ArenaItemsIntoIndices<I, T, A> {
470 fn next_back(&mut self) -> Option<Self::Item> {
471 loop {
472 match self.0.next_back() {
473 None => return None,
474 Some((_, Left(_))) => { },
475 Some((index, Right(_))) => return Some(I::try_from_usize(index).unwrap()),
476 }
477 }
478 }
479}
480
481impl<I: ArenaIndex, T, A: Allocator> FusedIterator for ArenaItemsIntoIndices<I, T, A> { }
482
483#[derive(Debug)]
487pub struct ArenaItemsIntoValues<I: ArenaIndex, T, A: Allocator = Global>(
488 vec::IntoIter<Either<I::Optional, T>, A>,
489);
490
491impl<I: ArenaIndex, T, A: Allocator> Iterator for ArenaItemsIntoValues<I, T, A> {
492 type Item = T;
493
494 fn next(&mut self) -> Option<Self::Item> {
495 loop {
496 match self.0.next() {
497 None => return None,
498 Some(Left(_)) => { },
499 Some(Right(item)) => return Some(item),
500 }
501 }
502 }
503
504 fn size_hint(&self) -> (usize, Option<usize>) {
505 (0, self.0.size_hint().1)
506 }
507}
508
509impl<I: ArenaIndex, T, A: Allocator> DoubleEndedIterator for ArenaItemsIntoValues<I, T, A> {
510 fn next_back(&mut self) -> Option<Self::Item> {
511 loop {
512 match self.0.next_back() {
513 None => return None,
514 Some(Left(_)) => { },
515 Some(Right(item)) => return Some(item),
516 }
517 }
518 }
519}
520
521impl<I: ArenaIndex, T, A: Allocator> FusedIterator for ArenaItemsIntoValues<I, T, A> { }
522
523#[derive(Debug, Clone)]
527pub struct ArenaItemsIntoIter<I: ArenaIndex, T, A: Allocator = Global>(
528 iter::Enumerate<vec::IntoIter<Either<I::Optional, T>, A>>,
529);
530
531impl<I: ArenaIndex, T, A: Allocator> Iterator for ArenaItemsIntoIter<I, T, A> {
532 type Item = (I, T);
533
534 fn next(&mut self) -> Option<Self::Item> {
535 loop {
536 match self.0.next() {
537 None => return None,
538 Some((_, Left(_))) => { },
539 Some((index, Right(item))) =>
540 return Some((I::try_from_usize(index).unwrap(), item)),
541 }
542 }
543 }
544
545 fn size_hint(&self) -> (usize, Option<usize>) {
546 (0, self.0.size_hint().1)
547 }
548}
549
550impl<I: ArenaIndex, T, A: Allocator> DoubleEndedIterator for ArenaItemsIntoIter<I, T, A> {
551 fn next_back(&mut self) -> Option<Self::Item> {
552 loop {
553 match self.0.next_back() {
554 None => return None,
555 Some((_, Left(_))) => { },
556 Some((index, Right(item))) =>
557 return Some((I::try_from_usize(index).unwrap(), item)),
558 }
559 }
560 }
561}
562
563impl<I: ArenaIndex, T, A: Allocator> FusedIterator for ArenaItemsIntoIter<I, T, A> { }
564
565impl<I: ArenaIndex, T, A: Allocator> IntoIterator for ArenaItems<I, T, A> {
566 type Item = (I, T);
567 type IntoIter = ArenaItemsIntoIter<I, T, A>;
568
569 fn into_iter(self) -> Self::IntoIter {
570 ArenaItemsIntoIter(self.vec.into_iter().enumerate())
571 }
572}
573
574impl<'a, I: ArenaIndex, T, A: Allocator> IntoIterator for &'a ArenaItems<I, T, A> {
575 type Item = (I, &'a T);
576 type IntoIter = ArenaItemsIter<'a, I, T>;
577
578 fn into_iter(self) -> Self::IntoIter { self.iter() }
579}
580
581mod forgettable_field {
582 use core::fmt::{self, Debug, Formatter};
583 use core::mem::{MaybeUninit, forget, replace};
584 use core::ops::{Deref, DerefMut};
585
586 pub struct ForgettableField<T>(MaybeUninit<T>);
587
588 impl<T> ForgettableField<T> {
589 pub const fn new(value: T) -> Self { ForgettableField(MaybeUninit::new(value)) }
590
591 pub fn into_inner(mut this: Self) -> T {
592 let inner = replace(&mut this.0, MaybeUninit::uninit());
593 forget(this);
594 unsafe { inner.assume_init() }
595 }
596
597 pub fn take_and_forget<Owner>(mut owner: Owner, f: impl FnOnce(&mut Owner) -> &mut Self) -> T {
598 let this = replace(f(&mut owner), ForgettableField(MaybeUninit::uninit()));
599 forget(owner);
600 Self::into_inner(this)
601 }
602 }
603
604 impl<T> Drop for ForgettableField<T> {
605 fn drop(&mut self) {
606 unsafe { self.0.assume_init_drop() }
607 }
608 }
609
610 impl<T> Deref for ForgettableField<T> {
611 type Target = T;
612
613 fn deref(&self) -> &T { unsafe { self.0.assume_init_ref() } }
614 }
615
616 impl<T> DerefMut for ForgettableField<T> {
617 fn deref_mut(&mut self) -> &mut T { unsafe { self.0.assume_init_mut() } }
618 }
619
620 impl<T: Default> Default for ForgettableField<T> {
621 fn default() -> Self { ForgettableField::new(T::default()) }
622 }
623
624 impl<T: Debug> Debug for ForgettableField<T> {
625 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
626 self.deref().fmt(f)
627 }
628 }
629}
630
631use forgettable_field::*;
632
633#[derive(Debug)]
635pub struct Arena<I: ArenaIndex, T: 'static, A: Allocator = Global> {
636 items: ForgettableField<ArenaItems<I, T, A>>,
637}
638
639impl<I: ArenaIndex, T, A: Allocator> Arena<I, T, A> {
640 pub const fn new() -> Self where A: ConstDefault {
642 Self::new_in(ConstDefault::DEFAULT)
643 }
644
645 pub fn with_capacity(capacity: usize) -> Self where A: ConstDefault {
647 Self::with_capacity_in(capacity, ConstDefault::DEFAULT)
648 }
649
650 pub const fn new_in(alloc: A) -> Self {
652 Arena {
653 items: ForgettableField::new(ArenaItems::new_in(alloc))
654 }
655 }
656
657 pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
659 Arena {
660 items: ForgettableField::new(ArenaItems::with_capacity_in(capacity, alloc))
661 }
662 }
663
664 pub fn into_items(#[allow(unused_mut)] mut self) -> ArenaItems<I, T, A> {
668 ForgettableField::take_and_forget(self, |x| &mut x.items)
669 }
670
671 pub fn items(&self) -> &ArenaItems<I, T, A> { &self.items }
675
676 pub fn items_mut(&mut self) -> &mut ArenaItems<I, T, A> { &mut self.items }
681
682 pub fn reserve(&mut self) {
691 if self.items().len_equals_to_min_capacity() {
692 self.items_mut().reserve(1);
693 assert!(I::try_from_usize(self.items().min_capacity()).is_some());
694 }
695 }
696
697 pub fn reserve_exact(&mut self) {
709 if self.items().len_equals_to_min_capacity() {
710 self.items_mut().reserve_exact(1);
711 assert!(I::try_from_usize(self.items().min_capacity()).is_some());
712 }
713 }
714
715 pub fn try_reserve(&mut self) -> Result<(), TryReserveError> {
724 if self.items().len_equals_to_min_capacity() {
725 self.items_mut().try_reserve(1)?;
726 I::try_from_usize(self.items().min_capacity())
727 .ok_or(TryReserveError::from(TryReserveErrorKind::CapacityOverflow)).map(|_| ())
728 } else {
729 Ok(())
730 }
731 }
732
733 pub fn try_reserve_exact(&mut self) -> Result<(), TryReserveError> {
746 if self.items().len_equals_to_min_capacity() {
747 self.items_mut().try_reserve_exact(1)?;
748 I::try_from_usize(self.items().min_capacity())
749 .ok_or(TryReserveError::from(TryReserveErrorKind::CapacityOverflow)).map(|_| ())
750 } else {
751 Ok(())
752 }
753 }
754
755 pub fn insert<R>(&mut self, item: impl FnOnce(I) -> (T, R)) -> R {
770 if let Some(index) = I::to_option(self.items.vacancy) {
771 let (item, result) = item(index);
772 self.items.vacancy = replace(&mut self.items.vec[I::try_to_usize(index).unwrap()], Right(item)).left()
773 .unwrap_or_else(|| unsafe { unreachable_unchecked() });
774 result
775 } else {
776 let index = I::try_from_usize(self.items.len()).expect("out of indices");
777 let (item, result) = item(index);
778 self.items.vec.push(Right(item));
779 result
780 }
781 }
782
783 pub fn remove(&mut self, index: I) -> T {
787 let vacancy = self.items.vacancy;
788 match replace(&mut self.items.vec[I::try_to_usize(index).expect("invalid index")], Left(vacancy)) {
789 Left(vacancy) => {
790 self.items.vec[I::try_to_usize(index).unwrap()] = Left(vacancy);
791 panic!("invalid index");
792 },
793 Right(item) => {
794 self.items.vacancy = I::some(index);
795 item
796 }
797 }
798 }
799}
800
801impl<I: ArenaIndex, T, A: Allocator> Default for Arena<I, T, A> where A: ConstDefault {
802 fn default() -> Self { Arena::new() }
803}
804
805impl<I: ArenaIndex, T, A: Allocator> Index<I> for Arena<I, T, A> {
806 type Output = T;
807
808 fn index(&self, index: I) -> &T {
809 self.items.vec[I::try_to_usize(index).expect("invalid index")].as_ref().right().expect("invalid index")
810 }
811}
812
813impl<I: ArenaIndex, T, A: Allocator> IndexMut<I> for Arena<I, T, A> {
814 fn index_mut(&mut self, index: I) -> &mut T {
815 self.items.vec[I::try_to_usize(index).expect("invalid index")].as_mut().right().expect("invalid index")
816 }
817}
818
819#[cfg(test)]
820mod test {
821 use quickcheck_macros::quickcheck;
822
823 use core::num::NonZeroU32;
824 use core::sync::atomic::{AtomicI8, Ordering};
825 use crate::*;
826
827 struct Test {
828 this: NonZeroU32,
829 value: i8
830 }
831
832 const fn _new_test_arena() -> Arena<NonZeroU32, Test> {
833 Arena::new()
834 }
835
836 struct TestWithDrop {
837 value: i8
838 }
839
840 static TEST_DROP: AtomicI8 = AtomicI8::new(-1);
841
842 const fn _new_test_with_drop_arena() -> Arena<NonZeroU32, TestWithDrop> {
843 Arena::new()
844 }
845
846 impl Drop for TestWithDrop {
847 fn drop(&mut self) {
848 TEST_DROP.store(self.value, Ordering::SeqCst);
849 }
850 }
851
852 #[quickcheck]
853 fn new_arena_min_capacity_is_zero(capacity: Option<u8>) -> bool {
854 let capacity = capacity.map(|capacity| capacity as usize);
855 capacity.map_or_else(
856 || <Arena::<NonZeroU32, Test>>::new(),
857 |capacity| <Arena::<NonZeroU32, Test>>::with_capacity(capacity)
858 ).items().min_capacity() == 0
859 }
860
861 #[quickcheck]
862 fn arena_contains_inserted_item(capacity: Option<u8>, value: i8) -> bool {
863 let capacity = capacity.map(|capacity| capacity as usize);
864 let mut arena = capacity.map_or_else(
865 || <Arena::<NonZeroU32, Test>>::new(),
866 |capacity| <Arena::<NonZeroU32, Test>>::with_capacity(capacity)
867 );
868 let index = arena.insert(|this| (Test { this, value }, this));
869 arena[index].this == index && arena[index].value == value
870 }
871
872 #[test]
873 fn drop_components() {
874 {
875 let mut arena: Arena<NonZeroU32, TestWithDrop> = Arena::new();
876 arena.insert(|this: NonZeroU32| (TestWithDrop { value: 7 }, this));
877 TEST_DROP.store(-1, Ordering::SeqCst);
878 }
879 assert_eq!(TEST_DROP.load(Ordering::SeqCst), 7);
880 }
881
882 #[test]
883 fn try_reserve() {
884 let mut arena: Arena<i8, ()> = Arena::new();
885 while arena.try_reserve().is_ok() {
886 arena.insert(|_| ((), ()));
887 }
888 }
889}