cl-arena: Add iterator allocation to TypedArena

This is copied almost verbatim from rustc-arena e82c861d7e/compiler/rustc_arena/src/lib.rs (L203)

TODO: Add unit tests, run unit tests in Miri
This commit is contained in:
John 2024-05-04 22:17:34 -05:00
parent edf175e53b
commit d2eb165759

View File

@ -77,12 +77,14 @@ mod chunk {
pub mod typed_arena { pub mod typed_arena {
//! A [TypedArena] can hold many instances of a single type, and will properly [Drop] them. //! 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 crate::{chunk::ArenaChunk, constants::*};
use alloc::vec::Vec; use alloc::vec::Vec;
use core::{ use core::{
cell::{Cell, RefCell}, cell::{Cell, RefCell},
marker::PhantomData, marker::PhantomData,
mem, ptr, mem, ptr, slice,
}; };
/// A [TypedArena] can hold many instances of a single type, and will properly [Drop] them when /// 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 { pub fn alloc(&'arena self, value: T) -> &'arena mut T {
if self.head == self.tail { if self.head == self.tail {
self.grow(1); 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::<T> == 0 || len == 0
#[inline]
fn alloc_raw_slice(&self, len: usize) -> *mut T {
assert!(mem::size_of::<T>() != 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<I>(&'arena self, iter: I) -> &'arena mut [T]
where I: IntoIterator<Item = T> {
// Collect them all into a buffer so they're allocated contiguously
let mut buf = iter.into_iter().collect::<Vec<_>>();
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::<T>() == 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] #[cold]
#[inline(never)] #[inline(never)]
fn grow(&self, len: usize) { fn grow(&self, len: usize) {