BE-260: HashQL: Switch to bump allocator with scoping support

Cargo.toml

@@ -109,7 +109,7 @@ axum-core                          = { version = "0.5.5" }
 bitvec                             = { version = "1.0.1", default-features = false }
 bon                                = { version = "3.8.1", default-features = false, features = ["implied-bounds"] }
 bstr                               = { version = "1.12.1" }
-bumpalo                            = { version = "3.19.0", default-features = false, features = ["allocator-api2"] }
+bump-scope                         = { version = "1.5.1", default-features = false, features = ["alloc", "nightly-clone-to-uninit", "panic-on-alloc", "std"] }
 bytes                              = { version = "1.10.1" }
 bytes-utils                        = { version = "0.1.4", default-features = false }
 camino                             = { version = "1.1.12", default-features = false }

libs/@local/hashql/core/Cargo.toml

@@ -22,10 +22,9 @@ text-size = { workspace = true, public = true }
 # Private workspace dependencies
 
 # Private third-party dependencies
-allocator-api2       = { workspace = true }
 anstyle-lossy        = { workspace = true }
 bitvec               = { workspace = true, features = ["alloc"] }
-bumpalo              = { workspace = true }
+bump-scope           = { workspace = true }
 derive_more          = { workspace = true, features = ["debug", "from"] }
 ena                  = { workspace = true }
 lexical              = { workspace = true, features = ["parse-integers", "parse-floats", "format"] }

libs/@local/hashql/core/benches/type_system.rs

@@ -4,7 +4,7 @@ use core::hint::black_box;
 
 use codspeed_criterion_compat::{BatchSize, Bencher, Criterion, criterion_group, criterion_main};
 use hashql_core::{
-    heap::{BumpAllocator as _, Heap},
+    heap::{Heap, ResetAllocator as _},
     r#type::{
         TypeId,
         builder::{TypeBuilder, lazy},

libs/@local/hashql/core/src/heap/allocator.rs

@@ -2,9 +2,9 @@
 
 use core::{alloc, ptr};
 
-use bumpalo::Bump;
+use bump_scope::{Bump, BumpBox, BumpScope};
 
-use super::BumpAllocator;
+use super::{BumpAllocator, bump::ResetAllocator};
 
 /// Internal arena allocator.
 #[derive(Debug)]
@@ -21,68 +21,149 @@ impl Allocator {
     /// Creates a new allocator with at least `capacity` bytes pre-allocated.
     #[inline]
     pub(crate) fn with_capacity(capacity: usize) -> Self {
-        Self(Bump::with_capacity(capacity))
+        Self(Bump::with_size(capacity))
     }
 
-    /// Sets the allocation limit for the allocator.
+    /// Allocates a value using a closure to avoid moving before allocation.
     #[inline]
-    pub(crate) fn set_allocation_limit(&self, capacity: Option<usize>) {
-        self.0.set_allocation_limit(capacity);
+    pub(crate) fn alloc_with<T>(&self, func: impl FnOnce() -> T) -> &mut T {
+        BumpBox::leak(self.0.alloc_with(func))
     }
+}
+
+impl BumpAllocator for Allocator {
+    type Scoped<'scope> = AllocatorScope<'scope>;
 
-    /// Allocates a value using a closure to avoid moving before allocation.
     #[inline]
-    pub(crate) fn alloc_with<T>(&self, func: impl FnOnce() -> T) -> &mut T {
-        self.0.alloc_with(func)
+    fn scoped<T>(&mut self, func: impl FnOnce(Self::Scoped<'_>) -> T) -> T {
+        self.0.scoped(|scope| func(AllocatorScope(scope)))
     }
 
-    /// Copies a slice into the arena.
     #[inline]
-    pub(crate) fn try_alloc_slice_copy<T>(&self, slice: &[T]) -> Result<&mut [T], alloc::AllocError>
-    where
-        T: Copy,
-    {
+    fn try_allocate_slice_copy<T: Copy>(&self, slice: &[T]) -> Result<&mut [T], alloc::AllocError> {
         self.0
             .try_alloc_slice_copy(slice)
+            .map(BumpBox::leak)
             .map_err(|_err| alloc::AllocError)
     }
 }
 
-impl BumpAllocator for Allocator {
-    #[inline]
-    fn allocate_slice_copy<T: Copy>(&self, slice: &[T]) -> Result<&mut [T], alloc::AllocError> {
-        self.try_alloc_slice_copy(slice)
-    }
-
+impl ResetAllocator for Allocator {
     #[inline]
     fn reset(&mut self) {
         self.0.reset();
     }
 }
 
-// SAFETY: Delegates to bumpalo::Bump via allocator_api2.
+// SAFETY: Delegates to bump_scope
 #[expect(unsafe_code, reason = "proxy to bump")]
 unsafe impl alloc::Allocator for Allocator {
     #[inline]
     fn allocate(&self, layout: alloc::Layout) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
-        allocator_api2::alloc::Allocator::allocate(&&self.0, layout)
+        bump_scope::alloc::Allocator::allocate(&self.0, layout).map_err(|_err| alloc::AllocError)
+    }
+
+    #[inline]
+    fn allocate_zeroed(
+        &self,
+        layout: alloc::Layout,
+    ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
+        bump_scope::alloc::Allocator::allocate_zeroed(&self.0, layout)
+            .map_err(|_err| alloc::AllocError)
+    }
+
+    #[inline]
+    unsafe fn deallocate(&self, ptr: ptr::NonNull<u8>, layout: alloc::Layout) {
+        // SAFETY: Caller upholds Allocator contract.
+        unsafe {
+            bump_scope::alloc::Allocator::deallocate(&self.0, ptr, layout);
+        }
+    }
+
+    #[inline]
+    unsafe fn grow(
+        &self,
+        ptr: ptr::NonNull<u8>,
+        old_layout: alloc::Layout,
+        new_layout: alloc::Layout,
+    ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
+        // SAFETY: Caller upholds Allocator contract.
+        unsafe {
+            bump_scope::alloc::Allocator::grow(&self.0, ptr, old_layout, new_layout)
+                .map_err(|_err| alloc::AllocError)
+        }
+    }
+
+    #[inline]
+    unsafe fn grow_zeroed(
+        &self,
+        ptr: ptr::NonNull<u8>,
+        old_layout: alloc::Layout,
+        new_layout: alloc::Layout,
+    ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
+        // SAFETY: Caller upholds Allocator contract.
+        unsafe {
+            bump_scope::alloc::Allocator::grow_zeroed(&self.0, ptr, old_layout, new_layout)
+                .map_err(|_err| alloc::AllocError)
+        }
+    }
+
+    #[inline]
+    unsafe fn shrink(
+        &self,
+        ptr: ptr::NonNull<u8>,
+        old_layout: alloc::Layout,
+        new_layout: alloc::Layout,
+    ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
+        // SAFETY: Caller upholds Allocator contract.
+        unsafe {
+            bump_scope::alloc::Allocator::shrink(&self.0, ptr, old_layout, new_layout)
+                .map_err(|_err| alloc::AllocError)
+        }
+    }
+}
+
+pub struct AllocatorScope<'scope>(BumpScope<'scope>);
+
+impl BumpAllocator for AllocatorScope<'_> {
+    type Scoped<'scope> = AllocatorScope<'scope>;
+
+    #[inline]
+    fn scoped<T>(&mut self, func: impl FnOnce(Self::Scoped<'_>) -> T) -> T {
+        self.0.scoped(|scope| func(AllocatorScope(scope)))
+    }
+
+    #[inline]
+    fn try_allocate_slice_copy<T: Copy>(&self, slice: &[T]) -> Result<&mut [T], alloc::AllocError> {
+        self.0
+            .try_alloc_slice_copy(slice)
+            .map(BumpBox::leak)
             .map_err(|_err| alloc::AllocError)
     }
+}
+
+// SAFETY: Delegates to bump_scope
+#[expect(unsafe_code, reason = "proxy to bump")]
+unsafe impl alloc::Allocator for AllocatorScope<'_> {
+    #[inline]
+    fn allocate(&self, layout: alloc::Layout) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
+        bump_scope::alloc::Allocator::allocate(&self.0, layout).map_err(|_err| alloc::AllocError)
+    }
 
     #[inline]
     fn allocate_zeroed(
         &self,
         layout: alloc::Layout,
     ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
-        allocator_api2::alloc::Allocator::allocate_zeroed(&&self.0, layout)
+        bump_scope::alloc::Allocator::allocate_zeroed(&self.0, layout)
             .map_err(|_err| alloc::AllocError)
     }
 
     #[inline]
     unsafe fn deallocate(&self, ptr: ptr::NonNull<u8>, layout: alloc::Layout) {
         // SAFETY: Caller upholds Allocator contract.
         unsafe {
-            allocator_api2::alloc::Allocator::deallocate(&&self.0, ptr, layout);
+            bump_scope::alloc::Allocator::deallocate(&self.0, ptr, layout);
         }
     }
 
@@ -95,7 +176,7 @@ unsafe impl alloc::Allocator for Allocator {
     ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
         // SAFETY: Caller upholds Allocator contract.
         unsafe {
-            allocator_api2::alloc::Allocator::grow(&&self.0, ptr, old_layout, new_layout)
+            bump_scope::alloc::Allocator::grow(&self.0, ptr, old_layout, new_layout)
                 .map_err(|_err| alloc::AllocError)
         }
     }
@@ -109,7 +190,7 @@ unsafe impl alloc::Allocator for Allocator {
     ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
         // SAFETY: Caller upholds Allocator contract.
         unsafe {
-            allocator_api2::alloc::Allocator::grow_zeroed(&&self.0, ptr, old_layout, new_layout)
+            bump_scope::alloc::Allocator::grow_zeroed(&self.0, ptr, old_layout, new_layout)
                 .map_err(|_err| alloc::AllocError)
         }
     }
@@ -123,7 +204,7 @@ unsafe impl alloc::Allocator for Allocator {
     ) -> Result<ptr::NonNull<[u8]>, alloc::AllocError> {
         // SAFETY: Caller upholds Allocator contract.
         unsafe {
-            allocator_api2::alloc::Allocator::shrink(&&self.0, ptr, old_layout, new_layout)
+            bump_scope::alloc::Allocator::shrink(&self.0, ptr, old_layout, new_layout)
                 .map_err(|_err| alloc::AllocError)
         }
     }