diff --git a/compiler/cl-arena/src/lib.rs b/compiler/cl-arena/src/lib.rs index 5bca8e0..1d19c80 100644 --- a/compiler/cl-arena/src/lib.rs +++ b/compiler/cl-arena/src/lib.rs @@ -77,12 +77,14 @@ mod chunk { pub mod typed_arena { //! A [TypedArena] can hold many instances of a single type, and will properly [Drop] them. + #![allow(clippy::mut_from_ref)] + use crate::{chunk::ArenaChunk, constants::*}; use alloc::vec::Vec; use core::{ cell::{Cell, RefCell}, marker::PhantomData, - mem, ptr, + mem, ptr, slice, }; /// A [TypedArena] can hold many instances of a single type, and will properly [Drop] them when @@ -112,7 +114,6 @@ pub mod typed_arena { } } - #[allow(clippy::mut_from_ref)] pub fn alloc(&'arena self, value: T) -> &'arena mut T { if self.head == self.tail { self.grow(1); @@ -134,6 +135,52 @@ pub mod typed_arena { } } + fn can_allocate(&self, len: usize) -> bool { + len <= unsafe { self.tail.get().offset_from(self.head.get()) as usize } + } + + /// # Panics + /// Panics if size_of:: == 0 || len == 0 + #[inline] + fn alloc_raw_slice(&self, len: usize) -> *mut T { + assert!(mem::size_of::() != 0); + assert!(len != 0); + + if !self.can_allocate(len) { + self.grow(len) + } + + let out = self.head.get(); + + unsafe { self.head.set(out.add(len)) }; + out + } + + pub fn alloc_from_iter(&'arena self, iter: I) -> &'arena mut [T] + where I: IntoIterator { + // Collect them all into a buffer so they're allocated contiguously + let mut buf = iter.into_iter().collect::>(); + if buf.is_empty() { + return &mut []; + } + + let len = buf.len(); + // If T is a ZST, calling alloc_raw_slice will panic + let slice = if mem::size_of::() == 0 { + self.head + .set(ptr::without_provenance_mut(self.head.get().addr() + len)); + ptr::NonNull::dangling().as_ptr() + } else { + self.alloc_raw_slice(len) + }; + + unsafe { + buf.as_ptr().copy_to_nonoverlapping(slice, len); + buf.set_len(0); + slice::from_raw_parts_mut(slice, len) + } + } + #[cold] #[inline(never)] fn grow(&self, len: usize) {