use std::{
alloc::Layout,
ptr::{self, NonNull},
slice, str,
};
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
use std::mem::offset_of;
use crate::bump::Bump;
use oxc_data_structures::assert_unchecked;
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
use crate::tracking::AllocationStats;
#[derive(Default)]
pub struct Allocator {
bump: Bump,
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
pub(crate) stats: AllocationStats,
}
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
#[expect(clippy::cast_possible_wrap)]
pub const STATS_FIELD_OFFSET: isize =
(offset_of!(Allocator, stats) as isize) - (offset_of!(Allocator, bump) as isize);
impl Allocator {
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn new() -> Self {
Self {
bump: Bump::new(),
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
stats: AllocationStats::default(),
}
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn with_capacity(capacity: usize) -> Self {
Self {
bump: Bump::with_capacity(capacity),
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
stats: AllocationStats::default(),
}
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn alloc<T>(&self, val: T) -> &mut T {
const { assert!(!std::mem::needs_drop::<T>(), "Cannot allocate Drop type in arena") };
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
self.stats.record_allocation();
self.bump.alloc(val)
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn alloc_str<'alloc>(&'alloc self, src: &str) -> &'alloc str {
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
self.stats.record_allocation();
self.bump.alloc_str(src)
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn alloc_slice_copy<T: Copy>(&self, src: &[T]) -> &mut [T] {
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
self.stats.record_allocation();
self.bump.alloc_slice_copy(src)
}
pub fn alloc_layout(&self, layout: Layout) -> NonNull<u8> {
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
self.stats.record_allocation();
self.bump.alloc_layout(layout)
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn alloc_concat_strs_array<'a, const N: usize>(&'a self, strings: [&str; N]) -> &'a str {
#[expect(clippy::checked_conversions)]
let total_len = strings.iter().fold(0usize, |total_len, s| {
let len = s.len();
unsafe { assert_unchecked!(len <= (isize::MAX as usize)) };
total_len.checked_add(len).unwrap()
});
assert!(
isize::try_from(total_len).is_ok(),
"attempted to create a string longer than `isize::MAX` bytes"
);
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
self.stats.record_allocation();
unsafe { self.alloc_concat_strs_array_with_total_len_in(strings, total_len) }
}
unsafe fn alloc_concat_strs_array_with_total_len_in<'a, const N: usize>(
&'a self,
strings: [&str; N],
total_len: usize,
) -> &'a str {
if total_len == 0 {
return "";
}
let layout = unsafe { Layout::from_size_align_unchecked(total_len, 1) };
let start_ptr = self.bump().alloc_layout(layout);
let mut end_ptr = start_ptr;
for str in strings {
let src_ptr = str.as_ptr();
let len = str.len();
unsafe { ptr::copy_nonoverlapping(src_ptr, end_ptr.as_ptr(), len) };
end_ptr = unsafe { end_ptr.add(len) };
}
debug_assert_eq!(end_ptr.as_ptr() as usize - start_ptr.as_ptr() as usize, total_len);
unsafe {
let slice = slice::from_raw_parts(start_ptr.as_ptr(), total_len);
str::from_utf8_unchecked(slice)
}
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn reset(&mut self) {
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
self.stats.reset();
self.bump.reset();
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub fn capacity(&self) -> usize {
self.bump.allocated_bytes()
}
pub fn used_bytes(&self) -> usize {
let mut bytes = 0;
let chunks_iter = unsafe { self.bump.iter_allocated_chunks_raw() };
for (_, size) in chunks_iter {
bytes += size;
}
bytes
}
#[expect(clippy::inline_always)]
#[inline(always)]
pub(crate) fn bump(&self) -> &Bump {
&self.bump
}
#[cfg(feature = "from_raw_parts")]
#[expect(clippy::inline_always)]
#[inline(always)]
pub(crate) fn from_bump(bump: Bump) -> Self {
Self {
bump,
#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
stats: AllocationStats::default(),
}
}
}
#[cfg(test)]
mod test {
use super::Allocator;
#[test]
fn test_api() {
let mut allocator = Allocator::default();
{
let array = allocator.alloc([123; 10]);
assert_eq!(array, &[123; 10]);
let str = allocator.alloc_str("hello");
assert_eq!(str, "hello");
}
allocator.reset();
}
#[test]
fn string_from_array_len_1() {
let allocator = Allocator::default();
let s = allocator.alloc_concat_strs_array(["hello"]);
assert_eq!(s, "hello");
}
#[test]
fn string_from_array_len_2() {
let allocator = Allocator::default();
let s = allocator.alloc_concat_strs_array(["hello", "world!"]);
assert_eq!(s, "helloworld!");
}
#[test]
fn string_from_array_len_3() {
let hello = "hello";
let world = std::string::String::from("world");
let allocator = Allocator::default();
let s = allocator.alloc_concat_strs_array([hello, &world, "!"]);
assert_eq!(s, "helloworld!");
}
#[test]
fn string_from_empty_array() {
let allocator = Allocator::default();
let s = allocator.alloc_concat_strs_array([]);
assert_eq!(s, "");
}
#[test]
fn string_from_array_of_empty_strs() {
let allocator = Allocator::default();
let s = allocator.alloc_concat_strs_array(["", "", ""]);
assert_eq!(s, "");
}
#[test]
fn string_from_array_containing_some_empty_strs() {
let allocator = Allocator::default();
let s = allocator.alloc_concat_strs_array(["", "hello", ""]);
assert_eq!(s, "hello");
}
}