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..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); - /// Default alignment for all buffers. + /// 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 2523587253c..966ac380828 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -4,16 +4,15 @@ //! Allocator-backed storage for Vortex buffers. use std::alloc::Layout; -use std::fmt; 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; @@ -27,130 +26,9 @@ 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(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. #[derive(Clone, Copy, Debug, Default)] @@ -224,7 +102,11 @@ unsafe impl Allocator for StaticBufferAllocator { } } -static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); +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, @@ -258,7 +140,7 @@ impl Allocation { Self { ptr, layout, - allocator: BufferAllocatorRef::statically_allocated(), + allocator: GLOBAL_ALLOCATOR_REF.clone(), } } @@ -314,7 +196,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)); @@ -335,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 @@ -346,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 { @@ -368,9 +244,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, } } } @@ -387,8 +263,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 { @@ -433,22 +312,41 @@ 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); @@ -466,7 +364,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)); @@ -486,10 +385,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); @@ -498,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 3b48889fa8d..ce08a8bbbc1 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -30,9 +30,14 @@ 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. 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>, } @@ -100,34 +105,18 @@ impl Buffer { } } - fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self { - 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") - }; - Self { - ptr, - length, - alignment, - backing: Some(Arc::new(BufferBacking::External { _owner: owner })), - } - } - - fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self { - let length = bytes.len() / size_of::(); + 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 ptr = - NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("Bytes pointer is null"); + NonNull::new(owner.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null"); Self { ptr, length, alignment, - backing: Some(Arc::new(BufferBacking::Bytes(bytes))), + backing: Some(Arc::new(BufferBacking::External { _owner: owner })), } } @@ -152,10 +141,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() } @@ -172,7 +158,7 @@ impl Buffer { /// /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned 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. @@ -204,7 +190,7 @@ impl Buffer { /// /// [`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)) + BufferMut::zeroed_aligned(len, alignment).freeze() } /// Create a new zeroed `Buffer` with the requested alignment. @@ -264,7 +250,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::()) @@ -275,8 +262,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,36 +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`. - pub fn from_bytes_aligned(bytes: Bytes, alignment: Alignment) -> Self { - 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 { @@ -396,7 +357,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, } } @@ -559,7 +520,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"); } @@ -631,7 +595,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::() @@ -697,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`. /// @@ -810,10 +806,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 @@ -824,7 +816,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) } @@ -892,7 +884,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 @@ -905,20 +897,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)) } } @@ -931,11 +924,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, } } } @@ -963,6 +956,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]; @@ -1082,6 +1089,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") @@ -1151,6 +1166,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); @@ -1177,6 +1193,89 @@ 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] 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 27c3dad3532..6093bd12944 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -19,16 +19,25 @@ 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. 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 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, } @@ -40,7 +49,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. @@ -55,11 +64,7 @@ impl BufferMut { /// /// [`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, - 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. @@ -89,7 +94,7 @@ impl BufferMut { capacity, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -132,7 +137,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::() }; @@ -148,7 +153,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. @@ -163,7 +168,7 @@ impl BufferMut { /// /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned 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. @@ -193,7 +198,7 @@ impl BufferMut { len, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -227,7 +232,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::() }; @@ -253,7 +258,7 @@ impl BufferMut { /// /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned 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. @@ -273,7 +278,7 @@ impl BufferMut { 0, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -282,7 +287,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. @@ -297,7 +302,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. @@ -316,7 +321,7 @@ 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_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. @@ -350,7 +355,7 @@ impl BufferMut { other, alignment, preferred_alignment, - BufferAllocatorRef::statically_allocated(), + DEFAULT_BUFFER_ALLOCATOR.clone(), ) } @@ -460,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, reserve additional + alignment bytes in case we need to realign the buffer. + // Otherwise, reserve additional alignment bytes in case we need to realign the buffer. self.reserve_allocate(additional); } @@ -498,41 +503,24 @@ impl BufferMut { .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. + 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_nonoverlapping( - self.ptr.cast::().as_ptr(), - allocation.ptr().as_ptr().add(new_offset), + 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; - 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::(); @@ -545,9 +533,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 @@ -730,7 +717,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 } } @@ -847,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 { @@ -867,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. @@ -893,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 @@ -906,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. @@ -983,10 +982,22 @@ impl FromIterator for BufferMut { #[cfg(test)] mod test { + use std::mem::size_of; + use crate::Alignment; 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::>(), 80); + } + #[test] fn capacity() { let mut n = 57; @@ -1029,15 +1040,74 @@ mod test { } #[test] - fn static_growth_copies_live_data() { + 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); + 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] + 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)?;