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/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index 3c77e146ad0..1d1216ee978 100644 --- a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::marker::PhantomData; use std::mem; use std::mem::MaybeUninit; use std::ops::Range; @@ -51,6 +52,8 @@ impl> UnpackStrategy for BitPackingStr /// /// The usual pattern of usage should follow /// ``` +/// use std::mem::MaybeUninit; +/// /// use lending_iterator::gat; /// use lending_iterator::prelude::Item; /// #[gat(Item)] @@ -65,17 +68,20 @@ impl> UnpackStrategy for BitPackingStr /// let mut ctx = vortex_array::array_session().create_execution_ctx(); /// let array = BitPackedData::encode(&buffer![2, 3, 4, 5].into_array(), 2, &mut ctx).unwrap(); /// let mut unpacked_chunks: BitUnpackedChunks = array.unpacked_chunks().unwrap(); +/// let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; /// -/// if let Some(header) = unpacked_chunks.initial() { +/// if let Some(header) = unpacked_chunks.initial(&mut scratch) { /// // handle partial initial chunk /// } /// -/// let mut chunks_iter = unpacked_chunks.full_chunks(); -/// while let Some(chunk) = chunks_iter.next() { -/// // handle full bitpacked chunks of 1024 elements +/// { +/// let mut chunks_iter = unpacked_chunks.full_chunks(&mut scratch); +/// while let Some(chunk) = chunks_iter.next() { +/// // handle full bitpacked chunks of 1024 elements +/// } /// } /// -/// if let Some(trailer) = unpacked_chunks.trailer() { +/// if let Some(trailer) = unpacked_chunks.trailer(&mut scratch) { /// // handle partial trailing chunk /// } /// ``` @@ -88,7 +94,7 @@ pub struct UnpackedChunks> { // 0 indicates full chunk of CHUNK_SIZE last_chunk_length: usize, packed: ByteBuffer, - buffer: [MaybeUninit; CHUNK_SIZE], + _marker: PhantomData, } pub type BitUnpackedChunks = UnpackedChunks; @@ -104,13 +110,16 @@ impl BitUnpackedChunks { ) } - pub fn full_chunks(&mut self) -> BitUnpackIterator<'_, T> { + pub fn full_chunks<'a>( + &'a self, + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], + ) -> BitUnpackIterator<'a, T> { let elems_per_chunk = self.elems_per_chunk(); let last_chunk_is_sliced = self.last_chunk_is_sliced() as usize; let first_chunk_is_sliced = self.first_chunk_is_sliced(); BitUnpackIterator::new( buffer_as_slice(&self.packed), - &mut self.buffer, + scratch, self.bit_width, elems_per_chunk, self.num_chunks - last_chunk_is_sliced, @@ -148,9 +157,9 @@ impl> UnpackedChunks { offset, len, packed, - buffer: [const { MaybeUninit::::uninit() }; CHUNK_SIZE], num_chunks, last_chunk_length, + _marker: PhantomData, }) } @@ -160,10 +169,13 @@ impl> UnpackedChunks { } /// Access first chunk of the array if the last chunk has fewer than 1024 due to slicing - pub fn initial(&mut self) -> Option<&mut [T]> { + pub fn initial<'a>( + &self, + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], + ) -> Option<&'a mut [T]> { (self.first_chunk_is_sliced() || self.num_chunks == 1).then(|| { let chunk: &[T::Physical] = &buffer_as_slice(&self.packed)[..self.elems_per_chunk()]; - let dst: &mut [MaybeUninit] = &mut self.buffer; + let dst: &mut [MaybeUninit] = scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; let header_end_slice = if self.num_chunks == 1 { @@ -176,7 +188,7 @@ impl> UnpackedChunks { // 2. buffer is exactly CHUNK_SIZE. unsafe { self.strategy.unpack_chunk(self.bit_width, chunk, dst); - mem::transmute(&mut self.buffer[self.offset..][..header_end_slice]) + mem::transmute(&mut scratch[self.offset..][..header_end_slice]) } }) } @@ -184,9 +196,10 @@ impl> UnpackedChunks { /// Decode all chunks (initial, full, and trailer) directly into the output range. pub fn decode_into(&mut self, output: &mut [MaybeUninit]) { debug_assert_eq!(output.len(), self.len); + let mut scratch = [const { MaybeUninit::::uninit() }; CHUNK_SIZE]; let mut local_idx = 0; - if let Some(initial) = self.initial() { + if let Some(initial) = self.initial(&mut scratch) { local_idx = initial.len(); // TODO(connor): use maybe_uninit_write_slice when it gets stabilized. @@ -197,7 +210,7 @@ impl> UnpackedChunks { local_idx = self.decode_full_chunks_into_at(output, local_idx); - if let Some(trailer) = self.trailer() { + if let Some(trailer) = self.trailer(&mut scratch) { // TODO(connor): use maybe_uninit_write_slice when it gets stabilized. // SAFETY: &[T] and &[MaybeUninit] have the same layout. let init_trailer: &[MaybeUninit] = unsafe { mem::transmute(trailer) }; @@ -226,9 +239,10 @@ impl> UnpackedChunks { where F: FnMut(&mut [T], Range), { + let mut scratch = [const { MaybeUninit::::uninit() }; CHUNK_SIZE]; let mut local_idx = 0; - if let Some(initial) = self.initial() { + if let Some(initial) = self.initial(&mut scratch) { let chunk_len = initial.len(); f(initial, local_idx..local_idx + chunk_len); local_idx += chunk_len; @@ -240,16 +254,16 @@ impl> UnpackedChunks { for i in self.full_chunks_range() { let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; unsafe { - let dst: &mut [T::Physical] = mem::transmute(&mut self.buffer[..]); + let dst: &mut [T::Physical] = mem::transmute(&mut scratch[..]); self.strategy.unpack_chunk(self.bit_width, chunk, dst); - let unpacked: &mut [T] = mem::transmute(&mut self.buffer[..]); + let unpacked: &mut [T] = mem::transmute(&mut scratch[..]); f(unpacked, local_idx..local_idx + CHUNK_SIZE); } local_idx += CHUNK_SIZE; } } - if let Some(trailer) = self.trailer() { + if let Some(trailer) = self.trailer(&mut scratch) { let chunk_len = trailer.len(); f(trailer, local_idx..local_idx + chunk_len); local_idx += chunk_len; @@ -320,18 +334,21 @@ impl> UnpackedChunks { } /// Access last chunk of the array if the last chunk has fewer than 1024 due to slicing - pub fn trailer(&mut self) -> Option<&mut [T]> { + pub fn trailer<'a>( + &self, + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], + ) -> Option<&'a mut [T]> { (self.last_chunk_is_sliced() && self.num_chunks > 1).then(|| { let chunk: &[T::Physical] = &buffer_as_slice(&self.packed) [(self.num_chunks - 1) * self.elems_per_chunk()..][..self.elems_per_chunk()]; - let dst: &mut [MaybeUninit] = &mut self.buffer; + let dst: &mut [MaybeUninit] = scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; // SAFETY: // 1. chunk is elems_per_chunk. // 2. buffer is exactly CHUNK_SIZE. unsafe { self.strategy.unpack_chunk(self.bit_width, chunk, dst); - mem::transmute(&mut self.buffer[..self.last_chunk_length]) + mem::transmute(&mut scratch[..self.last_chunk_length]) } }) } diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 3b4e2b9770a..a199eee56fc 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::ops::Range; use itertools::Itertools; @@ -55,7 +56,8 @@ fn bitpacked_is_constant( array: ArrayView<'_, BitPacked>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut bit_unpack_iterator = array.unpacked_chunks::()?; + let bit_unpack_iterator = array.unpacked_chunks::()?; + let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; let patches = array .patches() .map(|p| -> VortexResult<_> { @@ -68,7 +70,7 @@ fn bitpacked_is_constant( let mut header_constant_value = None; let mut current_idx = 0; - if let Some(header) = bit_unpack_iterator.initial() { + if let Some(header) = bit_unpack_iterator.initial(&mut scratch) { if let Some((indices, patches, offset)) = &patches { apply_patches( header, @@ -87,40 +89,42 @@ fn bitpacked_is_constant( } let mut first_chunk_value = None; - let mut chunks_iter = bit_unpack_iterator.full_chunks(); - while let Some(chunk) = chunks_iter.next() { - if let Some((indices, patches, offset)) = &patches { - let chunk_len = chunk.len(); - apply_patches( - chunk, - current_idx..current_idx + chunk_len, - indices, - patches.as_slice::(), - *offset, - ) - } - - if !compute_is_constant::<_, WIDTH>(chunk) { - return Ok(false); - } + { + let mut chunks_iter = bit_unpack_iterator.full_chunks(&mut scratch); + while let Some(chunk) = chunks_iter.next() { + if let Some((indices, patches, offset)) = &patches { + let chunk_len = chunk.len(); + apply_patches( + chunk, + current_idx..current_idx + chunk_len, + indices, + patches.as_slice::(), + *offset, + ) + } - if let Some(chunk_value) = first_chunk_value { - if chunk_value != chunk[0] { + if !compute_is_constant::<_, WIDTH>(chunk) { return Ok(false); } - } else { - if let Some(header_value) = header_constant_value - && header_value != chunk[0] - { - return Ok(false); + + if let Some(chunk_value) = first_chunk_value { + if chunk_value != chunk[0] { + return Ok(false); + } + } else { + if let Some(header_value) = header_constant_value + && header_value != chunk[0] + { + return Ok(false); + } + first_chunk_value = Some(chunk[0]); } - first_chunk_value = Some(chunk[0]); - } - current_idx += chunk.len(); + current_idx += chunk.len(); + } } - if let Some(trailer) = bit_unpack_iterator.trailer() { + if let Some(trailer) = bit_unpack_iterator.trailer(&mut scratch) { if let Some((indices, patches, offset)) = &patches { let chunk_len = trailer.len(); apply_patches( diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 7dc3a371bda..44a6a8e4045 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -47,7 +47,6 @@ use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -568,17 +567,17 @@ impl PcoData { } ); - let mut chunk_meta_buffer = ByteBufferMut::with_capacity(cc.meta_size_hint()); + let mut chunk_meta_buffer = Vec::with_capacity(cc.meta_size_hint()); cc.write_meta(&mut chunk_meta_buffer) .map_err(vortex_err_from_pco)?; - chunk_meta_buffers.push(chunk_meta_buffer.freeze()); + chunk_meta_buffers.push(ByteBuffer::from(chunk_meta_buffer)); let mut page_infos = vec![]; for (page_idx, page_n_values) in cc.n_per_page().into_iter().enumerate() { - let mut page = ByteBufferMut::with_capacity(cc.page_size_hint(page_idx)); + let mut page = Vec::with_capacity(cc.page_size_hint(page_idx)); cc.write_page(page_idx, &mut page) .map_err(vortex_err_from_pco)?; - page_buffers.push(page.freeze()); + page_buffers.push(ByteBuffer::from(page)); page_infos.push(PcoPageInfo { n_values: u32::try_from(page_n_values)?, }); diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index f3a488ea9ba..e18106f4ca4 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -48,7 +48,6 @@ use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_buffer::Buffer; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -217,7 +216,7 @@ impl VTable for Sparse { match idx { 0 => { let fill_value_buffer = - ScalarValue::to_proto_bytes::(array.fill_value.value()).freeze(); + ScalarValue::to_proto_bytes::>(array.fill_value.value()).into(); BufferHandle::new_host(fill_value_buffer) } _ => vortex_panic!("SparseArray buffer index {idx} out of bounds"), 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/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 38a9016e4b8..6d9c8906fc9 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -19,7 +19,6 @@ use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_btrblocks::BtrBlocksCompressorBuilder; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_file::OpenOptionsSessionExt; @@ -72,7 +71,7 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { ), }; - let mut full_buff = ByteBufferMut::empty(); + let mut full_buff = Vec::new(); let _footer = write_options .blocking(&*RUNTIME) .write(&mut full_buff, array_data.to_array_iterator()) diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 6917d979239..7e950bfe5b5 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -6,7 +6,6 @@ use std::hash::Hash; use std::hash::Hasher; use itertools::Itertools; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -107,7 +106,7 @@ impl VTable for Constant { fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { match idx { 0 => BufferHandle::new_host( - ScalarValue::to_proto_bytes::(array.scalar.value()).freeze(), + ScalarValue::to_proto_bytes::>(array.scalar.value()).into(), ), _ => vortex_panic!("ConstantArray buffer index {idx} out of bounds"), } 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-arrow/src/executor/byte_view.rs b/vortex-arrow/src/executor/byte_view.rs index 4406bce8ddc..3026317b511 100644 --- a/vortex-arrow/src/executor/byte_view.rs +++ b/vortex-arrow/src/executor/byte_view.rs @@ -6,12 +6,12 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::GenericByteViewArray; use arrow_array::types::ByteViewType; -use arrow_buffer::ScalarBuffer; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::Nullability; +use vortex_buffer::Buffer; use vortex_error::VortexResult; use crate::dtype::from_arrow_data_type; @@ -22,8 +22,8 @@ pub fn canonical_varbinview_to_arrow( array: &VarBinViewArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let views = - ScalarBuffer::::from(array.views_handle().as_host().clone().into_arrow_buffer()); + let views = Buffer::::from_byte_buffer(array.views_handle().as_host().clone()) + .into_arrow_scalar_buffer(); let buffers: Vec<_> = array .data_buffers() .iter() @@ -64,3 +64,23 @@ pub(super) fn to_arrow_byte_view( let varbinview = array.execute::(ctx)?; execute_varbinview_to_arrow::(&varbinview, ctx) } + +#[cfg(test)] +mod tests { + use arrow_array::types::StringViewType; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + + use super::*; + + #[test] + fn empty_views_are_aligned() -> VortexResult<()> { + let array = VarBinViewArray::from_iter_str(std::iter::empty::<&str>()); + let mut ctx = array_session().create_execution_ctx(); + + let arrow = canonical_varbinview_to_arrow::(&array, &mut ctx)?; + + assert!(arrow.is_empty()); + Ok(()) + } +} 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..7b405b300b6 --- /dev/null +++ b/vortex-buffer/src/allocation.rs @@ -0,0 +1,472 @@ +// 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 static_ref() -> &'static Self { + &STATIC_ALLOCATOR + } + + /// 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), + Bytes(bytes::Bytes), + #[cfg(feature = "arrow")] + Arrow(arrow_buffer::Buffer), + External { + _owner: Box, + }, +} + +impl BufferBacking { + #[inline(always)] + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + match self { + Self::Owned(allocation) => allocation.allocator(), + Self::Bytes(_) | Self::External { .. } => &STATIC_ALLOCATOR, + #[cfg(feature = "arrow")] + Self::Arrow(_) => &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..41e8dbbc733 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -3,7 +3,6 @@ use arrow_buffer::ArrowNativeType; use arrow_buffer::OffsetBuffer; -use bytes::Bytes; use vortex_error::vortex_panic; use crate::Alignment; @@ -13,7 +12,10 @@ use crate::ByteBuffer; impl Buffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_scalar_buffer(self) -> arrow_buffer::ScalarBuffer { - let buffer = arrow_buffer::Buffer::from(self.into_inner()); + if self.is_empty() { + return Vec::new().into(); + } + let buffer = self.into_byte_buffer().into_arrow_buffer(); arrow_buffer::ScalarBuffer::from(buffer) } @@ -25,22 +27,18 @@ impl Buffer { /// alignment is not sufficient for type T. pub fn from_arrow_scalar_buffer(arrow: arrow_buffer::ScalarBuffer) -> Self { let length = arrow.len(); - let bytes = Bytes::from_owner(ArrowWrapper(arrow.into_inner())); + let arrow = arrow.into_inner(); let alignment = Alignment::of::(); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if arrow.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, arrow.len() / size_of::()); + Self::from_arrow_owner(arrow, length, alignment) } /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`. @@ -55,6 +53,10 @@ impl Buffer { impl ByteBuffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_buffer(self) -> arrow_buffer::Buffer { + if let Some(crate::BufferBacking::Arrow(arrow)) = self.backing.as_deref() { + let offset = self.ptr.addr().get() - arrow.as_ptr().addr(); + return arrow.slice_with_length(offset, self.length); + } arrow_buffer::Buffer::from(self.into_inner()) } @@ -66,30 +68,14 @@ impl ByteBuffer { pub fn from_arrow_buffer(arrow: arrow_buffer::Buffer, alignment: Alignment) -> Self { let length = arrow.len(); - let bytes = Bytes::from_owner(ArrowWrapper(arrow)); - if bytes.as_ptr().align_offset(*alignment) != 0 { + if arrow.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(), - } - } -} - -/// A wrapper struct to allow `arrow_buffer::Buffer` to implement `AsRef<[u8]>` for -/// `Bytes::from_owner`. -struct ArrowWrapper(arrow_buffer::Buffer); - -impl AsRef<[u8]> for ArrowWrapper { - fn as_ref(&self) -> &[u8] { - self.0.as_slice() + Self::from_arrow_owner(arrow, length, alignment) } } @@ -118,11 +104,22 @@ mod test { assert_eq!(scalar.as_ptr(), buf.as_ptr(), "Conversion not zero-copy") } + #[test] + fn empty_into_arrow_scalar_buffer() { + let scalar = Buffer::::empty().into_arrow_scalar_buffer(); + + assert!(scalar.is_empty()); + assert_eq!(scalar.as_ptr().align_offset(align_of::()), 0); + } + #[test] fn from_arrow_buffer() { let arrow = ArrowBuffer::from_vec(vec![0i32, 1, 2]); let buf = Buffer::from_arrow_buffer(arrow.clone(), Alignment::of::()); assert_eq!(arrow.as_ref(), buf.as_slice(), "Buffer values differ"); assert_eq!(arrow.as_ptr(), buf.as_ptr(), "Conversion not zero-copy"); + + let round_trip = buf.into_arrow_buffer(); + assert_eq!(round_trip.as_ptr(), arrow.as_ptr()); } } 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..122c8892709 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -8,9 +8,10 @@ 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 bytes::Buf; use bytes::Bytes; @@ -18,6 +19,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 +30,141 @@ 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: Option>, } -/// 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] = { - 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) } -}; +// 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 {} 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: None, } } } -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: Some(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: Some(Arc::new(BufferBacking::External { _owner: owner })), + } + } + + fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self { + let length = bytes.len() / size_of::(); + if length == 0 { + return Self::empty_aligned(alignment); + } + let ptr = + NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("Bytes pointer is null"); + Self { + ptr, + length, + alignment, + physical_alignment: alignment, + overallocated: false, + backing: Some(Arc::new(BufferBacking::Bytes(bytes))), + } + } + + #[cfg(feature = "arrow")] + pub(crate) fn from_arrow_owner( + arrow: arrow_buffer::Buffer, + length: usize, + alignment: Alignment, + ) -> Self { + if length == 0 { + return Self::empty_aligned(alignment); + } + let ptr = NonNull::new(arrow.as_ptr().cast_mut().cast()) + .vortex_expect("Arrow buffer pointer is null"); + Self { + ptr, + length, + alignment, + physical_alignment: alignment, + overallocated: false, + backing: Some(Arc::new(BufferBacking::Arrow(arrow))), + } + } + /// 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 +175,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 +207,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 @@ -150,8 +241,7 @@ impl Buffer { /// Create a new empty `ByteBuffer` with the provided alignment. /// - /// This does not allocate: empty buffers are backed by a zero-length `Bytes` that is - /// aligned to [`Alignment::MAX`]. + /// This does not allocate. Empty buffers use an aligned dangling pointer. pub fn empty_aligned(alignment: Alignment) -> Self { if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( @@ -161,10 +251,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: None, } } @@ -176,6 +268,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 +294,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 +348,7 @@ impl Buffer { size_of::() ); } - let length = bytes.len() / size_of::(); - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + Self::from_bytes(bytes, alignment) } /// Create a buffer with values from the TrustedLen iterator. @@ -249,7 +367,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.allocator().clone(); + let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf .spare_capacity_mut() .iter_mut() @@ -266,7 +385,6 @@ impl Buffer { /// Clear the buffer, preserving existing capacity. pub fn clear(&mut self) { - self.bytes.clear(); self.length = 0; } @@ -288,17 +406,36 @@ impl Buffer { self.alignment } + /// Returns the allocator to use for derived buffers. + /// + /// External buffers use the static allocator. + pub fn allocator(&self) -> &BufferAllocatorRef { + match self.backing.as_deref() { + Some(backing) => backing.allocator(), + None => BufferAllocatorRef::static_ref(), + } + } + + /// 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 +509,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 +520,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: self.backing.clone(), } } @@ -427,74 +565,139 @@ 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: self.backing.clone(), } } - /// 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 + if let Some(backing) = self.backing.as_ref() + && let BufferBacking::Bytes(bytes) = backing.as_ref() + { + let offset = self.ptr.cast::().addr().get() - bytes.as_ptr().addr(); + let length = self.length * size_of::(); + if offset == 0 && length == bytes.len() && Arc::strong_count(backing) == 1 { + return match self.backing { + Some(backing) => match Arc::try_unwrap(backing) { + Ok(BufferBacking::Bytes(bytes)) => bytes, + _ => unreachable!(), + }, + None => unreachable!(), + }; + } + return bytes.slice(offset..offset + length); + } + match self.backing { + Some(backing) => Bytes::from_owner(BufferBytesOwner { + ptr: self.ptr.cast(), + length: self.length * size_of::(), + backing, + }), + None => Bytes::new(), + } } /// 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; + let Some(backing) = backing else { + return Ok(BufferMut::empty_aligned(alignment)); + }; + if !matches!(backing.as_ref(), BufferBacking::Owned(_)) { + return Err(Self { + ptr, + length, + alignment, + physical_alignment, + overallocated, + backing: Some(backing), + }); + } + 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()); + let capacity = if allocation.size() == 0 { + 0 + } else if overallocated { + (allocation.size() - physical_alignment.as_usize()) / size_of::() + } else { + (allocation.size() - offset) / size_of::() + }; + Ok(BufferMut { + allocation, + ptr, + length, + capacity, + alignment, + physical_alignment, + overallocated, + _marker: Default::default(), + }) + } + Ok(_) => unreachable!(), + Err(backing) => Err(Self { + ptr, + length, + alignment, + physical_alignment, + overallocated, + backing: Some(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.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 +713,8 @@ impl Buffer { "Buffer is not aligned to requested alignment {alignment}, copying: {bt}" ) } - Self::copy_from_aligned(self, alignment) + let allocator = self.allocator().clone(); + BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze() } } @@ -546,10 +750,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 +836,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_bytes(bytes, Alignment::of::()) } } @@ -695,11 +898,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,10 +990,17 @@ 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 bytes::Bytes; use crate::Alignment; use crate::Buffer; + use crate::BufferBacking; use crate::ByteBuffer; use crate::buffer; @@ -879,6 +1114,110 @@ 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 bytes_round_trip_reuses_owner() { + let bytes = Bytes::from_static(&[1, 2, 3, 4]); + let ptr = bytes.as_ptr(); + + let buffer = ByteBuffer::from(bytes); + assert!(matches!( + buffer.backing.as_deref(), + Some(BufferBacking::Bytes(_)) + )); + let bytes = buffer.into_inner(); + + assert_eq!(bytes.as_ptr(), ptr); + assert_eq!(bytes.as_ref(), &[1, 2, 3, 4]); + } + + #[test] + fn external_try_into_mut_preserves_backing() { + let buffer = ByteBuffer::from(Bytes::from_static(&[1, 2, 3, 4])); + let Some(original_backing) = buffer.backing.as_ref() else { + panic!("external buffer has no backing") + }; + let backing = Arc::as_ptr(original_backing); + + let Err(buffer) = buffer.try_into_mut() else { + panic!("external buffer became mutable") + }; + + let Some(new_backing) = buffer.backing.as_ref() else { + panic!("external buffer has no backing") + }; + assert_eq!(Arc::as_ptr(new_backing), backing); + } + + #[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. @@ -887,6 +1226,11 @@ mod test { assert!(buf.is_aligned(Alignment::MAX)); } + #[test] + fn empty_has_no_backing() { + assert!(Buffer::::empty().backing.is_none()); + } + #[test] fn empty_slice_preserves_alignment() { let buf = Buffer::::zeroed_aligned(8, Alignment::new(64)); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index e5cb03c558b..7887d8a5acd 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -2,41 +2,53 @@ // 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; use std::fmt::Formatter; -use std::io::Write; use std::ops::Deref; 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) capacity: 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 +58,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 +87,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 +116,44 @@ 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, + capacity, 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 +163,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 +188,52 @@ 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, + capacity: len, alignment, + physical_alignment: actual_alignment, + overallocated: true, _marker: Default::default(), } } @@ -136,7 +250,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 +266,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 +279,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 +313,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 +342,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 +379,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 +399,38 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.bytes.capacity() / size_of::() + self.capacity + } + + /// Returns a raw pointer to the buffer's data. + #[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 +452,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 +461,58 @@ 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(); + 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::(), + ); + } + } + // SAFETY: new_offset was computed within the allocation for physical_alignment. + self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; + self.capacity = logical_size / size_of::(); + self.physical_alignment = physical_alignment; + self.overallocated = true; } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -329,13 +552,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 +568,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 +587,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 +615,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 +624,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 +643,50 @@ 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 { + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + let capacity = if self.allocation.size() == 0 { + 0 + } else if self.overallocated { + self.allocation.size() - self.physical_alignment.as_usize() + } else { + self.allocation.size() - offset + }; ByteBufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length * size_of::(), + capacity, 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 +713,14 @@ impl BufferMut { /// /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { - if self.as_ptr().align_offset(*alignment) == 0 { - Self { - bytes: self.bytes, - length: self.length, - alignment, - _marker: std::marker::PhantomData, - } + if self.as_ptr().align_offset(alignment.as_usize()) == 0 { + Self { alignment, ..self } } else { - Self::copy_from_aligned(self, alignment) + let capacity = self.capacity(); + let allocator = self.allocation.allocator().clone(); + let mut aligned = Self::with_capacity_aligned_in(capacity, alignment, allocator); + aligned.extend_from_slice(&self); + aligned } } @@ -564,9 +744,13 @@ impl BufferMut { ); BufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length, + capacity: self.capacity, alignment: self.alignment, + physical_alignment: self.physical_alignment, + overallocated: self.overallocated, _marker: std::marker::PhantomData, } } @@ -574,14 +758,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 +839,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 +884,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,132 +965,17 @@ 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 } } -impl Buf for ByteBufferMut { - fn remaining(&self) -> usize { - self.len() - } - - fn chunk(&self) -> &[u8] { - self.as_slice() - } - - fn advance(&mut self, cnt: usize) { - if !self.alignment.is_offset_aligned(cnt) { - vortex_panic!( - "Cannot advance buffer by {} items, resulting alignment is not {}", - cnt, - self.alignment - ); - } - self.bytes.advance(cnt); - self.length -= cnt; - } -} - -/// As per the BufMut implementation, we must support internal resizing when -/// asked to extend the buffer. -/// See: -unsafe impl BufMut for ByteBufferMut { - #[inline] - fn remaining_mut(&self) -> usize { - usize::MAX - self.len() - } - - #[inline] - unsafe fn advance_mut(&mut self, cnt: usize) { - if !self.alignment.is_offset_aligned(cnt) { - vortex_panic!( - "Cannot advance buffer by {} items, resulting alignment is not {}", - cnt, - self.alignment - ); - } - unsafe { self.bytes.advance_mut(cnt) }; - self.length -= cnt; - } - - #[inline] - fn chunk_mut(&mut self) -> &mut UninitSlice { - self.bytes.chunk_mut() - } - - fn put(&mut self, mut src: T) - where - Self: Sized, - { - while src.has_remaining() { - let chunk = src.chunk(); - self.extend_from_slice(chunk); - src.advance(chunk.len()); - } - } - - #[inline] - fn put_slice(&mut self, src: &[u8]) { - self.extend_from_slice(src); - } - - #[inline] - fn put_bytes(&mut self, val: u8, cnt: usize) { - self.push_n(val, cnt) - } -} - -/// 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); - Ok(buf.len()) - } - - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } -} - #[cfg(test)] mod test { - use bytes::Buf; - use bytes::BufMut; - use crate::Alignment; use crate::BufferMut; - use crate::ByteBufferMut; use crate::buffer_mut; #[test] @@ -914,6 +993,46 @@ 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 raising_logical_alignment_preserves_capacity() { + let buffer = + BufferMut::::with_capacity_preferred_aligned(1, Alignment::of::(), None); + let capacity = buffer.capacity(); + + let mut buffer = buffer.aligned(Alignment::new(2)); + + assert_eq!(buffer.capacity(), capacity); + buffer.extend(0..100); + assert!(Alignment::new(2).is_ptr_aligned(buffer.as_ptr())); + assert_eq!(buffer.as_slice(), (0..100).collect::>()); + } + #[test] fn from_iter() { let buf = BufferMut::from_iter([0, 10, 20, 30]); @@ -990,34 +1109,16 @@ mod test { assert_eq!(buf.as_slice(), &[1u32, 2, 3]); } - #[test] - fn bytes_buf() { - let mut buf = ByteBufferMut::copy_from("helloworld".as_bytes()); - assert_eq!(buf.remaining(), 10); - assert_eq!(buf.chunk(), b"helloworld"); - - buf.advance(5); - assert_eq!(buf.remaining(), 5); - assert_eq!(buf.as_slice(), b"world"); - assert_eq!(buf.chunk(), b"world"); - } - - #[test] - fn bytes_buf_mut() { - let mut buf = ByteBufferMut::copy_from("hello".as_bytes()); - assert_eq!(BufMut::remaining_mut(&buf), usize::MAX - 5); - - buf.put_slice(b"world"); - assert_eq!(buf.as_slice(), b"helloworld"); - } - #[test] fn buffer_mut_zeroed() { const LEN: usize = 17; 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 +1132,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, }