arena/
typed.rs

1use core::mem;
2use core::{
3    cell::{Cell, RefCell},
4    marker::PhantomData,
5    ptr, slice,
6};
7
8use tiny_vec::TinyVec;
9
10use crate::chunk::ArenaChunk;
11use crate::{PAGE_SIZE, HUGE_PAGE};
12
13pub struct TypedArena<'ctx, T> {
14    elems: RefCell<Vec<ArenaChunk<T>>>,
15    start: Cell<*mut T>,
16    end: Cell<*mut T>,
17    _marker: PhantomData<&'ctx T>,
18}
19
20impl<'ctx, T> TypedArena<'ctx, T> {
21
22    fn reserve(&self, amount: usize) {
23        let mut chunks = self.elems.borrow_mut();
24        let last_chunk = chunks.last();
25
26        if last_chunk.is_some_and(|c| c.can_alloc(amount)) {
27            return
28        }
29
30        let elem_size = mem::size_of::<T>().max(1);
31
32        let mut new_cap =
33        if let Some(last) = last_chunk {
34            let len = last.capacity();
35            let max_size = HUGE_PAGE / elem_size;
36            usize::min(len, max_size / 2) * 2
37        } else {
38            PAGE_SIZE / elem_size
39        };
40
41        new_cap = usize::max(amount, new_cap);
42
43        let mut chunk = ArenaChunk::new(new_cap);
44        self.start.set(chunk.start());
45        self.end.set(chunk.end());
46        chunks.push(chunk);
47    }
48
49    /// Allocs an slice of elements from the given [Iterator]
50    ///
51    /// The iterator must be an [`ExactSizeIterator`]
52    #[allow(clippy::mut_from_ref)]
53    pub fn alloc_iter<I>(&self, values: I) -> &'ctx mut [T]
54    where
55        I: IntoIterator<Item = T>,
56    {
57        /* We can't pre-alloc the space like we do on the DroplessArena.
58         * This arena is reentrant, which means that the iterator may
59         * call alloc_iter again. If we pre-allocated the space we may
60         * have initialized elements after uninitialized ones.
61         *
62         * Example:
63         * An initial call to alloc_iter pre-allocates space for 4 elements.
64         * While filling the space, the iterator calls alloc_iter again,
65         * which causes it to re-alloc more space.
66         * On the third call to iterator.next(), a panic occurs.
67         *
68         * In that case, we would have the following structure in memory.
69         *
70         *     (alloc_iter_1) -- calls -- (alloc_iter2)
71         * [ INIT INIT UNINIT UNINIT ] [INIT INIT PANIC! ]
72         *                |------------------------^
73         *                We poll the iterator to get the third item,
74         *                but the iterator panics.
75         *
76         * Since we have mixed init and uninit data, we either drop uninitalized
77         * elements, or leak memory
78         *
79         * Collecting the elements beforehand takes care of panic safety and
80         * reentrancy. Also, this function is called less often that its dropless
81         * counterpart.
82         * The TinyVec can hold 8 elements on the stack, so small iterators won't
83         * even cause heap allocations
84         * */
85        let mut values: TinyVec<_, 8> = values.into_iter().collect();
86        let len = values.len();
87
88        self.reserve(len);
89
90        self.elems
91            .borrow_mut()
92            .last_mut()
93            .expect("We've grown the vector, so it's imposible\
94                it doesn't have, at least, one element")
95            .add_len(len);
96
97        let ptr = self.start.get();
98        unsafe {
99            values.as_ptr().copy_to_nonoverlapping(ptr, len);
100            values.set_len(0);
101
102            self.start.set(ptr.add(len));
103
104            slice::from_raw_parts_mut(ptr, len)
105        }
106    }
107
108    /// Allocs an element, and returns a reference to it
109    #[allow(clippy::mut_from_ref)]
110    pub fn alloc(&self, value: T) -> &'ctx mut T {
111        self.reserve(1);
112
113        self.elems
114            .borrow_mut()
115            .last_mut()
116            .expect("We've grown the vector, so it's imposible\
117                     it doesn't have, at least, one element")
118            .add_len(1);
119
120        let ptr = self.start.get();
121        unsafe {
122            ptr::write(ptr, value);
123            self.start.set(ptr.add(1));
124            &mut *ptr
125        }
126    }
127}
128
129impl<T> Default for TypedArena<'_, T> {
130    fn default() -> Self {
131        Self {
132            elems: RefCell::default(),
133            start: Cell::new(ptr::null_mut()),
134            end: Cell::new(ptr::null_mut()),
135            _marker: PhantomData,
136        }
137    }
138}