diff --git a/Cargo.lock b/Cargo.lock index 2448cb9aff8..0b59cb19243 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10693,6 +10693,7 @@ dependencies = [ name = "vortex-buffer" version = "0.1.0" dependencies = [ + "allocator-api2", "arrow-buffer 59.2.0", "bitvec", "bytes", diff --git a/Cargo.toml b/Cargo.toml index ae164db13c2..1bb45b099d1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -97,6 +97,7 @@ rust-version = "1.95" version = "0.1.0" [workspace.dependencies] +allocator-api2 = "0.2.21" alp = "0.0.2" anyhow = "1.0.100" arbitrary = "1.3.2" diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 0d831e58775..81dda5940fd 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1144,7 +1144,7 @@ impl ZstdData { let value_bytes = values.buffer_handle().try_to_host_sync()?; // Align frames to buffer alignment. This is necessary for overaligned buffers. - let alignment = *value_bytes.alignment(); + let alignment = value_bytes.alignment().as_usize(); let step_width = (values_per_frame * byte_width).div_ceil(alignment) * alignment; let frame_byte_starts = (0..n_values * byte_width) diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 5f7f785f71f..4e62bc67857 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -357,7 +357,7 @@ fn compute_output_layout( let mut total_size = 0usize; for (&size, &alignment) in output_sizes.iter().zip(output_alignments.iter()) { - total_size = total_size.next_multiple_of(*alignment); + total_size = total_size.next_multiple_of(alignment.as_usize()); offsets.push(total_size); total_size += size; } diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index f84ec182269..df82cebe0f7 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -84,7 +84,7 @@ impl ArrayRef { .unwrap_or_else(FlatBuffer::alignment); // Create a shared buffer of zeros we can use for padding - let zeros = ByteBuffer::zeroed(*max_alignment); + let zeros = ByteBuffer::zeroed(max_alignment.as_usize()); // We push an empty buffer with the maximum alignment, so then subsequent buffers // will be aligned. For subsequent buffers, we always push a 1-byte alignment. @@ -96,7 +96,7 @@ impl ArrayRef { // Push all the array buffers with padding as necessary. for buffer in array_buffers { let padding = if options.include_padding { - let padding = pos.next_multiple_of(*buffer.alignment()) - pos; + let padding = pos.next_multiple_of(buffer.alignment().as_usize()) - pos; if padding > 0 { pos += padding; buffers.push(zeros.slice(0..padding)); @@ -139,7 +139,7 @@ impl ArrayRef { let fb_length = fb_buffer.len(); if options.include_padding { - let padding = pos.next_multiple_of(*FlatBuffer::alignment()) - pos; + let padding = pos.next_multiple_of(FlatBuffer::alignment().as_usize()) - pos; if padding > 0 { buffers.push(zeros.slice(0..padding)); } diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index 705c992f87d..c77d07f253d 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -23,6 +23,7 @@ serde = ["dep:serde", "serde/serde_derive"] warn-copy = ["dep:tracing"] [dependencies] +allocator-api2 = { workspace = true } arrow-buffer = { workspace = true } bitvec = { workspace = true } bytes = { workspace = true } diff --git a/vortex-buffer/src/alignment.rs b/vortex-buffer/src/alignment.rs index 58f5d16f10c..f4f476c7691 100644 --- a/vortex-buffer/src/alignment.rs +++ b/vortex-buffer/src/alignment.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Display; -use std::ops::Deref; use vortex_error::VortexError; use vortex_error::VortexExpect; @@ -12,9 +11,9 @@ use vortex_error::vortex_err; /// The alignment of a buffer. /// -/// This type is a wrapper around `usize` that ensures the alignment is a non-zero power of 2. +/// This type stores the base-2 exponent of a non-zero power-of-two alignment. #[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct Alignment(usize); +pub struct Alignment(u8); impl Alignment { /// Largest alignment accepted from untrusted serialized input. @@ -41,10 +40,11 @@ impl Alignment { /// /// Panics if `align` is zero or is not a power of 2. #[inline] + #[expect(clippy::cast_possible_truncation, reason = "usize has at most 64 bits")] pub const fn new(align: usize) -> Self { assert!(align > 0, "Alignment must be greater than 0"); assert!(align.is_power_of_two(), "Alignment must be a power of 2"); - Self(align) + Self(align.trailing_zeros() as u8) } /// Create a new 1-byte alignment. @@ -104,7 +104,7 @@ impl Alignment { #[inline] pub const fn is_offset_aligned(&self, offset: usize) -> bool { // Alignment is always a power of 2, so a mask test is equivalent to `offset % self == 0`. - offset & (self.0 - 1) == 0 + offset & (self.as_usize() - 1) == 0 } /// Check if the given pointer is aligned to this alignment. @@ -115,8 +115,7 @@ impl Alignment { /// Returns the log2 of the alignment. pub fn exponent(&self) -> u8 { - u8::try_from(self.0.trailing_zeros()) - .vortex_expect("alignment is a power of 2 within usize, so its exponent fits in u8") + self.0 } /// Create from the log2 exponent of the alignment. @@ -131,7 +130,7 @@ impl Alignment { (exponent as u32) < usize::BITS, "Alignment exponent must fit in usize" ); - Self::new(1 << exponent) + Self(exponent) } /// Create from the log2 exponent of the alignment, returning an error rather than panicking if @@ -166,20 +165,17 @@ impl Alignment { } Ok(alignment) } -} -impl Display for Alignment { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.0) + /// Return the alignment in bytes. + #[inline] + pub const fn as_usize(self) -> usize { + 1 << self.0 } } -impl Deref for Alignment { - type Target = usize; - - #[inline] - fn deref(&self) -> &Self::Target { - &self.0 +impl Display for Alignment { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_usize()) } } @@ -200,14 +196,14 @@ impl From for Alignment { impl From for usize { #[inline] fn from(value: Alignment) -> Self { - value.0 + value.as_usize() } } impl From for u32 { #[inline] fn from(value: Alignment) -> Self { - u32::try_from(value.0).vortex_expect("Alignment must fit into u32") + u32::try_from(value.as_usize()).vortex_expect("Alignment must fit into u32") } } @@ -225,7 +221,7 @@ impl TryFrom for Alignment { return Err(vortex_err!("Alignment must be a power of 2, got {value}")); } - Ok(Self(value)) + Ok(Self::new(value)) } } @@ -243,7 +239,7 @@ mod test { fn alignment_above_u16() { // 64KiB alignment (one past `u16::MAX`) is valid — common on ARM with 64K pages. let alignment = Alignment::new(u16::MAX as usize + 1); - assert_eq!(*alignment, 1 << 16); + assert_eq!(alignment.as_usize(), 1 << 16); assert_eq!(alignment, Alignment::from_exponent(16)); } diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs new file mode 100644 index 00000000000..415ba4fc4b4 --- /dev/null +++ b/vortex-buffer/src/allocation.rs @@ -0,0 +1,465 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Allocator-backed storage for Vortex buffers. + +use std::alloc::Layout; +use std::fmt; +use std::fmt::Debug; +use std::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 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. +#[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 is_statically_allocated(&self) -> bool { + self.0.is_none() + } + + /// Create a mutable buffer with this allocator. + pub fn with_capacity(&self, capacity: usize) -> BufferMut { + BufferMut::with_capacity_in(capacity, self.clone()) + } + + /// 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) }, + } + } +} + +/// The allocator used by buffer APIs that do not take an allocator. +#[derive(Clone, Copy, Debug, Default)] +pub struct StaticBufferAllocator; + +impl StaticBufferAllocator { + /// Create a mutable buffer with the static allocator. + pub fn with_capacity(capacity: usize) -> BufferMut { + BufferMut::with_capacity(capacity) + } + + /// Create an aligned mutable buffer with the static allocator. + pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> BufferMut { + BufferMut::with_capacity_aligned(capacity, alignment) + } + + /// Create a zeroed mutable buffer with the static allocator. + pub fn zeroed(len: usize) -> BufferMut { + BufferMut::zeroed(len) + } + + /// Copy values into a mutable buffer made by the static allocator. + pub fn copy_from(values: impl AsRef<[T]>) -> BufferMut { + BufferMut::copy_from(values) + } +} + +// SAFETY: Global satisfies the Allocator contract and this type only forwards to it. +unsafe impl Allocator for StaticBufferAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + Global.allocate(layout) + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + Global.allocate_zeroed(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } + } + + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.shrink(ptr, old_layout, new_layout) } + } +} + +static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); + +pub(crate) struct Allocation { + ptr: NonNull, + layout: Layout, + allocator: BufferAllocatorRef, +} + +// SAFETY: Allocation owns its memory, and its allocator is Send + Sync. +unsafe impl Send for Allocation {} +// SAFETY: shared access to Allocation never permits mutation of the allocation. +unsafe impl Sync for Allocation {} + +impl Allocation { + pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, false) + } + + pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, true) + } + + pub(crate) fn from_vec(vec: Vec) -> Self { + assert!(!std::mem::needs_drop::()); + + let mut vec = ManuallyDrop::new(vec); + let layout = Layout::array::(vec.capacity()) + .unwrap_or_else(|_| unreachable!("a Vec capacity always has a valid layout")); + let ptr = NonNull::new(vec.as_mut_ptr().cast()) + .vortex_expect("a Vec always has a non-null pointer"); + + Self { + ptr, + layout, + allocator: BufferAllocatorRef::statically_allocated(), + } + } + + fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { + if layout.size() == 0 { + return Self { + ptr: layout.dangling_ptr(), + layout, + allocator, + }; + } + + let allocation = if zeroed { + allocator.allocate_zeroed(layout) + } else { + allocator.allocate(layout) + } + .unwrap_or_else(|_| handle_alloc_error(layout)); + + Self { + ptr: allocation.cast(), + layout, + allocator, + } + } + + #[inline(always)] + pub(crate) fn ptr(&self) -> NonNull { + self.ptr + } + + #[inline(always)] + pub(crate) fn size(&self) -> usize { + self.layout.size() + } + + #[inline(always)] + pub(crate) fn alignment(&self) -> usize { + self.layout.align() + } + + #[inline(always)] + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + &self.allocator + } + + pub(crate) fn grow(&mut self, new_layout: Layout) { + let allocation = if self.layout.size() == 0 { + self.allocator.allocate(new_layout) + } else { + // SAFETY: ptr denotes a live block owned by allocator, old_layout fits the block, and + // the caller only grows the allocation. + unsafe { self.allocator.grow(self.ptr, self.layout, new_layout) } + } + .unwrap_or_else(|_| handle_alloc_error(new_layout)); + self.ptr = allocation.cast(); + self.layout = new_layout; + } +} + +impl Drop for Allocation { + fn drop(&mut self) { + if self.layout.size() == 0 { + return; + } + // SAFETY: ptr and layout describe a live block allocated by self.allocator. + unsafe { self.allocator.deallocate(self.ptr, self.layout) } + } +} + +pub(crate) trait BufferOwner: Send + Sync + 'static { + fn as_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), + External { _owner: Box }, +} + +impl BufferBacking { + #[inline(always)] + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + match self { + Self::Owned(allocation) => allocation.allocator(), + Self::External { .. } => &STATIC_ALLOCATOR, + } + } +} + +#[cfg(test)] +mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use allocator_api2::alloc::AllocError; + use allocator_api2::alloc::Allocator; + use allocator_api2::alloc::Global; + + use crate::Alignment; + use crate::BufferAllocatorRef; + + #[derive(Clone, Debug, Default)] + struct TrackingAllocator { + state: Arc, + } + + #[derive(Debug, Default)] + struct TrackingState { + allocations: AtomicUsize, + deallocations: AtomicUsize, + grows: AtomicUsize, + alignment: AtomicUsize, + } + + // SAFETY: this forwards all memory operations to Global and only records call metadata. + unsafe impl Allocator for TrackingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.state.allocations.fetch_add(1, Ordering::Relaxed); + self.state + .alignment + .store(layout.align(), Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + self.state.deallocations.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller passes the pointer and layout returned by Global. + unsafe { Global.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.state.grows.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + } + + #[test] + fn allocation_lives_until_last_view() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let buffer = BufferAllocatorRef::new(allocator) + .copy_from([1u32, 2, 3, 4]) + .freeze(); + let view = buffer.slice(0..2); + + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!( + state.alignment.load(Ordering::Relaxed), + Alignment::of::().as_usize() + ); + drop(buffer); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + drop(view); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + } + + #[test] + 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 initial_capacity = buffer.capacity(); + buffer.extend(std::iter::repeat_n(7, initial_capacity)); + + buffer.push(u32::MAX); + + assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]); + assert_eq!(buffer[initial_capacity], u32::MAX); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + assert_eq!(state.grows.load(Ordering::Relaxed), 1); + + drop(buffer); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + } + + #[test] + fn zero_capacity_does_not_allocate() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(0); + + assert_eq!(buffer.capacity(), 0); + assert!(Alignment::DEFAULT_ALIGNMENT.is_offset_aligned(buffer.as_ptr().addr())); + assert_eq!(state.allocations.load(Ordering::Relaxed), 0); + + buffer.push(42); + + assert_eq!(buffer.as_slice(), [42]); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + } +} diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index 63eec6deb60..f63aaadaff0 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -28,19 +28,15 @@ impl Buffer { let bytes = Bytes::from_owner(ArrowWrapper(arrow.into_inner())); let alignment = Alignment::of::(); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if bytes.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + debug_assert_eq!(length, bytes.len() / size_of::()); + Self::from_bytes_aligned(bytes, alignment) } /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`. @@ -67,19 +63,15 @@ impl ByteBuffer { let length = arrow.len(); let bytes = Bytes::from_owner(ArrowWrapper(arrow)); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if bytes.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + debug_assert_eq!(length, bytes.len()); + Self::from_bytes_aligned(bytes, alignment) } } diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index a1bb0c82c04..734f65a08a6 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -599,25 +599,6 @@ impl BitBufferMut { self.len += bit_len; } - /// Absorbs a mutable buffer that was previously split off. - /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.append_buffer(&other). - pub fn unsplit(&mut self, other: Self) { - if (self.offset + self.len).is_multiple_of(8) && other.offset == 0 { - // We are aligned and can just append the buffers - self.buffer.unsplit(other.buffer); - self.len += other.len; - return; - } - - // Otherwise, we need to append the bits one by one - self.append_buffer(&other.freeze()) - } - /// Freeze the buffer in its current state into an immutable `BoolBuffer`. #[inline] pub fn freeze(self) -> BitBuffer { diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index a59fff825f8..114083e5699 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -8,9 +8,11 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use std::marker::PhantomData; use std::ops::Deref; use std::ops::RangeBounds; +use std::ptr::NonNull; +use std::sync::Arc; +use std::sync::LazyLock; use bytes::Buf; use bytes::Bytes; @@ -18,6 +20,9 @@ use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; +use crate::BufferAllocatorRef; +use crate::BufferBacking; use crate::BufferMut; use crate::ByteBuffer; use crate::debug::TruncatedDebug; @@ -26,64 +31,119 @@ use crate::trusted_len::TrustedLen; /// An immutable buffer of items of `T`. #[derive(Clone)] pub struct Buffer { - pub(crate) bytes: Bytes, + pub(crate) ptr: NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, - pub(crate) _marker: PhantomData, + pub(crate) physical_alignment: Alignment, + // One physical-alignment block is reserved outside the logical capacity. + pub(crate) overallocated: bool, + pub(crate) backing: Arc, } +// SAFETY: Buffer is an immutable view over backing memory. Its pointer remains valid while the +// backing is live, and sharing elements follows the same bounds as sharing a slice. +unsafe impl Send for Buffer {} +// SAFETY: see the Send implementation above. +unsafe impl Sync for Buffer {} + /// Zero-length backing for empty buffers, "aligned" to [`Alignment::MAX`] so it satisfies any /// valid alignment without allocating. A zero-length slice never reads memory, so it may use a /// dangling pointer as long as it is non-null and aligned. -const EMPTY_BACKING: &[u8] = { +const EMPTY_BYTES: &[u8] = { let addr = 1usize << (usize::BITS - 1); assert!(Alignment::MAX.is_offset_aligned(addr)); // SAFETY: the pointer is non-null and aligned, and the slice is zero-length. unsafe { std::slice::from_raw_parts(std::ptr::without_provenance(addr), 0) } }; +static EMPTY_BACKING: LazyLock> = LazyLock::new(|| { + Arc::new(BufferBacking::External { + _owner: Box::new(EMPTY_BYTES), + }) +}); + impl Default for Buffer { fn default() -> Self { Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment: Alignment::of::(), - _marker: PhantomData, + physical_alignment: Alignment::MAX, + overallocated: false, + backing: EMPTY_BACKING.clone(), } } } -impl PartialEq for Buffer { +impl PartialEq for Buffer { #[inline] fn eq(&self, other: &Self) -> bool { - self.bytes == other.bytes + self.as_slice() == other.as_slice() } } -impl Eq for Buffer {} +impl Eq for Buffer {} -impl Ord for Buffer { +impl Ord for Buffer { #[inline] fn cmp(&self, other: &Self) -> Ordering { - self.bytes.cmp(&other.bytes) + self.as_slice().cmp(other.as_slice()) } } -impl PartialOrd for Buffer { +impl PartialOrd for Buffer { #[inline] fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + self.as_slice().partial_cmp(other.as_slice()) } } -impl Hash for Buffer { +impl Hash for Buffer { #[inline] fn hash(&self, state: &mut H) { - self.bytes.as_ref().hash(state) + self.as_slice().hash(state) } } impl Buffer { + pub(crate) fn from_allocation( + allocation: Allocation, + offset: usize, + length: usize, + alignment: Alignment, + physical_alignment: Alignment, + overallocated: bool, + ) -> Self { + // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; + Self { + ptr, + length, + alignment, + physical_alignment, + overallocated, + backing: Arc::new(BufferBacking::Owned(allocation)), + } + } + + 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, + physical_alignment: alignment, + overallocated: false, + backing: Arc::new(BufferBacking::External { _owner: owner }), + } + } + /// Returns a new `Buffer` copied from the provided `Vec`, `&[T]`, etc. /// /// Due to our underlying usage of `bytes::Bytes`, we are unable to take zero-copy ownership @@ -94,6 +154,11 @@ impl Buffer { BufferMut::copy_from(values).freeze() } + /// Returns a new `Buffer` copied with the provided allocator. + pub fn copy_from_in(values: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + BufferMut::copy_from_in(values, allocator).freeze() + } + /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -121,6 +186,11 @@ impl Buffer { Self::zeroed_aligned(len, Alignment::of::()) } + /// Create a new zeroed `Buffer` with the provided allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + BufferMut::zeroed_in(len, allocator).freeze() + } + /// Create a new zeroed `Buffer` with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -161,10 +231,12 @@ impl Buffer { ); } Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment, - _marker: PhantomData, + physical_alignment: Alignment::MAX, + overallocated: false, + backing: EMPTY_BACKING.clone(), } } @@ -176,6 +248,14 @@ impl Buffer { BufferMut::full(item, len).freeze() } + /// Create a full `Buffer` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + BufferMut::full_in(item, len, allocator).freeze() + } + /// Create a `Buffer` zero-copy from a `ByteBuffer`. /// /// ## Panics @@ -194,7 +274,31 @@ impl Buffer { /// Panics if the buffer is not aligned to the given alignment, if the length is not a multiple /// of the size of `T`, or if the given alignment is not aligned to that of `T`. pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self { - Self::from_bytes_aligned(buffer.into_inner(), alignment) + if !alignment.is_aligned_to(Alignment::of::()) { + vortex_panic!( + "Alignment {} must be compatible with the scalar type's alignment {}", + alignment, + Alignment::of::(), + ); + } + if !alignment.is_ptr_aligned(buffer.as_ptr()) { + vortex_panic!("Buffer must align to the requested alignment {}", alignment); + } + if !buffer.len().is_multiple_of(size_of::()) { + vortex_panic!( + "Buffer length {} must be a multiple of the scalar type's size {}", + buffer.len(), + size_of::() + ); + } + Self { + ptr: buffer.ptr.cast(), + length: buffer.length / size_of::(), + alignment, + physical_alignment: buffer.physical_alignment, + overallocated: buffer.overallocated, + backing: buffer.backing, + } } /// Create a `Buffer` zero-copy from a `Bytes`. @@ -224,13 +328,7 @@ impl Buffer { size_of::() ); } - let length = bytes.len() / size_of::(); - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + Self::from_owner(bytes, alignment) } /// Create a buffer with values from the TrustedLen iterator. @@ -249,7 +347,8 @@ impl Buffer { Ok(mut_buf) => mut_buf.map_each_in_place(f), Err(buf) => { let len = buf.len(); - let mut out_buf = BufferMut::with_capacity(len); + let allocator = buf.backing.allocator().clone(); + let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf .spare_capacity_mut() .iter_mut() @@ -266,7 +365,6 @@ impl Buffer { /// Clear the buffer, preserving existing capacity. pub fn clear(&mut self) { - self.bytes.clear(); self.length = 0; } @@ -288,17 +386,33 @@ impl Buffer { self.alignment } + /// Returns the allocator to use for derived buffers. + /// + /// External buffers use the static allocator. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.backing.allocator() + } + + /// Returns a raw pointer to the buffer's data. + #[inline(always)] + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() + } + /// Returns a slice over the buffer of elements of type T. #[inline(always)] pub fn as_slice(&self) -> &[T] { - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(self.bytes.as_ptr().cast(), self.length) } + // SAFETY: ptr points into the live backing and construction checks its alignment. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } } /// Return a view over the buffer as an opaque byte slice. #[inline(always)] pub fn as_bytes(&self) -> &[u8] { - self.bytes.as_ref() + // SAFETY: the element range is initialized and remains live through backing. + unsafe { + std::slice::from_raw_parts(self.ptr.as_ptr().cast(), size_of_val(self.as_slice())) + } } /// Returns an iterator over the buffer of elements of type T. @@ -372,8 +486,6 @@ impl Buffer { } let begin_byte = begin * size_of::(); - let end_byte = end * size_of::(); - if !alignment.is_offset_aligned(begin_byte) { vortex_panic!( "range start must be aligned to {alignment:?}, byte {}", @@ -385,10 +497,13 @@ impl Buffer { } Self { - bytes: self.bytes.slice(begin_byte..end_byte), + // SAFETY: begin is in bounds and the alignment check applies to the new pointer. + ptr: unsafe { self.ptr.add(begin) }, length: end - begin, alignment, - _marker: Default::default(), + physical_alignment: self.physical_alignment, + overallocated: self.overallocated, + backing: Arc::clone(&self.backing), } } @@ -427,74 +542,106 @@ impl Buffer { vortex_panic!("slice_ref subset must be aligned to {:?}", alignment); } - let subset_u8 = - unsafe { std::slice::from_raw_parts(subset.as_ptr().cast(), size_of_val(subset)) }; + let start = self.as_ptr().addr(); + let end = start + size_of_val(self.as_slice()); + let subset_start = subset.as_ptr().addr(); + let subset_end = subset_start + .checked_add(size_of_val(subset)) + .vortex_expect("slice_ref address overflow"); + if subset_start < start || subset_end > end { + vortex_panic!("slice_ref subset must be contained in the buffer"); + } Self { - bytes: self.bytes.slice_ref(subset_u8), + ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"), length: subset.len(), alignment, - _marker: Default::default(), + physical_alignment: self.physical_alignment, + overallocated: self.overallocated, + backing: Arc::clone(&self.backing), } } - /// Returns the underlying aligned buffer. - pub fn inner(&self) -> &Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - &self.bytes - } - - /// Returns the underlying aligned buffer. + /// Returns the underlying bytes without copying. pub fn into_inner(self) -> Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - self.bytes + Bytes::from_owner(BufferBytesOwner { + ptr: self.ptr.cast(), + length: self.length * size_of::(), + backing: self.backing, + }) } /// Return the ByteBuffer for this `Buffer`. pub fn into_byte_buffer(self) -> ByteBuffer { ByteBuffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, - _marker: Default::default(), + physical_alignment: self.physical_alignment, + overallocated: self.overallocated, + backing: self.backing, } } /// Try to convert self into `BufferMut` if there is only a single strong reference. pub fn try_into_mut(self) -> Result, Self> { - self.bytes - .try_into_mut() - .map(|bytes| BufferMut { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) - .map_err(|bytes| Self { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) + let Self { + ptr, + length, + alignment, + physical_alignment, + overallocated, + backing, + } = self; + match Arc::try_unwrap(backing) { + Ok(BufferBacking::Owned(allocation)) => { + let offset = ptr.addr().get() - allocation.ptr().addr().get(); + let overallocated = overallocated + && offset + == allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + Ok(BufferMut { + allocation, + ptr, + length, + alignment, + physical_alignment, + overallocated, + _marker: Default::default(), + }) + } + Ok(backing) => Err(Self { + ptr, + length, + alignment, + physical_alignment, + overallocated, + backing: Arc::new(backing), + }), + Err(backing) => Err(Self { + ptr, + length, + alignment, + physical_alignment, + overallocated, + backing, + }), + } } /// Convert self into `BufferMut`, cloning the data if there are multiple strong references. pub fn into_mut(self) -> BufferMut { - self.try_into_mut() - .unwrap_or_else(|buffer| BufferMut::::copy_from_aligned(&buffer, buffer.alignment)) + self.try_into_mut().unwrap_or_else(|buffer| { + let allocator = buffer.backing.allocator().clone(); + BufferMut::::copy_from_aligned_in(&buffer, buffer.alignment, allocator) + }) } /// Returns whether a `Buffer` is aligned to the given alignment. pub fn is_aligned(&self, alignment: Alignment) -> bool { - alignment.is_ptr_aligned(self.bytes.as_ptr()) + alignment.is_ptr_aligned(self.as_ptr()) } /// Return a `Buffer` with the given alignment. Where possible, this will be zero-copy. @@ -510,7 +657,8 @@ impl Buffer { "Buffer is not aligned to requested alignment {alignment}, copying: {bt}" ) } - Self::copy_from_aligned(self, alignment) + let allocator = self.backing.allocator().clone(); + BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze() } } @@ -546,10 +694,12 @@ impl Buffer { ); Buffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, - _marker: PhantomData, + physical_alignment: self.physical_alignment, + overallocated: self.overallocated, + backing: self.backing, } } } @@ -630,48 +780,45 @@ impl FromIterator for Buffer { } } -// Helper struct to allow us to zero-copy any vec into a buffer +// Helper struct that preserves drop glue for non-native Vec elements. #[repr(transparent)] struct Wrapper(Vec); -impl AsRef<[u8]> for Wrapper { - fn as_ref(&self) -> &[u8] { - let data = self.0.as_ptr().cast::(); - let len = self.0.len() * size_of::(); - unsafe { std::slice::from_raw_parts(data, len) } +impl crate::BufferOwner for Wrapper { + fn as_ptr(&self) -> *const u8 { + self.0.as_ptr().cast() + } + + fn len(&self) -> usize { + self.0.len() * size_of::() } } impl From> for Buffer where - T: Send + 'static, + T: Send + Sync + 'static, { fn from(value: Vec) -> Self { - let original_len = value.len(); - let wrapped_vec = Wrapper(value); - - let bytes = Bytes::from_owner(wrapped_vec); - - assert_eq!(bytes.as_ptr().align_offset(align_of::()), 0); - - Self { - bytes, - length: original_len, - alignment: Alignment::of::(), - _marker: PhantomData, + let length = value.len(); + let alignment = Alignment::of::(); + if std::mem::needs_drop::() { + Self::from_owner(Wrapper(value), alignment) + } else { + Self::from_allocation( + Allocation::from_vec(value), + 0, + length, + alignment, + alignment, + false, + ) } } } impl From for ByteBuffer { fn from(bytes: Bytes) -> Self { - let length = bytes.len(); - Self { - bytes, - length, - alignment: Alignment::of::(), - _marker: Default::default(), - } + Self::from_owner(bytes, Alignment::of::()) } } @@ -695,11 +842,36 @@ impl Buf for ByteBuffer { self.alignment ); } - self.bytes.advance(cnt); + assert!(cnt <= self.length, "cannot advance past the buffer length"); + // SAFETY: cnt is within the initialized byte range. + self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; } } +struct BufferBytesOwner { + ptr: NonNull, + length: usize, + backing: Arc, +} + +// SAFETY: the owner exposes immutable initialized bytes and keeps their backing live. +unsafe impl Send for BufferBytesOwner {} +unsafe impl Sync for BufferBytesOwner {} + +impl AsRef<[u8]> for BufferBytesOwner { + fn as_ref(&self) -> &[u8] { + let _ = &self.backing; + // SAFETY: ptr and length came from a live Buffer. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } + } +} + +fn empty_ptr() -> NonNull { + let addr = 1usize << (usize::BITS - 1); + NonNull::new(std::ptr::without_provenance_mut(addr)).vortex_expect("empty pointer is non-null") +} + /// Owned iterator over a [`Buffer`]. pub struct BufferIterator { // Keep the buffer alive for the duration of the iteration. @@ -762,6 +934,11 @@ impl From> for Buffer { #[cfg(test)] mod test { + use std::mem::align_of; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use bytes::Buf; use crate::Alignment; @@ -879,6 +1056,76 @@ mod test { assert_eq!(vec, buff.as_ref()); } + #[test] + fn from_vec_adopts_allocation() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u32, 2, 3, 4, 5]); + let ptr = vec.as_ptr(); + let capacity = vec.capacity(); + + let buffer = Buffer::from(vec); + assert_eq!(buffer.as_ptr(), ptr); + + let Ok(mut buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), capacity); + assert_eq!(buffer.allocation.alignment(), align_of::()); + + buffer.extend(6..=32); + assert_eq!(buffer.as_slice(), (1..=32).collect::>()); + assert_eq!(buffer.allocation.alignment(), align_of::()); + } + + #[test] + fn from_u8_vec_preserves_capacity() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u8, 2, 3]); + + let buffer = Buffer::from(vec); + let Ok(buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), 16); + } + + #[test] + fn sliced_buffer_into_mut_has_safe_capacity() { + let mut original = crate::BufferMut::with_capacity(128); + original.extend(0u32..100); + let original = original.freeze(); + let sliced = original.slice(64..96); + drop(original); + + let Ok(mut sliced) = sliced.try_into_mut() else { + panic!("uniquely owned slice should become mutable") + }; + let capacity = sliced.capacity(); + sliced.push_n(0, capacity - sliced.len()); + assert_eq!(sliced.len(), capacity); + } + + #[test] + fn from_vec_preserves_drop_glue() { + struct DropValue(Arc); + + impl Drop for DropValue { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let values = (0..3) + .map(|_| DropValue(Arc::clone(&drops))) + .collect::>(); + let buffer = Buffer::from(values); + + assert_eq!(drops.load(Ordering::Relaxed), 0); + drop(buffer); + assert_eq!(drops.load(Ordering::Relaxed), 3); + } + #[test] fn empty_aligned_max_alignment() { // Empty buffers are backed by a static and must satisfy any valid alignment. diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index e5cb03c558b..c6b0cf9a252 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use core::mem::MaybeUninit; +use std::alloc::Layout; use std::any::type_name; use std::cmp::max; use std::fmt::Debug; @@ -12,31 +13,45 @@ use std::ops::DerefMut; use bytes::Buf; use bytes::BufMut; -use bytes::BytesMut; use bytes::buf::UninitSlice; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; use crate::Buffer; +use crate::BufferAllocatorRef; use crate::ByteBufferMut; use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. -#[derive(PartialEq, Eq)] pub struct BufferMut { - pub(crate) bytes: BytesMut, + pub(crate) allocation: Allocation, + pub(crate) ptr: std::ptr::NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, + pub(crate) physical_alignment: Alignment, + // One physical-alignment block is reserved outside the logical capacity. + pub(crate) overallocated: bool, pub(crate) _marker: std::marker::PhantomData, } +// SAFETY: BufferMut uniquely owns its allocation and only exposes T across threads. +unsafe impl Send for BufferMut {} +// SAFETY: shared access to BufferMut only exposes shared access to T. +unsafe impl Sync for BufferMut {} + impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. pub fn with_capacity(capacity: usize) -> Self { - Self::with_capacity_aligned(capacity, Alignment::of::()) + Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new `BufferMut` with the requested capacity and allocator. + pub fn with_capacity_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(capacity, Alignment::of::(), allocator) } /// Create a new `BufferMut` with the requested alignment and capacity. @@ -46,10 +61,24 @@ impl BufferMut { /// /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { - Self::with_capacity_preferred_aligned( + Self::with_capacity_aligned_in( + capacity, + alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new `BufferMut` with the requested alignment, capacity, and allocator. + pub fn with_capacity_aligned_in( + capacity: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::with_capacity_preferred_aligned_in( capacity, alignment, Some(Alignment::DEFAULT_ALIGNMENT), + allocator, ) } @@ -61,6 +90,21 @@ impl BufferMut { capacity: usize, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::with_capacity_preferred_aligned_in( + capacity, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new allocator-backed `BufferMut` with a requested and preferred alignment. + pub fn with_capacity_preferred_aligned_in( + capacity: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { let actual = max( alignment, @@ -75,20 +119,43 @@ impl BufferMut { ); } - let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + *actual); - bytes.align_empty(actual); - + let size = capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let layout = 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() }; Self { - bytes, + allocation, + ptr, length: 0, alignment, + physical_alignment: actual, + overallocated: true, _marker: Default::default(), } } /// Create a new zeroed `BufferMut`. pub fn zeroed(len: usize) -> Self { - Self::zeroed_aligned(len, Alignment::of::()) + Self::zeroed_in(len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new zeroed `BufferMut` with the requested allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + Self::zeroed_aligned_in(len, Alignment::of::(), allocator) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -98,7 +165,21 @@ impl BufferMut { /// /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::zeroed_aligned_in(len, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create a zeroed `BufferMut` with an alignment and allocator. + pub fn zeroed_aligned_in( + len: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -109,17 +190,51 @@ impl BufferMut { len: usize, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a zeroed allocator-backed buffer with a requested and preferred alignment. + pub fn zeroed_preferred_aligned_in( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { let preferred_alignment = preferred_alignment.unwrap_or(Alignment::of::()); let actual_alignment = max(preferred_alignment, alignment); - let mut bytes = BytesMut::zeroed((len * size_of::()) + *actual_alignment); - bytes.advance(bytes.as_ptr().align_offset(*actual_alignment)); - unsafe { bytes.set_len(len * size_of::()) }; - let actual_len = bytes.len().checked_div(size_of::()).unwrap_or(0); + let size = len + .checked_mul(size_of::()) + .vortex_expect("buffer length overflow"); + let layout = 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() }; Self { - bytes, - length: actual_len, + allocation, + ptr, + length: len, alignment, + physical_alignment: actual_alignment, + overallocated: true, _marker: Default::default(), } } @@ -136,7 +251,12 @@ impl BufferMut { /// /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned pub fn empty_aligned(alignment: Alignment) -> Self { - Self::empty_preferred_aligned(alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::empty_aligned_in(alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create an empty `BufferMut` with an alignment and allocator. + pub fn empty_aligned_in(alignment: Alignment, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(0, alignment, allocator) } /// Create a new empty `BufferMut` with the provided alignment. @@ -147,7 +267,12 @@ impl BufferMut { alignment: Alignment, preferred_alignment: Option, ) -> Self { - BufferMut::with_capacity_preferred_aligned(0, alignment, preferred_alignment) + BufferMut::with_capacity_preferred_aligned_in( + 0, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) } /// Create a new full `BufferMut` with the given value. @@ -155,14 +280,27 @@ impl BufferMut { where T: Copy, { - let mut buffer = BufferMut::::with_capacity(len); + Self::full_in(item, len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a full `BufferMut` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + let mut buffer = BufferMut::::with_capacity_in(len, allocator); buffer.push_n(item, len); buffer } /// Create a mutable scalar buffer by copying the contents of the slice. pub fn copy_from(other: impl AsRef<[T]>) -> Self { - Self::copy_from_aligned(other, Alignment::of::()) + Self::copy_from_in(other, BufferAllocatorRef::statically_allocated()) + } + + /// Create a mutable scalar buffer by copying with the given allocator. + pub fn copy_from_in(other: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + Self::copy_from_aligned_in(other, Alignment::of::(), allocator) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -176,7 +314,21 @@ impl BufferMut { /// /// Panics when the requested alignment isn't itself aligned to type T. pub fn copy_from_aligned(other: impl AsRef<[T]>, alignment: Alignment) -> Self { - Self::copy_from_preferred_aligned(other, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::copy_from_aligned_in(other, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Copy values into a mutable buffer with the given alignment and allocator. + pub fn copy_from_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -191,13 +343,32 @@ impl BufferMut { other: impl AsRef<[T]>, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Copy values with the given allocator, requested alignment, and preferred alignment. + pub fn copy_from_preferred_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!("Given alignment is not aligned to type T") } let other = other.as_ref(); - let mut buffer = - Self::with_capacity_preferred_aligned(other.len(), alignment, preferred_alignment); + let mut buffer = Self::with_capacity_preferred_aligned_in( + other.len(), + alignment, + preferred_alignment, + allocator, + ); buffer.extend_from_slice(other); debug_assert_eq!(buffer.alignment(), alignment); buffer @@ -209,10 +380,14 @@ impl BufferMut { self.alignment } + /// Returns the allocator that owns this buffer. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.allocation.allocator() + } + /// Returns the length of the buffer. #[inline(always)] pub fn len(&self) -> usize { - debug_assert_eq!(self.length, self.bytes.len() / size_of::()); self.length } @@ -225,29 +400,47 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.bytes.capacity() / size_of::() + if self.allocation.size() == 0 { + return 0; + } + + if !self.overallocated { + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + return (self.allocation.size() - offset) / size_of::(); + } + + (self.allocation.size() - self.physical_alignment.as_usize()) / size_of::() + } + + /// Returns a raw pointer to the buffer's data. + #[inline(always)] + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() + } + + /// Returns a mutable raw pointer to the buffer's data. + #[inline(always)] + pub fn as_mut_ptr(&mut self) -> *mut T { + self.ptr.as_ptr() } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_slice(&self) -> &[T] { - let raw_slice = self.bytes.as_ref(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(raw_slice.as_ptr().cast(), self.length) } + // SAFETY: ptr is in the live allocation and construction checks its alignment. + unsafe { std::slice::from_raw_parts(self.as_ptr(), self.length) } } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_mut_slice(&mut self) -> &mut [T] { - let raw_slice = self.bytes.as_mut(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts_mut(raw_slice.as_mut_ptr().cast(), self.length) } + // SAFETY: BufferMut uniquely owns the allocation and the initialized range is in bounds. + unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.length) } } /// Clear the buffer, retaining any existing capacity. #[inline] pub fn clear(&mut self) { - unsafe { self.bytes.set_len(0) } self.length = 0; } @@ -269,8 +462,7 @@ impl BufferMut { /// Reserves capacity for at least `additional` more elements to be inserted in the buffer. #[inline] pub fn reserve(&mut self, additional: usize) { - let additional_bytes = additional * size_of::(); - if additional_bytes <= self.bytes.capacity() - self.bytes.len() { + if additional <= self.capacity() - self.length { // We can fit the additional bytes in the remaining capacity. Nothing to do. return; } @@ -279,17 +471,77 @@ impl BufferMut { self.reserve_allocate(additional); } - /// A separate function so we can inline the reserve call's fast path. According to `BytesMut` - /// this has significant performance implications. + /// A separate function so we can inline the reserve call's fast path. fn reserve_allocate(&mut self, additional: usize) { - let new_capacity: usize = ((self.length + additional) * size_of::()) + *self.alignment; - // Make sure we at least double in size each time we re-allocate to amortize the cost - let new_capacity = new_capacity.max(self.bytes.capacity() * 2); - - let mut bytes = BytesMut::with_capacity(new_capacity); - bytes.align_empty(self.alignment); - bytes.extend_from_slice(&self.bytes); - self.bytes = bytes; + let required = self + .length + .checked_add(additional) + .vortex_expect("buffer capacity overflow"); + let required_size = required + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let physical_alignment = max(self.alignment, self.physical_alignment); + let current_size = self + .capacity() + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let logical_size = required_size + .max(current_size.saturating_mul(2)) + .max(physical_alignment.as_usize()); + let allocation_size = logical_size + .checked_add(physical_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(); + let new_offset = if self.allocation.allocator().is_statically_allocated() { + let allocation = + Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); + let new_offset = allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + // SAFETY: both allocations have room for the initialized elements and do not overlap. + unsafe { + std::ptr::copy_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(physical_alignment.as_usize()); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits + // overlap. + unsafe { + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } + } + new_offset + }; + // SAFETY: new_offset was computed within the allocation for physical_alignment. + self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; + self.physical_alignment = physical_alignment; + self.overallocated = true; } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -329,13 +581,9 @@ impl BufferMut { /// ``` #[inline] pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit] { - let dst = self.bytes.spare_capacity_mut().as_mut_ptr(); - unsafe { - std::slice::from_raw_parts_mut( - dst as *mut MaybeUninit, - self.capacity() - self.length, - ) - } + // SAFETY: offset + length is within the allocation and points at spare capacity. + let dst = unsafe { self.as_mut_ptr().add(self.length) }.cast::>(); + unsafe { std::slice::from_raw_parts_mut(dst, self.capacity() - self.length) } } /// Sets the length of the buffer. @@ -349,7 +597,6 @@ impl BufferMut { #[inline] pub unsafe fn set_len(&mut self, len: usize) { debug_assert!(len <= self.capacity()); - unsafe { self.bytes.set_len(len * size_of::()) }; self.length = len; } @@ -369,9 +616,8 @@ impl BufferMut { pub unsafe fn push_unchecked(&mut self, item: T) { // SAFETY: the caller ensures we have sufficient capacity unsafe { - let dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let dst = self.as_mut_ptr().add(self.length); dst.write(item); - self.bytes.set_len(self.bytes.len() + size_of::()) } self.length += 1; } @@ -398,7 +644,8 @@ impl BufferMut { where T: Copy, { - let mut dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + // SAFETY: the caller guarantees enough spare capacity. + let mut dst = unsafe { self.as_mut_ptr().add(self.length) }; // SAFETY: we checked the capacity in the reserve call unsafe { let end = dst.add(n); @@ -406,7 +653,6 @@ impl BufferMut { dst.write(item); dst = dst.add(1); } - self.bytes.set_len(self.bytes.len() + (n * size_of::())); } self.length += n; } @@ -426,86 +672,41 @@ impl BufferMut { #[inline] pub fn extend_from_slice(&mut self, slice: &[T]) { self.reserve(slice.len()); - let raw_slice = - unsafe { std::slice::from_raw_parts(slice.as_ptr().cast(), size_of_val(slice)) }; - self.bytes.extend_from_slice(raw_slice); - self.length += slice.len(); - } - - /// Splits the buffer into two at the given index. - /// - /// Afterward, self contains elements `[0, at)`, and the returned buffer contains elements - /// `[at, capacity)`. It’s guaranteed that the memory does not move, that is, the address of - /// self does not change, and the address of the returned slice is at bytes after that. - /// - /// This is an O(1) operation that just increases the reference count and sets a few indices. - /// - /// Panics if either half would have a length that is not a multiple of the alignment. - pub fn split_off(&mut self, at: usize) -> Self { - if at > self.capacity() { - vortex_panic!("Cannot split buffer of capacity {} at {}", self.len(), at); - } - - let bytes_at = at * size_of::(); - if !self.alignment.is_offset_aligned(bytes_at) { - vortex_panic!( - "Cannot split buffer at {}, resulting alignment is not {}", - at, - self.alignment - ); - } - - let new_bytes = self.bytes.split_off(bytes_at); - - // Adjust the lengths, given that length may be < at - let new_length = self.length.saturating_sub(at); - self.length = self.length.min(at); - - BufferMut { - bytes: new_bytes, - length: new_length, - alignment: self.alignment, - _marker: Default::default(), - } - } - - /// Absorbs a mutable buffer that was previously split off. - /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.extend_from_slice(other.as_ref()). - pub fn unsplit(&mut self, other: Self) { - if self.alignment != other.alignment { - vortex_panic!( - "Cannot unsplit buffers with different alignments: {} and {}", - self.alignment, - other.alignment + // SAFETY: reserve made the destination valid and non-overlapping for slice.len() values. + unsafe { + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + self.as_mut_ptr().add(self.length), + slice.len(), ); } - self.bytes.unsplit(other.bytes); - self.length += other.length; + self.length += slice.len(); } /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { ByteBufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, + physical_alignment: self.physical_alignment, + overallocated: self.overallocated, _marker: Default::default(), } } /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { - Buffer { - bytes: self.bytes.freeze(), - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - } + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + Buffer::from_allocation( + self.allocation, + offset, + self.length, + self.alignment, + self.physical_alignment, + self.overallocated, + ) } /// Map each element of the buffer with a closure. @@ -532,15 +733,15 @@ impl BufferMut { /// /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { - if self.as_ptr().align_offset(*alignment) == 0 { + if self.as_ptr().align_offset(alignment.as_usize()) == 0 { Self { - bytes: self.bytes, - length: self.length, alignment, - _marker: std::marker::PhantomData, + physical_alignment: max(self.physical_alignment, alignment), + ..self } } else { - Self::copy_from_aligned(self, alignment) + let allocator = self.allocation.allocator().clone(); + Self::copy_from_aligned_in(self, alignment, allocator) } } @@ -564,9 +765,12 @@ impl BufferMut { ); BufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, + physical_alignment: self.physical_alignment, + overallocated: self.overallocated, _marker: std::marker::PhantomData, } } @@ -574,14 +778,24 @@ impl BufferMut { impl Clone for BufferMut { fn clone(&self) -> Self { - // NOTE(ngates): we cannot derive Clone since BytesMut copies on clone and the alignment - // might be messed up. - let mut buffer = BufferMut::::with_capacity_aligned(self.capacity(), self.alignment); + let mut buffer = BufferMut::::with_capacity_aligned_in( + self.capacity(), + self.alignment, + self.allocation.allocator().clone(), + ); buffer.extend_from_slice(self.as_slice()); buffer } } +impl PartialEq for BufferMut { + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +} + +impl Eq for BufferMut {} + impl Debug for BufferMut { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct(&format!("BufferMut<{}>", type_name::())) @@ -645,7 +859,7 @@ impl BufferMut { let unwritten = self.capacity() - self.len(); // We store `begin` in the case that the lower bound hint is incorrect. - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); // As a first step, we manually iterate the iterator up to the known capacity. @@ -690,7 +904,7 @@ impl BufferMut { .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), ); - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); iter.for_each(|item| { @@ -771,8 +985,8 @@ where impl FromIterator for BufferMut { fn from_iter>(iter: I) -> Self { - // We don't infer the capacity here and just let the first call to `extend` do it for us. - let mut buffer = Self::with_capacity(0); + let iter = iter.into_iter(); + let mut buffer = Self::with_capacity(iter.size_hint().0); buffer.extend(iter); buffer } @@ -795,7 +1009,9 @@ impl Buf for ByteBufferMut { self.alignment ); } - self.bytes.advance(cnt); + assert!(cnt <= self.length, "advance out of bounds"); + // SAFETY: cnt is checked against the initialized length above. + self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; } } @@ -818,13 +1034,15 @@ unsafe impl BufMut for ByteBufferMut { self.alignment ); } - unsafe { self.bytes.advance_mut(cnt) }; - self.length -= cnt; + self.reserve(cnt); + self.length += cnt; } #[inline] fn chunk_mut(&mut self) -> &mut UninitSlice { - self.bytes.chunk_mut() + let spare = self.spare_capacity_mut(); + // SAFETY: spare points to valid uninitialized byte capacity owned by this buffer. + unsafe { UninitSlice::from_raw_parts_mut(spare.as_mut_ptr().cast(), spare.len()) } } fn put(&mut self, mut src: T) @@ -849,35 +1067,6 @@ unsafe impl BufMut for ByteBufferMut { } } -/// Extension trait for [`BytesMut`] that provides functions for aligning the buffer. -trait AlignedBytesMut { - /// Align an empty `BytesMut` to the specified alignment. - /// - /// ## Panics - /// - /// Panics if the buffer is not empty, or if there is not enough capacity to align the buffer. - fn align_empty(&mut self, alignment: Alignment); -} - -impl AlignedBytesMut for BytesMut { - fn align_empty(&mut self, alignment: Alignment) { - // TODO(joe): this is slow fixme - if !self.is_empty() { - vortex_panic!("ByteBufferMut must be empty"); - } - - let padding = self.as_ptr().align_offset(*alignment); - self.capacity() - .checked_sub(padding) - .vortex_expect("Not enough capacity to align buffer"); - - // SAFETY: We know the buffer is empty, and we know we have enough capacity, so we can - // safely set the length to the padding and advance the buffer to the aligned offset. - unsafe { self.set_len(padding) }; - self.advance(padding); - } -} - impl Write for ByteBufferMut { fn write(&mut self, buf: &[u8]) -> std::io::Result { self.extend_from_slice(buf); @@ -914,6 +1103,32 @@ mod test { assert_eq!(buf.alignment(), Alignment::new(1024)); } + #[test] + fn growth_preserves_alignment_and_values() { + let alignment = Alignment::new(4096); + let mut buffer = BufferMut::::with_capacity_aligned(1, alignment); + + for value in 0..10_000 { + buffer.push(value); + assert!(alignment.is_offset_aligned(buffer.as_ptr().addr())); + } + + assert_eq!(buffer.as_slice(), (0..10_000).collect::>()); + } + + #[test] + fn growth_seeds_and_doubles_logical_capacity() { + let alignment = Alignment::new(64); + let mut buffer = BufferMut::::empty_aligned(alignment); + + buffer.push(0); + let capacity = buffer.capacity(); + assert_eq!(capacity, Alignment::DEFAULT_ALIGNMENT.as_usize()); + + buffer.reserve(capacity); + assert_eq!(buffer.capacity(), capacity * 2); + } + #[test] fn from_iter() { let buf = BufferMut::from_iter([0, 10, 20, 30]); @@ -1017,7 +1232,10 @@ mod test { let mut buf = BufferMut::::zeroed(LEN); - assert_eq!(buf.as_ptr().align_offset(*Alignment::of::()), 0); + assert_eq!( + buf.as_ptr().align_offset(Alignment::of::().as_usize()), + 0 + ); assert_eq!(buf.as_slice(), &[0; LEN]); buf[3] = 7; @@ -1031,7 +1249,7 @@ mod test { let mut buf = BufferMut::::zeroed_aligned(LEN, alignment); - assert_eq!(buf.as_ptr().align_offset(*alignment), 0); + assert_eq!(buf.as_ptr().align_offset(alignment.as_usize()), 0); assert_eq!(buf.as_slice(), &[0; LEN]); buf[3] = 7; diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index ee113481353..99b1e7c3081 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -5,11 +5,10 @@ //! A library for working with custom aligned buffers of sized values. //! -//! The `vortex-buffer` crate is built around `bytes::Bytes` and therefore supports zero-copy -//! cloning and slicing, but differs in that it can define and maintain a custom alignment. +//! The `vortex-buffer` crate supports zero-copy cloning and slicing with a custom allocator and +//! runtime alignment. //! -//! * `Buffer` and `BufferMut` provide immutable and mutable wrappers around `bytes::Bytes` -//! and `bytes::BytesMut` respectively. +//! * `Buffer` and `BufferMut` provide immutable and mutable typed buffers. //! * `ByteBuffer` and `ByteBufferMut` are type aliases for `u8` buffers. //! * `BufferString` is a wrapper around a `ByteBuffer` that enforces utf-8 encoding. //! * `ConstBuffer` provides similar functionality to `Buffer` except with a @@ -47,6 +46,7 @@ //! `arrow_buffer::OffsetBuffer`. pub use alignment::*; +pub use allocation::*; pub use bit::*; pub use buffer::*; pub use buffer_mut::*; @@ -55,6 +55,7 @@ pub use r#const::*; pub use dispatch::*; pub use string::*; mod alignment; +mod allocation; #[cfg(feature = "arrow")] mod arrow; mod bit; diff --git a/vortex-buffer/src/serde.rs b/vortex-buffer/src/serde.rs index 9563236a50a..5d36f8da3d0 100644 --- a/vortex-buffer/src/serde.rs +++ b/vortex-buffer/src/serde.rs @@ -22,7 +22,7 @@ where where S: Serializer, { - serializer.serialize_bytes(self.inner().as_ref()) + serializer.serialize_bytes(self.as_bytes()) } } diff --git a/vortex-cuda/src/device_buffer.rs b/vortex-cuda/src/device_buffer.rs index 1c0068bf578..869f5c4dc95 100644 --- a/vortex-cuda/src/device_buffer.rs +++ b/vortex-cuda/src/device_buffer.rs @@ -474,7 +474,7 @@ impl DeviceBuffer for CudaDeviceBuffer { fn aligned(self: Arc, alignment: Alignment) -> VortexResult> { let effective_ptr = self.device_ptr + self.offset as u64; - if effective_ptr.is_multiple_of(*alignment as u64) { + if effective_ptr.is_multiple_of(alignment.as_usize() as u64) { Ok(Arc::new(CudaDeviceBuffer { allocation: Arc::clone(&self.allocation), offset: self.offset, diff --git a/vortex-file/src/footer/serializer.rs b/vortex-file/src/footer/serializer.rs index 5010dadb0d2..a0b846d6b61 100644 --- a/vortex-file/src/footer/serializer.rs +++ b/vortex-file/src/footer/serializer.rs @@ -246,7 +246,7 @@ fn write_buffer( .map_err(|_| vortex_err!("metadata segment length exceeds maximum u32"))?; let alignment = buffer.alignment(); - let padding = offset.next_multiple_of(*alignment as u64) - *offset; + let padding = offset.next_multiple_of(alignment.as_usize() as u64) - *offset; let segment_offset = *offset + padding; let segment = PostscriptSegment { diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 616e9a52606..9f8b7aedbd3 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -228,7 +228,7 @@ impl State { let mut requests = vec![first_req]; let mut current_start = requests[0].offset; let mut current_end = requests[0].offset + requests[0].length as u64; - let align = *self.coalesced_buffer_alignment as u64; + let align = self.coalesced_buffer_alignment.as_usize() as u64; // Track requests that we've already decided to remove (or that were cancelled) so that // we don't repeatedly process them during range scans. @@ -593,7 +593,7 @@ mod tests { assert_eq!(coalesced.alignment(), Alignment::new(4)); for req in coalesced.requests() { let rel = req.offset - coalesced.range().start; - assert_eq!(rel % *req.alignment as u64, 0); + assert_eq!(rel % req.alignment.as_usize() as u64, 0); } } _ => panic!("Expected coalesced request"), diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 1b69f06e7c2..1ab7af31064 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -309,7 +309,7 @@ impl FileSegmentSource { let coalesce_config = reader.coalesce_config().map(|mut config| { // Aligning the coalesced start down can add up to (alignment - 1) bytes. // Increase max_size to keep the effective payload window consistent. - let extra = (*max_alignment as u64).saturating_sub(1); + let extra = (max_alignment.as_usize() as u64).saturating_sub(1); config.max_size = config.max_size.saturating_add(extra); config }); diff --git a/vortex-file/src/segments/writer.rs b/vortex-file/src/segments/writer.rs index e163c2cc868..1673ce40101 100644 --- a/vortex-file/src/segments/writer.rs +++ b/vortex-file/src/segments/writer.rs @@ -70,7 +70,7 @@ impl SegmentSink for BufferedSegmentSink { // Add any padding required to align the segment. let byte_offset = self.byte_offset.load(Ordering::Relaxed); - let padding = byte_offset.next_multiple_of(*alignment as u64) - byte_offset; + let padding = byte_offset.next_multiple_of(alignment.as_usize() as u64) - byte_offset; let offset = byte_offset + padding; specs.push(SegmentSpec { offset, diff --git a/vortex-tui/src/inspect.rs b/vortex-tui/src/inspect.rs index b035036be99..122c3ceec9a 100644 --- a/vortex-tui/src/inspect.rs +++ b/vortex-tui/src/inspect.rs @@ -185,22 +185,22 @@ async fn exec_inspect_json( dtype: ps.dtype.map(|s| SegmentInfoJson { offset: s.offset, length: s.length, - alignment: *s.alignment, + alignment: s.alignment.as_usize(), }), layout: SegmentInfoJson { offset: ps.layout.offset, length: ps.layout.length, - alignment: *ps.layout.alignment, + alignment: ps.layout.alignment.as_usize(), }, statistics: ps.statistics.map(|s| SegmentInfoJson { offset: s.offset, length: s.length, - alignment: *s.alignment, + alignment: s.alignment.as_usize(), }), footer: SegmentInfoJson { offset: ps.footer.offset, length: ps.footer.length, - alignment: *ps.footer.alignment, + alignment: ps.footer.alignment.as_usize(), }, }) } else { @@ -239,7 +239,7 @@ async fn exec_inspect_json( offset: segment.offset, end_offset: segment.offset + segment.length as u64, length: segment.length, - alignment: *segment.alignment, + alignment: segment.alignment.as_usize(), path: segment_paths[i] .as_ref() .map(|p| p.iter().map(|s| s.as_ref()).collect::>().join(".")), @@ -611,7 +611,7 @@ impl FooterSegments { print!( "{:>length_w$} {:>align_w$} ", segment.length, - *segment.alignment, + segment.alignment.as_usize(), length_w = length_width, align_w = alignment_width, ); diff --git a/vortex-tui/src/segments.rs b/vortex-tui/src/segments.rs index 64be95e2f1e..e8ca3d94ab7 100644 --- a/vortex-tui/src/segments.rs +++ b/vortex-tui/src/segments.rs @@ -90,7 +90,7 @@ pub async fn exec_segments(session: &VortexSession, args: SegmentsArgs) -> Vorte row_count: seg.row_count, byte_offset: seg.spec.offset, byte_length: seg.spec.length, - alignment: *seg.spec.alignment, + alignment: seg.spec.alignment.as_usize(), byte_gap, } }) diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index b22440d88b0..344b3b3b244 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -244,7 +244,7 @@ impl VortexFileHandle { index: i, byte_offset: spec.offset, byte_length: spec.length, - alignment: *spec.alignment, + alignment: spec.alignment.as_usize(), column, layout_path, }