1use std::{
2 alloc,
3 cell::Cell,
4 ops::{Deref, DerefMut},
5 ptr::{self, NonNull},
6 rc::Rc,
7};
8
9struct ArenaElement {
10 value: NonNull<u8>,
11 drop: unsafe fn(NonNull<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: NonNull<u8>,
25 offset: usize,
26 elements: Vec<ArenaElement>,
27 valid: Rc<Cell<bool>>,
28}
29
30impl Arena {
31 pub fn new(size_in_bytes: usize) -> Self {
32 unsafe {
33 let layout = alloc::Layout::from_size_align(size_in_bytes, 1).unwrap();
34 let ptr = alloc::alloc(layout);
35 Self {
36 start: NonNull::new_unchecked(ptr),
37 offset: 0,
38 elements: Vec::new(),
39 valid: Rc::new(Cell::new(true)),
40 }
41 }
42 }
43
44 pub fn clear(&mut self) {
45 self.valid.set(false);
46 self.valid = Rc::new(Cell::new(true));
47 self.elements.clear();
48 self.offset = 0;
49 }
50
51 #[inline(always)]
52 pub fn alloc<T>(&mut self, f: impl FnOnce() -> T) -> ArenaRef<T> {
53 #[inline(always)]
54 unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
55 where
56 F: FnOnce() -> T,
57 {
58 ptr::write(ptr, f());
59 }
60
61 unsafe fn drop<T>(ptr: NonNull<u8>) {
62 std::ptr::drop_in_place(ptr.cast::<T>().as_ptr());
63 }
64
65 unsafe {
66 let layout = alloc::Layout::new::<T>().pad_to_align();
67 let ptr = NonNull::new_unchecked(self.start.as_ptr().add(self.offset).cast::<T>());
68 inner_writer(ptr.as_ptr(), f);
69
70 self.elements.push(ArenaElement {
71 value: ptr.cast(),
72 drop: drop::<T>,
73 });
74 self.offset += layout.size();
75 ArenaRef {
76 ptr,
77 valid: self.valid.clone(),
78 }
79 }
80 }
81}
82
83impl Drop for Arena {
84 fn drop(&mut self) {
85 self.clear();
86 }
87}
88
89pub struct ArenaRef<T: ?Sized> {
90 ptr: NonNull<T>,
91 valid: Rc<Cell<bool>>,
92}
93
94impl<T: ?Sized> Clone for ArenaRef<T> {
95 fn clone(&self) -> Self {
96 Self {
97 ptr: self.ptr,
98 valid: self.valid.clone(),
99 }
100 }
101}
102
103impl<T: ?Sized> ArenaRef<T> {
104 #[inline(always)]
105 pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaRef<U> {
106 ArenaRef {
107 ptr: unsafe { NonNull::new_unchecked(f(&mut *self)) },
108 valid: self.valid,
109 }
110 }
111
112 fn validate(&self) {
113 assert!(
114 self.valid.get(),
115 "attempted to dereference an ArenaRef after its Arena was cleared"
116 );
117 }
118}
119
120impl<T: ?Sized> Deref for ArenaRef<T> {
121 type Target = T;
122
123 #[inline(always)]
124 fn deref(&self) -> &Self::Target {
125 self.validate();
126 unsafe { self.ptr.as_ref() }
127 }
128}
129
130impl<T: ?Sized> DerefMut for ArenaRef<T> {
131 #[inline(always)]
132 fn deref_mut(&mut self) -> &mut Self::Target {
133 self.validate();
134 unsafe { self.ptr.as_mut() }
135 }
136}
137
138#[cfg(test)]
139mod tests {
140 use std::{cell::Cell, rc::Rc};
141
142 use super::*;
143
144 #[test]
145 fn test_arena() {
146 let mut arena = Arena::new(1024);
147 let a = arena.alloc(|| 1u64);
148 let b = arena.alloc(|| 2u32);
149 let c = arena.alloc(|| 3u16);
150 let d = arena.alloc(|| 4u8);
151 assert_eq!(*a, 1);
152 assert_eq!(*b, 2);
153 assert_eq!(*c, 3);
154 assert_eq!(*d, 4);
155
156 arena.clear();
157 let a = arena.alloc(|| 5u64);
158 let b = arena.alloc(|| 6u32);
159 let c = arena.alloc(|| 7u16);
160 let d = arena.alloc(|| 8u8);
161 assert_eq!(*a, 5);
162 assert_eq!(*b, 6);
163 assert_eq!(*c, 7);
164 assert_eq!(*d, 8);
165
166 // Ensure drop gets called.
167 let dropped = Rc::new(Cell::new(false));
168 struct DropGuard(Rc<Cell<bool>>);
169 impl Drop for DropGuard {
170 fn drop(&mut self) {
171 self.0.set(true);
172 }
173 }
174 arena.alloc(|| DropGuard(dropped.clone()));
175 arena.clear();
176 assert!(dropped.get());
177 }
178}