From a877c0d726a257990712a15395f30a5d9e15c351 Mon Sep 17 00:00:00 2001 From: John Date: Fri, 26 Apr 2024 00:02:50 -0500 Subject: [PATCH] cl-arena: Add arena allocator implementation based HEAVILY on rustc-arena It seems to pass the most basic of tests in miri (shame it needs the unstable `strict_provenance` feature to do so, but eh.) --- Cargo.toml | 1 + compiler/cl-arena/Cargo.toml | 10 + compiler/cl-arena/src/dropless_arena/tests.rs | 42 +++ compiler/cl-arena/src/lib.rs | 344 ++++++++++++++++++ compiler/cl-arena/src/typed_arena/tests.rs | 61 ++++ 5 files changed, 458 insertions(+) create mode 100644 compiler/cl-arena/Cargo.toml create mode 100644 compiler/cl-arena/src/dropless_arena/tests.rs create mode 100644 compiler/cl-arena/src/lib.rs create mode 100644 compiler/cl-arena/src/typed_arena/tests.rs diff --git a/Cargo.toml b/Cargo.toml index 96fe598..dfd4989 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ members = [ "compiler/cl-ast", "compiler/cl-parser", "compiler/cl-lexer", + "compiler/cl-arena", "repline", ] resolver = "2" diff --git a/compiler/cl-arena/Cargo.toml b/compiler/cl-arena/Cargo.toml new file mode 100644 index 0000000..276e0c5 --- /dev/null +++ b/compiler/cl-arena/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "cl-arena" +repository.workspace = true +version.workspace = true +authors.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true + +[dependencies] diff --git a/compiler/cl-arena/src/dropless_arena/tests.rs b/compiler/cl-arena/src/dropless_arena/tests.rs new file mode 100644 index 0000000..18e7b44 --- /dev/null +++ b/compiler/cl-arena/src/dropless_arena/tests.rs @@ -0,0 +1,42 @@ +use super::DroplessArena; + extern crate std; + use core::alloc::Layout; + use std::{prelude::rust_2021::*, vec}; + + #[test] + fn alloc_raw() { + let arena = DroplessArena::new(); + let bytes = arena.alloc_raw(Layout::for_value(&0u128)); + let byte2 = arena.alloc_raw(Layout::for_value(&0u128)); + + assert_ne!(bytes, byte2); + } + + #[test] + fn alloc() { + let arena = DroplessArena::new(); + let mut allocations = vec![]; + for i in 0..0x400 { + allocations.push(arena.alloc(i)); + } + } + + #[test] + fn alloc_strings() { + const KW: &[&str] = &["pub", "mut", "fn", "mod", "conlang", "sidon", "🦈"]; + let arena = DroplessArena::new(); + let mut allocations = vec![]; + for _ in 0..100 { + for kw in KW { + allocations.push(arena.alloc_str(kw)); + } + } + } + + #[test] + #[should_panic] + fn alloc_zsts() { + struct Zst; + let arena = DroplessArena::new(); + arena.alloc(Zst); + } diff --git a/compiler/cl-arena/src/lib.rs b/compiler/cl-arena/src/lib.rs new file mode 100644 index 0000000..e8eae66 --- /dev/null +++ b/compiler/cl-arena/src/lib.rs @@ -0,0 +1,344 @@ +//! Typed and dropless arena allocation, paraphrased from [the Rust Compiler's `rustc_arena`](https://github.com/rust-lang/rust/blob/master/compiler/rustc_arena/src/lib.rs). See [LICENSE][1]. +//! +//! An Arena Allocator is a type of allocator which provides stable locations for allocations within +//! itself for the entire duration of its lifetime. +//! +//! [1]: https://raw.githubusercontent.com/rust-lang/rust/master/LICENSE-MIT + +#![feature(dropck_eyepatch, new_uninit, strict_provenance)] +#![no_std] + +extern crate alloc; + +pub(crate) mod constants { + //! Size constants for arena chunk growth + pub(crate) const MIN_CHUNK: usize = 4096; + pub(crate) const MAX_CHUNK: usize = 2 * 1024 * 1024; +} + +mod chunk { + //! An [ArenaChunk] contains a block of raw memory for use in arena allocators. + use alloc::boxed::Box; + use core::{ + mem::{self, MaybeUninit}, + ptr::{self, NonNull}, + }; + + pub struct ArenaChunk { + pub(crate) mem: NonNull<[MaybeUninit]>, + pub(crate) filled: usize, + } + + impl ArenaChunk { + pub fn new(cap: usize) -> Self { + let slice = Box::new_uninit_slice(cap); + Self { mem: NonNull::from(Box::leak(slice)), filled: 0 } + } + + /// Drops all elements inside self, and resets the filled count to 0 + /// + /// # Safety + /// + /// The caller must ensure that `self.filled` elements of self are currently initialized + pub unsafe fn drop_elements(&mut self) { + if mem::needs_drop::() { + // Safety: the caller has ensured that `filled` elements are initialized + unsafe { + let slice = self.mem.as_mut(); + for t in slice[..self.filled].iter_mut() { + t.assume_init_drop(); + } + } + self.filled = 0; + } + } + + /// Gets a pointer to the start of the arena + pub fn start(&mut self) -> *mut T { + self.mem.as_ptr() as _ + } + + /// Gets a pointer to the end of the arena + pub fn end(&mut self) -> *mut T { + if mem::size_of::() == 0 { + ptr::without_provenance_mut(usize::MAX) // pointers to ZSTs must be unique + } else { + unsafe { self.start().add(self.mem.len()) } + } + } + } + + impl Drop for ArenaChunk { + fn drop(&mut self) { + let _ = unsafe { Box::from_raw(self.mem.as_ptr()) }; + } + } +} + +pub mod typed_arena { + //! A [TypedArena] can hold many instances of a single type, and will properly [Drop] them. + use crate::{chunk::ArenaChunk, constants::*}; + use alloc::vec::Vec; + use core::{ + cell::{Cell, RefCell}, + marker::PhantomData, + mem, ptr, + }; + + /// A [TypedArena] can hold many instances of a single type, and will properly [Drop] them when + /// it falls out of scope. + pub struct TypedArena { + _drops: PhantomData, + chunks: RefCell>>, + head: Cell<*mut T>, + tail: Cell<*mut T>, + } + + impl Default for TypedArena { + fn default() -> Self { + Self { + _drops: Default::default(), + chunks: Default::default(), + head: Cell::new(ptr::null_mut()), + tail: Cell::new(ptr::null_mut()), + } + } + } + + impl TypedArena { + pub fn new() -> Self { + Self::default() + } + + #[allow(clippy::mut_from_ref)] + pub fn alloc(&self, value: T) -> &mut T { + if self.head == self.tail { + self.grow(1); + } + + let out = if mem::size_of::() == 0 { + self.head + .set(ptr::without_provenance_mut(self.head.get().addr() + 1)); + ptr::NonNull::::dangling().as_ptr() + } else { + let out = self.head.get(); + self.head.set(unsafe { out.add(1) }); + out + }; + + unsafe { + ptr::write(out, value); + &mut *out + } + } + + #[cold] + #[inline(never)] + fn grow(&self, len: usize) { + let size = mem::size_of::().max(1); + + let mut chunks = self.chunks.borrow_mut(); + + let capacity = if let Some(last) = chunks.last_mut() { + last.filled = self.get_filled_of_chunk(last); + last.mem.len().min(MAX_CHUNK / size) * 2 + } else { + MIN_CHUNK / size + } + .max(len); + + let mut chunk = ArenaChunk::::new(capacity); + + self.head.set(chunk.start()); + self.tail.set(chunk.end()); + chunks.push(chunk); + } + + fn get_filled_of_chunk(&self, chunk: &mut ArenaChunk) -> usize { + let Self { head: tail, .. } = self; + let head = chunk.start(); + if mem::size_of::() == 0 { + tail.get().addr() - head.addr() + } else { + unsafe { tail.get().offset_from(head) as usize } + } + } + } + + unsafe impl Send for TypedArena {} + + unsafe impl<#[may_dangle] T> Drop for TypedArena { + fn drop(&mut self) { + let mut chunks = self.chunks.borrow_mut(); + + if let Some(last) = chunks.last_mut() { + last.filled = self.get_filled_of_chunk(last); + self.tail.set(self.head.get()); + } + + for chunk in chunks.iter_mut() { + unsafe { chunk.drop_elements() } + } + } + } + + #[cfg(test)] + mod tests; +} + +pub mod dropless_arena { + //! A [DroplessArena] can hold *any* combination of types as long as they don't implement + //! [Drop]. + use crate::{chunk::ArenaChunk, constants::*}; + use alloc::vec::Vec; + use core::{ + alloc::Layout, + cell::{Cell, RefCell}, + mem, ptr, slice, + }; + + pub struct DroplessArena { + chunks: RefCell>>, + head: Cell<*mut u8>, + tail: Cell<*mut u8>, + } + + impl Default for DroplessArena { + fn default() -> Self { + Self { + chunks: Default::default(), + head: Cell::new(ptr::null_mut()), + tail: Cell::new(ptr::null_mut()), + } + } + } + + impl DroplessArena { + pub fn new() -> Self { + Self::default() + } + + /// Allocates a `T` in the [DroplessArena], and returns a mutable reference to it. + /// + /// # Panics + /// - Panics if T implements [Drop] + /// - Panics if T is zero-sized + #[allow(clippy::mut_from_ref)] + pub fn alloc(&self, value: T) -> &mut T { + assert!(!mem::needs_drop::()); + assert!(mem::size_of::() != 0); + + let out = self.alloc_raw(Layout::new::()) as *mut T; + + unsafe { + ptr::write(out, value); + &mut *out + } + } + + /// Allocates a slice of `T`s`, copied from the given slice, returning a mutable reference + /// to it. + /// + /// # Panics + /// - Panics if T implements [Drop] + /// - Panics if T is zero-sized + /// - Panics if the slice is empty + #[allow(clippy::mut_from_ref)] + pub fn alloc_slice(&self, slice: &[T]) -> &mut [T] { + assert!(!mem::needs_drop::()); + assert!(mem::size_of::() != 0); + assert!(!slice.is_empty()); + + let mem = self.alloc_raw(Layout::for_value::<[T]>(slice)) as *mut T; + + unsafe { + mem.copy_from_nonoverlapping(slice.as_ptr(), slice.len()); + slice::from_raw_parts_mut(mem, slice.len()) + } + } + + /// Allocates a copy of the given [`&str`](str), returning a reference to the allocation. + /// + /// # Panics + /// Panics if the string is empty. + pub fn alloc_str(&self, string: &str) -> &str { + let slice = self.alloc_slice(string.as_bytes()); + + // Safety: This is a clone of the input string, which was valid + unsafe { core::str::from_utf8_unchecked(slice) } + } + + /// Allocates some [bytes](u8) based on the given [Layout]. + /// + /// # Panics + /// Panics if the provided [Layout] has size 0 + pub fn alloc_raw(&self, layout: Layout) -> *mut u8 { + /// Rounds the given size (or pointer value) *up* to the given alignment + fn align_up(size: usize, align: usize) -> usize { + (size + align - 1) & !(align - 1) + } + /// Rounds the given size (or pointer value) *down* to the given alignment + fn align_down(size: usize, align: usize) -> usize { + size & !(align - 1) + } + + assert!(layout.size() != 0); + loop { + let Self { head, tail, .. } = self; + let start = head.get().addr(); + let end = tail.get().addr(); + + let align = 8.max(layout.align()); + + let bytes = align_up(layout.size(), align); + + if let Some(end) = end.checked_sub(bytes) { + let end = align_down(end, layout.align()); + + if start <= end { + tail.set(tail.get().with_addr(end)); + return tail.get(); + } + } + + self.grow(layout.size()); + } + } + + /// Grows the allocator, doubling the chunk size until it reaches [MAX_CHUNK]. + #[cold] + #[inline(never)] + fn grow(&self, len: usize) { + let mut chunks = self.chunks.borrow_mut(); + + let capacity = if let Some(last) = chunks.last_mut() { + last.mem.len().min(MAX_CHUNK / 2) * 2 + } else { + MIN_CHUNK + } + .max(len); + + let mut chunk = ArenaChunk::::new(capacity); + + self.head.set(chunk.start()); + self.tail.set(chunk.end()); + chunks.push(chunk); + } + + /// Checks whether the given slice is allocated in this arena + pub fn contains_slice(&self, slice: &[T]) -> bool { + let ptr = slice.as_ptr().cast::().cast_mut(); + for chunk in self.chunks.borrow_mut().iter_mut() { + if chunk.start() <= ptr && ptr <= chunk.end() { + return true; + } + } + false + } + } + + unsafe impl Send for DroplessArena {} + + #[cfg(test)] + mod tests; +} diff --git a/compiler/cl-arena/src/typed_arena/tests.rs b/compiler/cl-arena/src/typed_arena/tests.rs new file mode 100644 index 0000000..74c8b74 --- /dev/null +++ b/compiler/cl-arena/src/typed_arena/tests.rs @@ -0,0 +1,61 @@ +use super::TypedArena; + extern crate std; + use std::{prelude::rust_2021::*, print, vec}; + #[test] + fn pushing_to_arena() { + let arena = TypedArena::new(); + let foo = arena.alloc("foo"); + let bar = arena.alloc("bar"); + let baz = arena.alloc("baz"); + + assert_eq!("foo", *foo); + assert_eq!("bar", *bar); + assert_eq!("baz", *baz); + } + + #[test] + fn pushing_vecs_to_arena() { + let arena = TypedArena::new(); + + let foo = arena.alloc(vec!["foo"]); + let bar = arena.alloc(vec!["bar"]); + let baz = arena.alloc(vec!["baz"]); + + assert_eq!("foo", foo[0]); + assert_eq!("bar", bar[0]); + assert_eq!("baz", baz[0]); + } + + #[test] + fn pushing_zsts() { + struct ZeroSized; + impl Drop for ZeroSized { + fn drop(&mut self) { + print!("") + } + } + + let arena = TypedArena::new(); + + for _ in 0..0x100 { + arena.alloc(ZeroSized); + } + } + + #[test] + fn pushing_nodrop_zsts() { + struct ZeroSized; + let arena = TypedArena::new(); + + for _ in 0..0x1000 { + arena.alloc(ZeroSized); + } + } + #[test] + fn resize() { + let arena = TypedArena::new(); + + for _ in 0..0x780 { + arena.alloc(0u128); + } + }