refactor(transformer): move common stack functionality into StackCommon trait

crates/oxc_transformer/src/helpers/stack/common.rs

@@ -0,0 +1,281 @@
+#![expect(clippy::unnecessary_safety_comment)]
+
+use std::{
+    alloc::{self, Layout},
+    mem::{align_of, size_of},
+    ptr::{self, NonNull},
+};
+
+use assert_unchecked::assert_unchecked;
+
+use super::StackCapacity;
+
+pub trait StackCommon<T>: StackCapacity<T> {
+    // Getter setter methods defined by implementer
+    fn start(&self) -> NonNull<T>;
+    fn end(&self) -> NonNull<T>;
+    fn cursor(&self) -> NonNull<T>;
+    fn set_start(&mut self, start: NonNull<T>);
+    fn set_end(&mut self, end: NonNull<T>);
+    fn set_cursor(&mut self, cursor: NonNull<T>);
+
+    /// Make allocation of `capacity_bytes` bytes, aligned for `T`.
+    ///
+    /// # Panics
+    /// Panics if out of memory (or may abort, depending on global allocator's behavior).
+    ///
+    /// # SAFETY
+    /// * `capacity_bytes` must not be 0.
+    /// * `capacity_bytes` must be a multiple of `mem::size_of::<T>()`.
+    /// * `capacity_bytes` must not exceed `Self::MAX_CAPACITY_BYTES`.
+    #[inline]
+    unsafe fn allocate(capacity_bytes: usize) -> (NonNull<T>, NonNull<T>) {
+        // SAFETY: Caller guarantees `capacity_bytes` satisfies requirements
+        let layout = Self::layout_for(capacity_bytes);
+        let (start, end) = allocate(layout);
+
+        // SAFETY: `start` and `end` are `NonNull` - just casting them
+        let start = NonNull::new_unchecked(start.as_ptr().cast::<T>());
+        let end = NonNull::new_unchecked(end.as_ptr().cast::<T>());
+
+        (start, end)
+    }
+
+    /// Grow allocation.
+    ///
+    /// `start` and `end` are set to the start and end of new allocation.
+    /// `current` is set so distance from `start` is old `capacity_bytes`.
+    /// This is where it should be if stack was previously full to capacity.
+    ///
+    /// # Panics
+    /// Panics if stack is already at maximum capacity.
+    ///
+    /// # SAFETY
+    /// Stack must have already allocated. i.e. `start` is not a dangling pointer.
+    #[inline]
+    unsafe fn grow(&mut self) {
+        // Grow allocation.
+        // SAFETY: Caller guarantees stack has allocated.
+        // `start` and `end` are boundaries of that allocation (`alloc` and `grow` ensure that).
+        // So `old_start_ptr` and `old_layout` accurately describe the current allocation.
+        // `grow` creates new allocation with byte size double what it currently is, or caps it
+        // at `MAX_CAPACITY_BYTES`.
+        // Old capacity in bytes was a multiple of `size_of::<T>()`, so double that must be too.
+        // `MAX_CAPACITY_BYTES` is also a multiple of `size_of::<T>()`.
+        // So new capacity in bytes must be a multiple of `size_of::<T>()`.
+        // `MAX_CAPACITY_BYTES <= isize::MAX`.
+        let old_start_ptr = NonNull::new_unchecked(self.start().as_ptr().cast::<u8>());
+        let old_layout = Self::layout_for(self.capacity_bytes());
+        let (start, end, current) = grow(old_start_ptr, old_layout, Self::MAX_CAPACITY_BYTES);
+
+        // Update pointers.
+        // SAFETY: `start`, `end`, and `current` are all `NonNull` - just casting them.
+        // All pointers returned from `grow` are aligned for `T`.
+        // Old capacity and new capacity in bytes are both multiples of `size_of::<T>()`,
+        // so distances `end - start` and `current - start` are both multiples of `size_of::<T>()`.
+        self.set_start(NonNull::new_unchecked(start.as_ptr().cast::<T>()));
+        self.set_end(NonNull::new_unchecked(end.as_ptr().cast::<T>()));
+        self.set_cursor(NonNull::new_unchecked(current.as_ptr().cast::<T>()));
+    }
+
+    /// Deallocate stack memory.
+    ///
+    /// Note: Does *not* drop the contents of the stack (the `T`s), only the memory allocated
+    /// by `allocate` / `grow`. If stack is not empty, also call `drop_contents()` before calling this.
+    ///
+    /// # SAFETY
+    /// Stack must have already allocated. i.e. `start` is not a dangling pointer.
+    #[inline]
+    unsafe fn deallocate(&self) {
+        // SAFETY: Caller guarantees stack is allocated.
+        // `start` and `end` are boundaries of that allocation (`allocate` and `grow` ensure that).
+        // So `start` and `layout` accurately describe the current allocation.
+        let layout = Self::layout_for(self.capacity_bytes());
+        alloc::dealloc(self.start().as_ptr().cast::<u8>(), layout);
+    }
+
+    /// Drop contents of stack.
+    ///
+    /// This function will be optimized out if `T` is non-drop, as `drop_in_place` calls
+    /// `std::mem::needs_drop` internally and is a no-op if it returns true.
+    ///
+    /// # SAFETY
+    /// * Stack must be allocated.
+    /// * Stack must contain `len` initialized entries, starting at `self.start()`.
+    #[inline]
+    unsafe fn drop_contents(&self, len: usize) {
+        // Drop contents. Next line copied from `std`'s `Vec`.
+        // SAFETY: Caller guarantees stack contains `len` initialized entries, starting at `start`.
+        ptr::drop_in_place(ptr::slice_from_raw_parts_mut(self.start().as_ptr(), len));
+    }
+
+    /// Get layout for allocation of `capacity_bytes` bytes.
+    ///
+    /// # SAFETY
+    /// * `capacity_bytes` must not be 0.
+    /// * `capacity_bytes` must be a multiple of `mem::size_of::<T>()`.
+    /// * `capacity_bytes` must not exceed `Self::MAX_CAPACITY_BYTES`.
+    #[inline]
+    unsafe fn layout_for(capacity_bytes: usize) -> Layout {
+        // `capacity_bytes` must not be 0 because cannot make 0-size allocations.
+        debug_assert!(capacity_bytes > 0);
+        // `capacity_bytes` must be a multiple of `size_of::<T>()` so that `cursor == end`
+        // checks in `push` methods accurately detects when full to capacity
+        debug_assert!(capacity_bytes % size_of::<T>() == 0);
+        // `capacity_bytes` must not exceed `Self::MAX_CAPACITY_BYTES` to prevent creating an allocation
+        // of illegal size
+        debug_assert!(capacity_bytes <= Self::MAX_CAPACITY_BYTES);
+
+        // SAFETY: `align_of::<T>()` trivially satisfies alignment requirements.
+        // Caller guarantees `capacity_bytes <= MAX_CAPACITY_BYTES`.
+        // `MAX_CAPACITY_BYTES` takes into account the rounding-up by alignment requirement.
+        Layout::from_size_align_unchecked(capacity_bytes, align_of::<T>())
+    }
+
+    /// Get offset of `cursor` in number of `T`s.
+    ///
+    /// # SAFETY
+    /// * `self.cursor()` and `self.start()` must be derived from same pointer.
+    /// * `self.cursor()` must be `>= self.start()`.
+    /// * Byte distance between `self.cursor()` and `self.start()` must be a multiple of `size_of::<T>()`.
+    unsafe fn cursor_offset(&self) -> usize {
+        // `offset_from` returns offset in units of `T`.
+        // SAFETY: Caller guarantees `cursor` and `start` are derived from same pointer.
+        // This implies that both pointers are always within bounds of a single allocation.
+        // Caller guarantees `cursor >= start`.
+        // Caller guarantees distance between pointers is a multiple of `size_of::<T>()`.
+        // `assert_unchecked!` is to help compiler to optimize.
+        // See: https://doc.rust-lang.org/std/primitive.pointer.html#method.sub_ptr
+        #[expect(clippy::cast_sign_loss)]
+        unsafe {
+            assert_unchecked!(self.cursor() >= self.start());
+            self.cursor().as_ptr().offset_from(self.start().as_ptr()) as usize
+        }
+    }
+
+    /// Get capacity.
+    #[inline]
+    fn capacity(&self) -> usize {
+        // SAFETY: `allocate` and `grow` both ensure:
+        // * `start` and `end` are both derived from same pointer
+        // * `start` and `end` are both within bounds of a single allocation.
+        // * `end` is always >= `start`.
+        // * Distance between `start` and `end` is always a multiple of `size_of::<T>()`.
+        // `assert_unchecked!` is to help compiler to optimize.
+        // See: https://doc.rust-lang.org/std/primitive.pointer.html#method.sub_ptr
+        #[expect(clippy::cast_sign_loss)]
+        unsafe {
+            assert_unchecked!(self.end() >= self.start());
+            self.end().as_ptr().offset_from(self.start().as_ptr()) as usize
+        }
+    }
+
+    /// Get capacity in bytes.
+    #[inline]
+    fn capacity_bytes(&self) -> usize {
+        // SAFETY: `allocate` and `grow` both ensure:
+        // * `start` and `end` are both derived from same pointer
+        // * `start` and `end` are both within bounds of a single allocation.
+        // * `end` is always >= `start`.
+        // * Distance between `start` and `end` is always a multiple of `size_of::<T>()`.
+        // `assert_unchecked!` is to help compiler to optimize.
+        // See: https://doc.rust-lang.org/std/primitive.pointer.html#method.sub_ptr
+        #[expect(clippy::cast_sign_loss)]
+        unsafe {
+            assert_unchecked!(self.end() >= self.start());
+            self.end().as_ptr().byte_offset_from(self.start().as_ptr()) as usize
+        }
+    }
+}
+
+/// Make allocation of with provided layout.
+///
+/// This is a separate non-generic function to improve compile time
+/// (same pattern as `std::vec::Vec` uses).
+///
+/// # Panics
+/// Panics if out of memory (or may abort, depending on global allocator's behavior).
+///
+/// # SAFETY
+/// `layout` must have non-zero size.
+#[inline]
+unsafe fn allocate(layout: Layout) -> (/* start */ NonNull<u8>, /* end */ NonNull<u8>) {
+    // SAFETY: Caller guarantees `layout` has non-zero-size
+    let ptr = alloc::alloc(layout);
+    if ptr.is_null() {
+        alloc::handle_alloc_error(layout);
+    }
+
+    // SAFETY: We checked `ptr` is non-null
+    let start = NonNull::new_unchecked(ptr);
+    // SAFETY: We allocated `layout.size()` bytes, so `end` is end of allocation
+    let end = NonNull::new_unchecked(ptr.add(layout.size()));
+
+    (start, end)
+}
+
+/// Grow existing allocation.
+///
+/// Grow by doubling size, with high bound of `max_capacity_bytes`.
+///
+/// # SAFETY
+/// * `old_start` and `old_layout` must describe an existing allocation.
+/// * `max_capacity_bytes` must be `>= old_layout.size()`.
+/// * `max_capacity_bytes` must be `<= isize::MAX`.
+unsafe fn grow(
+    old_start: NonNull<u8>,
+    old_layout: Layout,
+    max_capacity_bytes: usize,
+) -> (/* start */ NonNull<u8>, /* end */ NonNull<u8>, /* current */ NonNull<u8>) {
+    // Get new capacity
+    // Capacity in bytes cannot be larger than `isize::MAX`, so `* 2` cannot overflow
+    let old_capacity_bytes = old_layout.size();
+    let mut new_capacity_bytes = old_capacity_bytes * 2;
+    if new_capacity_bytes > max_capacity_bytes {
+        assert!(old_capacity_bytes < max_capacity_bytes, "Cannot grow beyond `Self::MAX_CAPACITY`");
+        new_capacity_bytes = max_capacity_bytes;
+    }
+    debug_assert!(new_capacity_bytes > old_capacity_bytes);
+
+    // Reallocate.
+    // SAFETY: Caller guarantees `old_start` and `old_layout` describe an existing allocation.
+    // Caller guarantees that `max_capacity_bytes <= isize::MAX`.
+    // `new_capacity_bytes` is capped above at `max_capacity_bytes`, so is a legal allocation size.
+
+    // `start` and `end` are boundaries of that allocation (`alloc` and `grow` ensure that).
+    // So `start` and `old_layout` accurately describe the current allocation.
+    // `old_capacity_bytes` was a multiple of `size_of::<T>()`, so double that must be too.
+    // `MAX_CAPACITY_BYTES` is also a multiple of `size_of::<T>()`.
+    // So `new_capacity_bytes` must be a multiple of `size_of::<T>()`.
+    // `new_capacity_bytes` is `<= MAX_CAPACITY_BYTES`, so is a legal allocation size.
+    // `layout_for` produces a layout with `T`'s alignment, so `new_ptr` is aligned for `T`.
+    let new_ptr = unsafe {
+        let old_ptr = old_start.as_ptr();
+        let new_ptr = alloc::realloc(old_ptr, old_layout, new_capacity_bytes);
+        if new_ptr.is_null() {
+            let new_layout =
+                Layout::from_size_align_unchecked(old_capacity_bytes, old_layout.align());
+            alloc::handle_alloc_error(new_layout);
+        }
+        new_ptr
+    };
+
+    // Update pointers.
+    //
+    // Stack was full to capacity, so new last index after push is the old capacity.
+    // i.e. `new_cursor - new_start == old_end - old_start`.
+    // Note: All pointers need to be updated even if allocation grew in place.
+    // From docs for `GlobalAlloc::realloc`:
+    // "Any access to the old `ptr` is Undefined Behavior, even if the allocation remained in-place."
+    // <https://doc.rust-lang.org/std/alloc/trait.GlobalAlloc.html#method.realloc>
+    // `end` changes whatever happens, so always need to be updated.
+    // `cursor` needs to be derived from `start` to make `offset_from` valid, so also needs updating.
+    //
+    // SAFETY: We checked that `new_ptr` is non-null.
+    // `old_capacity_bytes < new_capacity_bytes` (ensured above), so `new_cursor` must be in bounds.
+    let new_start = NonNull::new_unchecked(new_ptr);
+    let new_end = NonNull::new_unchecked(new_ptr.add(new_capacity_bytes));
+    let new_cursor = NonNull::new_unchecked(new_ptr.add(old_capacity_bytes));
+
+    (new_start, new_end, new_cursor)
+}

crates/oxc_transformer/src/helpers/stack/mod.rs

@@ -1,9 +1,11 @@
 mod capacity;
+mod common;
 mod non_empty;
 mod sparse;
 mod standard;
 
 use capacity::StackCapacity;
+use common::StackCommon;
 pub use non_empty::NonEmptyStack;
 pub use sparse::SparseStack;
 pub use standard::Stack;