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