refactor(allocator/pool): single AllocatorPool implementation

crates/oxc_allocator/src/pool/fixed_size.rs

@@ -15,38 +15,32 @@ use crate::{
     generated::fixed_size_constants::{BLOCK_ALIGN, BLOCK_SIZE, RAW_METADATA_SIZE},
 };
 
-use super::AllocatorGuard;
-
 const TWO_GIB: usize = 1 << 31;
 const FOUR_GIB: usize = 1 << 32;
 
-/// A thread-safe pool for reusing [`Allocator`] instances to reduce allocation overhead.
-///
-/// Internally uses a `Vec` protected by a `Mutex` to store available allocators.
-pub struct AllocatorPool {
+/// A thread-safe pool for reusing [`Allocator`] instances, that uses fixed-size allocators,
+/// suitable for use with raw transfer.
+pub struct FixedSizeAllocatorPool {
     /// 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`] for use across the specified number of threads.
-    pub fn new(thread_count: usize) -> AllocatorPool {
+impl FixedSizeAllocatorPool {
+    /// Create a new [`FixedSizeAllocatorPool`] for use across the specified number of threads.
+    pub fn new(thread_count: usize) -> FixedSizeAllocatorPool {
         // Each allocator consumes a large block of memory, so create them on demand instead of upfront,
         // in case not all threads end up being used (e.g. language server without `import` plugin)
         let allocators = Vec::with_capacity(thread_count);
-        AllocatorPool { allocators: Mutex::new(allocators), next_id: AtomicU32::new(0) }
+        FixedSizeAllocatorPool { 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.
-    ///
-    /// Returns an [`AllocatorGuard`] that gives access to the allocator.
+    /// Retrieve an [`Allocator`] from the pool, or create a new one if the pool is empty.
     ///
     /// # Panics
-    ///
     /// Panics if the underlying mutex is poisoned.
-    pub fn get(&self) -> AllocatorGuard<'_> {
+    pub fn get(&self) -> Allocator {
         let fixed_size_allocator = {
             let mut allocators = self.allocators.lock().unwrap();
             allocators.pop()
@@ -69,18 +63,19 @@ impl AllocatorPool {
         let allocator = unsafe {
             mem::transmute::<FixedSizeAllocator, ManuallyDrop<Allocator>>(fixed_size_allocator)
         };
-
-        AllocatorGuard { allocator, pool: self }
+        ManuallyDrop::into_inner(allocator)
     }
 
     /// Add an [`Allocator`] to the pool.
     ///
     /// The `Allocator` is reset by this method, so it's ready to be re-used.
     ///
-    /// # Panics
+    /// # SAFETY
+    /// The `Allocator` must have been created by a `FixedSizeAllocatorPool` (not `StandardAllocatorPool`).
     ///
+    /// # Panics
     /// Panics if the underlying mutex is poisoned.
-    pub(super) fn add(&self, allocator: Allocator) {
+    pub(super) unsafe fn add(&self, allocator: Allocator) {
         let mut fixed_size_allocator =
             FixedSizeAllocator { allocator: ManuallyDrop::new(allocator) };
         fixed_size_allocator.reset();

crates/oxc_allocator/src/pool/mod.rs

@@ -2,6 +2,9 @@ use std::{mem::ManuallyDrop, ops::Deref};
 
 use crate::Allocator;
 
+mod standard;
+use standard::StandardAllocatorPool;
+
 // Fixed size allocators are only supported on 64-bit little-endian platforms at present.
 // They are only enabled if `fixed_size` feature enabled, and `disable_fixed_size` feature is not enabled.
 //
@@ -19,22 +22,122 @@ mod fixed_size;
     target_pointer_width = "64",
     target_endian = "little"
 ))]
-pub use fixed_size::*;
-
-#[cfg(not(all(
+use fixed_size::FixedSizeAllocatorPool;
+#[cfg(all(
     feature = "fixed_size",
     not(feature = "disable_fixed_size"),
     target_pointer_width = "64",
     target_endian = "little"
-)))]
-mod standard;
+))]
+pub use fixed_size::{FixedSizeAllocatorMetadata, free_fixed_size_allocator};
+
+// Dummy implementations of interfaces from `fixed_size`, just to stop clippy complaining.
+// Seems to be necessary due to feature unification.
 #[cfg(not(all(
     feature = "fixed_size",
     not(feature = "disable_fixed_size"),
     target_pointer_width = "64",
     target_endian = "little"
 )))]
-pub use standard::*;
+pub use standard::{FixedSizeAllocatorMetadata, free_fixed_size_allocator};
+
+/// A thread-safe pool for reusing [`Allocator`] instances to reduce allocation overhead.
+///
+/// Uses either a standard or fixed-size allocator pool implementation, depending on Cargo features
+/// and platform support.
+#[repr(transparent)]
+pub struct AllocatorPool(AllocatorPoolInner);
+
+/// Inner type of [`AllocatorPool`], holding either a standard or fixed-size allocator pool.
+enum AllocatorPoolInner {
+    #[cfg_attr(all(feature = "fixed_size", not(feature = "disable_fixed_size")), expect(dead_code))]
+    Standard(StandardAllocatorPool),
+    #[cfg(all(
+        feature = "fixed_size",
+        not(feature = "disable_fixed_size"),
+        target_pointer_width = "64",
+        target_endian = "little"
+    ))]
+    FixedSize(FixedSizeAllocatorPool),
+}
+
+impl AllocatorPool {
+    /// Create a new [`AllocatorPool`] for use across the specified number of threads,
+    /// which uses either standard or fixed-size allocators depending on Cargo features.
+    pub fn new(thread_count: usize) -> AllocatorPool {
+        #[cfg(all(feature = "fixed_size", not(feature = "disable_fixed_size")))]
+        {
+            Self::new_fixed_size(thread_count)
+        }
+
+        #[cfg(not(all(feature = "fixed_size", not(feature = "disable_fixed_size"))))]
+        {
+            Self(AllocatorPoolInner::Standard(StandardAllocatorPool::new(thread_count)))
+        }
+    }
+
+    /// Create a new [`AllocatorPool`] for use across the specified number of threads,
+    /// which uses fixed-size allocators (suitable for raw transfer).
+    #[cfg(all(feature = "fixed_size", not(feature = "disable_fixed_size")))]
+    pub fn new_fixed_size(thread_count: usize) -> AllocatorPool {
+        #[cfg(all(target_pointer_width = "64", target_endian = "little"))]
+        {
+            Self(AllocatorPoolInner::FixedSize(FixedSizeAllocatorPool::new(thread_count)))
+        }
+
+        #[cfg(not(all(target_pointer_width = "64", target_endian = "little")))]
+        {
+            panic!("Fixed size allocators are only supported on 64-bit little-endian platforms");
+        }
+    }
+
+    /// Retrieve an [`Allocator`] from the pool, or create a new one if the pool is empty.
+    ///
+    /// Returns an [`AllocatorGuard`] that gives access to the allocator.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the underlying mutex is poisoned.
+    pub fn get(&self) -> AllocatorGuard<'_> {
+        let allocator = match &self.0 {
+            AllocatorPoolInner::Standard(pool) => pool.get(),
+            #[cfg(all(
+                feature = "fixed_size",
+                not(feature = "disable_fixed_size"),
+                target_pointer_width = "64",
+                target_endian = "little"
+            ))]
+            AllocatorPoolInner::FixedSize(pool) => pool.get(),
+        };
+
+        AllocatorGuard { allocator: ManuallyDrop::new(allocator), pool: self }
+    }
+
+    /// Add an [`Allocator`] to the pool.
+    ///
+    /// The `Allocator` is reset by this method, so it's ready to be re-used.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the underlying mutex is poisoned.
+    fn add(&self, allocator: Allocator) {
+        // SAFETY: This method is only called from `AllocatorGuard::drop`.
+        // `AllocatorGuard`s are only created by `AllocatorPool::get`, so the `Allocator` must have
+        // been created by this pool. Therefore, it is the correct type for the pool.
+        unsafe {
+            match &self.0 {
+                AllocatorPoolInner::Standard(pool) => pool.add(allocator),
+                #[cfg(all(
+                    feature = "fixed_size",
+                    not(feature = "disable_fixed_size"),
+                    target_pointer_width = "64",
+                    target_endian = "little"
+                ))]
+                AllocatorPoolInner::FixedSize(pool) => pool.add(allocator),
+            }
+        }
+    }
+}
 
 /// A guard object representing exclusive access to an [`Allocator`] from the pool.
 ///

crates/oxc_allocator/src/pool/standard.rs

@@ -1,48 +1,45 @@
-use std::{iter, mem::ManuallyDrop, sync::Mutex};
+use std::{iter, sync::Mutex};
 
 use crate::Allocator;
 
-use super::AllocatorGuard;
-
-/// A thread-safe pool for reusing [`Allocator`] instances to reduce allocation overhead.
+/// A thread-safe pool for reusing [`Allocator`] instances, that uses standard allocators.
 ///
-/// Internally uses a `Vec` protected by a `Mutex` to store available allocators.
-pub struct AllocatorPool {
+/// Unlike `FixedSizeAllocatorPool`, the `Allocator`s used in this pool are suitable for general use,
+/// but not for raw transfer.
+pub struct StandardAllocatorPool {
     allocators: Mutex<Vec<Allocator>>,
 }
 
-impl AllocatorPool {
-    /// Creates a new [`AllocatorPool`] for use across the specified number of threads.
-    pub fn new(thread_count: usize) -> AllocatorPool {
+impl StandardAllocatorPool {
+    /// Create a new [`StandardAllocatorPool`] for use across the specified number of threads.
+    #[cfg_attr(all(feature = "fixed_size", not(feature = "disable_fixed_size")), expect(dead_code))]
+    pub fn new(thread_count: usize) -> StandardAllocatorPool {
         let allocators = iter::repeat_with(Allocator::new).take(thread_count).collect();
-        AllocatorPool { allocators: Mutex::new(allocators) }
+        StandardAllocatorPool { allocators: Mutex::new(allocators) }
     }
 
-    /// Retrieves an [`Allocator`] from the pool, or creates a new one if the pool is empty.
-    ///
-    /// Returns an [`AllocatorGuard`] that gives access to the allocator.
+    /// Retrieve an [`Allocator`] from the pool, or create a new one if the pool is empty.
     ///
     /// # Panics
-    ///
     /// Panics if the underlying mutex is poisoned.
-    pub fn get(&self) -> AllocatorGuard<'_> {
+    pub fn get(&self) -> Allocator {
         let allocator = {
             let mut allocators = self.allocators.lock().unwrap();
             allocators.pop()
         };
-        let allocator = allocator.unwrap_or_else(Allocator::new);
-
-        AllocatorGuard { allocator: ManuallyDrop::new(allocator), pool: self }
+        allocator.unwrap_or_else(Allocator::new)
     }
 
     /// Add an [`Allocator`] to the pool.
     ///
     /// The `Allocator` is reset by this method, so it's ready to be re-used.
     ///
-    /// # Panics
+    /// # SAFETY
+    /// The `Allocator` must have been created by a `StandardAllocatorPool` (not `FixedSizeAllocatorPool`).
     ///
+    /// # Panics
     /// Panics if the underlying mutex is poisoned.
-    pub(super) fn add(&self, mut allocator: Allocator) {
+    pub(super) unsafe fn add(&self, mut allocator: Allocator) {
         allocator.reset();
         let mut allocators = self.allocators.lock().unwrap();
         allocators.push(allocator);
@@ -51,6 +48,12 @@ impl AllocatorPool {
 
 // Dummy implementations of interfaces from `fixed_size`, just to stop clippy complaining.
 // Seems to be necessary due to feature unification.
+#[cfg(not(all(
+    feature = "fixed_size",
+    not(feature = "disable_fixed_size"),
+    target_pointer_width = "64",
+    target_endian = "little"
+)))]
 #[allow(
     dead_code,
     missing_docs,
@@ -82,4 +85,10 @@ mod dummies {
         }
     }
 }
+#[cfg(not(all(
+    feature = "fixed_size",
+    not(feature = "disable_fixed_size"),
+    target_pointer_width = "64",
+    target_endian = "little"
+)))]
 pub use dummies::*;