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) -> ArenaRef<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>().pad_to_align();
70 let next_offset = self.offset.add(layout.size());
71 assert!(next_offset <= self.end);
72
73 let result = ArenaRef {
74 ptr: self.offset.cast(),
75 valid: self.valid.clone(),
76 };
77
78 inner_writer(result.ptr, f);
79 self.elements.push(ArenaElement {
80 value: self.offset,
81 drop: drop::<T>,
82 });
83 self.offset = next_offset;
84
85 result
86 }
87 }
88}
89
90impl Drop for Arena {
91 fn drop(&mut self) {
92 self.clear();
93 }
94}
95
96pub struct ArenaRef<T: ?Sized> {
97 ptr: *mut T,
98 valid: Rc<Cell<bool>>,
99}
100
101impl<T: ?Sized> Clone for ArenaRef<T> {
102 fn clone(&self) -> Self {
103 Self {
104 ptr: self.ptr,
105 valid: self.valid.clone(),
106 }
107 }
108}
109
110impl<T: ?Sized> ArenaRef<T> {
111 #[inline(always)]
112 pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaRef<U> {
113 ArenaRef {
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 ArenaRef<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 ArenaRef<T> {
138 #[inline(always)]
139 fn deref_mut(&mut self) -> &mut Self::Target {
140 self.validate();
141 unsafe { &mut *self.ptr }
142 }
143}
144
145#[cfg(test)]
146mod tests {
147 use std::{cell::Cell, rc::Rc};
148
149 use super::*;
150
151 #[test]
152 fn test_arena() {
153 let mut arena = Arena::new(1024);
154 let a = arena.alloc(|| 1u64);
155 let b = arena.alloc(|| 2u32);
156 let c = arena.alloc(|| 3u16);
157 let d = arena.alloc(|| 4u8);
158 assert_eq!(*a, 1);
159 assert_eq!(*b, 2);
160 assert_eq!(*c, 3);
161 assert_eq!(*d, 4);
162
163 arena.clear();
164 let a = arena.alloc(|| 5u64);
165 let b = arena.alloc(|| 6u32);
166 let c = arena.alloc(|| 7u16);
167 let d = arena.alloc(|| 8u8);
168 assert_eq!(*a, 5);
169 assert_eq!(*b, 6);
170 assert_eq!(*c, 7);
171 assert_eq!(*d, 8);
172
173 // Ensure drop gets called.
174 let dropped = Rc::new(Cell::new(false));
175 struct DropGuard(Rc<Cell<bool>>);
176 impl Drop for DropGuard {
177 fn drop(&mut self) {
178 self.0.set(true);
179 }
180 }
181 arena.alloc(|| DropGuard(dropped.clone()));
182 arena.clear();
183 assert!(dropped.get());
184 }
185}