From 23246828f9ec7222bb09aac7c47443575ef982b5 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 3 Sep 2026 07:18:58 -0400 Subject: [PATCH 1/4] fix(buffer): address allocator review follow-ups Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 9 ++++++++- vortex-buffer/src/buffer.rs | 8 +++++++- vortex-buffer/src/buffer_mut.rs | 19 +++++++++++++++++-- 3 files changed, 32 insertions(+), 4 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 2523587253c..24bcfdfbaca 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -27,8 +27,15 @@ 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. +/// +/// The static allocator does not need shared ownership, so it is stored without an [`Arc`]. This +/// makes cloning the common static allocator a simple value copy. #[derive(Clone)] -pub struct BufferAllocatorRef(Option>); +pub struct BufferAllocatorRef( + // `None` selects the static allocator without allocating or updating an Arc reference count. + // `Some` keeps a custom allocator alive for as long as its buffers need it. + Option>, +); impl BufferAllocatorRef { /// Wrap an allocator in a shared reference. diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 3b48889fa8d..cc3e2df7ec4 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -30,9 +30,13 @@ use crate::trusted_len::TrustedLen; /// An immutable buffer of items of `T`. #[derive(Clone)] pub struct Buffer { + /// The first element in this view, or an aligned dangling pointer when the buffer is empty. pub(crate) ptr: NonNull, + /// The number of initialized `T` values visible from `ptr`. pub(crate) length: usize, + /// The minimum alignment promised for `ptr` and preserved by aligned slices. pub(crate) alignment: Alignment, + /// Shared ownership of the storage containing `ptr`; empty buffers have no backing. pub(crate) backing: Option>, } @@ -631,7 +635,9 @@ impl Buffer { match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); - let capacity = if allocation.size() == 0 { + let capacity = if size_of::() == 0 { + usize::MAX + } else if allocation.size() == 0 { 0 } else { (allocation.size() - offset) / size_of::() diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 27c3dad3532..7f2fe206e4a 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -24,11 +24,17 @@ use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. pub struct BufferMut { + /// The owned allocation, including any bytes before `ptr` used for alignment. pub(crate) allocation: Allocation, + /// The first element, aligned to `alignment`; it may dangle for an empty or zero-sized buffer. pub(crate) ptr: std::ptr::NonNull, + /// The number of initialized `T` values starting at `ptr`. pub(crate) length: usize, + /// The number of `T` values that fit from `ptr`; this is `usize::MAX` for zero-sized `T`. pub(crate) capacity: usize, + /// The minimum alignment maintained for `ptr` across reallocations. pub(crate) alignment: Alignment, + /// Marks the buffer as logically owning values of `T` despite storing an erased allocation. pub(crate) _marker: std::marker::PhantomData, } @@ -132,7 +138,7 @@ impl BufferMut { // SAFETY: the allocation includes enough padding to reach this aligned pointer. let ptr = unsafe { allocation.ptr().add(offset).cast() }; let capacity = if size_of::() == 0 { - capacity + usize::MAX } else { (allocation.size() - offset) / size_of::() }; @@ -227,7 +233,7 @@ impl BufferMut { // SAFETY: the allocation includes enough padding to reach this aligned pointer. let ptr = unsafe { allocation.ptr().add(offset).cast() }; let capacity = if size_of::() == 0 { - len + usize::MAX } else { (allocation.size() - offset) / size_of::() }; @@ -1028,6 +1034,15 @@ mod test { assert_eq!(buffer.capacity(), capacity * 2); } + #[test] + fn zero_sized_elements_grow() { + let mut buffer = BufferMut::<()>::empty(); + assert_eq!(buffer.capacity(), usize::MAX); + buffer.push(()); + buffer.push(()); + assert_eq!(buffer.len(), 2); + } + #[test] fn static_growth_copies_live_data() { let mut buffer = BufferMut::::with_capacity(1); From 5404fa16fb20a45602bacb8e0f3d7482d1e6c770 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 3 Sep 2026 10:14:19 -0400 Subject: [PATCH 2/4] refactor(buffer): use ArcRef for allocators Signed-off-by: Nicholas Gates --- Cargo.lock | 1 + .../src/arrays/chunked/vtable/canonical.rs | 12 +- vortex-array/src/executor.rs | 12 +- vortex-array/src/memory.rs | 10 +- vortex-buffer/Cargo.toml | 1 + vortex-buffer/benches/allocation.rs | 30 ++- vortex-buffer/src/alignment.rs | 2 +- vortex-buffer/src/allocation.rs | 242 ++++++------------ vortex-buffer/src/buffer.rs | 55 ++-- vortex-buffer/src/buffer_mut.rs | 240 +++++------------ vortex-file/src/open.rs | 7 +- vortex-io/src/object_store/read_at.rs | 11 +- vortex-io/src/std_file/read_at.rs | 7 +- 13 files changed, 222 insertions(+), 408 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0a367df1534..35e339d2de6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10738,6 +10738,7 @@ name = "vortex-buffer" version = "0.1.0" dependencies = [ "allocator-api2", + "arcref", "arrow-buffer 59.2.0", "bitvec", "bytes", diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index df6e5bd3739..fd3be992e9e 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -179,8 +180,8 @@ fn swizzle_list_chunks( // We (somewhat arbitrarily) choose `u64` for our offsets and sizes here. These can always be // narrowed later by the compressor. let allocator = ctx.allocator(); - let mut offsets = allocator.zeroed::(len); - let mut sizes = allocator.zeroed::(len); + let mut offsets = BufferMut::::zeroed_in(len, allocator.clone()); + let mut sizes = BufferMut::::zeroed_in(len, allocator.clone()); let offsets_out = offsets.as_mut_slice(); let sizes_slice_out = sizes.as_mut_slice(); let mut next_list = 0usize; @@ -659,10 +660,11 @@ mod tests { #[test] fn list_canonicalize_uses_memory_session_allocator() { let allocations = Arc::new(AtomicUsize::new(0)); - let session = - crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator { + let session = crate::array_session().with_allocator(BufferAllocatorRef::new_arc(Arc::new( + CountingAllocator { allocations: Arc::clone(&allocations), - })); + }, + ))); let mut ctx = session.create_execution_ctx(); let l1 = ListArray::try_new( diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index dab9ccfcfb5..1b431d359a5 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -995,9 +995,9 @@ mod tests { #[test] fn execution_ctx_allocator_override() { - let first = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); - let second = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); - let third = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); + let first = BufferAllocatorRef::new_arc(Arc::new(vortex_buffer::StaticBufferAllocator)); + let second = BufferAllocatorRef::new_arc(Arc::new(vortex_buffer::StaticBufferAllocator)); + let third = BufferAllocatorRef::new_arc(Arc::new(vortex_buffer::StaticBufferAllocator)); let session = VortexSession::empty() .with::() .with_allocator(first.clone()); @@ -1007,12 +1007,12 @@ mod tests { .get_mut::() .set_allocator(third.clone()); - assert!(session.allocator().ptr_eq(&third)); - assert!(ctx.allocator().ptr_eq(&third)); + assert!(std::ptr::eq(session.allocator().as_ref(), third.as_ref())); + assert!(std::ptr::eq(ctx.allocator().as_ref(), third.as_ref())); let ctx = ctx.with_allocator(second.clone()); session.get_mut::().set_allocator(first); - assert!(ctx.allocator().ptr_eq(&second)); + assert!(std::ptr::eq(ctx.allocator().as_ref(), second.as_ref())); } } diff --git a/vortex-array/src/memory.rs b/vortex-array/src/memory.rs index 3d328915933..08ca7875765 100644 --- a/vortex-array/src/memory.rs +++ b/vortex-array/src/memory.rs @@ -7,6 +7,7 @@ use std::any::Any; pub use vortex_buffer::BufferAllocator; pub use vortex_buffer::BufferAllocatorRef; +pub use vortex_buffer::DEFAULT_BUFFER_ALLOCATOR; pub use vortex_buffer::StaticBufferAllocator; use vortex_session::SessionExt; use vortex_session::SessionGuard; @@ -38,7 +39,7 @@ impl MemorySession { impl Default for MemorySession { fn default() -> Self { - Self::new(BufferAllocatorRef::statically_allocated()) + Self::new(DEFAULT_BUFFER_ALLOCATOR.clone()) } } @@ -76,16 +77,17 @@ impl MemorySessionExt for S {} #[cfg(test)] mod tests { - use vortex_buffer::BufferAllocatorRef; + use vortex_buffer::BufferMut; + use vortex_buffer::DEFAULT_BUFFER_ALLOCATOR; use super::MemorySession; #[test] fn memory_session_replaces_allocator() { - let allocator = BufferAllocatorRef::statically_allocated(); + let allocator = DEFAULT_BUFFER_ALLOCATOR.clone(); let mut session = MemorySession::default(); session.set_allocator(allocator); - let buffer = session.allocator().copy_from([1u32, 2, 3]); + let buffer = BufferMut::copy_from_in([1u32, 2, 3], session.allocator()); assert_eq!(buffer.as_slice(), [1, 2, 3]); } } diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index 02cf4ce7782..d7469c9c0a3 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -24,6 +24,7 @@ warn-copy = ["dep:tracing"] [dependencies] allocator-api2 = { workspace = true } +arcref = { workspace = true } arrow-buffer = { workspace = true } bitvec = { workspace = true } bytes = { workspace = true } diff --git a/vortex-buffer/benches/allocation.rs b/vortex-buffer/benches/allocation.rs index 6cd923b5dfb..b71908216fe 100644 --- a/vortex-buffer/benches/allocation.rs +++ b/vortex-buffer/benches/allocation.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use allocator_api2::alloc::Global; use arrow_buffer::MutableBuffer; use bytes::BytesMut; @@ -9,8 +11,10 @@ use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; +use vortex_buffer::StaticBufferAllocator; const SIZES: &[usize] = &[0, 64, 256, 1024, 16_384, 65_536]; +static GLOBAL_ALLOCATOR: Global = Global; fn main() { divan::main(); @@ -22,19 +26,20 @@ fn allocate_drop_vortex(bencher: Bencher, size: usize) { } #[divan::bench(args = SIZES)] -fn allocate_drop_vortex_custom(bencher: Bencher, size: usize) { +fn allocate_drop_vortex_arc(bencher: Bencher, size: usize) { bencher - .with_inputs(|| BufferAllocatorRef::new(Global)) - .bench_refs(|allocator| drop(allocator.with_capacity::(size))); + .with_inputs(|| BufferAllocatorRef::new_arc(Arc::new(StaticBufferAllocator))) + .bench_refs(|allocator| drop(BufferMut::::with_capacity_in(size, allocator.clone()))); } #[divan::bench(args = SIZES)] fn allocate_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) { bencher.bench(|| { - drop(BufferMut::::with_capacity_preferred_aligned( + drop(BufferMut::::with_capacity_preferred_aligned_in( size, Alignment::of::(), None, + BufferAllocatorRef::new_ref(&GLOBAL_ALLOCATOR), )) }); } @@ -55,18 +60,25 @@ fn allocate_freeze_drop_vortex(bencher: Bencher, size: usize) { } #[divan::bench(args = SIZES)] -fn allocate_freeze_drop_vortex_custom(bencher: Bencher, size: usize) { +fn allocate_freeze_drop_vortex_arc(bencher: Bencher, size: usize) { bencher - .with_inputs(|| BufferAllocatorRef::new(Global)) - .bench_refs(|allocator| drop(allocator.with_capacity::(size).freeze())); + .with_inputs(|| BufferAllocatorRef::new_arc(Arc::new(StaticBufferAllocator))) + .bench_refs(|allocator| { + drop(BufferMut::::with_capacity_in(size, allocator.clone()).freeze()) + }); } #[divan::bench(args = SIZES)] fn allocate_freeze_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) { bencher.bench(|| { drop( - BufferMut::::with_capacity_preferred_aligned(size, Alignment::of::(), None) - .freeze(), + BufferMut::::with_capacity_preferred_aligned_in( + size, + Alignment::of::(), + None, + BufferAllocatorRef::new_ref(&GLOBAL_ALLOCATOR), + ) + .freeze(), ) }); } diff --git a/vortex-buffer/src/alignment.rs b/vortex-buffer/src/alignment.rs index f4f476c7691..fdbbd313403 100644 --- a/vortex-buffer/src/alignment.rs +++ b/vortex-buffer/src/alignment.rs @@ -25,7 +25,7 @@ impl Alignment { /// Default alignment for device-to-host buffer copies. pub const HOST_COPY: Self = Alignment::new(256); - /// Default alignment for all buffers. + /// Preferred alignment of the default Vortex allocator. /// /// Chosen to be larger than any SIMD register (e.g. AVX-512's 64-byte /// registers) so that buffers can be processed with vectorized loads/stores diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 24bcfdfbaca..e517f4bacc1 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -4,162 +4,34 @@ //! Allocator-backed storage for Vortex buffers. use std::alloc::Layout; -use std::fmt; +use std::cmp::max; use std::fmt::Debug; use std::mem::ManuallyDrop; use std::ptr::NonNull; -use std::sync::Arc; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; use allocator_api2::alloc::Global; use allocator_api2::alloc::handle_alloc_error; +use arcref::ArcRef; use vortex_error::VortexExpect; 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 {} /// A shared reference to a buffer allocator. /// -/// The static allocator does not need shared ownership, so it is stored without an [`Arc`]. This -/// makes cloning the common static allocator a simple value copy. -#[derive(Clone)] -pub struct BufferAllocatorRef( - // `None` selects the static allocator without allocating or updating an Arc reference count. - // `Some` keeps a custom allocator alive for as long as its buffers need it. - Option>, -); - -impl BufferAllocatorRef { - /// Wrap an allocator in a shared reference. - pub fn new(allocator: impl BufferAllocator) -> Self { - Self(Some(Arc::new(allocator))) - } - - /// Return a shared reference to the static allocator. - pub fn statically_allocated() -> Self { - Self(None) - } - - pub(crate) fn static_ref() -> &'static Self { - &STATIC_ALLOCATOR - } - - pub(crate) fn is_statically_allocated(&self) -> bool { - self.0.is_none() - } - - /// Returns true if both references point to the same allocator. - pub fn ptr_eq(&self, other: &Self) -> bool { - match (&self.0, &other.0) { - (None, None) => true, - (Some(lhs), Some(rhs)) => Arc::ptr_eq(lhs, rhs), - _ => false, - } - } - - /// 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 { - 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> { - match &self.0 { - Some(allocator) => allocator.allocate(layout), - None => Global.allocate(layout), - } - } - - fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - 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. - match &self.0 { - Some(allocator) => unsafe { allocator.deallocate(ptr, layout) }, - None => 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. - 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( - &self, - ptr: NonNull, - old_layout: Layout, - new_layout: Layout, - ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - 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( - &self, - ptr: NonNull, - old_layout: Layout, - new_layout: Layout, - ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - match &self.0 { - Some(allocator) => unsafe { allocator.shrink(ptr, old_layout, new_layout) }, - None => unsafe { Global.shrink(ptr, old_layout, new_layout) }, - } - } -} +/// Use [`ArcRef::new_ref`] for a static allocator or [`ArcRef::new_arc`] for an owned allocator. +pub type BufferAllocatorRef = ArcRef; /// The allocator used by buffer APIs that do not take an allocator. +/// +/// It uses the global allocator and prefers at least 256-byte alignment. #[derive(Clone, Copy, Debug, Default)] pub struct StaticBufferAllocator; @@ -188,16 +60,18 @@ impl StaticBufferAllocator { // 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) + Global.allocate(Self::preferred_layout(layout)?) } fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - Global.allocate_zeroed(layout) + Global.allocate_zeroed(Self::preferred_layout(layout)?) } unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + let preferred_layout = + Self::preferred_layout(layout).unwrap_or_else(|_| handle_alloc_error(layout)); // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.deallocate(ptr, layout) } + unsafe { Global.deallocate(ptr, preferred_layout) } } unsafe fn grow( @@ -206,7 +80,10 @@ unsafe impl Allocator for StaticBufferAllocator { old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. + let old_layout = Self::preferred_layout(old_layout)?; + let new_layout = Self::preferred_layout(new_layout)?; + // SAFETY: the caller upholds the Allocator contract and both layouts use the same promoted + // alignment. unsafe { Global.grow(ptr, old_layout, new_layout) } } @@ -216,7 +93,10 @@ unsafe impl Allocator for StaticBufferAllocator { old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. + let old_layout = Self::preferred_layout(old_layout)?; + let new_layout = Self::preferred_layout(new_layout)?; + // SAFETY: the caller upholds the Allocator contract and both layouts use the same promoted + // alignment. unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } } @@ -226,12 +106,29 @@ unsafe impl Allocator for StaticBufferAllocator { old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. + let old_layout = Self::preferred_layout(old_layout)?; + let new_layout = Self::preferred_layout(new_layout)?; + // SAFETY: the caller upholds the Allocator contract and both layouts use the same promoted + // alignment. unsafe { Global.shrink(ptr, old_layout, new_layout) } } } -static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); +impl StaticBufferAllocator { + fn preferred_layout(layout: Layout) -> Result { + Layout::from_size_align( + layout.size(), + max(layout.align(), Alignment::DEFAULT_ALIGNMENT.as_usize()), + ) + .map_err(|_| AllocError) + } +} + +static STATIC_BUFFER_ALLOCATOR: StaticBufferAllocator = StaticBufferAllocator; +static GLOBAL_ALLOCATOR: Global = Global; +/// The allocator used by buffer APIs that do not take an allocator. +pub static DEFAULT_BUFFER_ALLOCATOR: BufferAllocatorRef = ArcRef::new_ref(&STATIC_BUFFER_ALLOCATOR); +pub(crate) static GLOBAL_ALLOCATOR_REF: BufferAllocatorRef = ArcRef::new_ref(&GLOBAL_ALLOCATOR); pub(crate) struct Allocation { ptr: NonNull, @@ -265,7 +162,7 @@ impl Allocation { Self { ptr, layout, - allocator: BufferAllocatorRef::statically_allocated(), + allocator: GLOBAL_ALLOCATOR_REF.clone(), } } @@ -304,8 +201,7 @@ impl Allocation { self.layout.size() } - #[allow(clippy::inline_always)] - #[inline(always)] + #[cfg(test)] pub(crate) fn alignment(&self) -> usize { self.layout.align() } @@ -321,7 +217,7 @@ impl Allocation { 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. + // new_layout is larger. Allocator::grow permits a change in alignment. unsafe { self.allocator.grow(self.ptr, self.layout, new_layout) } } .unwrap_or_else(|_| handle_alloc_error(new_layout)); @@ -375,9 +271,9 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::Bytes(_) | Self::External { .. } => &STATIC_ALLOCATOR, + Self::Bytes(_) | Self::External { .. } => &DEFAULT_BUFFER_ALLOCATOR, #[cfg(feature = "arrow")] - Self::Arrow(_) => &STATIC_ALLOCATOR, + Self::Arrow(_) => &DEFAULT_BUFFER_ALLOCATOR, } } } @@ -394,8 +290,11 @@ mod tests { use allocator_api2::alloc::Allocator; use allocator_api2::alloc::Global; + use super::DEFAULT_BUFFER_ALLOCATOR; + use super::GLOBAL_ALLOCATOR_REF; use crate::Alignment; use crate::BufferAllocatorRef; + use crate::BufferMut; #[derive(Clone, Debug, Default)] struct TrackingAllocator { @@ -440,28 +339,47 @@ mod tests { #[test] fn allocator_identity() { - let static_allocator = BufferAllocatorRef::statically_allocated(); - assert!(static_allocator.ptr_eq(&BufferAllocatorRef::statically_allocated())); - - let custom_allocator = BufferAllocatorRef::new(TrackingAllocator::default()); - assert!(custom_allocator.ptr_eq(&custom_allocator.clone())); - assert!(!custom_allocator.ptr_eq(&static_allocator)); - assert!(!custom_allocator.ptr_eq(&BufferAllocatorRef::new(TrackingAllocator::default()))); + let static_allocator = DEFAULT_BUFFER_ALLOCATOR.clone(); + assert!(std::ptr::eq( + static_allocator.as_ref(), + DEFAULT_BUFFER_ALLOCATOR.as_ref() + )); + + let global_allocator = GLOBAL_ALLOCATOR_REF.clone(); + assert!(std::ptr::eq( + global_allocator.as_ref(), + GLOBAL_ALLOCATOR_REF.as_ref() + )); + assert!(!std::ptr::eq( + global_allocator.as_ref(), + static_allocator.as_ref() + )); + + let custom_allocator = BufferAllocatorRef::new_arc(Arc::new(TrackingAllocator::default())); + assert!(std::ptr::eq( + custom_allocator.as_ref(), + custom_allocator.clone().as_ref() + )); + assert!(!std::ptr::eq( + custom_allocator.as_ref(), + static_allocator.as_ref() + )); + let other = BufferAllocatorRef::new_arc(Arc::new(TrackingAllocator::default())); + assert!(!std::ptr::eq(custom_allocator.as_ref(), other.as_ref())); } #[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 allocator = BufferAllocatorRef::new_arc(Arc::new(allocator)); + let buffer = BufferMut::copy_from_in([1u32, 2, 3, 4], allocator).freeze(); let view = buffer.slice(0..2); assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!( state.alignment.load(Ordering::Relaxed), - Alignment::of::().as_usize() + Alignment::of::().as_usize() ); drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); @@ -473,7 +391,8 @@ mod tests { 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); + let allocator = BufferAllocatorRef::new_arc(Arc::new(allocator)); + let mut buffer = BufferMut::::with_capacity_in(1, allocator); let initial_capacity = buffer.capacity(); buffer.extend(std::iter::repeat_n(7, initial_capacity)); @@ -493,10 +412,11 @@ mod tests { 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); + let allocator = BufferAllocatorRef::new_arc(Arc::new(allocator)); + let mut buffer = BufferMut::::with_capacity_in(0, allocator); assert_eq!(buffer.capacity(), 0); - assert!(Alignment::DEFAULT_ALIGNMENT.is_offset_aligned(buffer.as_ptr().addr())); + assert!(Alignment::of::().is_offset_aligned(buffer.as_ptr().addr())); assert_eq!(state.allocations.load(Ordering::Relaxed), 0); buffer.push(42); diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index cc3e2df7ec4..34b46ac2313 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -156,10 +156,7 @@ impl Buffer { /// 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 - /// of the provided `Vec` while maintaining the ability to convert it back into a mutable - /// buffer. We could fix this by forking `Bytes`, or in many other complex ways, but for now - /// callers should prefer to construct `Buffer` from a `BufferMut`. + /// This always copies. Use [`Buffer::from`] to take ownership of a `Vec` without copying. pub fn copy_from(values: impl AsRef<[T]>) -> Self { BufferMut::copy_from(values).freeze() } @@ -171,18 +168,15 @@ impl Buffer { /// 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 - /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. - /// - /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned + /// The default allocator may use a larger physical alignment. pub fn copy_from_aligned(values: impl AsRef<[T]>, alignment: Alignment) -> Self { - Self::copy_from_preferred_aligned(values, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + BufferMut::copy_from_aligned(values, alignment).freeze() } /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. /// - /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger - /// of `alignment` and `preferred_alignment`. + /// `preferred_alignment` raises the alignment requested from the allocator without changing + /// the alignment reported by the buffer. The allocator may use a larger alignment. pub fn copy_from_preferred_aligned( values: impl AsRef<[T]>, alignment: Alignment, @@ -203,18 +197,15 @@ impl Buffer { /// Create a new zeroed `Buffer` with the requested alignment. /// - /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than - /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. - /// - /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned + /// The default allocator may use a larger physical alignment. pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + BufferMut::zeroed_aligned(len, alignment).freeze() } /// Create a new zeroed `Buffer` with the requested alignment. /// - /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger - /// of `alignment` and `preferred_alignment`. + /// `preferred_alignment` raises the alignment requested from the allocator without changing + /// the alignment reported by the buffer. The allocator may use a larger alignment. pub fn zeroed_preferred_aligned( len: usize, alignment: Alignment, @@ -400,7 +391,7 @@ impl Buffer { pub fn allocator(&self) -> &BufferAllocatorRef { match self.backing.as_deref() { Some(backing) => backing.allocator(), - None => BufferAllocatorRef::static_ref(), + None => &crate::allocation::DEFAULT_BUFFER_ALLOCATOR, } } @@ -635,18 +626,17 @@ impl Buffer { match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); - let capacity = if size_of::() == 0 { - usize::MAX - } else if allocation.size() == 0 { - 0 - } else { - (allocation.size() - offset) / size_of::() - }; + if offset != 0 { + // SAFETY: the allocation is uniquely owned, both ranges are within it, and + // copy permits overlap. Moving the visible values to the allocation base lets + // BufferMut keep its data pointer equal to its allocation pointer. + unsafe { + std::ptr::copy(ptr.as_ptr(), allocation.ptr().as_ptr().cast::(), length); + } + } Ok(BufferMut { allocation, - ptr, length, - capacity, alignment, _marker: Default::default(), }) @@ -1088,6 +1078,14 @@ mod test { let buffer = Buffer::from(vec); assert_eq!(buffer.as_ptr(), ptr); + assert!(std::ptr::eq( + buffer.allocator().as_ref(), + crate::allocation::GLOBAL_ALLOCATOR_REF.as_ref() + )); + assert!(!std::ptr::eq( + buffer.allocator().as_ref(), + crate::allocation::DEFAULT_BUFFER_ALLOCATOR.as_ref() + )); let Ok(mut buffer) = buffer.try_into_mut() else { panic!("Vec-backed buffer should be uniquely owned") @@ -1157,6 +1155,7 @@ mod test { let Ok(mut sliced) = sliced.try_into_mut() else { panic!("uniquely owned slice should become mutable") }; + assert_eq!(sliced.as_slice(), (64..96).collect::>()); let capacity = sliced.capacity(); sliced.push_n(0, capacity - sliced.len()); assert_eq!(sliced.len(), capacity); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 7f2fe206e4a..98bd868b9c5 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -19,20 +19,17 @@ use crate::Allocation; use crate::Buffer; use crate::BufferAllocatorRef; use crate::ByteBufferMut; +use crate::allocation::DEFAULT_BUFFER_ALLOCATOR; use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. pub struct BufferMut { - /// The owned allocation, including any bytes before `ptr` used for alignment. + /// The owned allocation whose base pointer is the first element. pub(crate) allocation: Allocation, - /// The first element, aligned to `alignment`; it may dangle for an empty or zero-sized buffer. - pub(crate) ptr: std::ptr::NonNull, - /// The number of initialized `T` values starting at `ptr`. + /// The number of initialized `T` values starting at the allocation pointer. pub(crate) length: usize, - /// The number of `T` values that fit from `ptr`; this is `usize::MAX` for zero-sized `T`. - pub(crate) capacity: usize, - /// The minimum alignment maintained for `ptr` across reallocations. + /// The minimum alignment maintained for the allocation pointer across reallocations. pub(crate) alignment: Alignment, /// Marks the buffer as logically owning values of `T` despite storing an erased allocation. pub(crate) _marker: std::marker::PhantomData, @@ -46,7 +43,7 @@ unsafe impl Sync for BufferMut {} impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. pub fn with_capacity(capacity: usize) -> Self { - Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + Self::with_capacity_in(capacity, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create a new `BufferMut` with the requested capacity and allocator. @@ -56,16 +53,9 @@ impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. /// - /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than - /// `alignment`. Use [`with_capacity_preferred_aligned`] to control the over-alignment. - /// - /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned + /// The default allocator may use a larger physical alignment. pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { - Self::with_capacity_aligned_in( - capacity, - alignment, - BufferAllocatorRef::statically_allocated(), - ) + Self::with_capacity_aligned_in(capacity, alignment, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create a new `BufferMut` with the requested alignment, capacity, and allocator. @@ -74,18 +64,13 @@ impl BufferMut { alignment: Alignment, allocator: BufferAllocatorRef, ) -> Self { - Self::with_capacity_preferred_aligned_in( - capacity, - alignment, - Some(Alignment::DEFAULT_ALIGNMENT), - allocator, - ) + Self::with_capacity_preferred_aligned_in(capacity, alignment, None, allocator) } /// Create a new `BufferMut` with the requested alignment and capacity. /// - /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger - /// of `alignment` and `preferred_alignment`. + /// `preferred_alignment` raises the alignment requested from the allocator without changing + /// the alignment reported by the buffer. The allocator may use a larger alignment. pub fn with_capacity_preferred_aligned( capacity: usize, alignment: Alignment, @@ -95,7 +80,7 @@ impl BufferMut { capacity, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -122,31 +107,12 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - 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 layout = Layout::from_size_align(size, actual.as_usize()) + .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()); - // SAFETY: the allocation includes enough padding to reach this aligned pointer. - let ptr = unsafe { allocation.ptr().add(offset).cast() }; - let capacity = if size_of::() == 0 { - usize::MAX - } else { - (allocation.size() - offset) / size_of::() - }; Self { allocation, - ptr, length: 0, - capacity, alignment, _marker: Default::default(), } @@ -154,7 +120,7 @@ impl BufferMut { /// Create a new zeroed `BufferMut`. pub fn zeroed(len: usize) -> Self { - Self::zeroed_in(len, BufferAllocatorRef::statically_allocated()) + Self::zeroed_in(len, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create a new zeroed `BufferMut` with the requested allocator. @@ -164,12 +130,9 @@ impl BufferMut { /// Create a new zeroed `BufferMut` with the requested alignment. /// - /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than - /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. - /// - /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned + /// The default allocator may use a larger physical alignment. pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - Self::zeroed_aligned_in(len, alignment, BufferAllocatorRef::statically_allocated()) + Self::zeroed_aligned_in(len, alignment, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create a zeroed `BufferMut` with an alignment and allocator. @@ -178,18 +141,13 @@ impl BufferMut { alignment: Alignment, allocator: BufferAllocatorRef, ) -> Self { - Self::zeroed_preferred_aligned_in( - len, - alignment, - Some(Alignment::DEFAULT_ALIGNMENT), - allocator, - ) + Self::zeroed_preferred_aligned_in(len, alignment, None, allocator) } /// Create a new zeroed `BufferMut` with the requested alignment. /// - /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger - /// of `alignment` and `preferred_alignment`. + /// `preferred_alignment` raises the alignment requested from the allocator without changing + /// the alignment reported by the buffer. The allocator may use a larger alignment. pub fn zeroed_preferred_aligned( len: usize, alignment: Alignment, @@ -199,7 +157,7 @@ impl BufferMut { len, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -215,33 +173,12 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - 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 layout = Layout::from_size_align(size, actual_alignment.as_usize()) + .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); let allocation = Allocation::allocate_zeroed(layout, allocator); - let offset = allocation - .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() }; - let capacity = if size_of::() == 0 { - usize::MAX - } else { - (allocation.size() - offset) / size_of::() - }; Self { allocation, - ptr, length: len, - capacity, alignment, _marker: Default::default(), } @@ -254,12 +191,9 @@ impl BufferMut { /// Create a new empty `BufferMut` with the provided alignment. /// - /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than - /// `alignment`. Use [`empty_preferred_aligned`] to control the over-alignment. - /// - /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned + /// The default allocator may use a larger physical alignment. pub fn empty_aligned(alignment: Alignment) -> Self { - Self::empty_aligned_in(alignment, BufferAllocatorRef::statically_allocated()) + Self::empty_aligned_in(alignment, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create an empty `BufferMut` with an alignment and allocator. @@ -269,8 +203,8 @@ impl BufferMut { /// Create a new empty `BufferMut` with the provided alignment. /// - /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger - /// of `alignment` and `preferred_alignment`. + /// `preferred_alignment` raises the alignment requested from the allocator without changing + /// the alignment reported by the buffer. The allocator may use a larger alignment. pub fn empty_preferred_aligned( alignment: Alignment, preferred_alignment: Option, @@ -279,7 +213,7 @@ impl BufferMut { 0, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -288,7 +222,7 @@ impl BufferMut { where T: Copy, { - Self::full_in(item, len, BufferAllocatorRef::statically_allocated()) + Self::full_in(item, len, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create a full `BufferMut` with the given value and allocator. @@ -303,7 +237,7 @@ impl BufferMut { /// Create a mutable scalar buffer by copying the contents of the slice. pub fn copy_from(other: impl AsRef<[T]>) -> Self { - Self::copy_from_in(other, BufferAllocatorRef::statically_allocated()) + Self::copy_from_in(other, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create a mutable scalar buffer by copying with the given allocator. @@ -313,16 +247,13 @@ impl BufferMut { /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. /// - /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than - /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. - /// - /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned + /// The default allocator may use a larger physical alignment. /// /// ## Panics /// /// 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_aligned_in(other, alignment, BufferAllocatorRef::statically_allocated()) + Self::copy_from_aligned_in(other, alignment, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Copy values into a mutable buffer with the given alignment and allocator. @@ -331,18 +262,13 @@ impl BufferMut { alignment: Alignment, allocator: BufferAllocatorRef, ) -> Self { - Self::copy_from_preferred_aligned_in( - other, - alignment, - Some(Alignment::DEFAULT_ALIGNMENT), - allocator, - ) + Self::copy_from_preferred_aligned_in(other, alignment, None, allocator) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. /// - /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger - /// of `alignment` and `preferred_alignment`. + /// `preferred_alignment` raises the alignment requested from the allocator without changing + /// the alignment reported by the buffer. The allocator may use a larger alignment. /// /// ## Panics /// @@ -356,7 +282,7 @@ impl BufferMut { other, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -411,21 +337,25 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.capacity + if size_of::() == 0 { + usize::MAX + } else { + self.allocation.size() / size_of::() + } } /// Returns a raw pointer to the buffer's data. #[allow(clippy::inline_always)] #[inline(always)] pub fn as_ptr(&self) -> *const T { - self.ptr.as_ptr() + self.allocation.ptr().as_ptr().cast() } /// Returns a mutable raw pointer to the buffer's data. #[allow(clippy::inline_always)] #[inline(always)] pub fn as_mut_ptr(&mut self) -> *mut T { - self.ptr.as_ptr() + self.allocation.ptr().as_ptr().cast() } /// Returns a slice over the buffer of elements of type T. @@ -471,7 +401,7 @@ impl BufferMut { return; } - // Otherwise, reserve additional + alignment bytes in case we need to realign the buffer. + // Otherwise, grow the allocation. self.reserve_allocate(additional); } @@ -484,64 +414,16 @@ impl BufferMut { let required_size = required .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let alignment = self.alignment; let current_size = self - .capacity + .capacity() .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let logical_size = required_size .max(current_size.saturating_mul(2)) .max(Alignment::DEFAULT_ALIGNMENT.as_usize()); - let allocation_size = logical_size - .checked_add(alignment.as_usize()) - .vortex_expect("buffer capacity overflow"); - let allocation_alignment = if self.allocation.size() == 0 { - 1 - } else { - self.allocation.alignment() - }; - let layout = Layout::from_size_align(allocation_size, allocation_alignment) + let layout = Layout::from_size_align(logical_size, self.alignment.as_usize()) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - - 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()); - let new_offset = allocation.ptr().as_ptr().align_offset(alignment.as_usize()); - // SAFETY: both allocations have room for the initialized elements and do not overlap. - unsafe { - std::ptr::copy_nonoverlapping( - self.ptr.cast::().as_ptr(), - 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(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 - }; - // SAFETY: new_offset was computed within the allocation for alignment. - self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; - self.capacity = logical_size / size_of::(); + self.allocation.grow(layout); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -551,9 +433,8 @@ impl BufferMut { /// reading from a file) before marking the data as initialized using the /// [`set_len`] method. /// - /// Note that the returned slice may be larger than the capacity requested at - /// construction, since the underlying allocation can be rounded up (e.g. to - /// satisfy alignment requirements). + /// The returned slice may be larger than the capacity requested at construction because the + /// buffer grows geometrically. /// /// [`set_len`]: BufferMut::set_len /// [`Vec::spare_capacity_mut`]: Vec::spare_capacity_mut @@ -685,15 +566,9 @@ impl BufferMut { /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { - let capacity = self - .capacity - .checked_mul(size_of::()) - .vortex_expect("buffer capacity overflow"); ByteBufferMut { allocation: self.allocation, - ptr: self.ptr.cast(), length: self.length * size_of::(), - capacity, alignment: self.alignment, _marker: Default::default(), } @@ -701,8 +576,7 @@ 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, offset, self.length, self.alignment) + Buffer::from_allocation(self.allocation, 0, self.length, self.alignment) } /// Map each element of the buffer with a closure. @@ -736,7 +610,6 @@ impl BufferMut { let allocator = self.allocation.allocator().clone(); let mut aligned = Self::with_capacity_aligned_in(capacity, alignment, allocator); aligned.extend_from_slice(&self); - aligned.capacity = capacity; aligned } } @@ -762,9 +635,7 @@ impl BufferMut { BufferMut { allocation: self.allocation, - ptr: self.ptr.cast(), length: self.length, - capacity: self.capacity, alignment: self.alignment, _marker: std::marker::PhantomData, } @@ -989,10 +860,18 @@ impl FromIterator for BufferMut { #[cfg(test)] mod test { + use std::mem::size_of; + use crate::Alignment; use crate::BufferMut; use crate::buffer_mut; + #[cfg(target_pointer_width = "64")] + #[test] + fn compact_size() { + assert_eq!(size_of::>(), 64); + } + #[test] fn capacity() { let mut n = 57; @@ -1044,15 +923,12 @@ mod test { } #[test] - fn static_growth_copies_live_data() { + fn static_growth_preserves_live_data() { let mut buffer = BufferMut::::with_capacity(1); let capacity = buffer.capacity(); buffer.extend(std::iter::repeat_n(7, capacity)); - let old_ptr = buffer.as_ptr(); - buffer.push(u32::MAX); - assert_ne!(buffer.as_ptr(), old_ptr); assert_eq!(&buffer[..capacity], vec![7; capacity]); assert_eq!(buffer[capacity], u32::MAX); } diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index b54da20e192..488d32f4c4d 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -792,9 +792,10 @@ mod tests { std::fs::write(&file_path, ByteBuffer::from(buf).as_slice()).unwrap(); let allocations = Arc::new(AtomicUsize::new(0)); - let session = session.with_allocator(BufferAllocatorRef::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); + let session = + session.with_allocator(BufferAllocatorRef::new_arc(Arc::new(CountingAllocator { + allocations: Arc::clone(&allocations), + }))); let _file = session.open_options().open_path(&file_path).await.unwrap(); std::fs::remove_file(&file_path).unwrap(); diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index ed36c2a02a1..9b4a3675fc5 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -19,6 +19,8 @@ use object_store::path::Path as ObjectPath; use vortex_array::buffer::BufferHandle; use vortex_array::memory::BufferAllocatorRef; use vortex_buffer::Alignment; +use vortex_buffer::BufferMut; +use vortex_buffer::DEFAULT_BUFFER_ALLOCATOR; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -48,12 +50,7 @@ pub struct ObjectStoreReadAt { impl ObjectStoreReadAt { /// Create a new object store source. pub fn new(store: Arc, path: ObjectPath, handle: Handle) -> Self { - Self::new_with_allocator( - store, - path, - handle, - BufferAllocatorRef::statically_allocated(), - ) + Self::new_with_allocator(store, path, handle, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Create a new object store source with a custom writable buffer allocator. @@ -101,7 +98,7 @@ async fn read_object_store_range( alignment, } = request; let range = offset..(offset + length as u64); - let mut buffer = allocator.with_capacity_aligned::(length, alignment); + let mut buffer = BufferMut::::with_capacity_aligned_in(length, alignment, allocator); // SAFETY: each return path checks that every byte was initialized. unsafe { buffer.set_len(length) }; diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 21aac88923f..24a79e59e09 100644 --- a/vortex-io/src/std_file/read_at.rs +++ b/vortex-io/src/std_file/read_at.rs @@ -19,6 +19,8 @@ use futures::future::BoxFuture; use vortex_array::buffer::BufferHandle; use vortex_array::memory::BufferAllocatorRef; use vortex_buffer::Alignment; +use vortex_buffer::BufferMut; +use vortex_buffer::DEFAULT_BUFFER_ALLOCATOR; use vortex_error::VortexResult; use crate::CoalesceConfig; @@ -71,7 +73,7 @@ pub struct FileReadAt { impl FileReadAt { /// Open a file for reading. pub fn open(path: impl AsRef, handle: Handle) -> VortexResult { - Self::open_with_allocator(path, handle, BufferAllocatorRef::statically_allocated()) + Self::open_with_allocator(path, handle, DEFAULT_BUFFER_ALLOCATOR.clone()) } /// Open a file for reading using a custom writable buffer allocator. @@ -126,7 +128,8 @@ impl VortexReadAt for FileReadAt { async move { handle .spawn_blocking(move || { - let mut buffer = allocator.with_capacity_aligned::(length, alignment); + let mut buffer = + BufferMut::::with_capacity_aligned_in(length, alignment, allocator); // SAFETY: read_exact_at initializes every byte before the buffer is frozen. unsafe { buffer.set_len(length) }; read_exact_at(&file, buffer.as_mut_slice(), offset)?; From f894ab949e5a50a01a5d878343391eee03968f8a Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 3 Sep 2026 14:28:36 -0400 Subject: [PATCH 3/4] fix(buffer): restore manual over-alignment Signed-off-by: Nicholas Gates --- vortex-buffer/src/alignment.rs | 2 +- vortex-buffer/src/allocation.rs | 74 +++++---- vortex-buffer/src/buffer.rs | 201 ++++++++++++++++++----- vortex-buffer/src/buffer_mut.rs | 279 ++++++++++++++++++++++++++------ 4 files changed, 429 insertions(+), 127 deletions(-) diff --git a/vortex-buffer/src/alignment.rs b/vortex-buffer/src/alignment.rs index fdbbd313403..4541eb2f850 100644 --- a/vortex-buffer/src/alignment.rs +++ b/vortex-buffer/src/alignment.rs @@ -25,7 +25,7 @@ impl Alignment { /// Default alignment for device-to-host buffer copies. pub const HOST_COPY: Self = Alignment::new(256); - /// Preferred alignment of the default Vortex allocator. + /// Default preferred alignment for Vortex buffers. /// /// Chosen to be larger than any SIMD register (e.g. AVX-512's 64-byte /// registers) so that buffers can be processed with vectorized loads/stores diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index e517f4bacc1..966ac380828 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -4,7 +4,6 @@ //! Allocator-backed storage for Vortex buffers. use std::alloc::Layout; -use std::cmp::max; use std::fmt::Debug; use std::mem::ManuallyDrop; use std::ptr::NonNull; @@ -20,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 {} @@ -30,8 +31,6 @@ impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static pub type BufferAllocatorRef = ArcRef; /// The allocator used by buffer APIs that do not take an allocator. -/// -/// It uses the global allocator and prefers at least 256-byte alignment. #[derive(Clone, Copy, Debug, Default)] pub struct StaticBufferAllocator; @@ -60,18 +59,16 @@ impl StaticBufferAllocator { // 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(Self::preferred_layout(layout)?) + Global.allocate(layout) } fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - Global.allocate_zeroed(Self::preferred_layout(layout)?) + Global.allocate_zeroed(layout) } unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { - let preferred_layout = - Self::preferred_layout(layout).unwrap_or_else(|_| handle_alloc_error(layout)); // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.deallocate(ptr, preferred_layout) } + unsafe { Global.deallocate(ptr, layout) } } unsafe fn grow( @@ -80,10 +77,7 @@ unsafe impl Allocator for StaticBufferAllocator { old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - let old_layout = Self::preferred_layout(old_layout)?; - let new_layout = Self::preferred_layout(new_layout)?; - // SAFETY: the caller upholds the Allocator contract and both layouts use the same promoted - // alignment. + // SAFETY: the caller upholds the Allocator contract. unsafe { Global.grow(ptr, old_layout, new_layout) } } @@ -93,10 +87,7 @@ unsafe impl Allocator for StaticBufferAllocator { old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - let old_layout = Self::preferred_layout(old_layout)?; - let new_layout = Self::preferred_layout(new_layout)?; - // SAFETY: the caller upholds the Allocator contract and both layouts use the same promoted - // alignment. + // SAFETY: the caller upholds the Allocator contract. unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } } @@ -106,24 +97,11 @@ unsafe impl Allocator for StaticBufferAllocator { old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - let old_layout = Self::preferred_layout(old_layout)?; - let new_layout = Self::preferred_layout(new_layout)?; - // SAFETY: the caller upholds the Allocator contract and both layouts use the same promoted - // alignment. + // SAFETY: the caller upholds the Allocator contract. unsafe { Global.shrink(ptr, old_layout, new_layout) } } } -impl StaticBufferAllocator { - fn preferred_layout(layout: Layout) -> Result { - Layout::from_size_align( - layout.size(), - max(layout.align(), Alignment::DEFAULT_ALIGNMENT.as_usize()), - ) - .map_err(|_| AllocError) - } -} - static STATIC_BUFFER_ALLOCATOR: StaticBufferAllocator = StaticBufferAllocator; static GLOBAL_ALLOCATOR: Global = Global; /// The allocator used by buffer APIs that do not take an allocator. @@ -201,7 +179,8 @@ impl Allocation { self.layout.size() } - #[cfg(test)] + #[allow(clippy::inline_always)] + #[inline(always)] pub(crate) fn alignment(&self) -> usize { self.layout.align() } @@ -238,8 +217,6 @@ impl Drop for Allocation { pub(crate) trait BufferOwner: Send + Sync + 'static { fn as_ptr(&self) -> *const u8; - - fn len(&self) -> usize; } impl BufferOwner for T @@ -249,10 +226,6 @@ where fn as_ptr(&self) -> *const u8 { self.as_ref().as_ptr() } - - fn len(&self) -> usize { - self.as_ref().len() - } } pub(crate) enum BufferBacking { @@ -379,7 +352,7 @@ mod tests { assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!( state.alignment.load(Ordering::Relaxed), - Alignment::of::().as_usize() + Alignment::of::().as_usize() ); drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); @@ -425,4 +398,29 @@ mod tests { assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!(state.grows.load(Ordering::Relaxed), 0); } + + #[test] + fn zero_sized_buffers_do_not_call_the_allocator() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let allocator = BufferAllocatorRef::new_arc(Arc::new(allocator)); + + let mut buffer = BufferMut::<()>::with_capacity_in(usize::MAX, allocator.clone()); + buffer.extend([(); 4]); + let buffer = buffer.freeze(); + assert_eq!(buffer.len(), 4); + assert!(std::ptr::eq( + buffer.allocator().as_ref(), + allocator.as_ref() + )); + drop(buffer); + + let buffer = BufferMut::<()>::zeroed_in(4, allocator); + assert_eq!(buffer.len(), 4); + drop(buffer); + + assert_eq!(state.allocations.load(Ordering::Relaxed), 0); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + } } diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 34b46ac2313..1a51eb16a56 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -34,7 +34,8 @@ pub struct Buffer { pub(crate) ptr: NonNull, /// The number of initialized `T` values visible from `ptr`. pub(crate) length: usize, - /// The minimum alignment promised for `ptr` and preserved by aligned slices. + /// The minimum alignment promised for `ptr` and preserved by aligned slices. It may be larger + /// than the native alignment of `T`. pub(crate) alignment: Alignment, /// Shared ownership of the storage containing `ptr`; empty buffers have no backing. pub(crate) backing: Option>, @@ -104,14 +105,13 @@ impl Buffer { } } - fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self { + fn from_owner(owner: impl crate::BufferOwner, length: usize, alignment: Alignment) -> Self { + if length == 0 { + return Self::empty_aligned(alignment); + } let owner: Box = Box::new(owner); - let length = owner.len() / size_of::(); - let ptr = if length == 0 { - empty_ptr() - } else { - NonNull::new(owner.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") - }; + let ptr = + NonNull::new(owner.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null"); Self { ptr, length, @@ -121,6 +121,9 @@ impl Buffer { } fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self { + if size_of::() == 0 { + vortex_panic!("cannot infer a zero-sized buffer length from bytes"); + } let length = bytes.len() / size_of::(); if length == 0 { return Self::empty_aligned(alignment); @@ -168,15 +171,18 @@ impl Buffer { /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. /// - /// The default allocator may use a larger physical alignment. + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. + /// + /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned pub fn copy_from_aligned(values: impl AsRef<[T]>, alignment: Alignment) -> Self { BufferMut::copy_from_aligned(values, alignment).freeze() } /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. /// - /// `preferred_alignment` raises the alignment requested from the allocator without changing - /// the alignment reported by the buffer. The allocator may use a larger alignment. + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. pub fn copy_from_preferred_aligned( values: impl AsRef<[T]>, alignment: Alignment, @@ -197,15 +203,18 @@ impl Buffer { /// Create a new zeroed `Buffer` with the requested alignment. /// - /// The default allocator may use a larger physical alignment. + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. + /// + /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { BufferMut::zeroed_aligned(len, alignment).freeze() } /// Create a new zeroed `Buffer` with the requested alignment. /// - /// `preferred_alignment` raises the alignment requested from the allocator without changing - /// the alignment reported by the buffer. The allocator may use a larger alignment. + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. pub fn zeroed_preferred_aligned( len: usize, alignment: Alignment, @@ -259,7 +268,8 @@ impl Buffer { /// ## Panics /// /// Panics if the buffer is not aligned to the size of `T`, or the length is not a multiple of - /// the size of `T`. + /// the size of `T`. Also panics if `T` is zero-sized because bytes do not contain an element + /// count. pub fn from_byte_buffer(buffer: ByteBuffer) -> Self { // TODO(ngates): should this preserve the current alignment of the buffer? Self::from_byte_buffer_aligned(buffer, Alignment::of::()) @@ -270,8 +280,12 @@ impl Buffer { /// ## Panics /// /// 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`. + /// of the size of `T`, if the given alignment is not aligned to that of `T`, or if `T` is + /// zero-sized. pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self { + if size_of::() == 0 { + vortex_panic!("cannot infer a zero-sized buffer length from bytes"); + } if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must be compatible with the scalar type's alignment {}", @@ -302,8 +316,12 @@ impl Buffer { /// ## Panics /// /// Panics if the buffer is not aligned to the size of `T`, or the length is not a multiple of - /// the size of `T`. + /// the size of `T`. Also panics if `T` is zero-sized because bytes do not contain an element + /// count. pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self { + if size_of::() == 0 { + vortex_panic!("cannot infer a zero-sized buffer length from bytes"); + } if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must be compatible with the scalar type's alignment {}", @@ -554,7 +572,10 @@ impl Buffer { let subset_end = subset_start .checked_add(size_of_val(subset)) .vortex_expect("slice_ref address overflow"); - if subset_start < start || subset_end > end { + if subset_start < start + || subset_end > end + || (size_of::() == 0 && subset.len() > self.len()) + { vortex_panic!("slice_ref subset must be contained in the buffer"); } @@ -626,17 +647,18 @@ impl Buffer { match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); - if offset != 0 { - // SAFETY: the allocation is uniquely owned, both ranges are within it, and - // copy permits overlap. Moving the visible values to the allocation base lets - // BufferMut keep its data pointer equal to its allocation pointer. - unsafe { - std::ptr::copy(ptr.as_ptr(), allocation.ptr().as_ptr().cast::(), length); - } - } + let capacity = if size_of::() == 0 { + usize::MAX + } else if allocation.size() == 0 { + 0 + } else { + (allocation.size() - offset) / size_of::() + }; Ok(BufferMut { allocation, + ptr, length, + capacity, alignment, _marker: Default::default(), }) @@ -806,10 +828,6 @@ 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::() - } } impl From> for Buffer @@ -820,7 +838,7 @@ where let length = value.len(); let alignment = Alignment::of::(); if std::mem::needs_drop::() { - Self::from_owner(Wrapper(value), alignment) + Self::from_owner(Wrapper(value), length, alignment) } else { Self::from_allocation(Allocation::from_vec(value), 0, length, alignment) } @@ -888,7 +906,7 @@ pub struct BufferIterator { // Keep the buffer alive for the duration of the iteration. _buffer: Buffer, ptr: *const T, - end: *const T, + remaining: usize, } // SAFETY: `BufferIterator` is a `Buffer` plus two cursors into it, so it can safely be @@ -901,20 +919,21 @@ impl Iterator for BufferIterator { #[inline] fn next(&mut self) -> Option { - if self.ptr == self.end { + if self.remaining == 0 { None } else { - // SAFETY: ptr is within the buffer and has not reached end. + // SAFETY: `remaining` proves another initialized value exists. For a ZST, `ptr` is a + // suitably aligned dangling pointer and reading it does not access memory. let value = unsafe { self.ptr.read() }; self.ptr = unsafe { self.ptr.add(1) }; + self.remaining -= 1; Some(value) } } #[inline] fn size_hint(&self) -> (usize, Option) { - let remaining = unsafe { self.end.offset_from(self.ptr) } as usize; - (remaining, Some(remaining)) + (self.remaining, Some(self.remaining)) } } @@ -927,11 +946,11 @@ impl IntoIterator for Buffer { #[inline] fn into_iter(self) -> Self::IntoIter { let ptr = self.as_slice().as_ptr(); - let end = unsafe { ptr.add(self.len()) }; + let remaining = self.len(); BufferIterator { _buffer: self, ptr, - end, + remaining, } } } @@ -959,6 +978,20 @@ mod test { use crate::ByteBuffer; use crate::buffer; + #[repr(align(64))] + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + struct AlignedZst; + + static ZST_DROPS: AtomicUsize = AtomicUsize::new(0); + + struct DropZst; + + impl Drop for DropZst { + fn drop(&mut self) { + ZST_DROPS.fetch_add(1, Ordering::Relaxed); + } + } + #[test] fn align() { let buf = buffer![0u8, 1, 2]; @@ -1182,6 +1215,98 @@ mod test { assert_eq!(drops.load(Ordering::Relaxed), 3); } + #[test] + fn zero_sized_buffer_views_and_iteration() { + let buffer = Buffer::full(AlignedZst, 8); + assert_eq!(buffer.len(), 8); + assert_eq!(buffer.as_slice(), &[AlignedZst; 8]); + assert!(buffer.as_bytes().is_empty()); + assert!(Alignment::of::().is_ptr_aligned(buffer.as_ptr())); + + let sliced = buffer.slice(2..6); + assert_eq!(sliced.len(), 4); + assert_eq!(sliced.as_slice(), &[AlignedZst; 4]); + + let subset = &buffer.as_slice()[3..5]; + let sliced_ref = buffer.slice_ref(subset); + assert_eq!(sliced_ref.len(), 2); + + let mut iter = buffer.clone().into_iter(); + assert_eq!(iter.len(), 8); + assert_eq!(iter.by_ref().take(3).count(), 3); + assert_eq!(iter.len(), 5); + assert_eq!(iter.count(), 5); + + let bytes = buffer.clone().into_bytes(); + assert!(bytes.is_empty()); + drop(bytes); + + drop(sliced); + drop(sliced_ref); + let mutable = buffer + .try_into_mut() + .unwrap_or_else(|_| panic!("uniquely owned ZST buffer should become mutable")); + assert_eq!(mutable.len(), 8); + assert_eq!(mutable.capacity(), usize::MAX); + } + + #[test] + #[should_panic(expected = "slice_ref subset must be contained in the buffer")] + fn zero_sized_slice_ref_rejects_oversized_subset() { + let buffer = Buffer::full(AlignedZst, 8); + drop(buffer.slice_ref(&[AlignedZst; 9])); + } + + #[test] + fn zero_sized_vec_is_adopted() { + let values = vec![AlignedZst; 5]; + let ptr = values.as_ptr(); + let buffer = Buffer::from(values); + + assert_eq!(buffer.as_ptr(), ptr); + assert_eq!(buffer.len(), 5); + assert!(std::ptr::eq( + buffer.allocator().as_ref(), + crate::allocation::GLOBAL_ALLOCATOR_REF.as_ref() + )); + + let mutable = buffer + .try_into_mut() + .unwrap_or_else(|_| panic!("uniquely owned ZST buffer should become mutable")); + assert_eq!(mutable.len(), 5); + assert_eq!(mutable.capacity(), usize::MAX); + } + + #[test] + fn zero_sized_vec_preserves_length_and_drop_glue() { + ZST_DROPS.store(0, Ordering::Relaxed); + let values = (0..4).map(|_| DropZst).collect::>(); + let buffer = Buffer::from(values); + assert_eq!(buffer.len(), 4); + assert_eq!(buffer.as_slice().len(), 4); + + let view = buffer.slice(1..3); + drop(buffer); + assert_eq!(ZST_DROPS.load(Ordering::Relaxed), 0); + drop(view); + assert_eq!(ZST_DROPS.load(Ordering::Relaxed), 4); + } + + #[test] + #[should_panic(expected = "zero-sized")] + fn zero_sized_from_byte_buffer_is_rejected() { + drop(Buffer::<()>::from_byte_buffer(ByteBuffer::empty())); + } + + #[test] + #[should_panic(expected = "zero-sized")] + fn zero_sized_from_bytes_is_rejected() { + drop(Buffer::<()>::from_bytes_aligned( + Bytes::new(), + Alignment::of::<()>(), + )); + } + #[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 98bd868b9c5..6093bd12944 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -25,11 +25,17 @@ use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. pub struct BufferMut { - /// The owned allocation whose base pointer is the first element. + /// The owned allocation, including any bytes before `ptr` used for alignment. pub(crate) allocation: Allocation, - /// The number of initialized `T` values starting at the allocation pointer. + /// The first element, aligned to `alignment`; it may dangle for an empty or zero-sized buffer. + pub(crate) ptr: std::ptr::NonNull, + /// The number of initialized `T` values starting at `ptr`. pub(crate) length: usize, - /// The minimum alignment maintained for the allocation pointer across reallocations. + /// The number of `T` values that fit between `ptr` and the end of the allocation. This is + /// `usize::MAX` for zero-sized `T`. + pub(crate) capacity: usize, + /// The minimum alignment maintained for `ptr` across reallocations. It may be larger than the + /// native alignment of `T`. pub(crate) alignment: Alignment, /// Marks the buffer as logically owning values of `T` despite storing an erased allocation. pub(crate) _marker: std::marker::PhantomData, @@ -53,7 +59,10 @@ impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. /// - /// The default allocator may use a larger physical alignment. + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`with_capacity_preferred_aligned`] to control the over-alignment. + /// + /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { Self::with_capacity_aligned_in(capacity, alignment, DEFAULT_BUFFER_ALLOCATOR.clone()) } @@ -64,13 +73,18 @@ impl BufferMut { alignment: Alignment, allocator: BufferAllocatorRef, ) -> Self { - Self::with_capacity_preferred_aligned_in(capacity, alignment, None, allocator) + Self::with_capacity_preferred_aligned_in( + capacity, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a new `BufferMut` with the requested alignment and capacity. /// - /// `preferred_alignment` raises the alignment requested from the allocator without changing - /// the alignment reported by the buffer. The allocator may use a larger alignment. + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. pub fn with_capacity_preferred_aligned( capacity: usize, alignment: Alignment, @@ -107,12 +121,31 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let layout = Layout::from_size_align(size, actual.as_usize()) - .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()); + // SAFETY: the allocation includes enough padding to reach this aligned pointer. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; + let capacity = if size_of::() == 0 { + usize::MAX + } else { + (allocation.size() - offset) / size_of::() + }; Self { allocation, + ptr, length: 0, + capacity, alignment, _marker: Default::default(), } @@ -130,7 +163,10 @@ impl BufferMut { /// Create a new zeroed `BufferMut` with the requested alignment. /// - /// The default allocator may use a larger physical alignment. + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. + /// + /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { Self::zeroed_aligned_in(len, alignment, DEFAULT_BUFFER_ALLOCATOR.clone()) } @@ -141,13 +177,18 @@ impl BufferMut { alignment: Alignment, allocator: BufferAllocatorRef, ) -> Self { - Self::zeroed_preferred_aligned_in(len, alignment, None, allocator) + Self::zeroed_preferred_aligned_in( + len, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a new zeroed `BufferMut` with the requested alignment. /// - /// `preferred_alignment` raises the alignment requested from the allocator without changing - /// the alignment reported by the buffer. The allocator may use a larger alignment. + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. pub fn zeroed_preferred_aligned( len: usize, alignment: Alignment, @@ -173,12 +214,33 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - let layout = Layout::from_size_align(size, actual_alignment.as_usize()) - .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() + .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() }; + let capacity = if size_of::() == 0 { + usize::MAX + } else { + (allocation.size() - offset) / size_of::() + }; Self { allocation, + ptr, length: len, + capacity, alignment, _marker: Default::default(), } @@ -191,7 +253,10 @@ impl BufferMut { /// Create a new empty `BufferMut` with the provided alignment. /// - /// The default allocator may use a larger physical alignment. + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`empty_preferred_aligned`] to control the over-alignment. + /// + /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned pub fn empty_aligned(alignment: Alignment) -> Self { Self::empty_aligned_in(alignment, DEFAULT_BUFFER_ALLOCATOR.clone()) } @@ -203,8 +268,8 @@ impl BufferMut { /// Create a new empty `BufferMut` with the provided alignment. /// - /// `preferred_alignment` raises the alignment requested from the allocator without changing - /// the alignment reported by the buffer. The allocator may use a larger alignment. + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. pub fn empty_preferred_aligned( alignment: Alignment, preferred_alignment: Option, @@ -247,7 +312,10 @@ impl BufferMut { /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. /// - /// The default allocator may use a larger physical alignment. + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. + /// + /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned /// /// ## Panics /// @@ -262,13 +330,18 @@ impl BufferMut { alignment: Alignment, allocator: BufferAllocatorRef, ) -> Self { - Self::copy_from_preferred_aligned_in(other, alignment, None, allocator) + 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. /// - /// `preferred_alignment` raises the alignment requested from the allocator without changing - /// the alignment reported by the buffer. The allocator may use a larger alignment. + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. /// /// ## Panics /// @@ -337,25 +410,21 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - if size_of::() == 0 { - usize::MAX - } else { - self.allocation.size() / size_of::() - } + self.capacity } /// Returns a raw pointer to the buffer's data. #[allow(clippy::inline_always)] #[inline(always)] pub fn as_ptr(&self) -> *const T { - self.allocation.ptr().as_ptr().cast() + self.ptr.as_ptr() } /// Returns a mutable raw pointer to the buffer's data. #[allow(clippy::inline_always)] #[inline(always)] pub fn as_mut_ptr(&mut self) -> *mut T { - self.allocation.ptr().as_ptr().cast() + self.ptr.as_ptr() } /// Returns a slice over the buffer of elements of type T. @@ -396,12 +465,12 @@ 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; } - // Otherwise, grow the allocation. + // Otherwise, reserve additional alignment bytes in case we need to realign the buffer. self.reserve_allocate(additional); } @@ -414,16 +483,47 @@ impl BufferMut { let required_size = required .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); + let alignment = self.alignment; let current_size = self - .capacity() + .capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let logical_size = required_size .max(current_size.saturating_mul(2)) .max(Alignment::DEFAULT_ALIGNMENT.as_usize()); - let layout = Layout::from_size_align(logical_size, self.alignment.as_usize()) + let allocation_size = logical_size + .checked_add(alignment.as_usize()) + .vortex_expect("buffer capacity overflow"); + 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.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); self.allocation.grow(layout); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(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::(), + ); + } + } + // SAFETY: new_offset was computed within the allocation for alignment. + self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; + self.capacity = logical_size / size_of::(); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -566,9 +666,15 @@ impl BufferMut { /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { + let capacity = self + .capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); ByteBufferMut { allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length * size_of::(), + capacity, alignment: self.alignment, _marker: Default::default(), } @@ -576,7 +682,8 @@ impl BufferMut { /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { - Buffer::from_allocation(self.allocation, 0, self.length, self.alignment) + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + Buffer::from_allocation(self.allocation, offset, self.length, self.alignment) } /// Map each element of the buffer with a closure. @@ -635,7 +742,9 @@ impl BufferMut { BufferMut { allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length, + capacity: self.capacity, alignment: self.alignment, _marker: std::marker::PhantomData, } @@ -724,9 +833,9 @@ 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.spare_capacity_mut().as_mut_ptr().cast(); - let mut dst: *mut T = begin.cast_mut(); + let begin = self.spare_capacity_mut().as_mut_ptr().cast::(); + let mut dst = begin; + let mut zst_items_written = 0; // As a first step, we manually iterate the iterator up to the known capacity. for _ in 0..unwritten { @@ -744,11 +853,17 @@ impl BufferMut { // SAFETY: The offsets fits in `isize`, and because we were able to reserve the memory // we know that `add` will not overflow. unsafe { dst = dst.add(1) }; + if size_of::() == 0 { + zst_items_written += 1; + } } - // SAFETY: `dst` was derived from `begin`, which were both valid references to byte data, - // and since the only operation that `dst` has is `add`, we know that `dst >= begin`. - let items_written = unsafe { dst.offset_from_unsigned(begin) }; + let items_written = if size_of::() == 0 { + zst_items_written + } else { + // SAFETY: `dst` starts at `begin` and only advances within the reserved allocation. + unsafe { dst.offset_from_unsigned(begin) } + }; let length = self.len() + items_written; // SAFETY: We have written valid items between the old length and the new length. @@ -770,8 +885,9 @@ impl BufferMut { .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), ); - let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); - let mut dst: *mut T = begin.cast_mut(); + let begin = self.spare_capacity_mut().as_mut_ptr().cast::(); + let mut dst = begin; + let mut zst_items_written = 0; iter.for_each(|item| { // SAFETY: We have reserved enough capacity to hold this item, and `dst` is a pointer @@ -783,11 +899,17 @@ impl BufferMut { // SAFETY: The offset fits in `isize`, and because we were able to reserve the memory // we know that `add` will not overflow. unsafe { dst = dst.add(1) }; + if size_of::() == 0 { + zst_items_written += 1; + } }); - // SAFETY: `dst` starts at `begin` and advances by one for each item, so both pointers refer - // to the same allocation and `dst` is at or after `begin`. - let items_written = unsafe { dst.offset_from_unsigned(begin) }; + let items_written = if size_of::() == 0 { + zst_items_written + } else { + // SAFETY: `dst` starts at `begin` and only advances within the reserved allocation. + unsafe { dst.offset_from_unsigned(begin) } + }; let length = self.len() + items_written; // SAFETY: We have written valid items between the old length and the new length. @@ -866,10 +988,14 @@ mod test { use crate::BufferMut; use crate::buffer_mut; + #[repr(align(4096))] + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + struct AlignedZst; + #[cfg(target_pointer_width = "64")] #[test] fn compact_size() { - assert_eq!(size_of::>(), 64); + assert_eq!(size_of::>(), 80); } #[test] @@ -914,12 +1040,65 @@ mod test { } #[test] - fn zero_sized_elements_grow() { - let mut buffer = BufferMut::<()>::empty(); + fn zero_sized_capacity_and_alignment() { + let alignment = Alignment::new(4096); + let mut buffer = BufferMut::::with_capacity_aligned(usize::MAX, alignment); + assert_eq!(buffer.capacity(), usize::MAX); - buffer.push(()); - buffer.push(()); + assert_eq!(buffer.len(), 0); + assert!(alignment.is_ptr_aligned(buffer.as_ptr())); + + buffer.push(AlignedZst); + buffer.push(AlignedZst); assert_eq!(buffer.len(), 2); + assert_eq!(buffer.as_slice(), &[AlignedZst; 2]); + assert!(alignment.is_ptr_aligned(buffer.as_ptr())); + } + + #[test] + fn zero_sized_extend_paths() { + let mut buffer = BufferMut::::empty(); + + buffer.extend([AlignedZst; 3]); + buffer.extend_from_slice(&[AlignedZst; 2]); + buffer.extend_trusted([AlignedZst; 4].into_iter()); + buffer.push_n(AlignedZst, 2); + + let result = + buffer.try_extend_trusted([Ok(AlignedZst), Err("stop"), Ok(AlignedZst)].into_iter()); + assert_eq!(result, Err("stop")); + + assert_eq!(buffer.len(), 12); + assert_eq!(buffer.as_slice(), &[AlignedZst; 12]); + + let mut mapped = 0; + let buffer = buffer.clone().map_each_in_place(|value| { + mapped += 1; + value + }); + assert_eq!(mapped, 12); + + let mut buffer = buffer; + buffer.truncate(5); + assert_eq!(buffer.len(), 5); + buffer.clear(); + assert!(buffer.is_empty()); + } + + #[test] + fn zero_sized_zeroed_and_from_iter() { + let zeroed = BufferMut::::zeroed(7); + assert_eq!(zeroed.len(), 7); + assert_eq!(zeroed.as_slice(), &[AlignedZst; 7]); + + let bytes = zeroed.into_byte_buffer(); + assert_eq!(bytes.len(), 0); + assert_eq!(bytes.capacity(), 0); + assert!(Alignment::of::().is_ptr_aligned(bytes.as_ptr())); + + let collected = [AlignedZst; 6].into_iter().collect::>(); + assert_eq!(collected.len(), 6); + assert_eq!(collected.as_slice(), &[AlignedZst; 6]); } #[test] From 8861f3e512f3dbd9a26968dbc0ad4d599e1a5fc0 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 3 Sep 2026 14:38:18 -0400 Subject: [PATCH 4/4] refactor(buffer): restrict bytes conversion to byte buffers Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer.rs | 91 ++++++++++++------------------------- 1 file changed, 30 insertions(+), 61 deletions(-) diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 1a51eb16a56..ce08a8bbbc1 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -120,24 +120,6 @@ impl Buffer { } } - fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self { - if size_of::() == 0 { - vortex_panic!("cannot infer a zero-sized buffer length from bytes"); - } - let length = bytes.len() / size_of::(); - if length == 0 { - return Self::empty_aligned(alignment); - } - let ptr = - NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("Bytes pointer is null"); - Self { - ptr, - length, - alignment, - backing: Some(Arc::new(BufferBacking::Bytes(bytes))), - } - } - #[cfg(feature = "arrow")] pub(crate) fn from_arrow_owner( arrow: arrow_buffer::Buffer, @@ -311,40 +293,6 @@ impl Buffer { } } - /// Create a `Buffer` zero-copy from a `Bytes`. - /// - /// ## Panics - /// - /// Panics if the buffer is not aligned to the size of `T`, or the length is not a multiple of - /// the size of `T`. Also panics if `T` is zero-sized because bytes do not contain an element - /// count. - pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self { - if size_of::() == 0 { - vortex_panic!("cannot infer a zero-sized buffer length from bytes"); - } - 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(bytes.as_ptr()) { - vortex_panic!( - "Bytes alignment must align to the requested alignment {}", - alignment, - ); - } - if !bytes.len().is_multiple_of(size_of::()) { - vortex_panic!( - "Bytes length {} must be a multiple of the scalar type's size {}", - bytes.len(), - size_of::() - ); - } - Self::from_bytes(bytes, alignment) - } - /// Create a buffer with values from the TrustedLen iterator. /// Should be preferred over `from_iter` when the iterator is known to be `TrustedLen`. pub fn from_trusted_len_iter>(iter: I) -> Self { @@ -715,6 +663,36 @@ impl Buffer { } } +impl Buffer { + fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self { + if bytes.is_empty() { + return Self::empty_aligned(alignment); + } + let ptr = NonNull::new(bytes.as_ptr().cast_mut()).vortex_expect("Bytes pointer is null"); + Self { + ptr, + length: bytes.len(), + alignment, + backing: Some(Arc::new(BufferBacking::Bytes(bytes))), + } + } + + /// Create a byte buffer zero-copy from [`Bytes`] with the requested alignment. + /// + /// ## Panics + /// + /// Panics if `bytes` is not aligned to `alignment`. + pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self { + if !alignment.is_ptr_aligned(bytes.as_ptr()) { + vortex_panic!( + "Bytes alignment must align to the requested alignment {}", + alignment, + ); + } + Self::from_bytes(bytes, alignment) + } +} + impl Buffer { /// Transmute a `Buffer` into a `Buffer`. /// @@ -1298,15 +1276,6 @@ mod test { drop(Buffer::<()>::from_byte_buffer(ByteBuffer::empty())); } - #[test] - #[should_panic(expected = "zero-sized")] - fn zero_sized_from_bytes_is_rejected() { - drop(Buffer::<()>::from_bytes_aligned( - Bytes::new(), - Alignment::of::<()>(), - )); - } - #[test] fn empty_aligned_max_alignment() { // Empty buffers are backed by a static and must satisfy any valid alignment.