176 lines
4 KiB
Rust
176 lines
4 KiB
Rust
use std::{
|
|
alloc,
|
|
cell::Cell,
|
|
ops::{Deref, DerefMut},
|
|
ptr,
|
|
rc::Rc,
|
|
};
|
|
|
|
struct ArenaElement {
|
|
value: *mut u8,
|
|
drop: unsafe fn(*mut u8),
|
|
}
|
|
|
|
impl Drop for ArenaElement {
|
|
#[inline(always)]
|
|
fn drop(&mut self) {
|
|
unsafe {
|
|
(self.drop)(self.value);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct Arena {
|
|
start: *mut u8,
|
|
end: *mut u8,
|
|
offset: *mut u8,
|
|
elements: Vec<ArenaElement>,
|
|
valid: Rc<Cell<bool>>,
|
|
}
|
|
|
|
impl Arena {
|
|
pub fn new(size_in_bytes: usize) -> Self {
|
|
unsafe {
|
|
let layout = alloc::Layout::from_size_align(size_in_bytes, 1).unwrap();
|
|
let start = alloc::alloc(layout);
|
|
let end = start.add(size_in_bytes);
|
|
Self {
|
|
start,
|
|
end,
|
|
offset: start,
|
|
elements: Vec::new(),
|
|
valid: Rc::new(Cell::new(true)),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn clear(&mut self) {
|
|
self.valid.set(false);
|
|
self.valid = Rc::new(Cell::new(true));
|
|
self.elements.clear();
|
|
self.offset = self.start;
|
|
}
|
|
|
|
#[inline(always)]
|
|
pub fn alloc<T>(&mut self, f: impl FnOnce() -> T) -> ArenaRef<T> {
|
|
#[inline(always)]
|
|
unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
|
|
where
|
|
F: FnOnce() -> T,
|
|
{
|
|
ptr::write(ptr, f());
|
|
}
|
|
|
|
unsafe fn drop<T>(ptr: *mut u8) {
|
|
std::ptr::drop_in_place(ptr.cast::<T>());
|
|
}
|
|
|
|
unsafe {
|
|
let layout = alloc::Layout::new::<T>().pad_to_align();
|
|
let next_offset = self.offset.add(layout.size());
|
|
assert!(next_offset <= self.end);
|
|
|
|
let result = ArenaRef {
|
|
ptr: self.offset.cast(),
|
|
valid: self.valid.clone(),
|
|
};
|
|
|
|
inner_writer(result.ptr, f);
|
|
self.elements.push(ArenaElement {
|
|
value: self.offset,
|
|
drop: drop::<T>,
|
|
});
|
|
self.offset = next_offset;
|
|
|
|
result
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for Arena {
|
|
fn drop(&mut self) {
|
|
self.clear();
|
|
}
|
|
}
|
|
|
|
pub struct ArenaRef<T: ?Sized> {
|
|
ptr: *mut T,
|
|
valid: Rc<Cell<bool>>,
|
|
}
|
|
|
|
impl<T: ?Sized> ArenaRef<T> {
|
|
#[inline(always)]
|
|
pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaRef<U> {
|
|
ArenaRef {
|
|
ptr: f(&mut self),
|
|
valid: self.valid,
|
|
}
|
|
}
|
|
|
|
fn validate(&self) {
|
|
assert!(
|
|
self.valid.get(),
|
|
"attempted to dereference an ArenaRef after its Arena was cleared"
|
|
);
|
|
}
|
|
}
|
|
|
|
impl<T: ?Sized> Deref for ArenaRef<T> {
|
|
type Target = T;
|
|
|
|
#[inline(always)]
|
|
fn deref(&self) -> &Self::Target {
|
|
self.validate();
|
|
unsafe { &*self.ptr }
|
|
}
|
|
}
|
|
|
|
impl<T: ?Sized> DerefMut for ArenaRef<T> {
|
|
#[inline(always)]
|
|
fn deref_mut(&mut self) -> &mut Self::Target {
|
|
self.validate();
|
|
unsafe { &mut *self.ptr }
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::{cell::Cell, rc::Rc};
|
|
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_arena() {
|
|
let mut arena = Arena::new(1024);
|
|
let a = arena.alloc(|| 1u64);
|
|
let b = arena.alloc(|| 2u32);
|
|
let c = arena.alloc(|| 3u16);
|
|
let d = arena.alloc(|| 4u8);
|
|
assert_eq!(*a, 1);
|
|
assert_eq!(*b, 2);
|
|
assert_eq!(*c, 3);
|
|
assert_eq!(*d, 4);
|
|
|
|
arena.clear();
|
|
let a = arena.alloc(|| 5u64);
|
|
let b = arena.alloc(|| 6u32);
|
|
let c = arena.alloc(|| 7u16);
|
|
let d = arena.alloc(|| 8u8);
|
|
assert_eq!(*a, 5);
|
|
assert_eq!(*b, 6);
|
|
assert_eq!(*c, 7);
|
|
assert_eq!(*d, 8);
|
|
|
|
// Ensure drop gets called.
|
|
let dropped = Rc::new(Cell::new(false));
|
|
struct DropGuard(Rc<Cell<bool>>);
|
|
impl Drop for DropGuard {
|
|
fn drop(&mut self) {
|
|
self.0.set(true);
|
|
}
|
|
}
|
|
arena.alloc(|| DropGuard(dropped.clone()));
|
|
arena.clear();
|
|
assert!(dropped.get());
|
|
}
|
|
}
|