feat(allocator): fixed size allocators

crates/oxc_allocator/Cargo.toml

@@ -35,5 +35,7 @@ serde = { workspace = true }
 serde_json = { workspace = true }
 
 [features]
+fixed_size = ["from_raw_parts"]
+disable_fixed_size = []
 from_raw_parts = []
 serialize = ["dep:serde", "oxc_estree/serialize"]

crates/oxc_allocator/src/fixed_size.rs

@@ -0,0 +1,123 @@
+use std::{
+    alloc::{self, GlobalAlloc, Layout, System},
+    mem::ManuallyDrop,
+    ops::{Deref, DerefMut},
+    ptr::NonNull,
+};
+
+use crate::Allocator;
+
+// Only 64-bit little-endian platforms are supported at present
+const IS_SUPPORTED_PLATFORM: bool =
+    cfg!(all(target_pointer_width = "64", target_endian = "little"));
+
+const TWO_GIB: usize = 1 << 31;
+// `1 << 32`.
+// We use `IS_SUPPORTED_PLATFORM as usize * 32` to avoid compilation failure on 32-bit platforms.
+const FOUR_GIB: usize = 1 << (IS_SUPPORTED_PLATFORM as usize * 32);
+
+// What we ideally want is an allocation 2 GiB in size, aligned on 4 GiB.
+// But system allocator on Mac OS refuses allocations with 4 GiB alignment.
+// https://github.com/rust-lang/rust/blob/556d20a834126d2d0ac20743b9792b8474d6d03c/library/std/src/sys/alloc/unix.rs#L16-L27
+// https://github.com/rust-lang/rust/issues/30170
+//
+// So we instead allocate 4 GiB with 2 GiB alignment, and then use either the 1st or 2nd half
+// of the allocation, one of which is guaranteed to be on a 4 GiB boundary.
+//
+// TODO: We could use this workaround only on Mac OS, and just allocate what we actually want on Linux.
+// Windows OS allocator also doesn't support high alignment allocations, so Rust contains a workaround
+// which over-allocates (6 GiB in this case).
+// https://github.com/rust-lang/rust/blob/556d20a834126d2d0ac20743b9792b8474d6d03c/library/std/src/sys/alloc/windows.rs#L120-L137
+// Could just use that built-in workaround, rather than implementing our own, or allocate a 6 GiB chunk
+// with alignment 16, to skip Rust's built-in workaround.
+// Note: Rust's workaround will likely commit a whole page of memory, just to store the real pointer.
+const ALLOC_SIZE: usize = FOUR_GIB;
+const ALLOC_ALIGN: usize = TWO_GIB;
+const CHUNK_SIZE: usize = TWO_GIB;
+const CHUNK_ALIGN: usize = FOUR_GIB;
+
+const ALLOC_LAYOUT: Layout = match Layout::from_size_align(ALLOC_SIZE, ALLOC_ALIGN) {
+    Ok(layout) => layout,
+    Err(_) => unreachable!(),
+};
+
+/// Structure which wraps an [`Allocator`] with fixed size of 2 GiB, and aligned on 4 GiB.
+///
+/// To achieve this, we manually allocate memory to back the `Allocator`'s single chunk.
+/// We over-allocate 4 GiB, and then use a part of that allocation to back the `Allocator`.
+/// Inner `Allocator` is wrapped in `ManuallyDrop` to prevent it freeing the memory itself,
+/// and `AllocatorWrapper` has a custom `Drop` impl which frees the whole of the original allocation.
+///
+/// We allocate via `System` allocator, bypassing any registered alternative global allocator
+/// (e.g. Mimalloc in linter). Mimalloc complains that it cannot serve allocations with high alignment,
+/// and presumably it's pointless to try to obtain such large allocations from a thread-local heap,
+/// so better to go direct to the system allocator anyway.
+pub struct FixedSizeAllocator {
+    /// `Allocator` which utilizes part of the original allocation
+    allocator: ManuallyDrop<Allocator>,
+    /// Pointer to start of original allocation
+    alloc_ptr: NonNull<u8>,
+}
+
+impl FixedSizeAllocator {
+    /// Create a new [`FixedSizeAllocator`].
+    ///
+    /// # Panics
+    /// Panics if not a 64-bit little-endian platform.
+    pub fn new() -> Self {
+        assert!(
+            IS_SUPPORTED_PLATFORM,
+            "Fixed size allocators are only supported on 64-bit little-endian platforms"
+        );
+
+        // Allocate block of memory.
+        // SAFETY: Layout does not have zero size.
+        let alloc_ptr = unsafe { System.alloc(ALLOC_LAYOUT) };
+        let Some(alloc_ptr) = NonNull::new(alloc_ptr) else {
+            alloc::handle_alloc_error(ALLOC_LAYOUT);
+        };
+
+        // Get pointer to use for allocator chunk, aligned to 4 GiB.
+        // SAFETY: `offset` is either 0 or `TWO_GIB`.
+        // We allocated 4 GiB of memory, so adding `offset` to `alloc_ptr` is in bounds.
+        let chunk_ptr = unsafe {
+            let offset = alloc_ptr.as_ptr() as usize % ALLOC_SIZE;
+            alloc_ptr.add(offset)
+        };
+
+        debug_assert!(chunk_ptr.as_ptr() as usize % CHUNK_ALIGN == 0);
+
+        // SAFETY: Memory region starting at `chunk_ptr` with `CHUNK_SIZE` bytes is within
+        // the allocation we just made.
+        // `chunk_ptr` has high alignment (4 GiB). `size` is large and a high power of 2 (2 GiB).
+        let allocator = unsafe { Allocator::from_raw_parts(chunk_ptr, CHUNK_SIZE) };
+
+        // Store pointer to original allocation, so it can be used to deallocate in `drop`
+        Self { allocator: ManuallyDrop::new(allocator), alloc_ptr }
+    }
+}
+
+impl Drop for FixedSizeAllocator {
+    fn drop(&mut self) {
+        // SAFETY: Originally allocated from `System` allocator at `alloc_ptr`, with layout `ALLOC_LAYOUT`
+        unsafe { System.dealloc(self.alloc_ptr.as_ptr(), ALLOC_LAYOUT) }
+    }
+}
+
+impl Deref for FixedSizeAllocator {
+    type Target = Allocator;
+
+    fn deref(&self) -> &Self::Target {
+        &self.allocator
+    }
+}
+
+impl DerefMut for FixedSizeAllocator {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.allocator
+    }
+}
+
+// SAFETY: `Allocator` is `Send`.
+// Moving `alloc_ptr: NonNull<u8>` across threads along with the `Allocator` is safe.
+unsafe impl Send for FixedSizeAllocator {}

crates/oxc_allocator/src/lib.rs

@@ -18,6 +18,13 @@
 //!
 //! * `from_raw_parts` - Adds [`Allocator::from_raw_parts`] method.
 //!   Usage of this feature is not advisable, and it will be removed as soon as we're able to.
+//!
+//! * `fixed_size` - Makes [`AllocatorPool`] create large fixed-size allocators, instead of
+//!   flexibly-sized ones.
+//!   Usage of this feature is not advisable, and it will be removed as soon as we're able to.
+//!
+//! * `disable_fixed_size` - Disables `fixed_size` feature.
+//!   Purpose is to prevent `--all-features` enabling fixed sized allocators.
 
 #![warn(missing_docs)]
 
@@ -29,6 +36,8 @@ mod allocator_api2;
 mod boxed;
 mod clone_in;
 mod convert;
+#[cfg(all(feature = "fixed_size", not(feature = "disable_fixed_size")))]
+mod fixed_size;
 #[cfg(feature = "from_raw_parts")]
 mod from_raw_parts;
 pub mod hash_map;

crates/oxc_allocator/src/pool.rs

@@ -73,13 +73,13 @@ impl Drop for AllocatorGuard<'_> {
     }
 }
 
+#[cfg(any(not(feature = "fixed_size"), feature = "disable_fixed_size"))]
 mod wrapper {
     use crate::Allocator;
 
     /// Structure which wraps an [`Allocator`].
     ///
-    /// This implementation adds no value to `Allocator`, but we can add support for fixed-size allocators
-    /// by providing a different implementation of `AllocatorWrapper` behind a feature flag.
+    /// Default implementation which is just a wrapper around an [`Allocator`].
     pub struct AllocatorWrapper(Allocator);
 
     impl AllocatorWrapper {
@@ -105,4 +105,37 @@ mod wrapper {
     }
 }
 
+#[cfg(all(feature = "fixed_size", not(feature = "disable_fixed_size")))]
+mod wrapper {
+    use crate::{Allocator, fixed_size::FixedSizeAllocator};
+
+    /// Structure which wraps an [`Allocator`] with fixed size of 2 GiB, and aligned on 4 GiB.
+    ///
+    /// See [`FixedSizeAllocator`] for more details.
+    pub struct AllocatorWrapper(FixedSizeAllocator);
+
+    impl AllocatorWrapper {
+        /// Create a new [`AllocatorWrapper`].
+        pub fn new() -> Self {
+            Self(FixedSizeAllocator::new())
+        }
+
+        /// Get reference to underlying [`Allocator`].
+        pub fn get(&self) -> &Allocator {
+            &self.0
+        }
+
+        /// Reset the [`Allocator`] in this [`AllocatorWrapper`].
+        pub fn reset(&mut self) {
+            self.0.reset();
+        }
+
+        /// Create a `Vec` of [`AllocatorWrapper`]s.
+        pub fn new_vec(size: usize) -> Vec<Self> {
+            // Each allocator consumes a large block of memory, so create them on demand instead of upfront
+            Vec::with_capacity(size)
+        }
+    }
+}
+
 use wrapper::AllocatorWrapper;