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 #[track_caller]
120 fn validate(&self) {
121 assert!(
122 self.valid.get(),
123 "attempted to dereference an ArenaRef after its Arena was cleared"
124 );
125 }
126}
127
128impl<T: ?Sized> Deref for ArenaBox<T> {
129 type Target = T;
130
131 #[inline(always)]
132 fn deref(&self) -> &Self::Target {
133 self.validate();
134 unsafe { &*self.ptr }
135 }
136}
137
138impl<T: ?Sized> DerefMut for ArenaBox<T> {
139 #[inline(always)]
140 fn deref_mut(&mut self) -> &mut Self::Target {
141 self.validate();
142 unsafe { &mut *self.ptr }
143 }
144}
145
146pub struct ArenaRef<T: ?Sized>(ArenaBox<T>);
147
148impl<T: ?Sized> From<ArenaBox<T>> for ArenaRef<T> {
149 fn from(value: ArenaBox<T>) -> Self {
150 ArenaRef(value)
151 }
152}
153
154impl<T: ?Sized> Clone for ArenaRef<T> {
155 fn clone(&self) -> Self {
156 Self(ArenaBox {
157 ptr: self.0.ptr,
158 valid: self.0.valid.clone(),
159 })
160 }
161}
162
163impl<T: ?Sized> Deref for ArenaRef<T> {
164 type Target = T;
165
166 #[inline(always)]
167 fn deref(&self) -> &Self::Target {
168 self.0.deref()
169 }
170}
171
172#[cfg(test)]
173mod tests {
174 use std::{cell::Cell, rc::Rc};
175
176 use super::*;
177
178 #[test]
179 fn test_arena() {
180 let mut arena = Arena::new(1024);
181 let a = arena.alloc(|| 1u64);
182 let b = arena.alloc(|| 2u32);
183 let c = arena.alloc(|| 3u16);
184 let d = arena.alloc(|| 4u8);
185 assert_eq!(*a, 1);
186 assert_eq!(*b, 2);
187 assert_eq!(*c, 3);
188 assert_eq!(*d, 4);
189
190 arena.clear();
191 let a = arena.alloc(|| 5u64);
192 let b = arena.alloc(|| 6u32);
193 let c = arena.alloc(|| 7u16);
194 let d = arena.alloc(|| 8u8);
195 assert_eq!(*a, 5);
196 assert_eq!(*b, 6);
197 assert_eq!(*c, 7);
198 assert_eq!(*d, 8);
199
200 // Ensure drop gets called.
201 let dropped = Rc::new(Cell::new(false));
202 struct DropGuard(Rc<Cell<bool>>);
203 impl Drop for DropGuard {
204 fn drop(&mut self) {
205 self.0.set(true);
206 }
207 }
208 arena.alloc(|| DropGuard(dropped.clone()));
209 arena.clear();
210 assert!(dropped.get());
211 }
212
213 #[test]
214 #[should_panic(expected = "not enough space in Arena")]
215 fn test_arena_overflow() {
216 let mut arena = Arena::new(16);
217 arena.alloc(|| 1u64);
218 arena.alloc(|| 2u64);
219 // This should panic.
220 arena.alloc(|| 3u64);
221 }
222
223 #[test]
224 fn test_arena_alignment() {
225 let mut arena = Arena::new(256);
226 let x1 = arena.alloc(|| 1u8);
227 let x2 = arena.alloc(|| 2u16);
228 let x3 = arena.alloc(|| 3u32);
229 let x4 = arena.alloc(|| 4u64);
230 let x5 = arena.alloc(|| 5u64);
231
232 assert_eq!(*x1, 1);
233 assert_eq!(*x2, 2);
234 assert_eq!(*x3, 3);
235 assert_eq!(*x4, 4);
236 assert_eq!(*x5, 5);
237
238 assert_eq!(x1.ptr.align_offset(std::mem::align_of_val(&*x1)), 0);
239 assert_eq!(x2.ptr.align_offset(std::mem::align_of_val(&*x2)), 0);
240 }
241
242 #[test]
243 #[should_panic(expected = "attempted to dereference an ArenaRef after its Arena was cleared")]
244 fn test_arena_use_after_clear() {
245 let mut arena = Arena::new(16);
246 let value = arena.alloc(|| 1u64);
247
248 arena.clear();
249 let _read_value = *value;
250 }
251}