1use std::{
2 alloc,
3 cell::Cell,
4 ops::{Deref, DerefMut},
5 ptr,
6 rc::Rc,
7};
8
9struct ArenaElement {
10 value: *mut u8,
11 drop: unsafe fn(*mut u8),
12}
13
14impl Drop for ArenaElement {
15 #[inline(always)]
16 fn drop(&mut self) {
17 unsafe {
18 (self.drop)(self.value);
19 }
20 }
21}
22
23pub struct Arena {
24 start: *mut u8,
25 end: *mut u8,
26 offset: *mut u8,
27 elements: Vec<ArenaElement>,
28 valid: Rc<Cell<bool>>,
29}
30
31impl Arena {
32 pub fn new(size_in_bytes: usize) -> Self {
33 unsafe {
34 let layout = alloc::Layout::from_size_align(size_in_bytes, 1).unwrap();
35 let start = alloc::alloc(layout);
36 let end = start.add(size_in_bytes);
37 Self {
38 start,
39 end,
40 offset: start,
41 elements: Vec::new(),
42 valid: Rc::new(Cell::new(true)),
43 }
44 }
45 }
46
47 pub fn clear(&mut self) {
48 self.valid.set(false);
49 self.valid = Rc::new(Cell::new(true));
50 self.elements.clear();
51 self.offset = self.start;
52 }
53
54 #[inline(always)]
55 pub fn alloc<T>(&mut self, f: impl FnOnce() -> T) -> ArenaBox<T> {
56 #[inline(always)]
57 unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
58 where
59 F: FnOnce() -> T,
60 {
61 ptr::write(ptr, f());
62 }
63
64 unsafe fn drop<T>(ptr: *mut u8) {
65 std::ptr::drop_in_place(ptr.cast::<T>());
66 }
67
68 unsafe {
69 let layout = alloc::Layout::new::<T>();
70 let offset = self.offset.add(self.offset.align_offset(layout.align()));
71 let next_offset = offset.add(layout.size());
72 assert!(next_offset <= self.end, "not enough space in Arena");
73
74 let result = ArenaBox {
75 ptr: offset.cast(),
76 valid: self.valid.clone(),
77 };
78
79 inner_writer(result.ptr, f);
80 self.elements.push(ArenaElement {
81 value: offset,
82 drop: drop::<T>,
83 });
84 self.offset = next_offset;
85
86 result
87 }
88 }
89}
90
91impl Drop for Arena {
92 fn drop(&mut self) {
93 self.clear();
94 }
95}
96
97pub struct ArenaBox<T: ?Sized> {
98 ptr: *mut T,
99 valid: Rc<Cell<bool>>,
100}
101
102impl<T: ?Sized> ArenaBox<T> {
103 #[inline(always)]
104 pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaBox<U> {
105 ArenaBox {
106 ptr: f(&mut self),
107 valid: self.valid,
108 }
109 }
110
111 fn validate(&self) {
112 assert!(
113 self.valid.get(),
114 "attempted to dereference an ArenaRef after its Arena was cleared"
115 );
116 }
117}
118
119impl<T: ?Sized> Deref for ArenaBox<T> {
120 type Target = T;
121
122 #[inline(always)]
123 fn deref(&self) -> &Self::Target {
124 self.validate();
125 unsafe { &*self.ptr }
126 }
127}
128
129impl<T: ?Sized> DerefMut for ArenaBox<T> {
130 #[inline(always)]
131 fn deref_mut(&mut self) -> &mut Self::Target {
132 self.validate();
133 unsafe { &mut *self.ptr }
134 }
135}
136
137pub struct ArenaRef<T: ?Sized>(ArenaBox<T>);
138
139impl<T: ?Sized> From<ArenaBox<T>> for ArenaRef<T> {
140 fn from(value: ArenaBox<T>) -> Self {
141 ArenaRef(value)
142 }
143}
144
145impl<T: ?Sized> Clone for ArenaRef<T> {
146 fn clone(&self) -> Self {
147 Self(ArenaBox {
148 ptr: self.0.ptr,
149 valid: self.0.valid.clone(),
150 })
151 }
152}
153
154impl<T: ?Sized> Deref for ArenaRef<T> {
155 type Target = T;
156
157 #[inline(always)]
158 fn deref(&self) -> &Self::Target {
159 self.0.deref()
160 }
161}
162
163#[cfg(test)]
164mod tests {
165 use std::{cell::Cell, rc::Rc};
166
167 use super::*;
168
169 #[test]
170 fn test_arena() {
171 let mut arena = Arena::new(1024);
172 let a = arena.alloc(|| 1u64);
173 let b = arena.alloc(|| 2u32);
174 let c = arena.alloc(|| 3u16);
175 let d = arena.alloc(|| 4u8);
176 assert_eq!(*a, 1);
177 assert_eq!(*b, 2);
178 assert_eq!(*c, 3);
179 assert_eq!(*d, 4);
180
181 arena.clear();
182 let a = arena.alloc(|| 5u64);
183 let b = arena.alloc(|| 6u32);
184 let c = arena.alloc(|| 7u16);
185 let d = arena.alloc(|| 8u8);
186 assert_eq!(*a, 5);
187 assert_eq!(*b, 6);
188 assert_eq!(*c, 7);
189 assert_eq!(*d, 8);
190
191 // Ensure drop gets called.
192 let dropped = Rc::new(Cell::new(false));
193 struct DropGuard(Rc<Cell<bool>>);
194 impl Drop for DropGuard {
195 fn drop(&mut self) {
196 self.0.set(true);
197 }
198 }
199 arena.alloc(|| DropGuard(dropped.clone()));
200 arena.clear();
201 assert!(dropped.get());
202 }
203
204 #[test]
205 #[should_panic(expected = "not enough space in Arena")]
206 fn test_arena_overflow() {
207 let mut arena = Arena::new(16);
208 arena.alloc(|| 1u64);
209 arena.alloc(|| 2u64);
210 // This should panic.
211 arena.alloc(|| 3u64);
212 }
213
214 #[test]
215 fn test_arena_alignment() {
216 let mut arena = Arena::new(256);
217 let x1 = arena.alloc(|| 1u8);
218 let x2 = arena.alloc(|| 2u16);
219 let x3 = arena.alloc(|| 3u32);
220 let x4 = arena.alloc(|| 4u64);
221 let x5 = arena.alloc(|| 5u64);
222
223 assert_eq!(*x1, 1);
224 assert_eq!(*x2, 2);
225 assert_eq!(*x3, 3);
226 assert_eq!(*x4, 4);
227 assert_eq!(*x5, 5);
228
229 assert_eq!(x1.ptr.align_offset(std::mem::align_of_val(&*x1)), 0);
230 assert_eq!(x2.ptr.align_offset(std::mem::align_of_val(&*x2)), 0);
231 }
232
233 #[test]
234 #[should_panic(expected = "attempted to dereference an ArenaRef after its Arena was cleared")]
235 fn test_arena_use_after_clear() {
236 let mut arena = Arena::new(16);
237 let value = arena.alloc(|| 1u64);
238
239 arena.clear();
240 let _read_value = *value;
241 }
242}