refactor(allocator): move allocation tracking code into own module

crates/oxc_allocator/README.md

@@ -27,3 +27,5 @@ This approach is significantly faster than using the system allocator for AST op
 
 - `serialize` - Enables serialization support for `Box` and `Vec` with `serde`
 - `from_raw_parts` - Adds unsafe `from_raw_parts` method (not recommended for general use)
+- `track_allocations` - For internal use only. The APIs provided by this feature are sketchy at best,
+  and possibly undefined behavior. Do not enable this feature under any circumstances in production code.

crates/oxc_allocator/src/alloc.rs

@@ -95,17 +95,11 @@ impl Alloc for Bump {
     /// Panics if reserving space for `layout` fails.
     #[inline(always)]
     fn alloc(&self, layout: Layout) -> NonNull<u8> {
-        // SAFETY: We only use `Bump` inside of `Allocator` in oxc, so the `self` reference should
-        // also be pointing to a valid `Allocator` struct, which we can use for finding the stats fields.
-        // This will go away when we add a custom allocator to oxc.
+        // SAFETY: This is UNSOUND (see comment on `get_stats_ref`). But usage is gated behind
+        // `track_allocations` feature, so should never be compiled in production code.
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
         unsafe {
-            use crate::allocator::NUM_ALLOC_FIELD_OFFSET;
-            use std::sync::atomic::{AtomicUsize, Ordering};
-            let num_alloc_ptr =
-                std::ptr::from_ref(self).byte_offset(NUM_ALLOC_FIELD_OFFSET).cast::<AtomicUsize>();
-            let num_alloc = num_alloc_ptr.as_ref().unwrap_unchecked();
-            num_alloc.fetch_add(1, Ordering::SeqCst);
+            crate::tracking::get_stats_ref(self).record_allocation();
         }
 
         self.alloc_layout(layout)
@@ -146,18 +140,11 @@ impl Alloc for Bump {
     /// Panics / aborts if reserving space for `new_layout` fails.
     #[inline(always)]
     unsafe fn grow(&self, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout) -> NonNull<u8> {
-        // SAFETY: We only use `Bump` inside of `Allocator` in oxc, so the `self` reference should
-        // also be pointing to a valid `Allocator` struct, which we can use for finding the stats fields.
-        // This will go away when we add a custom allocator to oxc.
+        // SAFETY: This is UNSOUND (see comment on `get_stats_ref`). But usage is gated behind
+        // `track_allocations` feature, so should never be compiled in production code.
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
         unsafe {
-            use crate::allocator::NUM_REALLOC_FIELD_OFFSET;
-            use std::sync::atomic::{AtomicUsize, Ordering};
-            let num_realloc_ptr = std::ptr::from_ref(self)
-                .byte_offset(NUM_REALLOC_FIELD_OFFSET)
-                .cast::<AtomicUsize>();
-            let num_realloc = num_realloc_ptr.as_ref().unwrap_unchecked();
-            num_realloc.fetch_add(1, Ordering::SeqCst);
+            crate::tracking::get_stats_ref(self).record_reallocation();
         }
 
         // SAFETY: Safety requirements of `Allocator::grow` are the same as for this method

crates/oxc_allocator/src/allocator.rs

@@ -11,6 +11,9 @@ use bumpalo::Bump;
 
 use oxc_data_structures::assert_unchecked;
 
+#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+use crate::tracking::AllocationStats;
+
 /// A bump-allocated memory arena.
 ///
 /// # Anatomy of an Allocator
@@ -218,25 +221,17 @@ use oxc_data_structures::assert_unchecked;
 #[derive(Default)]
 pub struct Allocator {
     bump: Bump,
-    /// Used to track the total number of allocations made in this allocator when the `track_allocations` feature is enabled.
-    #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-    #[doc(hidden)]
-    pub num_alloc: std::sync::atomic::AtomicUsize,
-    /// Used to track the total number of re-allocations made in this allocator when the `track_allocations` feature is enabled.
+    /// Used to track number of allocations made in this allocator when `track_allocations` feature is enabled
     #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-    #[doc(hidden)]
-    pub num_realloc: std::sync::atomic::AtomicUsize,
+    pub(crate) stats: AllocationStats,
 }
 
-// Consts used in `Alloc` trait for allocation tracking
-#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-#[expect(clippy::cast_possible_wrap)]
-pub const NUM_ALLOC_FIELD_OFFSET: isize =
-    (offset_of!(Allocator, num_alloc) as isize) - (offset_of!(Allocator, bump) as isize);
+/// Offset of `stats` field, relative to `bump` field.
+/// Used in `tracking` module for allocation tracking.
 #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
 #[expect(clippy::cast_possible_wrap)]
-pub const NUM_REALLOC_FIELD_OFFSET: isize =
-    (offset_of!(Allocator, num_realloc) as isize) - (offset_of!(Allocator, bump) as isize);
+pub const STATS_FIELD_OFFSET: isize =
+    (offset_of!(Allocator, stats) as isize) - (offset_of!(Allocator, bump) as isize);
 
 impl Allocator {
     /// Create a new [`Allocator`] with no initial capacity.
@@ -265,9 +260,7 @@ impl Allocator {
         Self {
             bump: Bump::new(),
             #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-            num_alloc: std::sync::atomic::AtomicUsize::new(0),
-            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-            num_realloc: std::sync::atomic::AtomicUsize::new(0),
+            stats: AllocationStats::default(),
         }
     }
 
@@ -282,9 +275,7 @@ impl Allocator {
         Self {
             bump: Bump::with_capacity(capacity),
             #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-            num_alloc: std::sync::atomic::AtomicUsize::new(0),
-            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-            num_realloc: std::sync::atomic::AtomicUsize::new(0),
+            stats: AllocationStats::default(),
         }
     }
 
@@ -310,7 +301,7 @@ impl Allocator {
         const { assert!(!std::mem::needs_drop::<T>(), "Cannot allocate Drop type in arena") };
 
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        self.stats.record_allocation();
 
         self.bump.alloc(val)
     }
@@ -334,7 +325,7 @@ impl Allocator {
     #[inline(always)]
     pub fn alloc_str<'alloc>(&'alloc self, src: &str) -> &'alloc str {
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        self.stats.record_allocation();
 
         self.bump.alloc_str(src)
     }
@@ -357,7 +348,7 @@ impl Allocator {
     #[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.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        self.stats.record_allocation();
 
         self.bump.alloc_slice_copy(src)
     }
@@ -372,7 +363,7 @@ impl Allocator {
     /// Panics if reserving space matching `layout` fails.
     pub fn alloc_layout(&self, layout: Layout) -> NonNull<u8> {
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        self.stats.record_allocation();
 
         self.bump.alloc_layout(layout)
     }
@@ -425,7 +416,7 @@ impl Allocator {
         );
 
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        self.stats.record_allocation();
 
         // Create actual `&str` in a separate function, to ensure that `alloc_concat_strs_array`
         // is inlined, so that compiler has knowledge to remove the overflow checks above.
@@ -520,10 +511,7 @@ impl Allocator {
     #[inline(always)]
     pub fn reset(&mut self) {
         #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-        {
-            self.num_alloc.store(0, std::sync::atomic::Ordering::SeqCst);
-            self.num_realloc.store(0, std::sync::atomic::Ordering::SeqCst);
-        }
+        self.stats.reset();
 
         self.bump.reset();
     }
@@ -649,9 +637,7 @@ impl Allocator {
         Self {
             bump,
             #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-            num_alloc: std::sync::atomic::AtomicUsize::new(0),
-            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
-            num_realloc: std::sync::atomic::AtomicUsize::new(0),
+            stats: AllocationStats::default(),
         }
     }
 }

crates/oxc_allocator/src/lib.rs

@@ -42,6 +42,8 @@ mod from_raw_parts;
 pub mod hash_map;
 mod string_builder;
 mod take_in;
+#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+mod tracking;
 mod vec;
 mod vec2;
 

crates/oxc_allocator/src/tracking.rs

@@ -0,0 +1,94 @@
+//! Allocation tracking.
+//!
+//! This module is only loaded when `track_allocations` feature is enabled.
+//! This feature is only used in `tasks/track_memory_allocations`.
+//!
+//! Current implementation is unsound - see comment on [`get_stats_ref`] below.
+//! It's OK to use it in our internal `tasks/track_memory_allocations` tool,
+//! but we must take great care that this is NEVER enabled in any other circumstances.
+//!
+//! Even without the unsoundness, we don't want this enabled outside of `tasks/track_memory_allocations`,
+//! as it imposes a performance cost on making allocations.
+//!
+//! The 2nd cargo feature `disable_track_allocations` is to ensure that compiling with `--all-features`
+//! will not load this module.
+//!
+//! As soon as we replace `bumpalo` with our own arena allocator, we'll remove the hack from `get_stats_ref`,
+//! and make this sound.
+
+use std::{
+    ptr,
+    sync::atomic::{AtomicUsize, Ordering::SeqCst},
+};
+
+use bumpalo::Bump;
+
+use crate::{Allocator, allocator::STATS_FIELD_OFFSET};
+
+/// Counters of allocations and reallocations made in an [`Allocator`].
+//
+// Note: These fields could be `Cell<usize>` instead of `AtomicUsize`, because `Allocator` should not
+// be `Sync`. But currently it is (which is unsound!) because of other terrible hacks.
+#[derive(Default)]
+pub struct AllocationStats {
+    /// Number of allocations
+    num_alloc: AtomicUsize,
+    /// Number of reallocations
+    num_realloc: AtomicUsize,
+}
+
+impl AllocationStats {
+    /// Record that an allocation was made.
+    pub(crate) fn record_allocation(&self) {
+        self.num_alloc.fetch_add(1, SeqCst);
+    }
+
+    /// Record that a reallocation was made.
+    pub(crate) fn record_reallocation(&self) {
+        self.num_realloc.fetch_add(1, SeqCst);
+    }
+
+    /// Reset allocation counters.
+    pub(crate) fn reset(&self) {
+        self.num_alloc.store(0, SeqCst);
+        self.num_realloc.store(0, SeqCst);
+    }
+}
+
+impl Allocator {
+    /// Get number of allocations and reallocations made in this [`Allocator`].
+    #[doc(hidden)]
+    pub fn get_allocation_stats(&self) -> (usize, usize) {
+        let num_alloc = self.stats.num_alloc.load(SeqCst);
+        let num_realloc = self.stats.num_realloc.load(SeqCst);
+        (num_alloc, num_realloc)
+    }
+}
+
+/// Get reference to [`AllocationStats`] for a [`Bump`].
+///
+/// # SAFETY
+///
+/// Caller must guarantee that the `Bump` provided to this function is wrapped in an [`Allocator`].
+///
+/// In Oxc, we never use `Bump` alone, without it being wrapped in an `Allocator`.
+/// However, we have no static guarantee of this relationship between `Bump` and `Allocator`,
+/// so it's usually impossible for callers to proveably satisfy the safety requirements of this method.
+///
+/// Even if the `Bump` *is* wrapped in an `Allocator`, this may still be UB, as we project beyond
+/// the bounds of the `&Bump`. Certainly stacked borrows memory model says this is UB, though it's unclear
+/// to me (@overlookmotel) whether stacked borrows is unnecessarily strict on this point.
+/// <https://github.com/rust-lang/unsafe-code-guidelines/issues/134>
+///
+/// This function (and the `track_allocations` feature in general) must only be used for internal tools,
+/// and must NEVER be compiled in production code.
+pub unsafe fn get_stats_ref(bump: &Bump) -> &AllocationStats {
+    // We assume the `Bump` is wrapped in an `Allocator`. We can therefore get a pointer to the `stats`
+    // field of `Allocator` from the memory location of the `Bump`.
+    // SAFETY: This is UNSOUND. See above.
+    unsafe {
+        let stats_ptr =
+            ptr::from_ref(bump).byte_offset(STATS_FIELD_OFFSET).cast::<AllocationStats>();
+        stats_ptr.as_ref().unwrap_unchecked()
+    }
+}

tasks/track_memory_allocations/src/lib.rs

@@ -119,13 +119,9 @@ pub fn run() -> Result<(), io::Error> {
         let sys_allocs = NUM_ALLOC.load(SeqCst);
         let sys_reallocs = NUM_REALLOC.load(SeqCst);
         #[cfg(not(feature = "is_all_features"))]
-        let arena_allocs = allocator.num_alloc.load(SeqCst);
+        let (arena_allocs, arena_reallocs) = allocator.get_allocation_stats();
         #[cfg(feature = "is_all_features")]
-        let arena_allocs = 0;
-        #[cfg(not(feature = "is_all_features"))]
-        let arena_reallocs = allocator.num_realloc.load(SeqCst);
-        #[cfg(feature = "is_all_features")]
-        let arena_reallocs = 0;
+        let (arena_allocs, arena_reallocs) = (0, 0);
         let arena_bytes = allocator.used_bytes();
 
         let s = format!(