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> ArenaRef<T> {
102 #[inline(always)]
103 pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaRef<U> {
104 ArenaRef {
105 ptr: f(&mut self),
106 valid: self.valid,
107 }
108 }
109
110 fn validate(&self) {
111 assert!(
112 self.valid.get(),
113 "attempted to dereference an ArenaRef after its Arena was cleared"
114 );
115 }
116}
117
118impl<T: ?Sized> Deref for ArenaRef<T> {
119 type Target = T;
120
121 #[inline(always)]
122 fn deref(&self) -> &Self::Target {
123 self.validate();
124 unsafe { &*self.ptr }
125 }
126}
127
128impl<T: ?Sized> DerefMut for ArenaRef<T> {
129 #[inline(always)]
130 fn deref_mut(&mut self) -> &mut Self::Target {
131 self.validate();
132 unsafe { &mut *self.ptr }
133 }
134}
135
136#[cfg(test)]
137mod tests {
138 use std::{cell::Cell, rc::Rc};
139
140 use super::*;
141
142 #[test]
143 fn test_arena() {
144 let mut arena = Arena::new(1024);
145 let a = arena.alloc(|| 1u64);
146 let b = arena.alloc(|| 2u32);
147 let c = arena.alloc(|| 3u16);
148 let d = arena.alloc(|| 4u8);
149 assert_eq!(*a, 1);
150 assert_eq!(*b, 2);
151 assert_eq!(*c, 3);
152 assert_eq!(*d, 4);
153
154 arena.clear();
155 let a = arena.alloc(|| 5u64);
156 let b = arena.alloc(|| 6u32);
157 let c = arena.alloc(|| 7u16);
158 let d = arena.alloc(|| 8u8);
159 assert_eq!(*a, 5);
160 assert_eq!(*b, 6);
161 assert_eq!(*c, 7);
162 assert_eq!(*d, 8);
163
164 // Ensure drop gets called.
165 let dropped = Rc::new(Cell::new(false));
166 struct DropGuard(Rc<Cell<bool>>);
167 impl Drop for DropGuard {
168 fn drop(&mut self) {
169 self.0.set(true);
170 }
171 }
172 arena.alloc(|| DropGuard(dropped.clone()));
173 arena.clear();
174 assert!(dropped.get());
175 }
176}