From ebc91e1fb5466d9932db810cd83b04de91817bbc Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 27 Aug 2026 12:34:32 -0400 Subject: [PATCH 01/17] feat(buffer): add allocator-backed storage Signed-off-by: Nicholas Gates --- Cargo.lock | 1 + Cargo.toml | 1 + vortex-buffer/Cargo.toml | 1 + vortex-buffer/src/allocation.rs | 368 +++++++++++++++++++++++++++++ vortex-buffer/src/arrow.rs | 16 +- vortex-buffer/src/buffer.rs | 325 +++++++++++++++++-------- vortex-buffer/src/buffer_mut.rs | 406 +++++++++++++++++++++----------- vortex-buffer/src/lib.rs | 9 +- vortex-buffer/src/serde.rs | 2 +- 9 files changed, 883 insertions(+), 246 deletions(-) create mode 100644 vortex-buffer/src/allocation.rs diff --git a/Cargo.lock b/Cargo.lock index 2448cb9aff8..0b59cb19243 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10693,6 +10693,7 @@ dependencies = [ name = "vortex-buffer" version = "0.1.0" dependencies = [ + "allocator-api2", "arrow-buffer 59.2.0", "bitvec", "bytes", diff --git a/Cargo.toml b/Cargo.toml index ae164db13c2..3736acb27ae 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -117,6 +117,7 @@ async-fs = "2.2.0" async-lock = "3.4" async-stream = "0.3.6" async-trait = "0.1.89" +allocator-api2 = "0.2.21" base16ct = "1.0.0" bigdecimal = "0.4.8" bindgen = "0.72.0" diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index 705c992f87d..c77d07f253d 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -23,6 +23,7 @@ serde = ["dep:serde", "serde/serde_derive"] warn-copy = ["dep:tracing"] [dependencies] +allocator-api2 = { workspace = true } arrow-buffer = { workspace = true } bitvec = { workspace = true } bytes = { workspace = true } diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs new file mode 100644 index 00000000000..54e62e0d58e --- /dev/null +++ b/vortex-buffer/src/allocation.rs @@ -0,0 +1,368 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Allocator-backed storage for Vortex buffers. + +use std::alloc::Layout; +use std::fmt; +use std::fmt::Debug; +use std::ptr::NonNull; +use std::sync::Arc; +use std::sync::LazyLock; + +use allocator_api2::alloc::AllocError; +use allocator_api2::alloc::Allocator; +use allocator_api2::alloc::Global; +use allocator_api2::alloc::handle_alloc_error; + +use crate::Alignment; +use crate::BufferMut; + +/// An allocator that can back a Vortex buffer. +pub trait BufferAllocator: Allocator + Debug + Send + Sync + 'static {} + +impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static {} + +/// A shared reference to a buffer allocator. +#[derive(Clone)] +pub struct BufferAllocatorRef(Arc); + +impl BufferAllocatorRef { + /// Wrap an allocator in a shared reference. + pub fn new(allocator: impl BufferAllocator) -> Self { + Self(Arc::new(allocator)) + } + + /// Return a shared reference to the static allocator. + pub fn statically_allocated() -> Self { + STATIC_ALLOCATOR.clone() + } + + /// Create a mutable buffer with this allocator. + pub fn with_capacity(&self, capacity: usize) -> BufferMut { + BufferMut::with_capacity_in(capacity, self.clone()) + } + + /// Create an aligned mutable buffer with this allocator. + pub fn with_capacity_aligned(&self, capacity: usize, alignment: Alignment) -> BufferMut { + BufferMut::with_capacity_aligned_in(capacity, alignment, self.clone()) + } + + /// Create a zeroed mutable buffer with this allocator. + pub fn zeroed(&self, len: usize) -> BufferMut { + BufferMut::zeroed_in(len, self.clone()) + } + + /// Copy values into a mutable buffer made by this allocator. + pub fn copy_from(&self, values: impl AsRef<[T]>) -> BufferMut { + BufferMut::copy_from_in(values, self.clone()) + } +} + +impl Debug for BufferAllocatorRef { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +// SAFETY: all calls are forwarded to the same allocator value held by the Arc. +unsafe impl Allocator for BufferAllocatorRef { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.0.allocate(layout) + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + self.0.allocate_zeroed(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow(ptr, old_layout, new_layout) } + } + + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } + } + + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.shrink(ptr, old_layout, new_layout) } + } +} + +/// The allocator used by buffer APIs that do not take an allocator. +#[derive(Clone, Copy, Debug, Default)] +pub struct StaticBufferAllocator; + +impl StaticBufferAllocator { + /// Create a mutable buffer with the static allocator. + pub fn with_capacity(capacity: usize) -> BufferMut { + BufferMut::with_capacity(capacity) + } + + /// Create an aligned mutable buffer with the static allocator. + pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> BufferMut { + BufferMut::with_capacity_aligned(capacity, alignment) + } + + /// Create a zeroed mutable buffer with the static allocator. + pub fn zeroed(len: usize) -> BufferMut { + BufferMut::zeroed(len) + } + + /// Copy values into a mutable buffer made by the static allocator. + pub fn copy_from(values: impl AsRef<[T]>) -> BufferMut { + BufferMut::copy_from(values) + } +} + +// SAFETY: Global satisfies the Allocator contract and this type only forwards to it. +unsafe impl Allocator for StaticBufferAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + Global.allocate(layout) + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + Global.allocate_zeroed(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } + } + + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.shrink(ptr, old_layout, new_layout) } + } +} + +static STATIC_ALLOCATOR: LazyLock = + LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); + +pub(crate) struct Allocation { + ptr: NonNull, + layout: Layout, + capacity: usize, + allocator: BufferAllocatorRef, +} + +// SAFETY: Allocation owns its memory, and its allocator is Send + Sync. +unsafe impl Send for Allocation {} +// SAFETY: shared access to Allocation never permits mutation of the allocation. +unsafe impl Sync for Allocation {} + +impl Allocation { + pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, false) + } + + pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, true) + } + + fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { + if layout.size() == 0 { + return Self { + ptr: layout.dangling_ptr(), + layout, + capacity: 0, + allocator, + }; + } + + let allocation = if zeroed { + allocator.allocate_zeroed(layout) + } else { + allocator.allocate(layout) + } + .unwrap_or_else(|_| handle_alloc_error(layout)); + + Self { + ptr: allocation.cast(), + layout, + capacity: allocation.len(), + allocator, + } + } + + pub(crate) fn ptr(&self) -> NonNull { + self.ptr + } + + pub(crate) fn size(&self) -> usize { + self.capacity + } + + pub(crate) fn alignment(&self) -> Alignment { + Alignment::new(self.layout.align()) + } + + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + &self.allocator + } + + pub(crate) fn grow(&mut self, layout: Layout) { + if self.layout.size() == 0 { + *self = Self::allocate(layout, self.allocator.clone()); + return; + } + + let allocation = + // SAFETY: ptr and layout describe a live block allocated by self.allocator. The new + // layout is at least as large as the old layout. + unsafe { self.allocator.grow(self.ptr, self.layout, layout) } + .unwrap_or_else(|_| handle_alloc_error(layout)); + self.ptr = allocation.cast(); + self.layout = layout; + self.capacity = allocation.len(); + } +} + +impl Drop for Allocation { + fn drop(&mut self) { + if self.layout.size() == 0 { + return; + } + // SAFETY: ptr and layout describe a live block allocated by self.allocator. + unsafe { self.allocator.deallocate(self.ptr, self.layout) } + } +} + +pub(crate) trait BufferOwner: Send + Sync + 'static { + fn as_slice(&self) -> &[u8]; +} + +impl BufferOwner for T +where + T: AsRef<[u8]> + Send + Sync + 'static, +{ + fn as_slice(&self) -> &[u8] { + self.as_ref() + } +} + +pub(crate) enum BufferBacking { + Owned(Allocation), + External { _owner: Box }, +} + +impl BufferBacking { + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + match self { + Self::Owned(allocation) => allocation.allocator(), + Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), + } + } +} + +#[cfg(test)] +mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use allocator_api2::alloc::AllocError; + use allocator_api2::alloc::Allocator; + use allocator_api2::alloc::Global; + + use crate::Alignment; + use crate::BufferAllocatorRef; + + #[derive(Clone, Debug, Default)] + struct TrackingAllocator { + state: Arc, + } + + #[derive(Debug, Default)] + struct TrackingState { + allocations: AtomicUsize, + deallocations: AtomicUsize, + alignment: AtomicUsize, + } + + // SAFETY: this forwards all memory operations to Global and only records call metadata. + unsafe impl Allocator for TrackingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.state.allocations.fetch_add(1, Ordering::Relaxed); + self.state + .alignment + .store(layout.align(), Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + self.state.deallocations.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller passes the pointer and layout returned by Global. + unsafe { Global.deallocate(ptr, layout) } + } + } + + #[test] + fn allocation_lives_until_last_view() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let buffer = BufferAllocatorRef::new(allocator) + .copy_from([1u32, 2, 3, 4]) + .freeze(); + let view = buffer.slice(0..2); + + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!( + state.alignment.load(Ordering::Relaxed), + *Alignment::DEFAULT_ALIGNMENT + ); + drop(buffer); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + drop(view); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + } +} diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index 63eec6deb60..efaddf6ee14 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -35,12 +35,8 @@ impl Buffer { ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + debug_assert_eq!(length, bytes.len() / size_of::()); + Self::from_bytes_aligned(bytes, alignment) } /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`. @@ -74,12 +70,8 @@ impl ByteBuffer { ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + debug_assert_eq!(length, bytes.len()); + Self::from_bytes_aligned(bytes, alignment) } } diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index a59fff825f8..13032c9ea6a 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -8,9 +8,11 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use std::marker::PhantomData; use std::ops::Deref; use std::ops::RangeBounds; +use std::ptr::NonNull; +use std::sync::Arc; +use std::sync::LazyLock; use bytes::Buf; use bytes::Bytes; @@ -18,6 +20,9 @@ use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; +use crate::BufferAllocatorRef; +use crate::BufferBacking; use crate::BufferMut; use crate::ByteBuffer; use crate::debug::TruncatedDebug; @@ -26,64 +31,110 @@ use crate::trusted_len::TrustedLen; /// An immutable buffer of items of `T`. #[derive(Clone)] pub struct Buffer { - pub(crate) bytes: Bytes, + pub(crate) ptr: NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, - pub(crate) _marker: PhantomData, + pub(crate) backing: Arc, } +// SAFETY: Buffer is an immutable view over backing memory. Its pointer remains valid while the +// backing is live, and sharing elements follows the same bounds as sharing a slice. +unsafe impl Send for Buffer {} +// SAFETY: see the Send implementation above. +unsafe impl Sync for Buffer {} + /// Zero-length backing for empty buffers, "aligned" to [`Alignment::MAX`] so it satisfies any /// valid alignment without allocating. A zero-length slice never reads memory, so it may use a /// dangling pointer as long as it is non-null and aligned. -const EMPTY_BACKING: &[u8] = { +const EMPTY_BYTES: &[u8] = { let addr = 1usize << (usize::BITS - 1); assert!(Alignment::MAX.is_offset_aligned(addr)); // SAFETY: the pointer is non-null and aligned, and the slice is zero-length. unsafe { std::slice::from_raw_parts(std::ptr::without_provenance(addr), 0) } }; +static EMPTY_BACKING: LazyLock> = LazyLock::new(|| { + Arc::new(BufferBacking::External { + _owner: Box::new(EMPTY_BYTES), + }) +}); + impl Default for Buffer { fn default() -> Self { Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment: Alignment::of::(), - _marker: PhantomData, + backing: EMPTY_BACKING.clone(), } } } -impl PartialEq for Buffer { +impl PartialEq for Buffer { #[inline] fn eq(&self, other: &Self) -> bool { - self.bytes == other.bytes + self.as_slice() == other.as_slice() } } -impl Eq for Buffer {} +impl Eq for Buffer {} -impl Ord for Buffer { +impl Ord for Buffer { #[inline] fn cmp(&self, other: &Self) -> Ordering { - self.bytes.cmp(&other.bytes) + self.as_slice().cmp(other.as_slice()) } } -impl PartialOrd for Buffer { +impl PartialOrd for Buffer { #[inline] fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + self.as_slice().partial_cmp(other.as_slice()) } } -impl Hash for Buffer { +impl Hash for Buffer { #[inline] fn hash(&self, state: &mut H) { - self.bytes.as_ref().hash(state) + self.as_slice().hash(state) } } impl Buffer { + pub(crate) fn from_allocation( + allocation: Allocation, + offset: usize, + length: usize, + alignment: Alignment, + ) -> Self { + // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; + Self { + ptr, + length, + alignment, + backing: Arc::new(BufferBacking::Owned(allocation)), + } + } + + fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self { + let bytes = owner.as_slice(); + let length = bytes.len() / size_of::(); + let ptr = if length == 0 { + empty_ptr() + } else { + NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") + }; + Self { + ptr, + length, + alignment, + backing: Arc::new(BufferBacking::External { + _owner: Box::new(owner), + }), + } + } + /// Returns a new `Buffer` copied from the provided `Vec`, `&[T]`, etc. /// /// Due to our underlying usage of `bytes::Bytes`, we are unable to take zero-copy ownership @@ -94,6 +145,11 @@ impl Buffer { BufferMut::copy_from(values).freeze() } + /// Returns a new `Buffer` copied with the provided allocator. + pub fn copy_from_in(values: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + BufferMut::copy_from_in(values, allocator).freeze() + } + /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -121,6 +177,11 @@ impl Buffer { Self::zeroed_aligned(len, Alignment::of::()) } + /// Create a new zeroed `Buffer` with the provided allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + BufferMut::zeroed_in(len, allocator).freeze() + } + /// Create a new zeroed `Buffer` with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -161,10 +222,10 @@ impl Buffer { ); } Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment, - _marker: PhantomData, + backing: EMPTY_BACKING.clone(), } } @@ -176,6 +237,14 @@ impl Buffer { BufferMut::full(item, len).freeze() } + /// Create a full `Buffer` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + BufferMut::full_in(item, len, allocator).freeze() + } + /// Create a `Buffer` zero-copy from a `ByteBuffer`. /// /// ## Panics @@ -194,7 +263,29 @@ impl Buffer { /// Panics if the buffer is not aligned to the given alignment, if the length is not a multiple /// of the size of `T`, or if the given alignment is not aligned to that of `T`. pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self { - Self::from_bytes_aligned(buffer.into_inner(), alignment) + if !alignment.is_aligned_to(Alignment::of::()) { + vortex_panic!( + "Alignment {} must be compatible with the scalar type's alignment {}", + alignment, + Alignment::of::(), + ); + } + if !alignment.is_ptr_aligned(buffer.as_ptr()) { + vortex_panic!("Buffer must align to the requested alignment {}", alignment); + } + if !buffer.len().is_multiple_of(size_of::()) { + vortex_panic!( + "Buffer length {} must be a multiple of the scalar type's size {}", + buffer.len(), + size_of::() + ); + } + Self { + ptr: buffer.ptr.cast(), + length: buffer.length / size_of::(), + alignment, + backing: buffer.backing, + } } /// Create a `Buffer` zero-copy from a `Bytes`. @@ -224,13 +315,7 @@ impl Buffer { size_of::() ); } - let length = bytes.len() / size_of::(); - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + Self::from_owner(bytes, alignment) } /// Create a buffer with values from the TrustedLen iterator. @@ -249,7 +334,8 @@ impl Buffer { Ok(mut_buf) => mut_buf.map_each_in_place(f), Err(buf) => { let len = buf.len(); - let mut out_buf = BufferMut::with_capacity(len); + let allocator = buf.backing.allocator().clone(); + let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf .spare_capacity_mut() .iter_mut() @@ -266,7 +352,6 @@ impl Buffer { /// Clear the buffer, preserving existing capacity. pub fn clear(&mut self) { - self.bytes.clear(); self.length = 0; } @@ -288,17 +373,33 @@ impl Buffer { self.alignment } + /// Returns the allocator to use for derived buffers. + /// + /// External buffers use the static allocator. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.backing.allocator() + } + + /// Returns a raw pointer to the buffer's data. + #[inline(always)] + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() + } + /// Returns a slice over the buffer of elements of type T. #[inline(always)] pub fn as_slice(&self) -> &[T] { - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(self.bytes.as_ptr().cast(), self.length) } + // SAFETY: ptr points into the live backing and construction checks its alignment. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } } /// Return a view over the buffer as an opaque byte slice. #[inline(always)] pub fn as_bytes(&self) -> &[u8] { - self.bytes.as_ref() + // SAFETY: the element range is initialized and remains live through backing. + unsafe { + std::slice::from_raw_parts(self.ptr.as_ptr().cast(), size_of_val(self.as_slice())) + } } /// Returns an iterator over the buffer of elements of type T. @@ -372,8 +473,6 @@ impl Buffer { } let begin_byte = begin * size_of::(); - let end_byte = end * size_of::(); - if !alignment.is_offset_aligned(begin_byte) { vortex_panic!( "range start must be aligned to {alignment:?}, byte {}", @@ -385,10 +484,11 @@ impl Buffer { } Self { - bytes: self.bytes.slice(begin_byte..end_byte), + // SAFETY: begin is in bounds and the alignment check applies to the new pointer. + ptr: unsafe { self.ptr.add(begin) }, length: end - begin, alignment, - _marker: Default::default(), + backing: Arc::clone(&self.backing), } } @@ -427,74 +527,90 @@ impl Buffer { vortex_panic!("slice_ref subset must be aligned to {:?}", alignment); } - let subset_u8 = - unsafe { std::slice::from_raw_parts(subset.as_ptr().cast(), size_of_val(subset)) }; + let start = self.as_ptr().addr(); + let end = start + size_of_val(self.as_slice()); + let subset_start = subset.as_ptr().addr(); + let subset_end = subset_start + .checked_add(size_of_val(subset)) + .vortex_expect("slice_ref address overflow"); + if subset_start < start || subset_end > end { + vortex_panic!("slice_ref subset must be contained in the buffer"); + } Self { - bytes: self.bytes.slice_ref(subset_u8), + ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"), length: subset.len(), alignment, - _marker: Default::default(), + backing: Arc::clone(&self.backing), } } - /// Returns the underlying aligned buffer. - pub fn inner(&self) -> &Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - &self.bytes - } - - /// Returns the underlying aligned buffer. + /// Returns the underlying bytes without copying. pub fn into_inner(self) -> Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - self.bytes + Bytes::from_owner(BufferBytesOwner { + ptr: self.ptr.cast(), + length: self.length * size_of::(), + backing: self.backing, + }) } /// Return the ByteBuffer for this `Buffer`. pub fn into_byte_buffer(self) -> ByteBuffer { ByteBuffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, - _marker: Default::default(), + backing: self.backing, } } /// Try to convert self into `BufferMut` if there is only a single strong reference. pub fn try_into_mut(self) -> Result, Self> { - self.bytes - .try_into_mut() - .map(|bytes| BufferMut { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) - .map_err(|bytes| Self { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) + let Self { + ptr, + length, + alignment, + backing, + } = self; + match Arc::try_unwrap(backing) { + Ok(BufferBacking::Owned(allocation)) => { + let offset = ptr.addr().get() - allocation.ptr().addr().get(); + let capacity = (allocation.size() - offset) / size_of::(); + Ok(BufferMut { + allocation, + offset, + length, + capacity, + alignment, + _marker: Default::default(), + }) + } + Ok(backing) => Err(Self { + ptr, + length, + alignment, + backing: Arc::new(backing), + }), + Err(backing) => Err(Self { + ptr, + length, + alignment, + backing, + }), + } } /// Convert self into `BufferMut`, cloning the data if there are multiple strong references. pub fn into_mut(self) -> BufferMut { - self.try_into_mut() - .unwrap_or_else(|buffer| BufferMut::::copy_from_aligned(&buffer, buffer.alignment)) + self.try_into_mut().unwrap_or_else(|buffer| { + let allocator = buffer.backing.allocator().clone(); + BufferMut::::copy_from_aligned_in(&buffer, buffer.alignment, allocator) + }) } /// Returns whether a `Buffer` is aligned to the given alignment. pub fn is_aligned(&self, alignment: Alignment) -> bool { - alignment.is_ptr_aligned(self.bytes.as_ptr()) + alignment.is_ptr_aligned(self.as_ptr()) } /// Return a `Buffer` with the given alignment. Where possible, this will be zero-copy. @@ -510,7 +626,8 @@ impl Buffer { "Buffer is not aligned to requested alignment {alignment}, copying: {bt}" ) } - Self::copy_from_aligned(self, alignment) + let allocator = self.backing.allocator().clone(); + BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze() } } @@ -546,10 +663,10 @@ impl Buffer { ); Buffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, - _marker: PhantomData, + backing: self.backing, } } } @@ -644,34 +761,21 @@ impl AsRef<[u8]> for Wrapper { impl From> for Buffer where - T: Send + 'static, + T: Send + Sync + 'static, { fn from(value: Vec) -> Self { - let original_len = value.len(); let wrapped_vec = Wrapper(value); - - let bytes = Bytes::from_owner(wrapped_vec); - - assert_eq!(bytes.as_ptr().align_offset(align_of::()), 0); - - Self { - bytes, - length: original_len, - alignment: Alignment::of::(), - _marker: PhantomData, - } + assert_eq!( + wrapped_vec.as_ref().as_ptr().align_offset(align_of::()), + 0 + ); + Self::from_owner(wrapped_vec, Alignment::of::()) } } impl From for ByteBuffer { fn from(bytes: Bytes) -> Self { - let length = bytes.len(); - Self { - bytes, - length, - alignment: Alignment::of::(), - _marker: Default::default(), - } + Self::from_owner(bytes, Alignment::of::()) } } @@ -695,11 +799,36 @@ impl Buf for ByteBuffer { self.alignment ); } - self.bytes.advance(cnt); + assert!(cnt <= self.length, "cannot advance past the buffer length"); + // SAFETY: cnt is within the initialized byte range. + self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; } } +struct BufferBytesOwner { + ptr: NonNull, + length: usize, + backing: Arc, +} + +// SAFETY: the owner exposes immutable initialized bytes and keeps their backing live. +unsafe impl Send for BufferBytesOwner {} +unsafe impl Sync for BufferBytesOwner {} + +impl AsRef<[u8]> for BufferBytesOwner { + fn as_ref(&self) -> &[u8] { + let _ = &self.backing; + // SAFETY: ptr and length came from a live Buffer. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } + } +} + +fn empty_ptr() -> NonNull { + let addr = 1usize << (usize::BITS - 1); + NonNull::new(std::ptr::without_provenance_mut(addr)).vortex_expect("empty pointer is non-null") +} + /// Owned iterator over a [`Buffer`]. pub struct BufferIterator { // Keep the buffer alive for the duration of the iteration. diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index e5cb03c558b..0fa715c29bd 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use core::mem::MaybeUninit; +use std::alloc::Layout; use std::any::type_name; use std::cmp::max; use std::fmt::Debug; @@ -12,23 +13,25 @@ use std::ops::DerefMut; use bytes::Buf; use bytes::BufMut; -use bytes::BytesMut; use bytes::buf::UninitSlice; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; use crate::Buffer; +use crate::BufferAllocatorRef; use crate::ByteBufferMut; use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. -#[derive(PartialEq, Eq)] pub struct BufferMut { - pub(crate) bytes: BytesMut, + pub(crate) allocation: Allocation, + pub(crate) offset: usize, pub(crate) length: usize, + pub(crate) capacity: usize, pub(crate) alignment: Alignment, pub(crate) _marker: std::marker::PhantomData, } @@ -36,7 +39,12 @@ pub struct BufferMut { impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. pub fn with_capacity(capacity: usize) -> Self { - Self::with_capacity_aligned(capacity, Alignment::of::()) + Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new `BufferMut` with the requested capacity and allocator. + pub fn with_capacity_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(capacity, Alignment::of::(), allocator) } /// Create a new `BufferMut` with the requested alignment and capacity. @@ -46,10 +54,24 @@ impl BufferMut { /// /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { - Self::with_capacity_preferred_aligned( + Self::with_capacity_aligned_in( + capacity, + alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new `BufferMut` with the requested alignment, capacity, and allocator. + pub fn with_capacity_aligned_in( + capacity: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::with_capacity_preferred_aligned_in( capacity, alignment, Some(Alignment::DEFAULT_ALIGNMENT), + allocator, ) } @@ -61,6 +83,21 @@ impl BufferMut { capacity: usize, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::with_capacity_preferred_aligned_in( + capacity, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new allocator-backed `BufferMut` with a requested and preferred alignment. + pub fn with_capacity_preferred_aligned_in( + capacity: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { let actual = max( alignment, @@ -75,12 +112,19 @@ impl BufferMut { ); } - let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + *actual); - bytes.align_empty(actual); + let size = capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let layout = Layout::from_size_align(size, *actual) + .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); + let allocation = Allocation::allocate(layout, allocator); + let capacity = allocation.size() / size_of::(); Self { - bytes, + allocation, + offset: 0, length: 0, + capacity, alignment, _marker: Default::default(), } @@ -88,7 +132,12 @@ impl BufferMut { /// Create a new zeroed `BufferMut`. pub fn zeroed(len: usize) -> Self { - Self::zeroed_aligned(len, Alignment::of::()) + Self::zeroed_in(len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new zeroed `BufferMut` with the requested allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + Self::zeroed_aligned_in(len, Alignment::of::(), allocator) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -98,7 +147,21 @@ impl BufferMut { /// /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::zeroed_aligned_in(len, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create a zeroed `BufferMut` with an alignment and allocator. + pub fn zeroed_aligned_in( + len: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -109,16 +172,36 @@ impl BufferMut { len: usize, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a zeroed allocator-backed buffer with a requested and preferred alignment. + pub fn zeroed_preferred_aligned_in( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { let preferred_alignment = preferred_alignment.unwrap_or(Alignment::of::()); let actual_alignment = max(preferred_alignment, alignment); - let mut bytes = BytesMut::zeroed((len * size_of::()) + *actual_alignment); - bytes.advance(bytes.as_ptr().align_offset(*actual_alignment)); - unsafe { bytes.set_len(len * size_of::()) }; - let actual_len = bytes.len().checked_div(size_of::()).unwrap_or(0); + let size = len + .checked_mul(size_of::()) + .vortex_expect("buffer length overflow"); + let layout = Layout::from_size_align(size, *actual_alignment) + .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); + let allocation = Allocation::allocate_zeroed(layout, allocator); + let capacity = allocation.size() / size_of::(); Self { - bytes, - length: actual_len, + allocation, + offset: 0, + length: len, + capacity, alignment, _marker: Default::default(), } @@ -136,7 +219,12 @@ impl BufferMut { /// /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned pub fn empty_aligned(alignment: Alignment) -> Self { - Self::empty_preferred_aligned(alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::empty_aligned_in(alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create an empty `BufferMut` with an alignment and allocator. + pub fn empty_aligned_in(alignment: Alignment, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(0, alignment, allocator) } /// Create a new empty `BufferMut` with the provided alignment. @@ -147,7 +235,12 @@ impl BufferMut { alignment: Alignment, preferred_alignment: Option, ) -> Self { - BufferMut::with_capacity_preferred_aligned(0, alignment, preferred_alignment) + BufferMut::with_capacity_preferred_aligned_in( + 0, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) } /// Create a new full `BufferMut` with the given value. @@ -155,14 +248,27 @@ impl BufferMut { where T: Copy, { - let mut buffer = BufferMut::::with_capacity(len); + Self::full_in(item, len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a full `BufferMut` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + let mut buffer = BufferMut::::with_capacity_in(len, allocator); buffer.push_n(item, len); buffer } /// Create a mutable scalar buffer by copying the contents of the slice. pub fn copy_from(other: impl AsRef<[T]>) -> Self { - Self::copy_from_aligned(other, Alignment::of::()) + Self::copy_from_in(other, BufferAllocatorRef::statically_allocated()) + } + + /// Create a mutable scalar buffer by copying with the given allocator. + pub fn copy_from_in(other: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + Self::copy_from_aligned_in(other, Alignment::of::(), allocator) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -176,7 +282,21 @@ impl BufferMut { /// /// Panics when the requested alignment isn't itself aligned to type T. pub fn copy_from_aligned(other: impl AsRef<[T]>, alignment: Alignment) -> Self { - Self::copy_from_preferred_aligned(other, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::copy_from_aligned_in(other, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Copy values into a mutable buffer with the given alignment and allocator. + pub fn copy_from_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -191,13 +311,32 @@ impl BufferMut { other: impl AsRef<[T]>, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Copy values with the given allocator, requested alignment, and preferred alignment. + pub fn copy_from_preferred_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!("Given alignment is not aligned to type T") } let other = other.as_ref(); - let mut buffer = - Self::with_capacity_preferred_aligned(other.len(), alignment, preferred_alignment); + let mut buffer = Self::with_capacity_preferred_aligned_in( + other.len(), + alignment, + preferred_alignment, + allocator, + ); buffer.extend_from_slice(other); debug_assert_eq!(buffer.alignment(), alignment); buffer @@ -209,10 +348,14 @@ impl BufferMut { self.alignment } + /// Returns the allocator that owns this buffer. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.allocation.allocator() + } + /// Returns the length of the buffer. #[inline(always)] pub fn len(&self) -> usize { - debug_assert_eq!(self.length, self.bytes.len() / size_of::()); self.length } @@ -225,29 +368,38 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.bytes.capacity() / size_of::() + self.capacity + } + + /// Returns a raw pointer to the buffer's data. + pub fn as_ptr(&self) -> *const T { + // SAFETY: offset always remains within the allocation. + unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } + } + + /// Returns a mutable raw pointer to the buffer's data. + pub fn as_mut_ptr(&mut self) -> *mut T { + // SAFETY: BufferMut uniquely owns the allocation and offset is in bounds. + unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_slice(&self) -> &[T] { - let raw_slice = self.bytes.as_ref(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(raw_slice.as_ptr().cast(), self.length) } + // SAFETY: the allocation is live, offset is in bounds, and construction checks alignment. + unsafe { std::slice::from_raw_parts(self.as_ptr(), self.length) } } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_mut_slice(&mut self) -> &mut [T] { - let raw_slice = self.bytes.as_mut(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts_mut(raw_slice.as_mut_ptr().cast(), self.length) } + // SAFETY: BufferMut uniquely owns the allocation and the initialized range is in bounds. + unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.length) } } /// Clear the buffer, retaining any existing capacity. #[inline] pub fn clear(&mut self) { - unsafe { self.bytes.set_len(0) } self.length = 0; } @@ -269,8 +421,7 @@ impl BufferMut { /// Reserves capacity for at least `additional` more elements to be inserted in the buffer. #[inline] pub fn reserve(&mut self, additional: usize) { - let additional_bytes = additional * size_of::(); - if additional_bytes <= self.bytes.capacity() - self.bytes.len() { + if additional <= self.capacity - self.length { // We can fit the additional bytes in the remaining capacity. Nothing to do. return; } @@ -279,17 +430,37 @@ impl BufferMut { self.reserve_allocate(additional); } - /// A separate function so we can inline the reserve call's fast path. According to `BytesMut` - /// this has significant performance implications. + /// A separate function so we can inline the reserve call's fast path. fn reserve_allocate(&mut self, additional: usize) { - let new_capacity: usize = ((self.length + additional) * size_of::()) + *self.alignment; + let required = self + .length + .checked_add(additional) + .vortex_expect("buffer capacity overflow"); // Make sure we at least double in size each time we re-allocate to amortize the cost - let new_capacity = new_capacity.max(self.bytes.capacity() * 2); - - let mut bytes = BytesMut::with_capacity(new_capacity); - bytes.align_empty(self.alignment); - bytes.extend_from_slice(&self.bytes); - self.bytes = bytes; + let new_capacity = required.max(self.capacity.saturating_mul(2)); + let new_size = new_capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let physical_alignment = max(self.alignment, self.allocation.alignment()); + let layout = Layout::from_size_align(new_size, *physical_alignment) + .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); + + if self.offset == 0 { + self.allocation.grow(layout); + } else { + let mut allocation = Allocation::allocate(layout, self.allocation.allocator().clone()); + // SAFETY: both ranges are valid for the initialized byte length and do not overlap. + unsafe { + std::ptr::copy_nonoverlapping( + self.as_ptr().cast::(), + allocation.ptr().as_ptr(), + self.length * size_of::(), + ); + } + std::mem::swap(&mut self.allocation, &mut allocation); + self.offset = 0; + } + self.capacity = self.allocation.size() / size_of::(); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -329,13 +500,9 @@ impl BufferMut { /// ``` #[inline] pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit] { - let dst = self.bytes.spare_capacity_mut().as_mut_ptr(); - unsafe { - std::slice::from_raw_parts_mut( - dst as *mut MaybeUninit, - self.capacity() - self.length, - ) - } + // SAFETY: offset + length is within the allocation and points at spare capacity. + let dst = unsafe { self.as_mut_ptr().add(self.length) }.cast::>(); + unsafe { std::slice::from_raw_parts_mut(dst, self.capacity() - self.length) } } /// Sets the length of the buffer. @@ -349,7 +516,6 @@ impl BufferMut { #[inline] pub unsafe fn set_len(&mut self, len: usize) { debug_assert!(len <= self.capacity()); - unsafe { self.bytes.set_len(len * size_of::()) }; self.length = len; } @@ -369,9 +535,8 @@ impl BufferMut { pub unsafe fn push_unchecked(&mut self, item: T) { // SAFETY: the caller ensures we have sufficient capacity unsafe { - let dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let dst = self.as_mut_ptr().add(self.length); dst.write(item); - self.bytes.set_len(self.bytes.len() + size_of::()) } self.length += 1; } @@ -398,7 +563,8 @@ impl BufferMut { where T: Copy, { - let mut dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + // SAFETY: the caller guarantees enough spare capacity. + let mut dst = unsafe { self.as_mut_ptr().add(self.length) }; // SAFETY: we checked the capacity in the reserve call unsafe { let end = dst.add(n); @@ -406,7 +572,6 @@ impl BufferMut { dst.write(item); dst = dst.add(1); } - self.bytes.set_len(self.bytes.len() + (n * size_of::())); } self.length += n; } @@ -426,19 +591,21 @@ impl BufferMut { #[inline] pub fn extend_from_slice(&mut self, slice: &[T]) { self.reserve(slice.len()); - let raw_slice = - unsafe { std::slice::from_raw_parts(slice.as_ptr().cast(), size_of_val(slice)) }; - self.bytes.extend_from_slice(raw_slice); + // SAFETY: reserve made the destination valid and non-overlapping for slice.len() values. + unsafe { + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + self.as_mut_ptr().add(self.length), + slice.len(), + ); + } self.length += slice.len(); } /// Splits the buffer into two at the given index. /// /// Afterward, self contains elements `[0, at)`, and the returned buffer contains elements - /// `[at, capacity)`. It’s guaranteed that the memory does not move, that is, the address of - /// self does not change, and the address of the returned slice is at bytes after that. - /// - /// This is an O(1) operation that just increases the reference count and sets a few indices. + /// `[at, capacity)`. The returned buffer uses a new allocation. /// /// Panics if either half would have a length that is not a multiple of the alignment. pub fn split_off(&mut self, at: usize) -> Self { @@ -455,27 +622,25 @@ impl BufferMut { ); } - let new_bytes = self.bytes.split_off(bytes_at); - - // Adjust the lengths, given that length may be < at let new_length = self.length.saturating_sub(at); + let new_capacity = self.capacity - at; + let mut other = Self::with_capacity_aligned_in( + new_capacity, + self.alignment, + self.allocation.allocator().clone(), + ); + if new_length > 0 { + other.extend_from_slice(&self.as_slice()[at..]); + } self.length = self.length.min(at); + self.capacity = at; - BufferMut { - bytes: new_bytes, - length: new_length, - alignment: self.alignment, - _marker: Default::default(), - } + other } /// Absorbs a mutable buffer that was previously split off. /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.extend_from_slice(other.as_ref()). + /// This appends the contents of `other` to this buffer. pub fn unsplit(&mut self, other: Self) { if self.alignment != other.alignment { vortex_panic!( @@ -484,15 +649,16 @@ impl BufferMut { other.alignment ); } - self.bytes.unsplit(other.bytes); - self.length += other.length; + self.extend_from_slice(other.as_slice()); } /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { ByteBufferMut { - bytes: self.bytes, + allocation: self.allocation, + offset: self.offset, length: self.length * size_of::(), + capacity: self.capacity * size_of::(), alignment: self.alignment, _marker: Default::default(), } @@ -500,12 +666,7 @@ impl BufferMut { /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { - Buffer { - bytes: self.bytes.freeze(), - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - } + Buffer::from_allocation(self.allocation, self.offset, self.length, self.alignment) } /// Map each element of the buffer with a closure. @@ -533,14 +694,10 @@ impl BufferMut { /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { if self.as_ptr().align_offset(*alignment) == 0 { - Self { - bytes: self.bytes, - length: self.length, - alignment, - _marker: std::marker::PhantomData, - } + Self { alignment, ..self } } else { - Self::copy_from_aligned(self, alignment) + let allocator = self.allocation.allocator().clone(); + Self::copy_from_aligned_in(self, alignment, allocator) } } @@ -564,8 +721,10 @@ impl BufferMut { ); BufferMut { - bytes: self.bytes, + allocation: self.allocation, + offset: self.offset, length: self.length, + capacity: self.capacity, alignment: self.alignment, _marker: std::marker::PhantomData, } @@ -574,14 +733,24 @@ impl BufferMut { impl Clone for BufferMut { fn clone(&self) -> Self { - // NOTE(ngates): we cannot derive Clone since BytesMut copies on clone and the alignment - // might be messed up. - let mut buffer = BufferMut::::with_capacity_aligned(self.capacity(), self.alignment); + let mut buffer = BufferMut::::with_capacity_aligned_in( + self.capacity(), + self.alignment, + self.allocation.allocator().clone(), + ); buffer.extend_from_slice(self.as_slice()); buffer } } +impl PartialEq for BufferMut { + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +} + +impl Eq for BufferMut {} + impl Debug for BufferMut { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct(&format!("BufferMut<{}>", type_name::())) @@ -645,7 +814,7 @@ impl BufferMut { let unwritten = self.capacity() - self.len(); // We store `begin` in the case that the lower bound hint is incorrect. - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); // As a first step, we manually iterate the iterator up to the known capacity. @@ -690,7 +859,7 @@ impl BufferMut { .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), ); - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); iter.for_each(|item| { @@ -795,8 +964,10 @@ impl Buf for ByteBufferMut { self.alignment ); } - self.bytes.advance(cnt); + assert!(cnt <= self.length, "advance out of bounds"); + self.offset += cnt; self.length -= cnt; + self.capacity -= cnt; } } @@ -818,13 +989,15 @@ unsafe impl BufMut for ByteBufferMut { self.alignment ); } - unsafe { self.bytes.advance_mut(cnt) }; - self.length -= cnt; + self.reserve(cnt); + self.length += cnt; } #[inline] fn chunk_mut(&mut self) -> &mut UninitSlice { - self.bytes.chunk_mut() + let spare = self.spare_capacity_mut(); + // SAFETY: spare points to valid uninitialized byte capacity owned by this buffer. + unsafe { UninitSlice::from_raw_parts_mut(spare.as_mut_ptr().cast(), spare.len()) } } fn put(&mut self, mut src: T) @@ -849,35 +1022,6 @@ unsafe impl BufMut for ByteBufferMut { } } -/// Extension trait for [`BytesMut`] that provides functions for aligning the buffer. -trait AlignedBytesMut { - /// Align an empty `BytesMut` to the specified alignment. - /// - /// ## Panics - /// - /// Panics if the buffer is not empty, or if there is not enough capacity to align the buffer. - fn align_empty(&mut self, alignment: Alignment); -} - -impl AlignedBytesMut for BytesMut { - fn align_empty(&mut self, alignment: Alignment) { - // TODO(joe): this is slow fixme - if !self.is_empty() { - vortex_panic!("ByteBufferMut must be empty"); - } - - let padding = self.as_ptr().align_offset(*alignment); - self.capacity() - .checked_sub(padding) - .vortex_expect("Not enough capacity to align buffer"); - - // SAFETY: We know the buffer is empty, and we know we have enough capacity, so we can - // safely set the length to the padding and advance the buffer to the aligned offset. - unsafe { self.set_len(padding) }; - self.advance(padding); - } -} - impl Write for ByteBufferMut { fn write(&mut self, buf: &[u8]) -> std::io::Result { self.extend_from_slice(buf); diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index ee113481353..99b1e7c3081 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -5,11 +5,10 @@ //! A library for working with custom aligned buffers of sized values. //! -//! The `vortex-buffer` crate is built around `bytes::Bytes` and therefore supports zero-copy -//! cloning and slicing, but differs in that it can define and maintain a custom alignment. +//! The `vortex-buffer` crate supports zero-copy cloning and slicing with a custom allocator and +//! runtime alignment. //! -//! * `Buffer` and `BufferMut` provide immutable and mutable wrappers around `bytes::Bytes` -//! and `bytes::BytesMut` respectively. +//! * `Buffer` and `BufferMut` provide immutable and mutable typed buffers. //! * `ByteBuffer` and `ByteBufferMut` are type aliases for `u8` buffers. //! * `BufferString` is a wrapper around a `ByteBuffer` that enforces utf-8 encoding. //! * `ConstBuffer` provides similar functionality to `Buffer` except with a @@ -47,6 +46,7 @@ //! `arrow_buffer::OffsetBuffer`. pub use alignment::*; +pub use allocation::*; pub use bit::*; pub use buffer::*; pub use buffer_mut::*; @@ -55,6 +55,7 @@ pub use r#const::*; pub use dispatch::*; pub use string::*; mod alignment; +mod allocation; #[cfg(feature = "arrow")] mod arrow; mod bit; diff --git a/vortex-buffer/src/serde.rs b/vortex-buffer/src/serde.rs index 9563236a50a..5d36f8da3d0 100644 --- a/vortex-buffer/src/serde.rs +++ b/vortex-buffer/src/serde.rs @@ -22,7 +22,7 @@ where where S: Serializer, { - serializer.serialize_bytes(self.inner().as_ref()) + serializer.serialize_bytes(self.as_bytes()) } } From 4acd035a3d374bdce26b8bfb5f82960f10ce62a5 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 27 Aug 2026 17:36:23 -0400 Subject: [PATCH 02/17] fix(buffer): sort allocator dependency Signed-off-by: Nicholas Gates --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 3736acb27ae..1bb45b099d1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -97,6 +97,7 @@ rust-version = "1.95" version = "0.1.0" [workspace.dependencies] +allocator-api2 = "0.2.21" alp = "0.0.2" anyhow = "1.0.100" arbitrary = "1.3.2" @@ -117,7 +118,6 @@ async-fs = "2.2.0" async-lock = "3.4" async-stream = "0.3.6" async-trait = "0.1.89" -allocator-api2 = "0.2.21" base16ct = "1.0.0" bigdecimal = "0.4.8" bindgen = "0.72.0" From 3adc04a3847f6a608560d566be1029690fc3aed9 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 10:40:13 -0400 Subject: [PATCH 03/17] fix(buffer): align within raw allocations Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 47 ++++++++++++++++-------- vortex-buffer/src/buffer_mut.rs | 64 ++++++++++++++++++++++----------- 2 files changed, 76 insertions(+), 35 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 54e62e0d58e..4a58340bb70 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -19,6 +19,8 @@ use crate::Alignment; use crate::BufferMut; /// An allocator that can back a Vortex buffer. +/// +/// Vortex over-allocates raw storage and aligns the buffer within it. pub trait BufferAllocator: Allocator + Debug + Send + Sync + 'static {} impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static {} @@ -190,6 +192,7 @@ pub(crate) struct Allocation { ptr: NonNull, layout: Layout, capacity: usize, + buffer_alignment: Alignment, allocator: BufferAllocatorRef, } @@ -199,20 +202,34 @@ unsafe impl Send for Allocation {} unsafe impl Sync for Allocation {} impl Allocation { - pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { - Self::allocate_impl(layout, allocator, false) - } - - pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { - Self::allocate_impl(layout, allocator, true) - } - - fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { + pub(crate) fn allocate( + layout: Layout, + buffer_alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::allocate_impl(layout, buffer_alignment, allocator, false) + } + + pub(crate) fn allocate_zeroed( + layout: Layout, + buffer_alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::allocate_impl(layout, buffer_alignment, allocator, true) + } + + fn allocate_impl( + layout: Layout, + buffer_alignment: Alignment, + allocator: BufferAllocatorRef, + zeroed: bool, + ) -> Self { if layout.size() == 0 { return Self { ptr: layout.dangling_ptr(), layout, capacity: 0, + buffer_alignment, allocator, }; } @@ -228,6 +245,7 @@ impl Allocation { ptr: allocation.cast(), layout, capacity: allocation.len(), + buffer_alignment, allocator, } } @@ -240,17 +258,17 @@ impl Allocation { self.capacity } - pub(crate) fn alignment(&self) -> Alignment { - Alignment::new(self.layout.align()) + pub(crate) fn buffer_alignment(&self) -> Alignment { + self.buffer_alignment } pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } - pub(crate) fn grow(&mut self, layout: Layout) { + pub(crate) fn grow(&mut self, layout: Layout, buffer_alignment: Alignment) { if self.layout.size() == 0 { - *self = Self::allocate(layout, self.allocator.clone()); + *self = Self::allocate(layout, buffer_alignment, self.allocator.clone()); return; } @@ -262,6 +280,7 @@ impl Allocation { self.ptr = allocation.cast(); self.layout = layout; self.capacity = allocation.len(); + self.buffer_alignment = buffer_alignment; } } @@ -358,7 +377,7 @@ mod tests { assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!( state.alignment.load(Ordering::Relaxed), - *Alignment::DEFAULT_ALIGNMENT + *Alignment::of::() ); drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 0fa715c29bd..91eada1b885 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -115,14 +115,22 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let layout = Layout::from_size_align(size, *actual) + let allocation_size = if size == 0 { + 0 + } else { + size.checked_add(*actual - 1) + .vortex_expect("buffer capacity overflow") + }; + let allocation_alignment = if size == 0 { *actual } else { 1 }; + let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let allocation = Allocation::allocate(layout, allocator); - let capacity = allocation.size() / size_of::(); + let allocation = Allocation::allocate(layout, actual, allocator); + let offset = allocation.ptr().as_ptr().align_offset(*actual); + let capacity = (allocation.size() - offset) / size_of::(); Self { allocation, - offset: 0, + offset, length: 0, capacity, alignment, @@ -193,13 +201,21 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - let layout = Layout::from_size_align(size, *actual_alignment) + let allocation_size = if size == 0 { + 0 + } else { + size.checked_add(*actual_alignment - 1) + .vortex_expect("buffer length overflow") + }; + let allocation_alignment = if size == 0 { *actual_alignment } else { 1 }; + let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); - let allocation = Allocation::allocate_zeroed(layout, allocator); - let capacity = allocation.size() / size_of::(); + let allocation = Allocation::allocate_zeroed(layout, actual_alignment, allocator); + let offset = allocation.ptr().as_ptr().align_offset(*actual_alignment); + let capacity = (allocation.size() - offset) / size_of::(); Self { allocation, - offset: 0, + offset, length: len, capacity, alignment, @@ -441,26 +457,32 @@ impl BufferMut { let new_size = new_capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let physical_alignment = max(self.alignment, self.allocation.alignment()); - let layout = Layout::from_size_align(new_size, *physical_alignment) + let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); + let allocation_size = new_size + .checked_add(*physical_alignment - 1) + .vortex_expect("buffer capacity overflow"); + let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - if self.offset == 0 { - self.allocation.grow(layout); - } else { - let mut allocation = Allocation::allocate(layout, self.allocation.allocator().clone()); - // SAFETY: both ranges are valid for the initialized byte length and do not overlap. + let old_offset = self.offset; + self.allocation.grow(layout, physical_alignment); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(*physical_alignment); + if old_offset != new_offset { + // SAFETY: both ranges are within the allocation and may overlap. unsafe { - std::ptr::copy_nonoverlapping( - self.as_ptr().cast::(), - allocation.ptr().as_ptr(), + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); } - std::mem::swap(&mut self.allocation, &mut allocation); - self.offset = 0; } - self.capacity = self.allocation.size() / size_of::(); + self.offset = new_offset; + self.capacity = (self.allocation.size() - new_offset) / size_of::(); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. From 5562097da239c198ff1b681b03ecf6d7a0c5fabf Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 12:05:39 -0400 Subject: [PATCH 04/17] fix(buffer): avoid realloc when growing buffers Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 46 +++++++++++++++++++++------------ vortex-buffer/src/buffer_mut.rs | 31 +++++++++++----------- 2 files changed, 44 insertions(+), 33 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 4a58340bb70..7ed655e114d 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -265,23 +265,6 @@ impl Allocation { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } - - pub(crate) fn grow(&mut self, layout: Layout, buffer_alignment: Alignment) { - if self.layout.size() == 0 { - *self = Self::allocate(layout, buffer_alignment, self.allocator.clone()); - return; - } - - let allocation = - // SAFETY: ptr and layout describe a live block allocated by self.allocator. The new - // layout is at least as large as the old layout. - unsafe { self.allocator.grow(self.ptr, self.layout, layout) } - .unwrap_or_else(|_| handle_alloc_error(layout)); - self.ptr = allocation.cast(); - self.layout = layout; - self.capacity = allocation.len(); - self.buffer_alignment = buffer_alignment; - } } impl Drop for Allocation { @@ -345,6 +328,7 @@ mod tests { struct TrackingState { allocations: AtomicUsize, deallocations: AtomicUsize, + grows: AtomicUsize, alignment: AtomicUsize, } @@ -363,6 +347,17 @@ mod tests { // SAFETY: the caller passes the pointer and layout returned by Global. unsafe { Global.deallocate(ptr, layout) } } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.state.grows.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } } #[test] @@ -384,4 +379,21 @@ mod tests { drop(view); assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); } + + #[test] + fn buffer_growth_allocates_and_copies() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(1); + let initial_capacity = buffer.capacity(); + buffer.extend(std::iter::repeat_n(7, initial_capacity)); + + buffer.push(u32::MAX); + + assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]); + assert_eq!(buffer[initial_capacity], u32::MAX); + assert_eq!(state.allocations.load(Ordering::Relaxed), 2); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + } } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 91eada1b885..64c7baf8082 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -464,23 +464,22 @@ impl BufferMut { let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let old_offset = self.offset; - self.allocation.grow(layout, physical_alignment); - let new_offset = self - .allocation - .ptr() - .as_ptr() - .align_offset(*physical_alignment); - if old_offset != new_offset { - // SAFETY: both ranges are within the allocation and may overlap. - unsafe { - std::ptr::copy( - self.allocation.ptr().as_ptr().add(old_offset), - self.allocation.ptr().as_ptr().add(new_offset), - self.length * size_of::(), - ); - } + let allocation = Allocation::allocate( + layout, + physical_alignment, + self.allocation.allocator().clone(), + ); + let new_offset = allocation.ptr().as_ptr().align_offset(*physical_alignment); + // SAFETY: the source contains `length` initialized elements and the fresh allocation has + // room for at least `new_capacity` elements. The allocations do not overlap. + unsafe { + std::ptr::copy_nonoverlapping( + self.allocation.ptr().as_ptr().add(self.offset), + allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); } + self.allocation = allocation; self.offset = new_offset; self.capacity = (self.allocation.size() - new_offset) / size_of::(); } From 5598a2c8a5720e09a23e80afeaa415e9047c098e Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 13:02:20 -0400 Subject: [PATCH 05/17] perf(buffer): avoid indirection for static allocator Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 84 +++++++++++++++++++++++++-------- 1 file changed, 65 insertions(+), 19 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 7ed655e114d..2519241e4a6 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -8,7 +8,6 @@ use std::fmt; use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; -use std::sync::LazyLock; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; @@ -27,17 +26,18 @@ impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static /// A shared reference to a buffer allocator. #[derive(Clone)] -pub struct BufferAllocatorRef(Arc); +pub struct BufferAllocatorRef(Option>); impl BufferAllocatorRef { /// Wrap an allocator in a shared reference. pub fn new(allocator: impl BufferAllocator) -> Self { - Self(Arc::new(allocator)) + Self(Some(Arc::new(allocator))) } /// Return a shared reference to the static allocator. + #[inline] pub fn statically_allocated() -> Self { - STATIC_ALLOCATOR.clone() + Self(None) } /// Create a mutable buffer with this allocator. @@ -63,53 +63,100 @@ impl BufferAllocatorRef { impl Debug for BufferAllocatorRef { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt(f) + match &self.0 { + Some(allocator) => allocator.fmt(f), + None => StaticBufferAllocator.fmt(f), + } } } -// SAFETY: all calls are forwarded to the same allocator value held by the Arc. +// SAFETY: all calls are forwarded to the allocator represented by this value. unsafe impl Allocator for BufferAllocatorRef { + #[inline] fn allocate(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate(layout) + match &self.0 { + Some(allocator) => allocator.allocate(layout), + None => Global.allocate(layout), + } } + #[inline] fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate_zeroed(layout) + match &self.0 { + Some(allocator) => allocator.allocate_zeroed(layout), + None => Global.allocate_zeroed(layout), + } } + #[inline] unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.deallocate(ptr, layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.deallocate(ptr, layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.deallocate(ptr, layout) } + } + } } + #[inline] unsafe fn grow( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.grow(ptr, old_layout, new_layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + } } + #[inline] unsafe fn grow_zeroed( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.grow_zeroed(ptr, old_layout, new_layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } + } + } } + #[inline] unsafe fn shrink( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.shrink(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.shrink(ptr, old_layout, new_layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.shrink(ptr, old_layout, new_layout) } + } + } } } @@ -185,8 +232,7 @@ unsafe impl Allocator for StaticBufferAllocator { } } -static STATIC_ALLOCATOR: LazyLock = - LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); +static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); pub(crate) struct Allocation { ptr: NonNull, @@ -299,7 +345,7 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), + Self::External { .. } => &STATIC_ALLOCATOR, } } } From 12d50a78f0c9d7d25979e03e4054e3c791be74d4 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 13:33:23 -0400 Subject: [PATCH 06/17] Revert "perf(buffer): avoid indirection for static allocator" This reverts commit 1186945ab2f8c0b4d9f6e6dbbcdcf82cf67ac40b. Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 84 ++++++++------------------------- 1 file changed, 19 insertions(+), 65 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 2519241e4a6..7ed655e114d 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -8,6 +8,7 @@ use std::fmt; use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; +use std::sync::LazyLock; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; @@ -26,18 +27,17 @@ impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static /// A shared reference to a buffer allocator. #[derive(Clone)] -pub struct BufferAllocatorRef(Option>); +pub struct BufferAllocatorRef(Arc); impl BufferAllocatorRef { /// Wrap an allocator in a shared reference. pub fn new(allocator: impl BufferAllocator) -> Self { - Self(Some(Arc::new(allocator))) + Self(Arc::new(allocator)) } /// Return a shared reference to the static allocator. - #[inline] pub fn statically_allocated() -> Self { - Self(None) + STATIC_ALLOCATOR.clone() } /// Create a mutable buffer with this allocator. @@ -63,100 +63,53 @@ impl BufferAllocatorRef { impl Debug for BufferAllocatorRef { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match &self.0 { - Some(allocator) => allocator.fmt(f), - None => StaticBufferAllocator.fmt(f), - } + self.0.fmt(f) } } -// SAFETY: all calls are forwarded to the allocator represented by this value. +// SAFETY: all calls are forwarded to the same allocator value held by the Arc. unsafe impl Allocator for BufferAllocatorRef { - #[inline] fn allocate(&self, layout: Layout) -> Result, AllocError> { - match &self.0 { - Some(allocator) => allocator.allocate(layout), - None => Global.allocate(layout), - } + self.0.allocate(layout) } - #[inline] fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - match &self.0 { - Some(allocator) => allocator.allocate_zeroed(layout), - None => Global.allocate_zeroed(layout), - } + self.0.allocate_zeroed(layout) } - #[inline] unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.deallocate(ptr, layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.deallocate(ptr, layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.deallocate(ptr, layout) } } - #[inline] unsafe fn grow( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.grow(ptr, old_layout, new_layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.grow(ptr, old_layout, new_layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow(ptr, old_layout, new_layout) } } - #[inline] unsafe fn grow_zeroed( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.grow_zeroed(ptr, old_layout, new_layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } } - #[inline] unsafe fn shrink( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.shrink(ptr, old_layout, new_layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.shrink(ptr, old_layout, new_layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.shrink(ptr, old_layout, new_layout) } } } @@ -232,7 +185,8 @@ unsafe impl Allocator for StaticBufferAllocator { } } -static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); +static STATIC_ALLOCATOR: LazyLock = + LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); pub(crate) struct Allocation { ptr: NonNull, @@ -345,7 +299,7 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::External { .. } => &STATIC_ALLOCATOR, + Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), } } } From 73c12a619a0a29f7d5ba1dba4a42601209ac2dbc Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 13:35:05 -0400 Subject: [PATCH 07/17] perf(buffer): inline hot buffer growth paths Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 5 +++++ vortex-buffer/src/buffer_mut.rs | 6 ++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 7ed655e114d..c1eb4a32139 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -250,18 +250,22 @@ impl Allocation { } } + #[inline(always)] pub(crate) fn ptr(&self) -> NonNull { self.ptr } + #[inline(always)] pub(crate) fn size(&self) -> usize { self.capacity } + #[inline(always)] pub(crate) fn buffer_alignment(&self) -> Alignment { self.buffer_alignment } + #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } @@ -296,6 +300,7 @@ pub(crate) enum BufferBacking { } impl BufferBacking { + #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 64c7baf8082..b9b38e5f729 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -388,12 +388,14 @@ impl BufferMut { } /// Returns a raw pointer to the buffer's data. + #[inline(always)] pub fn as_ptr(&self) -> *const T { // SAFETY: offset always remains within the allocation. unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } } /// Returns a mutable raw pointer to the buffer's data. + #[inline(always)] pub fn as_mut_ptr(&mut self) -> *mut T { // SAFETY: BufferMut uniquely owns the allocation and offset is in bounds. unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } @@ -961,8 +963,8 @@ where impl FromIterator for BufferMut { fn from_iter>(iter: I) -> Self { - // We don't infer the capacity here and just let the first call to `extend` do it for us. - let mut buffer = Self::with_capacity(0); + let iter = iter.into_iter(); + let mut buffer = Self::with_capacity(iter.size_hint().0); buffer.extend(iter); buffer } From b01ecde734ee6a9a840a561b4cc2d95a7f49082f Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 14:02:12 -0400 Subject: [PATCH 08/17] perf(buffer): preserve aligned seed capacity Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 24 ++++++++---------------- 1 file changed, 8 insertions(+), 16 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index b9b38e5f729..4c6cb016603 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -115,14 +115,10 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let allocation_size = if size == 0 { - 0 - } else { - size.checked_add(*actual - 1) - .vortex_expect("buffer capacity overflow") - }; - let allocation_alignment = if size == 0 { *actual } else { 1 }; - let layout = Layout::from_size_align(allocation_size, allocation_alignment) + let allocation_size = size + .checked_add(*actual) + .vortex_expect("buffer capacity overflow"); + let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let allocation = Allocation::allocate(layout, actual, allocator); let offset = allocation.ptr().as_ptr().align_offset(*actual); @@ -201,14 +197,10 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - let allocation_size = if size == 0 { - 0 - } else { - size.checked_add(*actual_alignment - 1) - .vortex_expect("buffer length overflow") - }; - let allocation_alignment = if size == 0 { *actual_alignment } else { 1 }; - let layout = Layout::from_size_align(allocation_size, allocation_alignment) + let allocation_size = size + .checked_add(*actual_alignment) + .vortex_expect("buffer length overflow"); + let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); let allocation = Allocation::allocate_zeroed(layout, actual_alignment, allocator); let offset = allocation.ptr().as_ptr().align_offset(*actual_alignment); From c9c8f35e0cc0ab182267db85b7c1d17a1029c78d Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 14:22:41 -0400 Subject: [PATCH 09/17] perf(buffer): restore byte-based growth Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 4c6cb016603..043381e634c 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -446,15 +446,15 @@ impl BufferMut { .length .checked_add(additional) .vortex_expect("buffer capacity overflow"); - // Make sure we at least double in size each time we re-allocate to amortize the cost - let new_capacity = required.max(self.capacity.saturating_mul(2)); - let new_size = new_capacity + let required_size = required .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); - let allocation_size = new_size - .checked_add(*physical_alignment - 1) + let required_size = required_size + .checked_add(*physical_alignment) .vortex_expect("buffer capacity overflow"); + let current_size = self.allocation.size() - self.offset; + let allocation_size = required_size.max(current_size.saturating_mul(2)); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); @@ -465,7 +465,7 @@ impl BufferMut { ); let new_offset = allocation.ptr().as_ptr().align_offset(*physical_alignment); // SAFETY: the source contains `length` initialized elements and the fresh allocation has - // room for at least `new_capacity` elements. The allocations do not overlap. + // room for at least `required` elements. The allocations do not overlap. unsafe { std::ptr::copy_nonoverlapping( self.allocation.ptr().as_ptr().add(self.offset), From a38f8f1010d6922106412d33813b994d77a06118 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 14:58:33 -0400 Subject: [PATCH 10/17] fix(buffer): account for alignment when doubling Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 043381e634c..8822ab21140 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -450,11 +450,16 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); + let padding = *physical_alignment - 1; let required_size = required_size - .checked_add(*physical_alignment) + .checked_add(padding) .vortex_expect("buffer capacity overflow"); let current_size = self.allocation.size() - self.offset; - let allocation_size = required_size.max(current_size.saturating_mul(2)); + let doubled_size = current_size + .checked_mul(2) + .and_then(|size| size.checked_add(padding)) + .unwrap_or(usize::MAX); + let allocation_size = required_size.max(doubled_size); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); From 473d53744579c0f9e492a6f42640b785f1d418c2 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 15:06:35 -0400 Subject: [PATCH 11/17] Revert "fix(buffer): account for alignment when doubling" This reverts commit b55382807f5bf40accd7d456cb6ff25b0443b233. Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 8822ab21140..043381e634c 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -450,16 +450,11 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); - let padding = *physical_alignment - 1; let required_size = required_size - .checked_add(padding) + .checked_add(*physical_alignment) .vortex_expect("buffer capacity overflow"); let current_size = self.allocation.size() - self.offset; - let doubled_size = current_size - .checked_mul(2) - .and_then(|size| size.checked_add(padding)) - .unwrap_or(usize::MAX); - let allocation_size = required_size.max(doubled_size); + let allocation_size = required_size.max(current_size.saturating_mul(2)); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); From 498c5b3a0e42b4356faf45698f77da6f96e12251 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 16:40:52 -0400 Subject: [PATCH 12/17] perf(buffer): compact allocator-backed storage Signed-off-by: Nicholas Gates --- encodings/zstd/src/array.rs | 2 +- encodings/zstd/src/zstd_buffers.rs | 2 +- vortex-array/src/serde.rs | 6 +- vortex-buffer/src/alignment.rs | 40 +++---- vortex-buffer/src/allocation.rs | 144 +++++++++++++++--------- vortex-buffer/src/arrow.rs | 4 +- vortex-buffer/src/bit/buf_mut.rs | 19 ---- vortex-buffer/src/buffer.rs | 101 +++++++++++++---- vortex-buffer/src/buffer_mut.rs | 159 ++++++++++++--------------- vortex-cuda/src/device_buffer.rs | 2 +- vortex-file/src/footer/serializer.rs | 2 +- vortex-file/src/read/driver.rs | 4 +- vortex-file/src/segments/source.rs | 2 +- vortex-file/src/segments/writer.rs | 2 +- vortex-tui/src/inspect.rs | 12 +- vortex-tui/src/segments.rs | 2 +- vortex-web/crate/src/wasm.rs | 2 +- 17 files changed, 280 insertions(+), 225 deletions(-) diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 0d831e58775..81dda5940fd 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1144,7 +1144,7 @@ impl ZstdData { let value_bytes = values.buffer_handle().try_to_host_sync()?; // Align frames to buffer alignment. This is necessary for overaligned buffers. - let alignment = *value_bytes.alignment(); + let alignment = value_bytes.alignment().as_usize(); let step_width = (values_per_frame * byte_width).div_ceil(alignment) * alignment; let frame_byte_starts = (0..n_values * byte_width) diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 5f7f785f71f..4e62bc67857 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -357,7 +357,7 @@ fn compute_output_layout( let mut total_size = 0usize; for (&size, &alignment) in output_sizes.iter().zip(output_alignments.iter()) { - total_size = total_size.next_multiple_of(*alignment); + total_size = total_size.next_multiple_of(alignment.as_usize()); offsets.push(total_size); total_size += size; } diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index f84ec182269..df82cebe0f7 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -84,7 +84,7 @@ impl ArrayRef { .unwrap_or_else(FlatBuffer::alignment); // Create a shared buffer of zeros we can use for padding - let zeros = ByteBuffer::zeroed(*max_alignment); + let zeros = ByteBuffer::zeroed(max_alignment.as_usize()); // We push an empty buffer with the maximum alignment, so then subsequent buffers // will be aligned. For subsequent buffers, we always push a 1-byte alignment. @@ -96,7 +96,7 @@ impl ArrayRef { // Push all the array buffers with padding as necessary. for buffer in array_buffers { let padding = if options.include_padding { - let padding = pos.next_multiple_of(*buffer.alignment()) - pos; + let padding = pos.next_multiple_of(buffer.alignment().as_usize()) - pos; if padding > 0 { pos += padding; buffers.push(zeros.slice(0..padding)); @@ -139,7 +139,7 @@ impl ArrayRef { let fb_length = fb_buffer.len(); if options.include_padding { - let padding = pos.next_multiple_of(*FlatBuffer::alignment()) - pos; + let padding = pos.next_multiple_of(FlatBuffer::alignment().as_usize()) - pos; if padding > 0 { buffers.push(zeros.slice(0..padding)); } diff --git a/vortex-buffer/src/alignment.rs b/vortex-buffer/src/alignment.rs index 58f5d16f10c..f4f476c7691 100644 --- a/vortex-buffer/src/alignment.rs +++ b/vortex-buffer/src/alignment.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Display; -use std::ops::Deref; use vortex_error::VortexError; use vortex_error::VortexExpect; @@ -12,9 +11,9 @@ use vortex_error::vortex_err; /// The alignment of a buffer. /// -/// This type is a wrapper around `usize` that ensures the alignment is a non-zero power of 2. +/// This type stores the base-2 exponent of a non-zero power-of-two alignment. #[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct Alignment(usize); +pub struct Alignment(u8); impl Alignment { /// Largest alignment accepted from untrusted serialized input. @@ -41,10 +40,11 @@ impl Alignment { /// /// Panics if `align` is zero or is not a power of 2. #[inline] + #[expect(clippy::cast_possible_truncation, reason = "usize has at most 64 bits")] pub const fn new(align: usize) -> Self { assert!(align > 0, "Alignment must be greater than 0"); assert!(align.is_power_of_two(), "Alignment must be a power of 2"); - Self(align) + Self(align.trailing_zeros() as u8) } /// Create a new 1-byte alignment. @@ -104,7 +104,7 @@ impl Alignment { #[inline] pub const fn is_offset_aligned(&self, offset: usize) -> bool { // Alignment is always a power of 2, so a mask test is equivalent to `offset % self == 0`. - offset & (self.0 - 1) == 0 + offset & (self.as_usize() - 1) == 0 } /// Check if the given pointer is aligned to this alignment. @@ -115,8 +115,7 @@ impl Alignment { /// Returns the log2 of the alignment. pub fn exponent(&self) -> u8 { - u8::try_from(self.0.trailing_zeros()) - .vortex_expect("alignment is a power of 2 within usize, so its exponent fits in u8") + self.0 } /// Create from the log2 exponent of the alignment. @@ -131,7 +130,7 @@ impl Alignment { (exponent as u32) < usize::BITS, "Alignment exponent must fit in usize" ); - Self::new(1 << exponent) + Self(exponent) } /// Create from the log2 exponent of the alignment, returning an error rather than panicking if @@ -166,20 +165,17 @@ impl Alignment { } Ok(alignment) } -} -impl Display for Alignment { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.0) + /// Return the alignment in bytes. + #[inline] + pub const fn as_usize(self) -> usize { + 1 << self.0 } } -impl Deref for Alignment { - type Target = usize; - - #[inline] - fn deref(&self) -> &Self::Target { - &self.0 +impl Display for Alignment { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_usize()) } } @@ -200,14 +196,14 @@ impl From for Alignment { impl From for usize { #[inline] fn from(value: Alignment) -> Self { - value.0 + value.as_usize() } } impl From for u32 { #[inline] fn from(value: Alignment) -> Self { - u32::try_from(value.0).vortex_expect("Alignment must fit into u32") + u32::try_from(value.as_usize()).vortex_expect("Alignment must fit into u32") } } @@ -225,7 +221,7 @@ impl TryFrom for Alignment { return Err(vortex_err!("Alignment must be a power of 2, got {value}")); } - Ok(Self(value)) + Ok(Self::new(value)) } } @@ -243,7 +239,7 @@ mod test { fn alignment_above_u16() { // 64KiB alignment (one past `u16::MAX`) is valid — common on ARM with 64K pages. let alignment = Alignment::new(u16::MAX as usize + 1); - assert_eq!(*alignment, 1 << 16); + assert_eq!(alignment.as_usize(), 1 << 16); assert_eq!(alignment, Alignment::from_exponent(16)); } diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index c1eb4a32139..5ba2e54b44c 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -6,14 +6,15 @@ use std::alloc::Layout; use std::fmt; use std::fmt::Debug; +use std::mem::ManuallyDrop; use std::ptr::NonNull; use std::sync::Arc; -use std::sync::LazyLock; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; use allocator_api2::alloc::Global; use allocator_api2::alloc::handle_alloc_error; +use vortex_error::VortexExpect; use crate::Alignment; use crate::BufferMut; @@ -27,17 +28,17 @@ impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static /// A shared reference to a buffer allocator. #[derive(Clone)] -pub struct BufferAllocatorRef(Arc); +pub struct BufferAllocatorRef(Option>); impl BufferAllocatorRef { /// Wrap an allocator in a shared reference. pub fn new(allocator: impl BufferAllocator) -> Self { - Self(Arc::new(allocator)) + Self(Some(Arc::new(allocator))) } /// Return a shared reference to the static allocator. pub fn statically_allocated() -> Self { - STATIC_ALLOCATOR.clone() + Self(None) } /// Create a mutable buffer with this allocator. @@ -63,23 +64,35 @@ impl BufferAllocatorRef { impl Debug for BufferAllocatorRef { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt(f) + match &self.0 { + Some(allocator) => allocator.fmt(f), + None => StaticBufferAllocator.fmt(f), + } } } // SAFETY: all calls are forwarded to the same allocator value held by the Arc. unsafe impl Allocator for BufferAllocatorRef { fn allocate(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate(layout) + match &self.0 { + Some(allocator) => allocator.allocate(layout), + None => Global.allocate(layout), + } } fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate_zeroed(layout) + match &self.0 { + Some(allocator) => allocator.allocate_zeroed(layout), + None => Global.allocate_zeroed(layout), + } } unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.deallocate(ptr, layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.deallocate(ptr, layout) }, + None => unsafe { Global.deallocate(ptr, layout) }, + } } unsafe fn grow( @@ -89,7 +102,10 @@ unsafe impl Allocator for BufferAllocatorRef { new_layout: Layout, ) -> Result, AllocError> { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.grow(ptr, old_layout, new_layout) }, + None => unsafe { Global.grow(ptr, old_layout, new_layout) }, + } } unsafe fn grow_zeroed( @@ -99,7 +115,10 @@ unsafe impl Allocator for BufferAllocatorRef { new_layout: Layout, ) -> Result, AllocError> { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.grow_zeroed(ptr, old_layout, new_layout) }, + None => unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) }, + } } unsafe fn shrink( @@ -109,7 +128,10 @@ unsafe impl Allocator for BufferAllocatorRef { new_layout: Layout, ) -> Result, AllocError> { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.shrink(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.shrink(ptr, old_layout, new_layout) }, + None => unsafe { Global.shrink(ptr, old_layout, new_layout) }, + } } } @@ -185,14 +207,11 @@ unsafe impl Allocator for StaticBufferAllocator { } } -static STATIC_ALLOCATOR: LazyLock = - LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); +static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); pub(crate) struct Allocation { ptr: NonNull, layout: Layout, - capacity: usize, - buffer_alignment: Alignment, allocator: BufferAllocatorRef, } @@ -202,34 +221,35 @@ unsafe impl Send for Allocation {} unsafe impl Sync for Allocation {} impl Allocation { - pub(crate) fn allocate( - layout: Layout, - buffer_alignment: Alignment, - allocator: BufferAllocatorRef, - ) -> Self { - Self::allocate_impl(layout, buffer_alignment, allocator, false) - } - - pub(crate) fn allocate_zeroed( - layout: Layout, - buffer_alignment: Alignment, - allocator: BufferAllocatorRef, - ) -> Self { - Self::allocate_impl(layout, buffer_alignment, allocator, true) - } - - fn allocate_impl( - layout: Layout, - buffer_alignment: Alignment, - allocator: BufferAllocatorRef, - zeroed: bool, - ) -> Self { + pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, false) + } + + pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, true) + } + + pub(crate) fn from_vec(vec: Vec) -> Self { + assert!(!std::mem::needs_drop::()); + + let mut vec = ManuallyDrop::new(vec); + let layout = Layout::array::(vec.capacity()) + .unwrap_or_else(|_| unreachable!("a Vec capacity always has a valid layout")); + let ptr = NonNull::new(vec.as_mut_ptr().cast()) + .vortex_expect("a Vec always has a non-null pointer"); + + Self { + ptr, + layout, + allocator: BufferAllocatorRef::statically_allocated(), + } + } + + fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { if layout.size() == 0 { return Self { ptr: layout.dangling_ptr(), layout, - capacity: 0, - buffer_alignment, allocator, }; } @@ -244,8 +264,6 @@ impl Allocation { Self { ptr: allocation.cast(), layout, - capacity: allocation.len(), - buffer_alignment, allocator, } } @@ -257,18 +275,31 @@ impl Allocation { #[inline(always)] pub(crate) fn size(&self) -> usize { - self.capacity + self.layout.size() } #[inline(always)] - pub(crate) fn buffer_alignment(&self) -> Alignment { - self.buffer_alignment + pub(crate) fn alignment(&self) -> usize { + self.layout.align() } #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } + + pub(crate) fn grow(&mut self, new_layout: Layout) { + let allocation = if self.layout.size() == 0 { + self.allocator.allocate(new_layout) + } else { + // SAFETY: ptr denotes a live block owned by allocator, old_layout fits the block, and + // the caller only grows the allocation. + unsafe { self.allocator.grow(self.ptr, self.layout, new_layout) } + } + .unwrap_or_else(|_| handle_alloc_error(new_layout)); + self.ptr = allocation.cast(); + self.layout = new_layout; + } } impl Drop for Allocation { @@ -282,15 +313,21 @@ impl Drop for Allocation { } pub(crate) trait BufferOwner: Send + Sync + 'static { - fn as_slice(&self) -> &[u8]; + fn as_ptr(&self) -> *const u8; + + fn len(&self) -> usize; } impl BufferOwner for T where T: AsRef<[u8]> + Send + Sync + 'static, { - fn as_slice(&self) -> &[u8] { - self.as_ref() + fn as_ptr(&self) -> *const u8 { + self.as_ref().as_ptr() + } + + fn len(&self) -> usize { + self.as_ref().len() } } @@ -304,7 +341,7 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), + Self::External { .. } => &STATIC_ALLOCATOR, } } } @@ -377,7 +414,7 @@ mod tests { assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!( state.alignment.load(Ordering::Relaxed), - *Alignment::of::() + Alignment::of::().as_usize() ); drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); @@ -386,7 +423,7 @@ mod tests { } #[test] - fn buffer_growth_allocates_and_copies() { + fn buffer_growth_uses_allocator_grow() { let allocator = TrackingAllocator::default(); let state = Arc::clone(&allocator.state); let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(1); @@ -397,8 +434,11 @@ mod tests { assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]); assert_eq!(buffer[initial_capacity], u32::MAX); - assert_eq!(state.allocations.load(Ordering::Relaxed), 2); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + assert_eq!(state.grows.load(Ordering::Relaxed), 1); + + drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); - assert_eq!(state.grows.load(Ordering::Relaxed), 0); } } diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index efaddf6ee14..f63aaadaff0 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -28,7 +28,7 @@ impl Buffer { let bytes = Bytes::from_owner(ArrowWrapper(arrow.into_inner())); let alignment = Alignment::of::(); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if bytes.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment @@ -63,7 +63,7 @@ impl ByteBuffer { let length = arrow.len(); let bytes = Bytes::from_owner(ArrowWrapper(arrow)); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if bytes.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index a1bb0c82c04..734f65a08a6 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -599,25 +599,6 @@ impl BitBufferMut { self.len += bit_len; } - /// Absorbs a mutable buffer that was previously split off. - /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.append_buffer(&other). - pub fn unsplit(&mut self, other: Self) { - if (self.offset + self.len).is_multiple_of(8) && other.offset == 0 { - // We are aligned and can just append the buffers - self.buffer.unsplit(other.buffer); - self.len += other.len; - return; - } - - // Otherwise, we need to append the bits one by one - self.append_buffer(&other.freeze()) - } - /// Freeze the buffer in its current state into an immutable `BoolBuffer`. #[inline] pub fn freeze(self) -> BitBuffer { diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 13032c9ea6a..7c7f24285c9 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -34,6 +34,7 @@ pub struct Buffer { pub(crate) ptr: NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, + pub(crate) physical_alignment: Alignment, pub(crate) backing: Arc, } @@ -65,6 +66,7 @@ impl Default for Buffer { ptr: empty_ptr(), length: 0, alignment: Alignment::of::(), + physical_alignment: Alignment::MAX, backing: EMPTY_BACKING.clone(), } } @@ -106,6 +108,7 @@ impl Buffer { offset: usize, length: usize, alignment: Alignment, + physical_alignment: Alignment, ) -> Self { // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. let ptr = unsafe { allocation.ptr().add(offset).cast() }; @@ -113,25 +116,25 @@ impl Buffer { ptr, length, alignment, + physical_alignment, backing: Arc::new(BufferBacking::Owned(allocation)), } } fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self { - let bytes = owner.as_slice(); - let length = bytes.len() / size_of::(); + let owner: Box = Box::new(owner); + let length = owner.len() / size_of::(); let ptr = if length == 0 { empty_ptr() } else { - NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") + NonNull::new(owner.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") }; Self { ptr, length, alignment, - backing: Arc::new(BufferBacking::External { - _owner: Box::new(owner), - }), + physical_alignment: alignment, + backing: Arc::new(BufferBacking::External { _owner: owner }), } } @@ -225,6 +228,7 @@ impl Buffer { ptr: empty_ptr(), length: 0, alignment, + physical_alignment: Alignment::MAX, backing: EMPTY_BACKING.clone(), } } @@ -284,6 +288,7 @@ impl Buffer { ptr: buffer.ptr.cast(), length: buffer.length / size_of::(), alignment, + physical_alignment: buffer.physical_alignment, backing: buffer.backing, } } @@ -488,6 +493,7 @@ impl Buffer { ptr: unsafe { self.ptr.add(begin) }, length: end - begin, alignment, + physical_alignment: self.physical_alignment, backing: Arc::clone(&self.backing), } } @@ -541,6 +547,7 @@ impl Buffer { ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"), length: subset.len(), alignment, + physical_alignment: self.physical_alignment, backing: Arc::clone(&self.backing), } } @@ -560,6 +567,7 @@ impl Buffer { ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, + physical_alignment: self.physical_alignment, backing: self.backing, } } @@ -570,18 +578,18 @@ impl Buffer { ptr, length, alignment, + physical_alignment, backing, } = self; match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); - let capacity = (allocation.size() - offset) / size_of::(); Ok(BufferMut { allocation, offset, length, - capacity, alignment, + physical_alignment, _marker: Default::default(), }) } @@ -589,12 +597,14 @@ impl Buffer { ptr, length, alignment, + physical_alignment, backing: Arc::new(backing), }), Err(backing) => Err(Self { ptr, length, alignment, + physical_alignment, backing, }), } @@ -666,6 +676,7 @@ impl Buffer { ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, + physical_alignment: self.physical_alignment, backing: self.backing, } } @@ -747,15 +758,17 @@ impl FromIterator for Buffer { } } -// Helper struct to allow us to zero-copy any vec into a buffer +// Helper struct that preserves drop glue for non-native Vec elements. #[repr(transparent)] struct Wrapper(Vec); -impl AsRef<[u8]> for Wrapper { - fn as_ref(&self) -> &[u8] { - let data = self.0.as_ptr().cast::(); - let len = self.0.len() * size_of::(); - unsafe { std::slice::from_raw_parts(data, len) } +impl crate::BufferOwner for Wrapper { + fn as_ptr(&self) -> *const u8 { + self.0.as_ptr().cast() + } + + fn len(&self) -> usize { + self.0.len() * size_of::() } } @@ -764,12 +777,13 @@ where T: Send + Sync + 'static, { fn from(value: Vec) -> Self { - let wrapped_vec = Wrapper(value); - assert_eq!( - wrapped_vec.as_ref().as_ptr().align_offset(align_of::()), - 0 - ); - Self::from_owner(wrapped_vec, Alignment::of::()) + let length = value.len(); + let alignment = Alignment::of::(); + if std::mem::needs_drop::() { + Self::from_owner(Wrapper(value), alignment) + } else { + Self::from_allocation(Allocation::from_vec(value), 0, length, alignment, alignment) + } } } @@ -891,6 +905,11 @@ impl From> for Buffer { #[cfg(test)] mod test { + use std::mem::align_of; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use bytes::Buf; use crate::Alignment; @@ -1008,6 +1027,48 @@ mod test { assert_eq!(vec, buff.as_ref()); } + #[test] + fn from_vec_adopts_allocation() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u32, 2, 3, 4, 5]); + let ptr = vec.as_ptr(); + let capacity = vec.capacity(); + + let buffer = Buffer::from(vec); + assert_eq!(buffer.as_ptr(), ptr); + + let Ok(mut buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), capacity); + assert_eq!(buffer.allocation.alignment(), align_of::()); + + buffer.extend(6..=32); + assert_eq!(buffer.as_slice(), (1..=32).collect::>()); + assert_eq!(buffer.allocation.alignment(), align_of::()); + } + + #[test] + fn from_vec_preserves_drop_glue() { + struct DropValue(Arc); + + impl Drop for DropValue { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let values = (0..3) + .map(|_| DropValue(Arc::clone(&drops))) + .collect::>(); + let buffer = Buffer::from(values); + + assert_eq!(drops.load(Ordering::Relaxed), 0); + drop(buffer); + assert_eq!(drops.load(Ordering::Relaxed), 3); + } + #[test] fn empty_aligned_max_alignment() { // Empty buffers are backed by a static and must satisfy any valid alignment. diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 043381e634c..1c6864fd1c1 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -31,8 +31,8 @@ pub struct BufferMut { pub(crate) allocation: Allocation, pub(crate) offset: usize, pub(crate) length: usize, - pub(crate) capacity: usize, pub(crate) alignment: Alignment, + pub(crate) physical_alignment: Alignment, pub(crate) _marker: std::marker::PhantomData, } @@ -116,20 +116,18 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let allocation_size = size - .checked_add(*actual) + .checked_add(actual.as_usize()) .vortex_expect("buffer capacity overflow"); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let allocation = Allocation::allocate(layout, actual, allocator); - let offset = allocation.ptr().as_ptr().align_offset(*actual); - let capacity = (allocation.size() - offset) / size_of::(); - + let allocation = Allocation::allocate(layout, allocator); + let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); Self { allocation, offset, length: 0, - capacity, alignment, + physical_alignment: actual, _marker: Default::default(), } } @@ -198,19 +196,21 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); let allocation_size = size - .checked_add(*actual_alignment) + .checked_add(actual_alignment.as_usize()) .vortex_expect("buffer length overflow"); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); - let allocation = Allocation::allocate_zeroed(layout, actual_alignment, allocator); - let offset = allocation.ptr().as_ptr().align_offset(*actual_alignment); - let capacity = (allocation.size() - offset) / size_of::(); + let allocation = Allocation::allocate_zeroed(layout, allocator); + let offset = allocation + .ptr() + .as_ptr() + .align_offset(actual_alignment.as_usize()); Self { allocation, offset, length: len, - capacity, alignment, + physical_alignment: actual_alignment, _marker: Default::default(), } } @@ -376,7 +376,7 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.capacity + (self.allocation.size() - self.offset) / size_of::() } /// Returns a raw pointer to the buffer's data. @@ -431,7 +431,7 @@ impl BufferMut { /// Reserves capacity for at least `additional` more elements to be inserted in the buffer. #[inline] pub fn reserve(&mut self, additional: usize) { - if additional <= self.capacity - self.length { + if additional <= self.capacity() - self.length { // We can fit the additional bytes in the remaining capacity. Nothing to do. return; } @@ -449,33 +449,35 @@ impl BufferMut { let required_size = required .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); + let physical_alignment = max(self.alignment, self.physical_alignment); let required_size = required_size - .checked_add(*physical_alignment) + .checked_add(physical_alignment.as_usize()) .vortex_expect("buffer capacity overflow"); let current_size = self.allocation.size() - self.offset; let allocation_size = required_size.max(current_size.saturating_mul(2)); - let layout = Layout::from_size_align(allocation_size, 1) + let layout = Layout::from_size_align(allocation_size, self.allocation.alignment()) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let allocation = Allocation::allocate( - layout, - physical_alignment, - self.allocation.allocator().clone(), - ); - let new_offset = allocation.ptr().as_ptr().align_offset(*physical_alignment); - // SAFETY: the source contains `length` initialized elements and the fresh allocation has - // room for at least `required` elements. The allocations do not overlap. - unsafe { - std::ptr::copy_nonoverlapping( - self.allocation.ptr().as_ptr().add(self.offset), - allocation.ptr().as_ptr().add(new_offset), - self.length * size_of::(), - ); + let old_offset = self.offset; + self.allocation.grow(layout); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits overlap. + unsafe { + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } } - self.allocation = allocation; self.offset = new_offset; - self.capacity = (self.allocation.size() - new_offset) / size_of::(); + self.physical_alignment = physical_alignment; } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -617,71 +619,27 @@ impl BufferMut { self.length += slice.len(); } - /// Splits the buffer into two at the given index. - /// - /// Afterward, self contains elements `[0, at)`, and the returned buffer contains elements - /// `[at, capacity)`. The returned buffer uses a new allocation. - /// - /// Panics if either half would have a length that is not a multiple of the alignment. - pub fn split_off(&mut self, at: usize) -> Self { - if at > self.capacity() { - vortex_panic!("Cannot split buffer of capacity {} at {}", self.len(), at); - } - - let bytes_at = at * size_of::(); - if !self.alignment.is_offset_aligned(bytes_at) { - vortex_panic!( - "Cannot split buffer at {}, resulting alignment is not {}", - at, - self.alignment - ); - } - - let new_length = self.length.saturating_sub(at); - let new_capacity = self.capacity - at; - let mut other = Self::with_capacity_aligned_in( - new_capacity, - self.alignment, - self.allocation.allocator().clone(), - ); - if new_length > 0 { - other.extend_from_slice(&self.as_slice()[at..]); - } - self.length = self.length.min(at); - self.capacity = at; - - other - } - - /// Absorbs a mutable buffer that was previously split off. - /// - /// This appends the contents of `other` to this buffer. - pub fn unsplit(&mut self, other: Self) { - if self.alignment != other.alignment { - vortex_panic!( - "Cannot unsplit buffers with different alignments: {} and {}", - self.alignment, - other.alignment - ); - } - self.extend_from_slice(other.as_slice()); - } - /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { ByteBufferMut { allocation: self.allocation, offset: self.offset, length: self.length * size_of::(), - capacity: self.capacity * size_of::(), alignment: self.alignment, + physical_alignment: self.physical_alignment, _marker: Default::default(), } } /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { - Buffer::from_allocation(self.allocation, self.offset, self.length, self.alignment) + Buffer::from_allocation( + self.allocation, + self.offset, + self.length, + self.alignment, + self.physical_alignment, + ) } /// Map each element of the buffer with a closure. @@ -708,8 +666,12 @@ impl BufferMut { /// /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { - if self.as_ptr().align_offset(*alignment) == 0 { - Self { alignment, ..self } + if self.as_ptr().align_offset(alignment.as_usize()) == 0 { + Self { + alignment, + physical_alignment: max(self.physical_alignment, alignment), + ..self + } } else { let allocator = self.allocation.allocator().clone(); Self::copy_from_aligned_in(self, alignment, allocator) @@ -739,8 +701,8 @@ impl BufferMut { allocation: self.allocation, offset: self.offset, length: self.length, - capacity: self.capacity, alignment: self.alignment, + physical_alignment: self.physical_alignment, _marker: std::marker::PhantomData, } } @@ -982,7 +944,6 @@ impl Buf for ByteBufferMut { assert!(cnt <= self.length, "advance out of bounds"); self.offset += cnt; self.length -= cnt; - self.capacity -= cnt; } } @@ -1073,6 +1034,19 @@ mod test { assert_eq!(buf.alignment(), Alignment::new(1024)); } + #[test] + fn growth_preserves_alignment_and_values() { + let alignment = Alignment::new(4096); + let mut buffer = BufferMut::::with_capacity_aligned(1, alignment); + + for value in 0..10_000 { + buffer.push(value); + assert!(alignment.is_offset_aligned(buffer.as_ptr().addr())); + } + + assert_eq!(buffer.as_slice(), (0..10_000).collect::>()); + } + #[test] fn from_iter() { let buf = BufferMut::from_iter([0, 10, 20, 30]); @@ -1176,7 +1150,10 @@ mod test { let mut buf = BufferMut::::zeroed(LEN); - assert_eq!(buf.as_ptr().align_offset(*Alignment::of::()), 0); + assert_eq!( + buf.as_ptr().align_offset(Alignment::of::().as_usize()), + 0 + ); assert_eq!(buf.as_slice(), &[0; LEN]); buf[3] = 7; @@ -1190,7 +1167,7 @@ mod test { let mut buf = BufferMut::::zeroed_aligned(LEN, alignment); - assert_eq!(buf.as_ptr().align_offset(*alignment), 0); + assert_eq!(buf.as_ptr().align_offset(alignment.as_usize()), 0); assert_eq!(buf.as_slice(), &[0; LEN]); buf[3] = 7; diff --git a/vortex-cuda/src/device_buffer.rs b/vortex-cuda/src/device_buffer.rs index 1c0068bf578..869f5c4dc95 100644 --- a/vortex-cuda/src/device_buffer.rs +++ b/vortex-cuda/src/device_buffer.rs @@ -474,7 +474,7 @@ impl DeviceBuffer for CudaDeviceBuffer { fn aligned(self: Arc, alignment: Alignment) -> VortexResult> { let effective_ptr = self.device_ptr + self.offset as u64; - if effective_ptr.is_multiple_of(*alignment as u64) { + if effective_ptr.is_multiple_of(alignment.as_usize() as u64) { Ok(Arc::new(CudaDeviceBuffer { allocation: Arc::clone(&self.allocation), offset: self.offset, diff --git a/vortex-file/src/footer/serializer.rs b/vortex-file/src/footer/serializer.rs index 5010dadb0d2..a0b846d6b61 100644 --- a/vortex-file/src/footer/serializer.rs +++ b/vortex-file/src/footer/serializer.rs @@ -246,7 +246,7 @@ fn write_buffer( .map_err(|_| vortex_err!("metadata segment length exceeds maximum u32"))?; let alignment = buffer.alignment(); - let padding = offset.next_multiple_of(*alignment as u64) - *offset; + let padding = offset.next_multiple_of(alignment.as_usize() as u64) - *offset; let segment_offset = *offset + padding; let segment = PostscriptSegment { diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 616e9a52606..9f8b7aedbd3 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -228,7 +228,7 @@ impl State { let mut requests = vec![first_req]; let mut current_start = requests[0].offset; let mut current_end = requests[0].offset + requests[0].length as u64; - let align = *self.coalesced_buffer_alignment as u64; + let align = self.coalesced_buffer_alignment.as_usize() as u64; // Track requests that we've already decided to remove (or that were cancelled) so that // we don't repeatedly process them during range scans. @@ -593,7 +593,7 @@ mod tests { assert_eq!(coalesced.alignment(), Alignment::new(4)); for req in coalesced.requests() { let rel = req.offset - coalesced.range().start; - assert_eq!(rel % *req.alignment as u64, 0); + assert_eq!(rel % req.alignment.as_usize() as u64, 0); } } _ => panic!("Expected coalesced request"), diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 1b69f06e7c2..1ab7af31064 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -309,7 +309,7 @@ impl FileSegmentSource { let coalesce_config = reader.coalesce_config().map(|mut config| { // Aligning the coalesced start down can add up to (alignment - 1) bytes. // Increase max_size to keep the effective payload window consistent. - let extra = (*max_alignment as u64).saturating_sub(1); + let extra = (max_alignment.as_usize() as u64).saturating_sub(1); config.max_size = config.max_size.saturating_add(extra); config }); diff --git a/vortex-file/src/segments/writer.rs b/vortex-file/src/segments/writer.rs index e163c2cc868..1673ce40101 100644 --- a/vortex-file/src/segments/writer.rs +++ b/vortex-file/src/segments/writer.rs @@ -70,7 +70,7 @@ impl SegmentSink for BufferedSegmentSink { // Add any padding required to align the segment. let byte_offset = self.byte_offset.load(Ordering::Relaxed); - let padding = byte_offset.next_multiple_of(*alignment as u64) - byte_offset; + let padding = byte_offset.next_multiple_of(alignment.as_usize() as u64) - byte_offset; let offset = byte_offset + padding; specs.push(SegmentSpec { offset, diff --git a/vortex-tui/src/inspect.rs b/vortex-tui/src/inspect.rs index b035036be99..122c3ceec9a 100644 --- a/vortex-tui/src/inspect.rs +++ b/vortex-tui/src/inspect.rs @@ -185,22 +185,22 @@ async fn exec_inspect_json( dtype: ps.dtype.map(|s| SegmentInfoJson { offset: s.offset, length: s.length, - alignment: *s.alignment, + alignment: s.alignment.as_usize(), }), layout: SegmentInfoJson { offset: ps.layout.offset, length: ps.layout.length, - alignment: *ps.layout.alignment, + alignment: ps.layout.alignment.as_usize(), }, statistics: ps.statistics.map(|s| SegmentInfoJson { offset: s.offset, length: s.length, - alignment: *s.alignment, + alignment: s.alignment.as_usize(), }), footer: SegmentInfoJson { offset: ps.footer.offset, length: ps.footer.length, - alignment: *ps.footer.alignment, + alignment: ps.footer.alignment.as_usize(), }, }) } else { @@ -239,7 +239,7 @@ async fn exec_inspect_json( offset: segment.offset, end_offset: segment.offset + segment.length as u64, length: segment.length, - alignment: *segment.alignment, + alignment: segment.alignment.as_usize(), path: segment_paths[i] .as_ref() .map(|p| p.iter().map(|s| s.as_ref()).collect::>().join(".")), @@ -611,7 +611,7 @@ impl FooterSegments { print!( "{:>length_w$} {:>align_w$} ", segment.length, - *segment.alignment, + segment.alignment.as_usize(), length_w = length_width, align_w = alignment_width, ); diff --git a/vortex-tui/src/segments.rs b/vortex-tui/src/segments.rs index 64be95e2f1e..e8ca3d94ab7 100644 --- a/vortex-tui/src/segments.rs +++ b/vortex-tui/src/segments.rs @@ -90,7 +90,7 @@ pub async fn exec_segments(session: &VortexSession, args: SegmentsArgs) -> Vorte row_count: seg.row_count, byte_offset: seg.spec.offset, byte_length: seg.spec.length, - alignment: *seg.spec.alignment, + alignment: seg.spec.alignment.as_usize(), byte_gap, } }) diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index b22440d88b0..344b3b3b244 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -244,7 +244,7 @@ impl VortexFileHandle { index: i, byte_offset: spec.offset, byte_length: spec.length, - alignment: *spec.alignment, + alignment: spec.alignment.as_usize(), column, layout_path, } From 2cb6a2ba915746f62ad8b99976d929355e9cba4c Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 20:05:57 -0400 Subject: [PATCH 13/17] perf(buffer): avoid empty data allocations Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 17 +++++++++++++++ vortex-buffer/src/buffer_mut.rs | 38 +++++++++++++++++++++++---------- 2 files changed, 44 insertions(+), 11 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 5ba2e54b44c..107c7271999 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -441,4 +441,21 @@ mod tests { drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); } + + #[test] + fn zero_capacity_does_not_allocate() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(0); + + assert_eq!(buffer.capacity(), 0); + assert!(Alignment::DEFAULT_ALIGNMENT.is_offset_aligned(buffer.as_ptr().addr())); + assert_eq!(state.allocations.load(Ordering::Relaxed), 0); + + buffer.push(42); + + assert_eq!(buffer.as_slice(), [42]); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + } } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 1c6864fd1c1..588babcc39b 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -115,11 +115,17 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let allocation_size = size - .checked_add(actual.as_usize()) - .vortex_expect("buffer capacity overflow"); - let layout = Layout::from_size_align(allocation_size, 1) - .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); + let layout = if size == 0 { + Layout::from_size_align(0, actual.as_usize()) + .unwrap_or_else(|_| vortex_panic!("invalid empty buffer alignment")) + } else { + let allocation_size = size + .checked_add(actual.as_usize()) + .vortex_expect("buffer capacity overflow"); + Layout::from_size_align(allocation_size, 1).unwrap_or_else(|_| { + vortex_panic!("buffer capacity exceeds maximum allocation size") + }) + }; let allocation = Allocation::allocate(layout, allocator); let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); Self { @@ -195,11 +201,16 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - let allocation_size = size - .checked_add(actual_alignment.as_usize()) - .vortex_expect("buffer length overflow"); - let layout = Layout::from_size_align(allocation_size, 1) - .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); + let layout = if size == 0 { + Layout::from_size_align(0, actual_alignment.as_usize()) + .unwrap_or_else(|_| vortex_panic!("invalid empty buffer alignment")) + } else { + let allocation_size = size + .checked_add(actual_alignment.as_usize()) + .vortex_expect("buffer length overflow"); + Layout::from_size_align(allocation_size, 1) + .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")) + }; let allocation = Allocation::allocate_zeroed(layout, allocator); let offset = allocation .ptr() @@ -455,7 +466,12 @@ impl BufferMut { .vortex_expect("buffer capacity overflow"); let current_size = self.allocation.size() - self.offset; let allocation_size = required_size.max(current_size.saturating_mul(2)); - let layout = Layout::from_size_align(allocation_size, self.allocation.alignment()) + let allocation_alignment = if self.allocation.size() == 0 { + 1 + } else { + self.allocation.alignment() + }; + let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let old_offset = self.offset; From 3086a788c61bf6d6c2246131529985e42e5fa3f4 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 21:25:23 -0400 Subject: [PATCH 14/17] perf(buffer): copy live data for static growth Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 4 +++ vortex-buffer/src/buffer_mut.rs | 45 ++++++++++++++++++++++++--------- 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 107c7271999..415ba4fc4b4 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -41,6 +41,10 @@ impl BufferAllocatorRef { Self(None) } + pub(crate) fn is_statically_allocated(&self) -> bool { + self.0.is_none() + } + /// Create a mutable buffer with this allocator. pub fn with_capacity(&self, capacity: usize) -> BufferMut { BufferMut::with_capacity_in(capacity, self.clone()) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 588babcc39b..463f648befa 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -475,23 +475,44 @@ impl BufferMut { .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let old_offset = self.offset; - self.allocation.grow(layout); - let new_offset = self - .allocation - .ptr() - .as_ptr() - .align_offset(physical_alignment.as_usize()); - if new_offset != old_offset { - // SAFETY: grow preserved the initialized elements at old_offset. The new allocation - // has room for the requested elements plus alignment padding, and copy permits overlap. + let new_offset = if self.allocation.allocator().is_statically_allocated() { + let allocation = + Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); + let new_offset = allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + // SAFETY: both allocations have room for the initialized elements and do not overlap. unsafe { - std::ptr::copy( + std::ptr::copy_nonoverlapping( self.allocation.ptr().as_ptr().add(old_offset), - self.allocation.ptr().as_ptr().add(new_offset), + allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); } - } + self.allocation = allocation; + new_offset + } else { + self.allocation.grow(layout); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits + // overlap. + unsafe { + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } + } + new_offset + }; self.offset = new_offset; self.physical_alignment = physical_alignment; } From ef98cb98ffd63b6857b3ec03b4d419566aabef6d Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 21:59:21 -0400 Subject: [PATCH 15/17] perf(buffer): double logical growth capacity Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 463f648befa..93ed9476e26 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -461,11 +461,13 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.physical_alignment); - let required_size = required_size + let current_size = self.allocation.size() - self.offset; + let logical_size = required_size + .max(current_size.saturating_mul(2)) + .max(physical_alignment.as_usize()); + let allocation_size = logical_size .checked_add(physical_alignment.as_usize()) .vortex_expect("buffer capacity overflow"); - let current_size = self.allocation.size() - self.offset; - let allocation_size = required_size.max(current_size.saturating_mul(2)); let allocation_alignment = if self.allocation.size() == 0 { 1 } else { @@ -1084,6 +1086,19 @@ mod test { assert_eq!(buffer.as_slice(), (0..10_000).collect::>()); } + #[test] + fn growth_seeds_and_doubles_logical_capacity() { + let alignment = Alignment::new(64); + let mut buffer = BufferMut::::empty_aligned(alignment); + + buffer.push(0); + let capacity = buffer.capacity(); + assert!(capacity >= alignment.as_usize()); + + buffer.reserve(capacity); + assert!(buffer.capacity() >= capacity * 2); + } + #[test] fn from_iter() { let buf = BufferMut::from_iter([0, 10, 20, 30]); From f200f3a5e4c816c21a748f7e6124eeb7750ec71f Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 22:37:46 -0400 Subject: [PATCH 16/17] perf(buffer): exclude alignment slack from growth Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer.rs | 59 ++++++++++++++++++++++++++++++++- vortex-buffer/src/buffer_mut.rs | 27 ++++++++++++--- 2 files changed, 81 insertions(+), 5 deletions(-) diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 7c7f24285c9..70918ecbf37 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -35,6 +35,8 @@ pub struct Buffer { pub(crate) length: usize, pub(crate) alignment: Alignment, pub(crate) physical_alignment: Alignment, + // One physical-alignment block is reserved outside the logical capacity. + pub(crate) overallocated: bool, pub(crate) backing: Arc, } @@ -67,6 +69,7 @@ impl Default for Buffer { length: 0, alignment: Alignment::of::(), physical_alignment: Alignment::MAX, + overallocated: false, backing: EMPTY_BACKING.clone(), } } @@ -109,6 +112,7 @@ impl Buffer { length: usize, alignment: Alignment, physical_alignment: Alignment, + overallocated: bool, ) -> Self { // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. let ptr = unsafe { allocation.ptr().add(offset).cast() }; @@ -117,6 +121,7 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing: Arc::new(BufferBacking::Owned(allocation)), } } @@ -134,6 +139,7 @@ impl Buffer { length, alignment, physical_alignment: alignment, + overallocated: false, backing: Arc::new(BufferBacking::External { _owner: owner }), } } @@ -229,6 +235,7 @@ impl Buffer { length: 0, alignment, physical_alignment: Alignment::MAX, + overallocated: false, backing: EMPTY_BACKING.clone(), } } @@ -289,6 +296,7 @@ impl Buffer { length: buffer.length / size_of::(), alignment, physical_alignment: buffer.physical_alignment, + overallocated: buffer.overallocated, backing: buffer.backing, } } @@ -494,6 +502,7 @@ impl Buffer { length: end - begin, alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: Arc::clone(&self.backing), } } @@ -548,6 +557,7 @@ impl Buffer { length: subset.len(), alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: Arc::clone(&self.backing), } } @@ -568,6 +578,7 @@ impl Buffer { length: self.length * size_of::(), alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: self.backing, } } @@ -579,17 +590,25 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing, } = self; match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); + let overallocated = overallocated + && offset + == allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); Ok(BufferMut { allocation, offset, length, alignment, physical_alignment, + overallocated, _marker: Default::default(), }) } @@ -598,6 +617,7 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing: Arc::new(backing), }), Err(backing) => Err(Self { @@ -605,6 +625,7 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing, }), } @@ -677,6 +698,7 @@ impl Buffer { length: self.length, alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: self.backing, } } @@ -782,7 +804,14 @@ where if std::mem::needs_drop::() { Self::from_owner(Wrapper(value), alignment) } else { - Self::from_allocation(Allocation::from_vec(value), 0, length, alignment, alignment) + Self::from_allocation( + Allocation::from_vec(value), + 0, + length, + alignment, + alignment, + false, + ) } } } @@ -1048,6 +1077,34 @@ mod test { assert_eq!(buffer.allocation.alignment(), align_of::()); } + #[test] + fn from_u8_vec_preserves_capacity() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u8, 2, 3]); + + let buffer = Buffer::from(vec); + let Ok(buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), 16); + } + + #[test] + fn sliced_buffer_into_mut_has_safe_capacity() { + let mut original = crate::BufferMut::with_capacity(128); + original.extend(0u32..100); + let original = original.freeze(); + let sliced = original.slice(64..96); + drop(original); + + let Ok(mut sliced) = sliced.try_into_mut() else { + panic!("uniquely owned slice should become mutable") + }; + let capacity = sliced.capacity(); + sliced.push_n(0, capacity - sliced.len()); + assert_eq!(sliced.len(), capacity); + } + #[test] fn from_vec_preserves_drop_glue() { struct DropValue(Arc); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 93ed9476e26..18680c43434 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -33,6 +33,8 @@ pub struct BufferMut { pub(crate) length: usize, pub(crate) alignment: Alignment, pub(crate) physical_alignment: Alignment, + // One physical-alignment block is reserved outside the logical capacity. + pub(crate) overallocated: bool, pub(crate) _marker: std::marker::PhantomData, } @@ -134,6 +136,7 @@ impl BufferMut { length: 0, alignment, physical_alignment: actual, + overallocated: true, _marker: Default::default(), } } @@ -222,6 +225,7 @@ impl BufferMut { length: len, alignment, physical_alignment: actual_alignment, + overallocated: true, _marker: Default::default(), } } @@ -387,7 +391,15 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - (self.allocation.size() - self.offset) / size_of::() + if self.allocation.size() == 0 { + return 0; + } + + if !self.overallocated { + return (self.allocation.size() - self.offset) / size_of::(); + } + + (self.allocation.size() - self.physical_alignment.as_usize()) / size_of::() } /// Returns a raw pointer to the buffer's data. @@ -461,7 +473,10 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.physical_alignment); - let current_size = self.allocation.size() - self.offset; + let current_size = self + .capacity() + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); let logical_size = required_size .max(current_size.saturating_mul(2)) .max(physical_alignment.as_usize()); @@ -517,6 +532,7 @@ impl BufferMut { }; self.offset = new_offset; self.physical_alignment = physical_alignment; + self.overallocated = true; } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -666,6 +682,7 @@ impl BufferMut { length: self.length * size_of::(), alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, _marker: Default::default(), } } @@ -678,6 +695,7 @@ impl BufferMut { self.length, self.alignment, self.physical_alignment, + self.overallocated, ) } @@ -742,6 +760,7 @@ impl BufferMut { length: self.length, alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, _marker: std::marker::PhantomData, } } @@ -1093,10 +1112,10 @@ mod test { buffer.push(0); let capacity = buffer.capacity(); - assert!(capacity >= alignment.as_usize()); + assert_eq!(capacity, Alignment::DEFAULT_ALIGNMENT.as_usize()); buffer.reserve(capacity); - assert!(buffer.capacity() >= capacity * 2); + assert_eq!(buffer.capacity(), capacity * 2); } #[test] From 235c4c8ca275536e1e1ba75c447c5c5fe2c06a06 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 22:44:59 -0400 Subject: [PATCH 17/17] perf(buffer): store aligned mutable pointer Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer.rs | 2 +- vortex-buffer/src/buffer_mut.rs | 43 +++++++++++++++++++++------------ 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 70918ecbf37..114083e5699 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -604,7 +604,7 @@ impl Buffer { .align_offset(physical_alignment.as_usize()); Ok(BufferMut { allocation, - offset, + ptr, length, alignment, physical_alignment, diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 18680c43434..c6b0cf9a252 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -29,7 +29,7 @@ use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. pub struct BufferMut { pub(crate) allocation: Allocation, - pub(crate) offset: usize, + pub(crate) ptr: std::ptr::NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, pub(crate) physical_alignment: Alignment, @@ -38,6 +38,11 @@ pub struct BufferMut { pub(crate) _marker: std::marker::PhantomData, } +// SAFETY: BufferMut uniquely owns its allocation and only exposes T across threads. +unsafe impl Send for BufferMut {} +// SAFETY: shared access to BufferMut only exposes shared access to T. +unsafe impl Sync for BufferMut {} + impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. pub fn with_capacity(capacity: usize) -> Self { @@ -130,9 +135,11 @@ impl BufferMut { }; let allocation = Allocation::allocate(layout, allocator); let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); + // SAFETY: the allocation includes enough padding to reach this aligned pointer. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; Self { allocation, - offset, + ptr, length: 0, alignment, physical_alignment: actual, @@ -219,9 +226,11 @@ impl BufferMut { .ptr() .as_ptr() .align_offset(actual_alignment.as_usize()); + // SAFETY: the allocation includes enough padding to reach this aligned pointer. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; Self { allocation, - offset, + ptr, length: len, alignment, physical_alignment: actual_alignment, @@ -396,7 +405,8 @@ impl BufferMut { } if !self.overallocated { - return (self.allocation.size() - self.offset) / size_of::(); + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + return (self.allocation.size() - offset) / size_of::(); } (self.allocation.size() - self.physical_alignment.as_usize()) / size_of::() @@ -405,21 +415,19 @@ impl BufferMut { /// Returns a raw pointer to the buffer's data. #[inline(always)] pub fn as_ptr(&self) -> *const T { - // SAFETY: offset always remains within the allocation. - unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } + self.ptr.as_ptr() } /// Returns a mutable raw pointer to the buffer's data. #[inline(always)] pub fn as_mut_ptr(&mut self) -> *mut T { - // SAFETY: BufferMut uniquely owns the allocation and offset is in bounds. - unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } + self.ptr.as_ptr() } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_slice(&self) -> &[T] { - // SAFETY: the allocation is live, offset is in bounds, and construction checks alignment. + // SAFETY: ptr is in the live allocation and construction checks its alignment. unsafe { std::slice::from_raw_parts(self.as_ptr(), self.length) } } @@ -491,7 +499,7 @@ impl BufferMut { let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let old_offset = self.offset; + let old_offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); let new_offset = if self.allocation.allocator().is_statically_allocated() { let allocation = Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); @@ -502,7 +510,7 @@ impl BufferMut { // SAFETY: both allocations have room for the initialized elements and do not overlap. unsafe { std::ptr::copy_nonoverlapping( - self.allocation.ptr().as_ptr().add(old_offset), + self.ptr.cast::().as_ptr(), allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); @@ -530,7 +538,8 @@ impl BufferMut { } new_offset }; - self.offset = new_offset; + // SAFETY: new_offset was computed within the allocation for physical_alignment. + self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; self.physical_alignment = physical_alignment; self.overallocated = true; } @@ -678,7 +687,7 @@ impl BufferMut { pub fn into_byte_buffer(self) -> ByteBufferMut { ByteBufferMut { allocation: self.allocation, - offset: self.offset, + ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, physical_alignment: self.physical_alignment, @@ -689,9 +698,10 @@ impl BufferMut { /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); Buffer::from_allocation( self.allocation, - self.offset, + offset, self.length, self.alignment, self.physical_alignment, @@ -756,7 +766,7 @@ impl BufferMut { BufferMut { allocation: self.allocation, - offset: self.offset, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, physical_alignment: self.physical_alignment, @@ -1000,7 +1010,8 @@ impl Buf for ByteBufferMut { ); } assert!(cnt <= self.length, "advance out of bounds"); - self.offset += cnt; + // SAFETY: cnt is checked against the initialized length above. + self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; } }