diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index b6e6a70c6cba..b30a932b35a7 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::alloc::{Layout, handle_alloc_error}; +use std::alloc::Layout; use std::mem; use std::ptr::NonNull; @@ -33,6 +33,31 @@ use std::sync::Mutex; use super::Buffer; +/// Error returned by fallible [`MutableBuffer`] operations. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MutableBufferError { + /// Arithmetic overflow when computing the required buffer length or capacity. + LengthOverflow, + /// The requested capacity cannot be represented as a valid allocation layout. + LayoutError, + /// An allocation failed due to insufficient memory. + AllocationError(Layout), +} + +impl std::fmt::Display for MutableBufferError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::LengthOverflow => write!(f, "buffer length overflow"), + Self::LayoutError => write!(f, "invalid allocation layout for requested capacity"), + Self::AllocationError(layout) => { + write!(f, "failed to allocate memory for layout {layout:?}") + } + } + } +} + +impl std::error::Error for MutableBufferError {} + /// A [`MutableBuffer`] is a wrapper over memory regions, used to build /// [`Buffer`]s out of items or slices of items. /// @@ -118,7 +143,7 @@ impl MutableBuffer { /// See [`MutableBuffer::with_capacity`]. #[inline] pub fn new(capacity: usize) -> Self { - Self::with_capacity(capacity) + Self::try_with_capacity(capacity).unwrap_or_else(|e| panic!("{e}")) } /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`. @@ -129,24 +154,35 @@ impl MutableBuffer { /// then `isize::MAX`, then this function will panic. #[inline] pub fn with_capacity(capacity: usize) -> Self { - let capacity = bit_util::round_upto_multiple_of_64(capacity); + Self::try_with_capacity(capacity).unwrap_or_else(|e| panic!("{e}")) + } + + /// Fallible version of [`MutableBuffer::with_capacity`]. + #[inline] + pub fn try_with_capacity(capacity: usize) -> Result { + let capacity = capacity + .checked_next_multiple_of(64) + .ok_or(MutableBufferError::LayoutError)?; let layout = Layout::from_size_align(capacity, ALIGNMENT) - .expect("failed to create layout for MutableBuffer"); + .map_err(|_| MutableBufferError::LayoutError)?; let data = match layout.size() { 0 => dangling_ptr(), _ => { // Safety: Verified size != 0 let raw_ptr = unsafe { std::alloc::alloc(layout) }; - NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout)) + match NonNull::new(raw_ptr) { + Some(ptr) => ptr, + None => return Err(MutableBufferError::AllocationError(layout)), + } } }; - Self { + Ok(Self { data, len: 0, layout, #[cfg(feature = "pool")] reservation: std::sync::Mutex::new(None), - } + }) } /// Allocates a new [MutableBuffer] with `len` and capacity to be at least `len` where @@ -165,22 +201,31 @@ impl MutableBuffer { /// /// Panics if `len` is too large to construct a valid allocation [`Layout`] pub fn from_len_zeroed(len: usize) -> Self { - let layout = Layout::from_size_align(len, ALIGNMENT).unwrap(); + Self::try_from_len_zeroed(len).unwrap_or_else(|e| panic!("{e}")) + } + + /// Fallible version of [`MutableBuffer::from_len_zeroed`]. + pub fn try_from_len_zeroed(len: usize) -> Result { + let layout = + Layout::from_size_align(len, ALIGNMENT).map_err(|_| MutableBufferError::LayoutError)?; let data = match layout.size() { 0 => dangling_ptr(), _ => { // Safety: Verified size != 0 let raw_ptr = unsafe { std::alloc::alloc_zeroed(layout) }; - NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout)) + match NonNull::new(raw_ptr) { + Some(ptr) => ptr, + None => return Err(MutableBufferError::AllocationError(layout)), + } } }; - Self { + Ok(Self { data, len, layout, #[cfg(feature = "pool")] reservation: std::sync::Mutex::new(None), - } + }) } /// Allocates a new [MutableBuffer] from given `Bytes`. @@ -259,6 +304,22 @@ impl MutableBuffer { } } + /// Fallible version of [`MutableBuffer::reserve`]. + #[inline] + pub fn try_reserve(&mut self, additional: usize) -> Result<(), MutableBufferError> { + let required_cap = self + .len + .checked_add(additional) + .ok_or(MutableBufferError::LengthOverflow)?; + if required_cap > self.layout.size() { + let new_capacity = required_cap + .checked_next_multiple_of(64) + .ok_or(MutableBufferError::LayoutError)?; + let new_capacity = std::cmp::max(new_capacity, self.layout.size().saturating_mul(2)); + self.try_reallocate(new_capacity)?; + } + Ok(()) + } /// Ensures that this buffer has at least `self.len + additional` bytes. This re-allocates iff /// `self.len + additional > capacity`. /// # Example @@ -279,108 +340,98 @@ impl MutableBuffer { // exits. #[inline(always)] pub fn reserve(&mut self, additional: usize) { - let required_cap = self - .len - .checked_add(additional) - .expect("buffer length overflow"); - if required_cap > self.layout.size() { - let new_capacity = bit_util::round_upto_multiple_of_64(required_cap); - let new_capacity = std::cmp::max(new_capacity, self.layout.size() * 2); - self.reallocate(new_capacity) - } + self.try_reserve(additional) + .unwrap_or_else(|e| panic!("{e}")) } - /// Adding to this mutable buffer `slice_to_repeat` repeated `repeat_count` times. - /// - /// # Example - /// - /// ## Repeat the same string bytes multiple times - /// ``` - /// # use arrow_buffer::buffer::MutableBuffer; - /// let mut buffer = MutableBuffer::new(0); - /// let bytes_to_repeat = b"ab"; - /// buffer.repeat_slice_n_times(bytes_to_repeat, 3); - /// assert_eq!(buffer.as_slice(), b"ababab"); - /// ``` - /// - /// # Panics - /// - /// Panics if the repeated slice byte length overflows `usize`, if the resulting buffer - /// length overflows `usize`, or if reserving the required capacity fails for the same - /// reasons as [`MutableBuffer::reserve`]. - pub fn repeat_slice_n_times( + /// Fallible version of [`MutableBuffer::repeat_slice_n_times`]. + pub fn try_repeat_slice_n_times( &mut self, slice_to_repeat: &[T], repeat_count: usize, - ) { + ) -> Result<(), MutableBufferError> { if repeat_count == 0 || slice_to_repeat.is_empty() { - return; + return Ok(()); } - - let bytes_to_repeat = size_of_val(slice_to_repeat); - let repeated_bytes = repeat_count - .checked_mul(bytes_to_repeat) - .expect("repeated slice byte length overflow"); + let bytes_per_copy = size_of_val(slice_to_repeat); + let total_bytes = repeat_count + .checked_mul(bytes_per_copy) + .ok_or(MutableBufferError::LengthOverflow)?; self.len - .checked_add(repeated_bytes) - .expect("mutable buffer length overflow"); + .checked_add(total_bytes) + .ok_or(MutableBufferError::LengthOverflow)?; // Ensure capacity - self.reserve(repeated_bytes); + self.try_reserve(total_bytes)?; // Save the length before we do all the copies to know where to start from let length_before = self.len; // Copy the initial slice once so we can use doubling strategy on it - self.extend_from_slice(slice_to_repeat); - - // This tracks how much bytes we have added by repeating so far - let added_repeats_length = bytes_to_repeat; - assert_eq!( - self.len - length_before, - added_repeats_length, - "should copy exactly the same number of bytes" - ); + self.try_extend_from_slice(slice_to_repeat)?; // Number of times the slice was repeated - let mut already_repeated_times = 1; + let mut already_repeated = 1usize; // We will use doubling strategy to fill the buffer in log(repeat_count) steps - while already_repeated_times < repeat_count { + while already_repeated < repeat_count { // How many slices can we copy in this iteration // (either double what we have, or just the remaining ones) - let number_of_slices_to_copy = - already_repeated_times.min(repeat_count - already_repeated_times); - let number_of_bytes_to_copy = number_of_slices_to_copy * bytes_to_repeat; - + let to_copy = already_repeated.min(repeat_count - already_repeated); + let byte_count = to_copy * bytes_per_copy; unsafe { // Get to the start of the data before we started copying anything let src = self.data.as_ptr().add(length_before) as *const u8; - // Go to the current location to copy to (end of current data) let dst = self.data.as_ptr().add(self.len); - - // SAFETY: the pointers are not overlapping as there is `number_of_bytes_to_copy` or less between them - std::ptr::copy_nonoverlapping(src, dst, number_of_bytes_to_copy) + // SAFETY: the pointers are not overlapping as there is `byte_count` or less between them + std::ptr::copy_nonoverlapping(src, dst, byte_count); } - // Advance the length by the amount of data we just copied (doubled) - self.len += number_of_bytes_to_copy; - - already_repeated_times += number_of_slices_to_copy; + self.len += byte_count; + already_repeated += to_copy; } + Ok(()) + } + /// Adding to this mutable buffer `slice_to_repeat` repeated `repeat_count` times. + /// + /// # Example + /// + /// ## Repeat the same string bytes multiple times + /// ``` + /// # use arrow_buffer::buffer::MutableBuffer; + /// let mut buffer = MutableBuffer::new(0); + /// let bytes_to_repeat = b"ab"; + /// buffer.repeat_slice_n_times(bytes_to_repeat, 3); + /// assert_eq!(buffer.as_slice(), b"ababab"); + /// ``` + /// + /// # Panics + /// + /// Panics if the repeated slice byte length overflows `usize`, if the resulting buffer + /// length overflows `usize`, or if reserving the required capacity fails for the same + /// reasons as [`MutableBuffer::reserve`]. + pub fn repeat_slice_n_times( + &mut self, + slice_to_repeat: &[T], + repeat_count: usize, + ) { + self.try_repeat_slice_n_times(slice_to_repeat, repeat_count) + .unwrap_or_else(|e| panic!("{e}")) } #[cold] - fn reallocate(&mut self, capacity: usize) { - let new_layout = Layout::from_size_align(capacity, self.layout.align()).unwrap(); + fn try_reallocate(&mut self, capacity: usize) -> Result<(), MutableBufferError> { + let new_layout = Layout::from_size_align(capacity, self.layout.align()) + .map_err(|_| MutableBufferError::LayoutError)?; + if new_layout.size() == 0 { if self.layout.size() != 0 { // Safety: data was allocated with layout unsafe { std::alloc::dealloc(self.as_mut_ptr(), self.layout) }; - self.layout = new_layout + self.layout = new_layout; } - return; + return Ok(()); } let data = match self.layout.size() { @@ -389,7 +440,10 @@ impl MutableBuffer { // Safety: verified new layout is valid and not empty _ => unsafe { std::alloc::realloc(self.as_mut_ptr(), self.layout, capacity) }, }; - self.data = NonNull::new(data).unwrap_or_else(|| handle_alloc_error(new_layout)); + self.data = match NonNull::new(data) { + Some(ptr) => ptr, + None => return Err(MutableBufferError::AllocationError(new_layout)), + }; self.layout = new_layout; #[cfg(feature = "pool")] { @@ -397,8 +451,8 @@ impl MutableBuffer { reservation.resize(self.layout.size()); } } + Ok(()) } - /// Truncates this buffer to `len` bytes /// /// If `len` is greater than the buffer's current length, this has no effect @@ -416,6 +470,26 @@ impl MutableBuffer { } } + /// Fallible version of [`MutableBuffer::resize`]. + #[inline] + pub fn try_resize(&mut self, new_len: usize, value: u8) -> Result<(), MutableBufferError> { + if new_len > self.len { + let diff = new_len - self.len; + self.try_reserve(diff)?; + // Safety: try_reserve ensured capacity >= new_len. + // write the value + unsafe { self.data.as_ptr().add(self.len).write_bytes(value, diff) }; + } + // this truncates the buffer when new_len < self.len + self.len = new_len; + #[cfg(feature = "pool")] + { + if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { + reservation.resize(self.len); + } + } + Ok(()) + } /// Resizes the buffer, either truncating its contents (with no change in capacity), or /// growing it (potentially reallocating it) and writing `value` in the newly available bytes. /// # Example @@ -434,22 +508,21 @@ impl MutableBuffer { // exits. #[inline(always)] pub fn resize(&mut self, new_len: usize, value: u8) { - if new_len > self.len { - let diff = new_len - self.len; - self.reserve(diff); - // write the value - unsafe { self.data.as_ptr().add(self.len).write_bytes(value, diff) }; - } - // this truncates the buffer when new_len < self.len - self.len = new_len; - #[cfg(feature = "pool")] - { - if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { - reservation.resize(self.len); - } - } + self.try_resize(new_len, value) + .unwrap_or_else(|e| panic!("{e}")) } + /// Fallible version of [`MutableBuffer::shrink_to_fit`]. + pub fn try_shrink_to_fit(&mut self) -> Result<(), MutableBufferError> { + let new_capacity = self + .len + .checked_next_multiple_of(64) + .ok_or(MutableBufferError::LayoutError)?; + if new_capacity < self.layout.size() { + self.try_reallocate(new_capacity)?; + } + Ok(()) + } /// Shrinks the capacity of the buffer as much as possible. /// The new capacity will aligned to the nearest 64 bit alignment. /// @@ -471,10 +544,7 @@ impl MutableBuffer { /// Panics if the current length is too large to round up to the next 64-byte boundary and /// construct a valid allocation layout. pub fn shrink_to_fit(&mut self) { - let new_capacity = bit_util::round_upto_multiple_of_64(self.len); - if new_capacity < self.layout.size() { - self.reallocate(new_capacity) - } + self.try_shrink_to_fit().unwrap_or_else(|e| panic!("{e}")) } /// Returns whether this buffer is empty or not. @@ -575,6 +645,25 @@ impl MutableBuffer { offsets } + /// Fallible version of [`MutableBuffer::extend_from_slice`]. + #[inline] + pub fn try_extend_from_slice( + &mut self, + items: &[T], + ) -> Result<(), MutableBufferError> { + let additional = mem::size_of_val(items); + self.try_reserve(additional)?; + unsafe { + // this assumes that `[ToByteSlice]` can be copied directly + // without calling `to_byte_slice` for each element, + // which is correct for all ArrowNativeType implementations. + let src = items.as_ptr() as *const u8; + let dst = self.data.as_ptr().add(self.len); + std::ptr::copy_nonoverlapping(src, dst, additional); + } + self.len += additional; + Ok(()) + } /// Extends this buffer from a slice of items that can be represented in bytes, increasing its capacity if needed. /// # Example /// ``` @@ -590,17 +679,8 @@ impl MutableBuffer { /// reasons as [`MutableBuffer::reserve`]. #[inline] pub fn extend_from_slice(&mut self, items: &[T]) { - let additional = mem::size_of_val(items); - self.reserve(additional); - unsafe { - // this assumes that `[ToByteSlice]` can be copied directly - // without calling `to_byte_slice` for each element, - // which is correct for all ArrowNativeType implementations. - let src = items.as_ptr() as *const u8; - let dst = self.data.as_ptr().add(self.len); - std::ptr::copy_nonoverlapping(src, dst, additional) - } - self.len += additional; + self.try_extend_from_slice(items) + .unwrap_or_else(|e| panic!("{e}")) } /// Extends the buffer with a new item, increasing its capacity if needed. @@ -640,6 +720,15 @@ impl MutableBuffer { self.len += additional; } + /// Fallible version of [`MutableBuffer::extend_zeros`]. + #[inline] + pub fn try_extend_zeros(&mut self, additional: usize) -> Result<(), MutableBufferError> { + let new_len = self + .len + .checked_add(additional) + .ok_or(MutableBufferError::LengthOverflow)?; + self.try_resize(new_len, 0) + } /// Extends the buffer by `additional` bytes equal to `0u8`, incrementing its capacity if needed. /// /// # Panics @@ -648,11 +737,8 @@ impl MutableBuffer { /// reserving a capacity that fails for the same reasons as [`MutableBuffer::reserve`]. #[inline] pub fn extend_zeros(&mut self, additional: usize) { - let new_len = self - .len - .checked_add(additional) - .expect("buffer length overflow"); - self.resize(new_len, 0); + self.try_extend_zeros(additional) + .unwrap_or_else(|e| panic!("{e}")) } /// # Safety @@ -1424,7 +1510,7 @@ mod tests { } #[test] - #[should_panic(expected = "failed to create layout for MutableBuffer: LayoutError")] + #[should_panic(expected = "invalid allocation layout for requested capacity")] fn test_with_capacity_panics_above_max_capacity() { let max_capacity = isize::MAX as usize - (isize::MAX as usize % ALIGNMENT); let _ = MutableBuffer::with_capacity(max_capacity + 1); @@ -1446,14 +1532,14 @@ mod tests { assert_eq!(pool.used(), 128); // Reallocate to a larger size - buffer.reallocate(200); + buffer.try_reallocate(200).unwrap(); // The capacity is exactly the requested size, not rounded up assert_eq!(buffer.capacity(), 200); assert_eq!(pool.used(), 200); // Reallocate to a smaller size - buffer.reallocate(50); + buffer.try_reallocate(50).unwrap(); // The capacity is exactly the requested size, not rounded up assert_eq!(buffer.capacity(), 50); @@ -1576,14 +1662,14 @@ mod tests { } #[test] - #[should_panic(expected = "repeated slice byte length overflow")] + #[should_panic(expected = "buffer length overflow")] fn test_repeat_slice_count_multiply_overflow() { let mut buffer = MutableBuffer::new(0); buffer.repeat_slice_n_times(&[0_u64], usize::MAX / mem::size_of::() + 1); } #[test] - #[should_panic(expected = "mutable buffer length overflow")] + #[should_panic(expected = "buffer length overflow")] fn test_repeat_slice_count_len_overflow() { let mut buffer = MutableBuffer::new(0); buffer.push(0_u8); @@ -1671,7 +1757,7 @@ mod tests { } #[test] - #[should_panic(expected = "failed to round upto multiple of 64")] + #[should_panic(expected = "invalid allocation layout for requested capacity")] fn test_mutable_new_capacity_overflow() { // Tests overflow during initial allocation let _ = MutableBuffer::new(usize::MAX - 10);