pratt_gen_core/
arena.rs

1use std::{alloc::*, cell::UnsafeCell};
2
3const N: usize = 4096;
4
5pub struct Arena {
6    buffer: UnsafeCell<Vec<Box<[u8; N]>>>,
7    size: UnsafeCell<usize>,
8}
9
10impl Arena {
11    pub fn new() -> Self {
12        Arena { buffer: UnsafeCell::new(vec![]), size: UnsafeCell::new(0) }
13    }
14    pub unsafe fn alloc<V>(&self, value: V) -> &V {
15        use std::mem::*;
16        use std::slice::*;
17        // aligned start
18        let mut a = (self.size() + align_of::<V>() - 1) / align_of::<V>() * align_of::<V>();
19        let mut b = a + size_of::<V>();
20        if (a + 1) / N < b / N {
21            b += N - a % N;
22            a += N - a % N;
23        }
24        // current capacity
25        let capacity = (*self.buffer.get()).len() * N;
26        if b > capacity {
27            let slice = alloc_zeroed(Layout::from_size_align_unchecked(N, 8)) as *mut _;
28            (*self.buffer.get()).push(Box::from_raw(slice));
29        }
30        *self.size.get() = b;
31        // allocate new slice
32        let slice = &mut (*self.buffer.get())[a / N][a % N .. b % N];
33        let v_ref = from_raw_parts(&value as *const _ as *const u8, size_of::<V>());
34        slice.copy_from_slice(v_ref);
35        &*(slice.as_ptr() as *const V)
36    }
37    pub unsafe fn alloc_str(&self, value: &str) -> &str {
38        if value.len() == 0 { return "" }
39        // aligned start
40        let mut a = self.size();
41        let mut b = a + value.len();
42        if (a + 1) / N < b / N {
43            b += N - a % N;
44            a += N - a % N;
45        }
46        // current capacity
47        let capacity = (*self.buffer.get()).len() * N;
48        if b > capacity {
49            let slice = alloc_zeroed(Layout::from_size_align_unchecked(N, 8)) as *mut _;
50            (*self.buffer.get()).push(Box::from_raw(slice));
51        }
52        *self.size.get() = b;
53        // allocate new slice
54        let slice = &mut (*self.buffer.get())[a / N][a % N .. b % N];
55        slice.copy_from_slice(value.as_bytes());
56        std::str::from_utf8(slice).unwrap()
57    }
58    pub unsafe fn size(&self) -> usize {
59        *self.size.get()
60    }
61    pub unsafe fn pop(&self, len: usize) {
62        *self.size.get() = len
63    }
64}
65
66#[cfg(test)]
67mod test {
68    use super::*;
69    use rand::Rng;
70    use rand_xoshiro::*;
71    use rand_core::*;
72
73    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
74    pub struct A(u8, u64);
75
76    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
77    pub struct B(u8);
78
79    #[test]
80    fn alloc_test() {
81        let mut rng = Xoroshiro64Star::from_seed(15u64.to_be_bytes());
82        let bump = Arena::new();
83        let mut a = (Vec::new(), Vec::new());
84        let mut b = (Vec::new(), Vec::new());
85        for _ in 0..4096 {
86            let v = A(rng.gen(), rng.gen());
87            a.1.push(v);
88            let v_ref = unsafe { bump.alloc(v) };
89            a.0.push(v_ref);
90            assert!(v_ref == &v);
91            let v = B(rng.gen());
92            b.1.push(v);
93            let v_ref = unsafe { bump.alloc(v) };
94            b.0.push(v_ref);
95            assert!(v_ref == &v);
96        }
97        for i in 0..4096 {
98            assert!(a.0[i] == &a.1[i], "a[{i}] {:?} {:?} ({:?})", a.0[i], a.1[i], a.0[i] as *const _);
99            println!("a[{i}] {:?} {:?} ({:?})", a.0[i], a.1[i], a.0[i] as *const _);
100            assert!(b.0[i] == &b.1[i], "b[{i}] {:?} {:?} ({:?})", b.0[i], b.1[i], b.0[i] as *const _);
101            println!("b[{i}] {:?} {:?} ({:?})", b.0[i], b.1[i], b.0[i] as *const _);
102        }
103    }
104}