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 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 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 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 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 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 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}