From 8d710a2b335718544b6cbf0a884c777790545986 Mon Sep 17 00:00:00 2001 From: overlookmotel <557937+overlookmotel@users.noreply.github.com> Date: Wed, 9 Jul 2025 08:56:19 +0000 Subject: [PATCH] feat(allocator): add `Allocator::alloc_bytes_start` method (#12083) Add an unsafe `Allocator::alloc_bytes_start` method. This method allows allocating space at the start of an `Allocator`'s data chunk (instead of at the end as usually happens). We'll use this in linter JS plugins prototype for writing source text into start of arena, which is what raw transfer needs. --- crates/oxc_allocator/src/from_raw_parts.rs | 167 ++++++++++++++++++++- 1 file changed, 166 insertions(+), 1 deletion(-) diff --git a/crates/oxc_allocator/src/from_raw_parts.rs b/crates/oxc_allocator/src/from_raw_parts.rs index 6ea8b3d9e73ba..dfbb696626d33 100644 --- a/crates/oxc_allocator/src/from_raw_parts.rs +++ b/crates/oxc_allocator/src/from_raw_parts.rs @@ -1,4 +1,9 @@ -//! Define additional [`Allocator::from_raw_parts`] method, used only by raw transfer. +//! Define additional methods, used only by raw transfer: +//! +//! * [`Allocator::from_raw_parts`] +//! * [`Allocator::alloc_bytes_start`] +//! * [`Allocator::data_ptr`] +//! * [`Allocator::set_data_ptr`] use std::{ alloc::Layout, @@ -9,6 +14,8 @@ use std::{ use bumpalo::Bump; +use oxc_data_structures::pointer_ext::PointerExt; + use crate::Allocator; /// Minimum alignment for allocator chunks. This is hard-coded on `bumpalo`. @@ -117,6 +124,164 @@ impl Allocator { Self::from_bump(bump) } + + /// Allocate space for `bytes` bytes at start of [`Allocator`]'s current chunk. + /// + /// Returns a pointer to the start of an uninitialized section of `bytes` bytes. + /// + /// Note: [`alloc_layout`] allocates at *end* of the current chunk, because `bumpalo` bumps downwards, + /// hence the need for this method, to allocate at *start* of current chunk. + /// + /// This method is dangerous, and should not ordinarily be used. + /// + /// This method moves the pointer to start of the current chunk forwards, so it no longer correctly + /// describes the start of the allocation obtained from system allocator. + /// + /// The `Allocator` **must not be allowed to be dropped** or it would be UB. + /// Only use this method if you prevent that possibililty. e.g.: + /// + /// 1. Set the data pointer back to its correct value before it is dropped, using [`set_data_ptr`]. + /// 2. Wrap the `Allocator` in `ManuallyDrop`, and taking care of deallocating it manually + /// with the correct pointer. + /// + /// # Panics + /// + /// Panics if insufficient capacity for `bytes` + /// (after rounding up to nearest multiple of [`RAW_MIN_ALIGN`]). + /// + /// # SAFETY + /// + /// `Allocator` must not be dropped after calling this method (see above). + /// + /// [`alloc_layout`]: Self::alloc_layout + /// [`set_data_ptr`]: Self::set_data_ptr + /// [`RAW_MIN_ALIGN`]: Self::RAW_MIN_ALIGN + pub unsafe fn alloc_bytes_start(&self, bytes: usize) -> NonNull { + // Round up number of bytes to reserve to multiple of `MIN_ALIGN`, + // so data pointer remains aligned on `MIN_ALIGN` + let alloc_bytes = (bytes + MIN_ALIGN - 1) & !(MIN_ALIGN - 1); + + let data_ptr = self.data_ptr(); + let cursor_ptr = self.cursor_ptr(); + // SAFETY: Cursor pointer is always `>=` data pointer. + // Both pointers are within same allocation, and derived from the same original pointer. + let free_capacity = unsafe { cursor_ptr.offset_from_usize(data_ptr) }; + + // Check sufficient capacity to write `alloc_bytes` bytes, without overwriting data already + // stored in allocator. + // Could use `>=` here and it would be sufficient capacity, but use `>` instead so this assertion + // fails if current chunk is the empty chunk and `bytes` is 0. + assert!(free_capacity > alloc_bytes); + + // Calculate new data pointer. + // SAFETY: We checked above that distance between data pointer and cursor is `>= alloc_bytes`, + // so moving data pointer forwards by `alloc_bytes` cannot place it after cursor pointer. + let new_data_ptr = unsafe { data_ptr.add(alloc_bytes) }; + + // Set new data pointer. + // SAFETY: `Allocator` must have at least 1 allocated chunk or check for sufficient capacity + // above would have failed. + // Data pointer is always aligned on `MIN_ALIGN`, and we rounded `alloc_bytes` up to a multiple + // of `MIN_ALIGN`, so that remains the case. + unsafe { self.set_data_ptr(new_data_ptr) }; + + // Return original data pointer + data_ptr + } + + /// Get data pointer for this [`Allocator`]'s current chunk. + pub fn data_ptr(&self) -> NonNull { + // SAFETY: We don't take any action with the `Allocator` while the `&ChunkFooter` + // reference is alive + let chunk_footer = unsafe { self.chunk_footer() }; + chunk_footer.data + } + + /// Set data pointer for this [`Allocator`]'s current chunk. + /// + /// This is dangerous, and this method should not ordinarily be used. + /// It is only here for manually writing data to start of the allocator chunk, + /// and then adjusting the start pointer to after it. + /// + /// # SAFETY + /// + /// * Allocator must have at least 1 allocated chunk. + /// It is UB to call this method on an `Allocator` which has not allocated + /// i.e. fresh from `Allocator::new`. + /// * `ptr` must point to within the allocation underlying this allocator. + /// * `ptr` must be aligned on [`RAW_MIN_ALIGN`]. + /// + /// [`RAW_MIN_ALIGN`]: Self::RAW_MIN_ALIGN + pub unsafe fn set_data_ptr(&self, ptr: NonNull) { + debug_assert!(is_multiple_of(ptr.as_ptr() as usize, MIN_ALIGN)); + + // SAFETY: Caller guarantees `Allocator` has at least 1 allocated chunk. + // We don't take any action with the `Allocator` while the `&mut ChunkFooter` reference + // is alive, beyond setting the data pointer. + let chunk_footer = unsafe { self.chunk_footer_mut() }; + chunk_footer.data = ptr; + } + + /// Get cursor pointer for this [`Allocator`]'s current chunk. + fn cursor_ptr(&self) -> NonNull { + // SAFETY: We don't take any action with the `Allocator` while the `&ChunkFooter` + // reference is alive + let chunk_footer = unsafe { self.chunk_footer() }; + chunk_footer.ptr.get() + } + + /// Get reference to current [`ChunkFooter`]. + /// + /// # SAFETY + /// + /// Caller must not allocate into this `Allocator`, or perform any other action which would create + /// a `&mut ChunkFooter` reference, while the `&ChunkFooter` reference returned by this method is alive. + unsafe fn chunk_footer(&self) -> &ChunkFooter { + let chunk_footer_ptr = self.chunk_footer_ptr(); + // SAFETY: Caller promises not to take any other action which would generate a mutable reference + // to the `ChunkFooter` while this reference is alive. + unsafe { chunk_footer_ptr.as_ref() } + } + + /// Get mutable reference to current [`ChunkFooter`]. + /// + /// It would be safer if this method took a `&mut self`, but that would preclude using this method + /// while any references to data in the arena exist, which is too restrictive. + /// So we just need to be careful how we use this method. + /// + /// # SAFETY + /// + /// * Allocator must have at least 1 allocated chunk. + /// It is UB to call this method on an `Allocator` which has not allocated + /// i.e. fresh from `Allocator::new`. + /// * Caller must not allocate into this `Allocator`, or perform any other action which would + /// read or alter the `ChunkFooter`, or create another reference to it, while the `&mut ChunkFooter` + /// reference returned by this method is alive. + #[expect(clippy::mut_from_ref)] + unsafe fn chunk_footer_mut(&self) -> &mut ChunkFooter { + let mut chunk_footer_ptr = self.chunk_footer_ptr(); + // SAFETY: Caller guarantees `Allocator` has an allocated chunk, so this isn't `EmptyChunkFooter`, + // which it'd be UB to obtain a mutable reference to. + // Caller promises not to take any other action which would generate another reference to the + // `ChunkFooter` while this reference is alive. + unsafe { chunk_footer_ptr.as_mut() } + } + + /// Get pointer to current chunk's [`ChunkFooter`]. + fn chunk_footer_ptr(&self) -> NonNull { + let current_chunk_footer_field_offset = get_current_chunk_footer_field_offset(); + + // Get pointer to current `ChunkFooter`. + // SAFETY: We've established the offset of the `current_chunk_footer` field above. + let current_chunk_footer_field = unsafe { + let bump = self.bump(); + let field_ptr = ptr::from_ref(bump) + .cast::>>() + .add(current_chunk_footer_field_offset); + &*field_ptr + }; + current_chunk_footer_field.get() + } } /// Allocator chunk footer.