feat(allocator): AllocatorPool store IDs in Allocators

crates/oxc_allocator/src/generated/assert_layouts.rs

@@ -9,17 +9,19 @@ use crate::*;
 
 #[cfg(target_pointer_width = "64")]
 const _: () = {
-    // Padding: 0 bytes
-    assert!(size_of::<FixedSizeAllocatorMetadata>() == 8);
+    // Padding: 4 bytes
+    assert!(size_of::<FixedSizeAllocatorMetadata>() == 16);
     assert!(align_of::<FixedSizeAllocatorMetadata>() == 8);
+    assert!(offset_of!(FixedSizeAllocatorMetadata, id) == 8);
     assert!(offset_of!(FixedSizeAllocatorMetadata, alloc_ptr) == 0);
 };
 
 #[cfg(target_pointer_width = "32")]
 const _: () = {
     // Padding: 0 bytes
-    assert!(size_of::<FixedSizeAllocatorMetadata>() == 4);
+    assert!(size_of::<FixedSizeAllocatorMetadata>() == 8);
     assert!(align_of::<FixedSizeAllocatorMetadata>() == 4);
+    assert!(offset_of!(FixedSizeAllocatorMetadata, id) == 4);
     assert!(offset_of!(FixedSizeAllocatorMetadata, alloc_ptr) == 0);
 };
 

crates/oxc_allocator/src/pool_fixed_size.rs

@@ -3,7 +3,10 @@ use std::{
     mem::ManuallyDrop,
     ops::Deref,
     ptr::NonNull,
-    sync::Mutex,
+    sync::{
+        Mutex,
+        atomic::{AtomicU32, Ordering},
+    },
 };
 
 use oxc_ast_macros::ast;
@@ -21,15 +24,18 @@ const FOUR_GIB: usize = 1 << 32;
 /// Internally uses a `Vec` protected by a `Mutex` to store available allocators.
 #[derive(Default)]
 pub struct AllocatorPool {
+    /// Allocators in the pool
     allocators: Mutex<Vec<FixedSizeAllocator>>,
+    /// ID to assign to next `Allocator` that's created
+    next_id: AtomicU32,
 }
 
 impl AllocatorPool {
     /// Creates a new [`AllocatorPool`] with capacity for the given number of `FixedSizeAllocator` instances.
     pub fn new(size: usize) -> AllocatorPool {
         // Each allocator consumes a large block of memory, so create them on demand instead of upfront
         let allocators = Vec::with_capacity(size);
-        AllocatorPool { allocators: Mutex::new(allocators) }
+        AllocatorPool { allocators: Mutex::new(allocators), next_id: AtomicU32::new(0) }
     }
 
     /// Retrieves an [`Allocator`] from the pool, or creates a new one if the pool is empty.
@@ -44,7 +50,16 @@ impl AllocatorPool {
             let mut allocators = self.allocators.lock().unwrap();
             allocators.pop()
         };
-        let allocator = allocator.unwrap_or_else(FixedSizeAllocator::new);
+
+        let allocator = allocator.unwrap_or_else(|| {
+            // Each allocator needs to have a unique ID, but the order those IDs are assigned in
+            // doesn't matter, so `Ordering::Relaxed` is fine
+            let id = self.next_id.fetch_add(1, Ordering::Relaxed);
+            // Protect against IDs wrapping around.
+            // TODO: Does this work? Do we need it anyway?
+            assert!(id < u32::MAX, "Created too many allocators");
+            FixedSizeAllocator::new(id)
+        });
 
         AllocatorGuard { allocator: ManuallyDrop::new(allocator), pool: self }
     }
@@ -94,6 +109,8 @@ impl Drop for AllocatorGuard<'_> {
 /// which is after the section of the allocation which [`Allocator`] uses for its chunk.
 #[ast]
 pub struct FixedSizeAllocatorMetadata {
+    /// ID of this allocator
+    pub id: u32,
     /// Pointer to start of original allocation backing the `FixedSizeAllocator`
     pub alloc_ptr: NonNull<u8>,
 }
@@ -179,7 +196,7 @@ pub struct FixedSizeAllocator {
 impl FixedSizeAllocator {
     /// Create a new [`FixedSizeAllocator`].
     #[expect(clippy::items_after_statements)]
-    pub fn new() -> Self {
+    pub fn new(id: u32) -> Self {
         // Allocate block of memory.
         // SAFETY: `ALLOC_LAYOUT` does not have zero size.
         let alloc_ptr = unsafe { System.alloc(ALLOC_LAYOUT) };
@@ -219,7 +236,7 @@ impl FixedSizeAllocator {
 
         // Write `FixedSizeAllocatorMetadata` to after space reserved for `RawTransferMetadata`,
         // which is after the end of the allocator chunk
-        let metadata = FixedSizeAllocatorMetadata { alloc_ptr };
+        let metadata = FixedSizeAllocatorMetadata { alloc_ptr, id };
         // SAFETY: `FIXED_METADATA_OFFSET` is `FIXED_METADATA_SIZE_ROUNDED` bytes before end of
         // the allocation, so there's space for `FixedSizeAllocatorMetadata`.
         // It's sufficiently aligned for `FixedSizeAllocatorMetadata`.

crates/oxc_ast_macros/src/generated/structs.rs

@@ -254,7 +254,7 @@ pub static STRUCTS: phf::Map<&'static str, StructDetails> = ::phf::Map {
         ),
         ("Pattern", StructDetails { field_order: None }),
         ("EmptyStatement", StructDetails { field_order: None }),
-        ("FixedSizeAllocatorMetadata", StructDetails { field_order: None }),
+        ("FixedSizeAllocatorMetadata", StructDetails { field_order: Some(&[1, 0]) }),
         ("Hashbang", StructDetails { field_order: None }),
         ("UnaryExpression", StructDetails { field_order: Some(&[0, 2, 1]) }),
         ("ThrowStatement", StructDetails { field_order: None }),

napi/parser/generated/deserialize/js.js

@@ -4020,6 +4020,10 @@ function deserializeErrorSeverity(pos) {
   }
 }
 
+function deserializeU32(pos) {
+  return uint32[pos >> 2];
+}
+
 function deserializeU8(pos) {
   return uint8[pos];
 }
@@ -5234,10 +5238,6 @@ function deserializeBoxTSExternalModuleReference(pos) {
   return deserializeTSExternalModuleReference(uint32[pos >> 2]);
 }
 
-function deserializeU32(pos) {
-  return uint32[pos >> 2];
-}
-
 function deserializeOptionNameSpan(pos) {
   if (uint32[(pos + 8) >> 2] === 0 && uint32[(pos + 12) >> 2] === 0) return null;
   return deserializeNameSpan(pos);

napi/parser/generated/deserialize/ts.js

@@ -4151,6 +4151,10 @@ function deserializeErrorSeverity(pos) {
   }
 }
 
+function deserializeU32(pos) {
+  return uint32[pos >> 2];
+}
+
 function deserializeU8(pos) {
   return uint8[pos];
 }
@@ -5365,10 +5369,6 @@ function deserializeBoxTSExternalModuleReference(pos) {
   return deserializeTSExternalModuleReference(uint32[pos >> 2]);
 }
 
-function deserializeU32(pos) {
-  return uint32[pos >> 2];
-}
-
 function deserializeOptionNameSpan(pos) {
   if (uint32[(pos + 8) >> 2] === 0 && uint32[(pos + 12) >> 2] === 0) return null;
   return deserializeNameSpan(pos);

napi/parser/generated/lazy/constructors.js

@@ -12472,6 +12472,10 @@ class StaticExport {
 
 const DebugStaticExport = class StaticExport {};
 
+function constructU32(pos, ast) {
+  return ast.buffer.uint32[pos >> 2];
+}
+
 function constructU8(pos, ast) {
   return ast.buffer[pos];
 }
@@ -13748,10 +13752,6 @@ function constructBoxTSExternalModuleReference(pos, ast) {
   return new TSExternalModuleReference(ast.buffer.uint32[pos >> 2], ast);
 }
 
-function constructU32(pos, ast) {
-  return ast.buffer.uint32[pos >> 2];
-}
-
 function constructOptionNameSpan(pos, ast) {
   if (ast.buffer.uint32[(pos + 8) >> 2] === 0 && ast.buffer.uint32[(pos + 12) >> 2] === 0) return null;
   return new NameSpan(pos, ast);