diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 2523587253c..69d9d4c0021 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -4,6 +4,7 @@ //! Allocator-backed storage for Vortex buffers. use std::alloc::Layout; +use std::any::Any; use std::fmt; use std::fmt::Debug; use std::mem::ManuallyDrop; @@ -27,8 +28,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. @@ -333,32 +341,13 @@ 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 -where - T: AsRef<[u8]> + Send + Sync + 'static, -{ - fn as_ptr(&self) -> *const u8 { - self.as_ref().as_ptr() - } - - fn len(&self) -> usize { - self.as_ref().len() - } -} - pub(crate) enum BufferBacking { Owned(Allocation), Bytes(bytes::Bytes), #[cfg(feature = "arrow")] Arrow(arrow_buffer::Buffer), External { - _owner: Box, + _owner: Box, }, } @@ -386,9 +375,13 @@ mod tests { use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; use allocator_api2::alloc::Global; + use rstest::rstest; + use vortex_error::VortexResult; + use vortex_error::vortex_err; use crate::Alignment; use crate::BufferAllocatorRef; + use crate::BufferMut; #[derive(Clone, Debug, Default)] struct TrackingAllocator { @@ -462,15 +455,18 @@ mod tests { assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); } - #[test] - fn buffer_growth_uses_allocator_grow() { + #[rstest] + fn buffer_growth_uses_allocator_grow(#[values(4, 64, 4096)] alignment: usize) { let allocator = TrackingAllocator::default(); let state = Arc::clone(&allocator.state); - let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(1); + let alignment = Alignment::new(alignment); + let mut buffer = + BufferAllocatorRef::new(allocator).with_capacity_aligned::(1, alignment); let initial_capacity = buffer.capacity(); buffer.extend(std::iter::repeat_n(7, initial_capacity)); buffer.push(u32::MAX); + assert!(alignment.is_ptr_aligned(buffer.as_ptr())); assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]); assert_eq!(buffer[initial_capacity], u32::MAX); @@ -498,4 +494,44 @@ mod tests { assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!(state.grows.load(Ordering::Relaxed), 0); } + + #[test] + fn empty_buffers_preserve_allocator_without_allocating() -> VortexResult<()> { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let allocator = BufferAllocatorRef::new(allocator); + let buffer = BufferMut::::zeroed_in(0, allocator.clone()); + let buffer = buffer.freeze(); + let copy = buffer.clone().into_mut(); + assert!(copy.allocator().ptr_eq(&allocator)); + let mut buffer = buffer + .try_into_mut() + .map_err(|_| vortex_err!("unique buffer"))?; + buffer.reserve(0); + assert!(buffer.is_empty()); + assert!(buffer.allocator().ptr_eq(&allocator)); + drop((copy, 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); + Ok(()) + } + + #[test] + fn shared_into_mut_preserves_allocator() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let allocator = BufferAllocatorRef::new(allocator); + let original = allocator.copy_from([1u32, 2, 3]).freeze(); + let mut copy = original.clone().into_mut(); + assert!(copy.allocator().ptr_eq(&allocator)); + copy[0] = 42; + assert_eq!(original.as_slice(), [1, 2, 3]); + assert_eq!(copy.as_slice(), [42, 2, 3]); + assert_eq!(state.allocations.load(Ordering::Relaxed), 2); + drop(copy); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + drop(original); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 2); + } } diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 3b48889fa8d..73d15991f0b 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -28,11 +28,38 @@ use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// An immutable buffer of items of `T`. +/// +/// Zero-sized element types are rejected at compile time when constructing a buffer. +/// +/// ```compile_fail +/// use vortex_buffer::Buffer; +/// let _ = Buffer::<()>::empty(); +/// ``` +/// +/// ```compile_fail +/// use vortex_buffer::Buffer; +/// let _ = Buffer::from(vec![(); 3]); +/// ``` +/// +/// ```compile_fail +/// use vortex_buffer::{Buffer, ByteBuffer}; +/// let _ = Buffer::<()>::from_byte_buffer(ByteBuffer::empty()); +/// ``` +/// +/// ```compile_fail +/// use bytes::Bytes; +/// use vortex_buffer::{Alignment, Buffer}; +/// let _ = Buffer::<()>::from_bytes_aligned(Bytes::new(), Alignment::none()); +/// ``` #[derive(Clone)] pub struct Buffer { + /// The first element in this view; may dangle for an empty 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. pub(crate) alignment: Alignment, + /// Shared ownership of the storage containing `ptr`, if any; `Buffer::empty` has no backing. pub(crate) backing: Option>, } @@ -100,22 +127,6 @@ 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::(); if length == 0 { @@ -228,6 +239,7 @@ impl Buffer { /// /// This does not allocate. Empty buffers use an aligned dangling pointer. pub fn empty_aligned(alignment: Alignment) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must align to the scalar type's alignment {}", @@ -277,6 +289,7 @@ impl Buffer { /// Panics if the buffer is not aligned to the given alignment, if the length is not a multiple /// of the size of `T`, or if the given alignment is not aligned to that of `T`. pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must be compatible with the scalar type's alignment {}", @@ -309,6 +322,7 @@ impl Buffer { /// 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 { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must be compatible with the scalar type's alignment {}", @@ -802,29 +816,25 @@ impl FromIterator for Buffer { } } -// Helper struct that preserves drop glue for non-native Vec elements. -#[repr(transparent)] -struct Wrapper(Vec); - -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 where T: Send + Sync + 'static, { fn from(value: Vec) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; let length = value.len(); let alignment = Alignment::of::(); if std::mem::needs_drop::() { - Self::from_owner(Wrapper(value), alignment) + // Keep the typed owner so its elements are dropped. + Self { + ptr: NonNull::new(value.as_ptr().cast_mut()) + .vortex_expect("a Vec always has a non-null pointer"), + length, + alignment, + backing: Some(Arc::new(BufferBacking::External { + _owner: Box::new(value), + })), + } } else { Self::from_allocation(Allocation::from_vec(value), 0, length, alignment) } @@ -1094,6 +1104,42 @@ mod test { assert_eq!(buffer.allocation.alignment(), align_of::()); } + #[test] + fn byte_owner_preserves_slice_and_lifetime() { + struct Owner { + values: Vec, + drops: Arc, + } + + impl AsRef<[u8]> for Owner { + fn as_ref(&self) -> &[u8] { + &self.values[1..4] + } + } + + impl Drop for Owner { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::Relaxed); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let owner = Owner { + values: vec![0, 1, 2, 3, 4], + drops: Arc::clone(&drops), + }; + let ptr = owner.as_ref().as_ptr(); + let buffer = ByteBuffer::from(Bytes::from_owner(owner)); + assert_eq!(buffer.as_ptr(), ptr); + assert_eq!(buffer.as_slice(), [1, 2, 3]); + let view = buffer.slice(1..); + drop(buffer); + assert_eq!(drops.load(Ordering::Relaxed), 0); + assert_eq!(view.as_slice(), [2, 3]); + drop(view); + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + #[test] fn bytes_round_trip_reuses_owner() { let bytes = Bytes::from_static(&[1, 2, 3, 4]); @@ -1151,9 +1197,15 @@ mod test { let Ok(mut sliced) = sliced.try_into_mut() else { panic!("uniquely owned slice should become mutable") }; + let ptr = sliced.as_ptr(); let capacity = sliced.capacity(); sliced.push_n(0, capacity - sliced.len()); assert_eq!(sliced.len(), capacity); + assert_eq!(sliced.as_ptr(), ptr); + sliced.push(42); + assert_eq!(&sliced[..32], (64u32..96).collect::>()); + assert_eq!(&sliced[32..capacity], vec![0; capacity - 32]); + assert_eq!(sliced[capacity], 42); } #[test] diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 27c3dad3532..71e27c4f8ce 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -23,12 +23,30 @@ use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. +/// +/// Zero-sized element types are rejected at compile time when constructing a buffer. +/// +/// ```compile_fail +/// use vortex_buffer::BufferMut; +/// let _ = BufferMut::<()>::empty(); +/// ``` +/// +/// ```compile_fail +/// use vortex_buffer::BufferMut; +/// let _ = BufferMut::<()>::zeroed(3); +/// ``` 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 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`. 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, } @@ -100,6 +118,7 @@ impl BufferMut { preferred_alignment: Option, allocator: BufferAllocatorRef, ) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; let actual = max( alignment, preferred_alignment.unwrap_or(Alignment::of::()), @@ -131,11 +150,7 @@ impl BufferMut { 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 { - capacity - } else { - (allocation.size() - offset) / size_of::() - }; + let capacity = (allocation.size() - offset) / size_of::(); Self { allocation, ptr, @@ -204,6 +219,7 @@ impl BufferMut { preferred_alignment: Option, allocator: BufferAllocatorRef, ) -> Self { + const { assert!(size_of::() != 0, "ZSTs are not supported") }; let preferred_alignment = preferred_alignment.unwrap_or(Alignment::of::()); let actual_alignment = max(preferred_alignment, alignment); let size = len @@ -226,11 +242,7 @@ impl BufferMut { .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 { - len - } else { - (allocation.size() - offset) / size_of::() - }; + let capacity = (allocation.size() - offset) / size_of::(); Self { allocation, ptr,