arrow_buffer/buffer/
mutable.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::alloc::Layout;
19use std::mem;
20use std::ptr::NonNull;
21
22use crate::alloc::{ALIGNMENT, Deallocation};
23use crate::{
24    bytes::Bytes,
25    native::{ArrowNativeType, ToByteSlice},
26    util::bit_util,
27};
28
29#[cfg(feature = "pool")]
30use crate::pool::{MemoryPool, MemoryReservation, lock_reservation};
31#[cfg(feature = "pool")]
32use std::sync::Mutex;
33
34use super::Buffer;
35
36/// Error returned by fallible [`MutableBuffer`] operations.
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub enum MutableBufferError {
39    /// Arithmetic overflow when computing the required buffer length or capacity.
40    LengthOverflow,
41    /// The requested capacity cannot be represented as a valid allocation layout.
42    LayoutError,
43    /// An allocation failed due to insufficient memory.
44    AllocationError(Layout),
45}
46
47impl std::fmt::Display for MutableBufferError {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        match self {
50            Self::LengthOverflow => write!(f, "buffer length overflow"),
51            Self::LayoutError => write!(f, "invalid allocation layout for requested capacity"),
52            Self::AllocationError(layout) => {
53                write!(f, "failed to allocate memory for layout {layout:?}")
54            }
55        }
56    }
57}
58
59impl std::error::Error for MutableBufferError {}
60
61/// A [`MutableBuffer`] is a wrapper over memory regions, used to build
62/// [`Buffer`]s out of items or slices of items.
63///
64/// [`Buffer`]s created from [`MutableBuffer`] (via `into`) are guaranteed to be
65/// aligned along cache lines and in multiples of 64 bytes.
66///
67/// Use [MutableBuffer::push] to insert an item, [MutableBuffer::extend_from_slice]
68/// to insert many items, and `into` to convert it to [`Buffer`]. For typed data,
69/// it is often more efficient to use [`Vec`] and convert it to [`Buffer`] rather
70/// than using [`MutableBuffer`] (see examples below).
71///
72/// # See Also
73/// * For a safe, strongly typed API consider using [`Vec`] and [`ScalarBuffer`](crate::ScalarBuffer)
74/// * To apply bitwise operations, see [`apply_bitwise_binary_op`] and [`apply_bitwise_unary_op`]
75///
76/// [`apply_bitwise_binary_op`]: crate::bit_util::apply_bitwise_binary_op
77/// [`apply_bitwise_unary_op`]: crate::bit_util::apply_bitwise_unary_op
78///
79/// # Example: Creating a [`Buffer`] from a [`MutableBuffer`]
80/// ```
81/// # use arrow_buffer::buffer::{Buffer, MutableBuffer};
82/// let mut buffer = MutableBuffer::new(0);
83/// buffer.push(256u32);
84/// buffer.extend_from_slice(&[1u32]);
85/// let buffer = Buffer::from(buffer);
86/// assert_eq!(buffer.as_slice(), &[0u8, 1, 0, 0, 1, 0, 0, 0])
87/// ```
88///
89/// The same can be achieved more efficiently by using a `Vec<u32>`
90/// ```
91/// # use arrow_buffer::buffer::Buffer;
92/// let mut vec = Vec::new();
93/// vec.push(256u32);
94/// vec.extend_from_slice(&[1u32]);
95/// let buffer = Buffer::from(vec);
96/// assert_eq!(buffer.as_slice(), &[0u8, 1, 0, 0, 1, 0, 0, 0]);
97/// ```
98///
99/// # Example: Creating a [`MutableBuffer`] from a `Vec<T>`
100/// ```
101/// # use arrow_buffer::buffer::MutableBuffer;
102/// let vec = vec![1u32, 2, 3];
103/// let mutable_buffer = MutableBuffer::from(vec); // reuses the allocation from vec
104/// assert_eq!(mutable_buffer.len(), 12); // 3 * 4 bytes
105/// ```
106///
107/// # Example: Creating a [`MutableBuffer`] from a [`Buffer`]
108/// ```
109/// # use arrow_buffer::buffer::{Buffer, MutableBuffer};
110/// let buffer: Buffer = Buffer::from(&[1u8, 2, 3, 4][..]);
111/// // Only possible to convert a Buffer into a MutableBuffer if uniquely owned
112/// // (i.e., there are no other references to it).
113/// let mut mutable_buffer = match buffer.into_mutable() {
114///    Ok(mutable) => mutable,
115///    Err(orig_buffer) => {
116///      panic!("buffer was not uniquely owned");
117///    }
118/// };
119/// mutable_buffer.push(5u8);
120/// let buffer = Buffer::from(mutable_buffer);
121/// assert_eq!(buffer.as_slice(), &[1u8, 2, 3, 4, 5])
122/// ```
123#[derive(Debug)]
124pub struct MutableBuffer {
125    // dangling iff capacity = 0
126    data: NonNull<u8>,
127    // invariant: len <= capacity
128    len: usize,
129    layout: Layout,
130
131    /// Memory reservation for tracking memory usage
132    #[cfg(feature = "pool")]
133    reservation: Mutex<Option<Box<dyn MemoryReservation>>>,
134}
135
136impl MutableBuffer {
137    /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`.
138    ///
139    /// See [`MutableBuffer::with_capacity`].
140    ///
141    /// # Panics
142    ///
143    /// See [`MutableBuffer::with_capacity`].
144    #[inline]
145    pub fn new(capacity: usize) -> Self {
146        Self::try_with_capacity(capacity).unwrap_or_else(|e| panic!("{e}"))
147    }
148
149    /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`.
150    ///
151    /// # Panics
152    ///
153    /// If `capacity`, when rounded up to the nearest multiple of [`ALIGNMENT`], is greater
154    /// then `isize::MAX`, then this function will panic.
155    #[inline]
156    pub fn with_capacity(capacity: usize) -> Self {
157        Self::try_with_capacity(capacity).unwrap_or_else(|e| panic!("{e}"))
158    }
159
160    /// Fallible version of [`MutableBuffer::with_capacity`].
161    #[inline]
162    pub fn try_with_capacity(capacity: usize) -> Result<Self, MutableBufferError> {
163        let capacity = capacity
164            .checked_next_multiple_of(64)
165            .ok_or(MutableBufferError::LayoutError)?;
166        let layout = Layout::from_size_align(capacity, ALIGNMENT)
167            .map_err(|_| MutableBufferError::LayoutError)?;
168        let data = match layout.size() {
169            0 => dangling_ptr(),
170            _ => {
171                // Safety: Verified size != 0
172                let raw_ptr = unsafe { std::alloc::alloc(layout) };
173                match NonNull::new(raw_ptr) {
174                    Some(ptr) => ptr,
175                    None => return Err(MutableBufferError::AllocationError(layout)),
176                }
177            }
178        };
179        Ok(Self {
180            data,
181            len: 0,
182            layout,
183            #[cfg(feature = "pool")]
184            reservation: std::sync::Mutex::new(None),
185        })
186    }
187
188    /// Allocates a new [MutableBuffer] with `len` and capacity to be at least `len` where
189    /// all bytes are guaranteed to be `0u8`.
190    /// # Example
191    /// ```
192    /// # use arrow_buffer::buffer::{Buffer, MutableBuffer};
193    /// let mut buffer = MutableBuffer::from_len_zeroed(127);
194    /// assert_eq!(buffer.len(), 127);
195    /// assert!(buffer.capacity() >= 127);
196    /// let data = buffer.as_slice_mut();
197    /// assert_eq!(data[126], 0u8);
198    /// ```
199    ///
200    /// # Panics
201    ///
202    /// Panics if `len` is too large to construct a valid allocation [`Layout`]
203    pub fn from_len_zeroed(len: usize) -> Self {
204        Self::try_from_len_zeroed(len).unwrap_or_else(|e| panic!("{e}"))
205    }
206
207    /// Fallible version of [`MutableBuffer::from_len_zeroed`].
208    pub fn try_from_len_zeroed(len: usize) -> Result<Self, MutableBufferError> {
209        let layout =
210            Layout::from_size_align(len, ALIGNMENT).map_err(|_| MutableBufferError::LayoutError)?;
211        let data = match layout.size() {
212            0 => dangling_ptr(),
213            _ => {
214                // Safety: Verified size != 0
215                let raw_ptr = unsafe { std::alloc::alloc_zeroed(layout) };
216                match NonNull::new(raw_ptr) {
217                    Some(ptr) => ptr,
218                    None => return Err(MutableBufferError::AllocationError(layout)),
219                }
220            }
221        };
222        Ok(Self {
223            data,
224            len,
225            layout,
226            #[cfg(feature = "pool")]
227            reservation: std::sync::Mutex::new(None),
228        })
229    }
230
231    /// Allocates a new [MutableBuffer] from given `Bytes`.
232    pub(crate) fn from_bytes(bytes: Bytes) -> Result<Self, Bytes> {
233        let layout = match bytes.deallocation() {
234            Deallocation::Standard(layout) => *layout,
235            Deallocation::Custom(..) => return Err(bytes),
236        };
237
238        let len = bytes.len();
239        let data = bytes.ptr();
240        #[cfg(feature = "pool")]
241        let reservation = lock_reservation(&bytes.reservation).take();
242        mem::forget(bytes);
243
244        Ok(Self {
245            data,
246            len,
247            layout,
248            #[cfg(feature = "pool")]
249            reservation: Mutex::new(reservation),
250        })
251    }
252
253    /// creates a new [MutableBuffer] with capacity and length capable of holding `len` bits.
254    /// This is useful to create a buffer for packed bitmaps.
255    ///
256    /// # Panics
257    ///
258    /// See [`MutableBuffer::from_len_zeroed`].
259    pub fn new_null(len: usize) -> Self {
260        let num_bytes = bit_util::ceil(len, 8);
261        MutableBuffer::from_len_zeroed(num_bytes)
262    }
263
264    /// Set the bits in the range of `[0, end)` to 0 (if `val` is false), or 1 (if `val`
265    /// is true). Also extend the length of this buffer to be `end`.
266    ///
267    /// This is useful when one wants to clear (or set) the bits and then manipulate
268    /// the buffer directly (e.g., modifying the buffer by holding a mutable reference
269    /// from `data_mut()`).
270    ///
271    /// # Panics
272    ///
273    /// Panics if `end` exceeds the buffer capacity.
274    pub fn with_bitset(mut self, end: usize, val: bool) -> Self {
275        assert!(end <= self.layout.size());
276        let v = if val { 255 } else { 0 };
277        unsafe {
278            std::ptr::write_bytes(self.data.as_ptr(), v, end);
279            self.len = end;
280        }
281        self
282    }
283
284    /// Ensure that `count` bytes from `start` contain zero bits
285    ///
286    /// This is used to initialize the bits in a buffer, however, it has no impact on the
287    /// `len` of the buffer and so can be used to initialize the memory region from
288    /// `len` to `capacity`.
289    ///
290    /// # Panics
291    ///
292    /// Panics if the byte range `start..start + count` exceeds the buffer capacity.
293    pub fn set_null_bits(&mut self, start: usize, count: usize) {
294        assert!(
295            start.saturating_add(count) <= self.layout.size(),
296            "range start index {start} and count {count} out of bounds for \
297            buffer of length {}",
298            self.layout.size(),
299        );
300
301        // Safety: `self.data[start..][..count]` is in-bounds and well-aligned for `u8`
302        unsafe {
303            std::ptr::write_bytes(self.data.as_ptr().add(start), 0, count);
304        }
305    }
306
307    /// Fallible version of [`MutableBuffer::reserve`].
308    #[inline]
309    pub fn try_reserve(&mut self, additional: usize) -> Result<(), MutableBufferError> {
310        let required_cap = self
311            .len
312            .checked_add(additional)
313            .ok_or(MutableBufferError::LengthOverflow)?;
314        if required_cap > self.layout.size() {
315            let new_capacity = required_cap
316                .checked_next_multiple_of(64)
317                .ok_or(MutableBufferError::LayoutError)?;
318            let new_capacity = std::cmp::max(new_capacity, self.layout.size().saturating_mul(2));
319            self.try_reallocate(new_capacity)?;
320        }
321        Ok(())
322    }
323    /// Ensures that this buffer has at least `self.len + additional` bytes. This re-allocates iff
324    /// `self.len + additional > capacity`.
325    /// # Example
326    /// ```
327    /// # use arrow_buffer::buffer::{Buffer, MutableBuffer};
328    /// let mut buffer = MutableBuffer::new(0);
329    /// buffer.reserve(253); // allocates for the first time
330    /// (0..253u8).for_each(|i| buffer.push(i)); // no reallocation
331    /// let buffer: Buffer = buffer.into();
332    /// assert_eq!(buffer.len(), 253);
333    /// ```
334    ///
335    /// # Panics
336    ///
337    /// Panics if `self.len + additional` overflows `usize`, or if the required capacity is too
338    /// large to round up to the next 64-byte boundary and construct a valid allocation layout.
339    // For performance reasons, this must be inlined so that the `if` is executed inside the caller, and not as an extra call that just
340    // exits.
341    #[inline(always)]
342    pub fn reserve(&mut self, additional: usize) {
343        self.try_reserve(additional)
344            .unwrap_or_else(|e| panic!("{e}"))
345    }
346
347    /// Fallible version of [`MutableBuffer::repeat_slice_n_times`].
348    pub fn try_repeat_slice_n_times<T: ArrowNativeType>(
349        &mut self,
350        slice_to_repeat: &[T],
351        repeat_count: usize,
352    ) -> Result<(), MutableBufferError> {
353        if repeat_count == 0 || slice_to_repeat.is_empty() {
354            return Ok(());
355        }
356        let bytes_per_copy = size_of_val(slice_to_repeat);
357        let total_bytes = repeat_count
358            .checked_mul(bytes_per_copy)
359            .ok_or(MutableBufferError::LengthOverflow)?;
360        self.len
361            .checked_add(total_bytes)
362            .ok_or(MutableBufferError::LengthOverflow)?;
363
364        // Ensure capacity
365        self.try_reserve(total_bytes)?;
366
367        // Save the length before we do all the copies to know where to start from
368        let length_before = self.len;
369
370        // Copy the initial slice once so we can use doubling strategy on it
371        self.try_extend_from_slice(slice_to_repeat)?;
372
373        // Number of times the slice was repeated
374        let mut already_repeated = 1usize;
375
376        // We will use doubling strategy to fill the buffer in log(repeat_count) steps
377        while already_repeated < repeat_count {
378            // How many slices can we copy in this iteration
379            // (either double what we have, or just the remaining ones)
380            let to_copy = already_repeated.min(repeat_count - already_repeated);
381            let byte_count = to_copy * bytes_per_copy;
382            unsafe {
383                // Get to the start of the data before we started copying anything
384                let src = self.data.as_ptr().add(length_before).cast_const();
385                // Go to the current location to copy to (end of current data)
386                let dst = self.data.as_ptr().add(self.len);
387                // SAFETY: the pointers are not overlapping as there is `byte_count` or less between them
388                std::ptr::copy_nonoverlapping(src, dst, byte_count);
389            }
390            // Advance the length by the amount of data we just copied (doubled)
391            self.len += byte_count;
392            already_repeated += to_copy;
393        }
394        Ok(())
395    }
396    /// Adding to this mutable buffer `slice_to_repeat` repeated `repeat_count` times.
397    ///
398    /// # Example
399    ///
400    /// ## Repeat the same string bytes multiple times
401    /// ```
402    /// # use arrow_buffer::buffer::MutableBuffer;
403    /// let mut buffer = MutableBuffer::new(0);
404    /// let bytes_to_repeat = b"ab";
405    /// buffer.repeat_slice_n_times(bytes_to_repeat, 3);
406    /// assert_eq!(buffer.as_slice(), b"ababab");
407    /// ```
408    ///
409    /// # Panics
410    ///
411    /// Panics if the repeated slice byte length overflows `usize`, if the resulting buffer
412    /// length overflows `usize`, or if reserving the required capacity fails for the same
413    /// reasons as [`MutableBuffer::reserve`].
414    pub fn repeat_slice_n_times<T: ArrowNativeType>(
415        &mut self,
416        slice_to_repeat: &[T],
417        repeat_count: usize,
418    ) {
419        self.try_repeat_slice_n_times(slice_to_repeat, repeat_count)
420            .unwrap_or_else(|e| panic!("{e}"))
421    }
422
423    #[cold]
424    fn try_reallocate(&mut self, capacity: usize) -> Result<(), MutableBufferError> {
425        let new_layout = Layout::from_size_align(capacity, self.layout.align())
426            .map_err(|_| MutableBufferError::LayoutError)?;
427
428        if new_layout.size() == 0 {
429            if self.layout.size() != 0 {
430                // Safety: data was allocated with layout
431                unsafe { std::alloc::dealloc(self.as_mut_ptr(), self.layout) };
432                self.layout = new_layout;
433            }
434            return Ok(());
435        }
436
437        let data = match self.layout.size() {
438            // Safety: new_layout is not empty
439            0 => unsafe { std::alloc::alloc(new_layout) },
440            // Safety: verified new layout is valid and not empty
441            _ => unsafe { std::alloc::realloc(self.as_mut_ptr(), self.layout, capacity) },
442        };
443        self.data = match NonNull::new(data) {
444            Some(ptr) => ptr,
445            None => return Err(MutableBufferError::AllocationError(new_layout)),
446        };
447        self.layout = new_layout;
448        #[cfg(feature = "pool")]
449        {
450            if let Some(reservation) = lock_reservation(&self.reservation).as_mut() {
451                reservation.resize(self.layout.size());
452            }
453        }
454        Ok(())
455    }
456    /// Truncates this buffer to `len` bytes
457    ///
458    /// If `len` is greater than the buffer's current length, this has no effect
459    #[inline(always)]
460    pub fn truncate(&mut self, len: usize) {
461        if len > self.len {
462            return;
463        }
464        self.len = len;
465        #[cfg(feature = "pool")]
466        {
467            if let Some(reservation) = lock_reservation(&self.reservation).as_mut() {
468                reservation.resize(self.len);
469            }
470        }
471    }
472
473    /// Fallible version of [`MutableBuffer::resize`].
474    #[inline]
475    pub fn try_resize(&mut self, new_len: usize, value: u8) -> Result<(), MutableBufferError> {
476        if new_len > self.len {
477            let diff = new_len - self.len;
478            self.try_reserve(diff)?;
479            // Safety: try_reserve ensured capacity >= new_len.
480            // write the value
481            unsafe { self.data.as_ptr().add(self.len).write_bytes(value, diff) };
482        }
483        // this truncates the buffer when new_len < self.len
484        self.len = new_len;
485        #[cfg(feature = "pool")]
486        {
487            if let Some(reservation) = lock_reservation(&self.reservation).as_mut() {
488                reservation.resize(self.len);
489            }
490        }
491        Ok(())
492    }
493    /// Resizes the buffer, either truncating its contents (with no change in capacity), or
494    /// growing it (potentially reallocating it) and writing `value` in the newly available bytes.
495    /// # Example
496    /// ```
497    /// # use arrow_buffer::buffer::{Buffer, MutableBuffer};
498    /// let mut buffer = MutableBuffer::new(0);
499    /// buffer.resize(253, 2); // allocates for the first time
500    /// assert_eq!(buffer.as_slice()[252], 2u8);
501    /// ```
502    ///
503    /// # Panics
504    ///
505    /// Panics if growing the buffer requires reserving a capacity that fails for the same
506    /// reasons as [`MutableBuffer::reserve`].
507    // For performance reasons, this must be inlined so that the `if` is executed inside the caller, and not as an extra call that just
508    // exits.
509    #[inline(always)]
510    pub fn resize(&mut self, new_len: usize, value: u8) {
511        self.try_resize(new_len, value)
512            .unwrap_or_else(|e| panic!("{e}"))
513    }
514
515    /// Fallible version of [`MutableBuffer::shrink_to_fit`].
516    pub fn try_shrink_to_fit(&mut self) -> Result<(), MutableBufferError> {
517        let new_capacity = self
518            .len
519            .checked_next_multiple_of(64)
520            .ok_or(MutableBufferError::LayoutError)?;
521        if new_capacity < self.layout.size() {
522            self.try_reallocate(new_capacity)?;
523        }
524        Ok(())
525    }
526    /// Shrinks the capacity of the buffer as much as possible.
527    /// The new capacity will aligned to the nearest 64 bit alignment.
528    ///
529    /// # Example
530    /// ```
531    /// # use arrow_buffer::buffer::{Buffer, MutableBuffer};
532    /// // 2 cache lines
533    /// let mut buffer = MutableBuffer::new(128);
534    /// assert_eq!(buffer.capacity(), 128);
535    /// buffer.push(1);
536    /// buffer.push(2);
537    ///
538    /// buffer.shrink_to_fit();
539    /// assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
540    /// ```
541    ///
542    /// # Panics
543    ///
544    /// Panics if the current length is too large to round up to the next 64-byte boundary and
545    /// construct a valid allocation layout.
546    pub fn shrink_to_fit(&mut self) {
547        self.try_shrink_to_fit().unwrap_or_else(|e| panic!("{e}"))
548    }
549
550    /// Returns whether this buffer is empty or not.
551    #[inline]
552    pub const fn is_empty(&self) -> bool {
553        self.len == 0
554    }
555
556    /// Returns the length (the number of bytes written) in this buffer.
557    /// The invariant `buffer.len() <= buffer.capacity()` is always upheld.
558    #[inline]
559    pub const fn len(&self) -> usize {
560        self.len
561    }
562
563    /// Returns the total capacity in this buffer, in bytes.
564    ///
565    /// The invariant `buffer.len() <= buffer.capacity()` is always upheld.
566    #[inline]
567    pub const fn capacity(&self) -> usize {
568        self.layout.size()
569    }
570
571    /// Clear all existing data from this buffer.
572    pub fn clear(&mut self) {
573        self.len = 0;
574        #[cfg(feature = "pool")]
575        {
576            if let Some(reservation) = lock_reservation(&self.reservation).as_mut() {
577                reservation.resize(self.len);
578            }
579        }
580    }
581
582    /// Returns the data stored in this buffer as a slice.
583    pub fn as_slice(&self) -> &[u8] {
584        self
585    }
586
587    /// Returns the data stored in this buffer as a mutable slice.
588    pub fn as_slice_mut(&mut self) -> &mut [u8] {
589        self
590    }
591
592    /// Returns a raw pointer to this buffer's internal memory
593    /// This pointer is guaranteed to be aligned along cache-lines.
594    #[inline]
595    pub const fn as_ptr(&self) -> *const u8 {
596        self.data.as_ptr()
597    }
598
599    /// Returns a mutable raw pointer to this buffer's internal memory
600    /// This pointer is guaranteed to be aligned along cache-lines.
601    #[inline]
602    pub fn as_mut_ptr(&mut self) -> *mut u8 {
603        self.data.as_ptr()
604    }
605
606    #[inline]
607    pub(super) fn into_buffer(self) -> Buffer {
608        let bytes = unsafe { Bytes::new(self.data, self.len, Deallocation::Standard(self.layout)) };
609        #[cfg(feature = "pool")]
610        {
611            let reservation = lock_reservation(&self.reservation).take();
612            *lock_reservation(&bytes.reservation) = reservation;
613        }
614        std::mem::forget(self);
615        Buffer::from(bytes)
616    }
617
618    /// View this buffer as a mutable slice of a specific type.
619    ///
620    /// # Panics
621    ///
622    /// This function panics if the underlying buffer is not aligned correctly for type `T`, or
623    /// if its length is not a multiple of `size_of::<T>()`.
624    pub fn typed_data_mut<T: ArrowNativeType>(&mut self) -> &mut [T] {
625        // SAFETY
626        // ArrowNativeType is trivially transmutable, is sealed to prevent potentially incorrect
627        // implementation outside this crate, and this method checks alignment
628        let (prefix, offsets, suffix) = unsafe { self.as_slice_mut().align_to_mut::<T>() };
629        assert!(prefix.is_empty() && suffix.is_empty());
630        offsets
631    }
632
633    /// View buffer as a immutable slice of a specific type.
634    ///
635    /// # Panics
636    ///
637    /// This function panics if the underlying buffer is not aligned correctly for type `T`, or
638    /// if its length is not a multiple of `size_of::<T>()`.
639    pub fn typed_data<T: ArrowNativeType>(&self) -> &[T] {
640        // SAFETY
641        // ArrowNativeType is trivially transmutable, is sealed to prevent potentially incorrect
642        // implementation outside this crate, and this method checks alignment
643        let (prefix, offsets, suffix) = unsafe { self.as_slice().align_to::<T>() };
644        assert!(prefix.is_empty() && suffix.is_empty());
645        offsets
646    }
647
648    /// Fallible version of [`MutableBuffer::extend_from_slice`].
649    #[inline]
650    pub fn try_extend_from_slice<T: ArrowNativeType>(
651        &mut self,
652        items: &[T],
653    ) -> Result<(), MutableBufferError> {
654        let additional = mem::size_of_val(items);
655        self.try_reserve(additional)?;
656        unsafe {
657            // this assumes that `[ToByteSlice]` can be copied directly
658            // without calling `to_byte_slice` for each element,
659            // which is correct for all ArrowNativeType implementations.
660            let src = items.as_ptr().cast::<u8>();
661            let dst = self.data.as_ptr().add(self.len);
662            std::ptr::copy_nonoverlapping(src, dst, additional);
663        }
664        self.len += additional;
665        Ok(())
666    }
667    /// Extends this buffer from a slice of items that can be represented in bytes, increasing its capacity if needed.
668    /// # Example
669    /// ```
670    /// # use arrow_buffer::buffer::MutableBuffer;
671    /// let mut buffer = MutableBuffer::new(0);
672    /// buffer.extend_from_slice(&[2u32, 0]);
673    /// assert_eq!(buffer.len(), 8) // u32 has 4 bytes
674    /// ```
675    ///
676    /// # Panics
677    ///
678    /// Panics if extending the buffer requires reserving a capacity that fails for the same
679    /// reasons as [`MutableBuffer::reserve`].
680    #[inline]
681    pub fn extend_from_slice<T: ArrowNativeType>(&mut self, items: &[T]) {
682        self.try_extend_from_slice(items)
683            .unwrap_or_else(|e| panic!("{e}"))
684    }
685
686    /// Extends the buffer with a new item, increasing its capacity if needed.
687    /// # Example
688    /// ```
689    /// # use arrow_buffer::buffer::MutableBuffer;
690    /// let mut buffer = MutableBuffer::new(0);
691    /// buffer.push(256u32);
692    /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
693    /// ```
694    ///
695    /// # Panics
696    ///
697    /// Panics if extending the buffer requires reserving a capacity that fails for the same
698    /// reasons as [`MutableBuffer::reserve`].
699    #[inline]
700    pub fn push<T: ToByteSlice>(&mut self, item: T) {
701        let additional = std::mem::size_of::<T>();
702        self.reserve(additional);
703        unsafe {
704            let src = item.to_byte_slice().as_ptr();
705            let dst = self.data.as_ptr().add(self.len);
706            std::ptr::copy_nonoverlapping(src, dst, additional);
707        }
708        self.len += additional;
709    }
710
711    /// Extends the buffer with a new item, without checking for sufficient capacity
712    /// # Safety
713    /// Caller must ensure that the capacity()-len()>=`size_of<T>`()
714    #[inline]
715    pub unsafe fn push_unchecked<T: ToByteSlice>(&mut self, item: T) {
716        let additional = std::mem::size_of::<T>();
717        let src = item.to_byte_slice().as_ptr();
718        let dst = unsafe { self.data.as_ptr().add(self.len) };
719        unsafe { std::ptr::copy_nonoverlapping(src, dst, additional) };
720        self.len += additional;
721    }
722
723    /// Fallible version of [`MutableBuffer::extend_zeros`].
724    #[inline]
725    pub fn try_extend_zeros(&mut self, additional: usize) -> Result<(), MutableBufferError> {
726        let new_len = self
727            .len
728            .checked_add(additional)
729            .ok_or(MutableBufferError::LengthOverflow)?;
730        self.try_resize(new_len, 0)
731    }
732    /// Extends the buffer by `additional` bytes equal to `0u8`, incrementing its capacity if needed.
733    ///
734    /// # Panics
735    ///
736    /// Panics if `self.len + additional` overflows `usize`, or if growing the buffer requires
737    /// reserving a capacity that fails for the same reasons as [`MutableBuffer::reserve`].
738    #[inline]
739    pub fn extend_zeros(&mut self, additional: usize) {
740        self.try_extend_zeros(additional)
741            .unwrap_or_else(|e| panic!("{e}"))
742    }
743
744    /// # Safety
745    /// The caller must ensure that the buffer was properly initialized up to `len`.
746    ///
747    /// # Panics
748    ///
749    /// Panics if `len` exceeds the buffer capacity.
750    #[inline]
751    pub unsafe fn set_len(&mut self, len: usize) {
752        assert!(len <= self.capacity());
753        self.len = len;
754    }
755
756    /// Invokes `f` with values `0..len` collecting the boolean results into a new `MutableBuffer`
757    ///
758    /// This is similar to `from_trusted_len_iter_bool`, however, can be significantly faster
759    /// as it eliminates the conditional `Iterator::next`
760    #[inline]
761    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, mut f: F) -> Self {
762        let mut buffer: Vec<u64> = Vec::with_capacity(bit_util::ceil(len, 64));
763
764        let chunks = len / 64;
765        let remainder = len % 64;
766        buffer.extend((0..chunks).map(|chunk| {
767            let mut packed = 0;
768            for bit_idx in 0..64 {
769                let i = bit_idx + chunk * 64;
770                packed |= (f(i) as u64) << bit_idx;
771            }
772
773            packed
774        }));
775
776        if remainder != 0 {
777            let mut packed = 0;
778            for bit_idx in 0..remainder {
779                let i = bit_idx + chunks * 64;
780                packed |= (f(i) as u64) << bit_idx;
781            }
782
783            buffer.push(packed)
784        }
785
786        let mut buffer: MutableBuffer = buffer.into();
787        buffer.truncate(bit_util::ceil(len, 8));
788        buffer
789    }
790
791    /// Extends this buffer with boolean values.
792    ///
793    /// This requires `iter` to report an exact size via `size_hint`.
794    /// `offset` indicates the starting offset in bits in this buffer to begin writing to
795    /// and must be less than or equal to the current length of this buffer.
796    /// All bits not written to (but readable due to byte alignment) will be zeroed out.
797    ///
798    /// # Panics
799    ///
800    /// Panics if `iter` does not report an exact size via `size_hint`, or if it yields fewer
801    /// items than reported, or if extending the buffer requires reserving a capacity that fails
802    /// for the same reasons as [`MutableBuffer::reserve`].
803    ///
804    /// # Safety
805    /// Callers must ensure that `iter` reports an exact size via `size_hint`
806    /// and that `I::next()` does not panic, or `set_len` will leave the buffer
807    /// in an inconsistent state, exposing uninitialized/stale bytes as though
808    /// they were valid.
809    #[inline]
810    pub unsafe fn extend_bool_trusted_len<I: Iterator<Item = bool>>(
811        &mut self,
812        mut iter: I,
813        offset: usize,
814    ) {
815        let (lower, upper) = iter.size_hint();
816        let len = upper.expect("Iterator must have exact size_hint");
817        assert_eq!(lower, len, "Iterator must have exact size_hint");
818        debug_assert!(
819            offset <= self.len * 8,
820            "offset must be <= buffer length in bits"
821        );
822
823        if len == 0 {
824            return;
825        }
826
827        let start_len = offset;
828        let end_bit = start_len + len;
829
830        // SAFETY: we will initialize all newly exposed bytes before they are read
831        let new_len_bytes = bit_util::ceil(end_bit, 8);
832        if new_len_bytes > self.len {
833            self.reserve(new_len_bytes - self.len);
834            // SAFETY: caller will initialize all newly exposed bytes before they are read
835            unsafe { self.set_len(new_len_bytes) };
836        }
837
838        let slice = self.as_slice_mut();
839
840        let mut bit_idx = start_len;
841
842        // ---- Unaligned prefix: advance to the next 64-bit boundary ----
843        let misalignment = bit_idx & 63;
844        let prefix_bits = if misalignment == 0 {
845            0
846        } else {
847            (64 - misalignment).min(end_bit - bit_idx)
848        };
849
850        if prefix_bits != 0 {
851            let byte_start = bit_idx / 8;
852            let byte_end = bit_util::ceil(bit_idx + prefix_bits, 8);
853            let bit_offset = bit_idx % 8;
854
855            // Clear any newly-visible bits in the existing partial byte
856            if bit_offset != 0 {
857                let keep_mask = (1u8 << bit_offset).wrapping_sub(1);
858                slice[byte_start] &= keep_mask;
859            }
860
861            // Zero any new bytes we will partially fill in this prefix
862            let zero_from = if bit_offset == 0 {
863                byte_start
864            } else {
865                byte_start + 1
866            };
867            if byte_end > zero_from {
868                slice[zero_from..byte_end].fill(0);
869            }
870
871            for _ in 0..prefix_bits {
872                let v = iter.next().unwrap();
873                if v {
874                    let byte_idx = bit_idx / 8;
875                    let bit = bit_idx % 8;
876                    slice[byte_idx] |= 1 << bit;
877                }
878                bit_idx += 1;
879            }
880        }
881
882        if bit_idx < end_bit {
883            // ---- Aligned middle: write u64 chunks ----
884            debug_assert_eq!(bit_idx & 63, 0);
885            let remaining_bits = end_bit - bit_idx;
886            let chunks = remaining_bits / 64;
887
888            let words_start = bit_idx / 8;
889            let words_end = words_start + chunks * 8;
890            for dst in slice[words_start..words_end].chunks_exact_mut(8) {
891                let mut packed: u64 = 0;
892                for i in 0..64 {
893                    packed |= (iter.next().unwrap() as u64) << i;
894                }
895                dst.copy_from_slice(&packed.to_le_bytes());
896                bit_idx += 64;
897            }
898
899            // ---- Unaligned suffix: remaining < 64 bits ----
900            let suffix_bits = end_bit - bit_idx;
901            if suffix_bits != 0 {
902                debug_assert_eq!(bit_idx % 8, 0);
903                let byte_start = bit_idx / 8;
904                let byte_end = bit_util::ceil(end_bit, 8);
905                slice[byte_start..byte_end].fill(0);
906
907                for _ in 0..suffix_bits {
908                    let v = iter.next().unwrap();
909                    if v {
910                        let byte_idx = bit_idx / 8;
911                        let bit = bit_idx % 8;
912                        slice[byte_idx] |= 1 << bit;
913                    }
914                    bit_idx += 1;
915                }
916            }
917        }
918
919        // Clear any unused bits in the last byte
920        let remainder = end_bit % 8;
921        if remainder != 0 {
922            let mask = (1u8 << remainder).wrapping_sub(1);
923            slice[bit_util::ceil(end_bit, 8) - 1] &= mask;
924        }
925
926        debug_assert_eq!(bit_idx, end_bit);
927    }
928
929    /// Register this [`MutableBuffer`] with the provided [`MemoryPool`]
930    ///
931    /// This claims the memory used by this buffer in the pool, allowing for
932    /// accurate accounting of memory usage. Any prior reservation will be
933    /// released so this works well when the buffer is being shared among
934    /// multiple arrays.
935    #[cfg(feature = "pool")]
936    pub fn claim(&self, pool: &dyn MemoryPool) {
937        *lock_reservation(&self.reservation) = Some(pool.reserve(self.capacity()));
938    }
939}
940
941/// Creates a non-null pointer with alignment of [`ALIGNMENT`]
942///
943/// This is similar to [`NonNull::dangling`]
944#[inline]
945pub(crate) fn dangling_ptr() -> NonNull<u8> {
946    // SAFETY: ALIGNMENT is a non-zero usize which is then cast
947    // to a *mut u8. Therefore, `ptr` is not null and the conditions for
948    // calling new_unchecked() are respected.
949    #[cfg(miri)]
950    {
951        // Since miri implies a nightly rust version we can use the unstable strict_provenance feature
952        unsafe { NonNull::new_unchecked(std::ptr::without_provenance_mut(ALIGNMENT)) }
953    }
954    #[cfg(not(miri))]
955    {
956        unsafe { NonNull::new_unchecked(ALIGNMENT as *mut u8) }
957    }
958}
959
960impl<A: ArrowNativeType> Extend<A> for MutableBuffer {
961    #[inline]
962    fn extend<T: IntoIterator<Item = A>>(&mut self, iter: T) {
963        let iterator = iter.into_iter();
964        self.extend_from_iter(iterator)
965    }
966}
967
968impl<T: ArrowNativeType> From<Vec<T>> for MutableBuffer {
969    fn from(mut value: Vec<T>) -> Self {
970        // Safety
971        // Vec::as_mut_ptr guaranteed to not be null and ArrowNativeType are trivially transmutable
972        let data = unsafe { NonNull::new_unchecked(value.as_mut_ptr().cast()) };
973        let len = value.len() * mem::size_of::<T>();
974        // Safety
975        // Vec guaranteed to have a valid layout matching that of `Layout::array`
976        // This is based on `RawVec::current_memory`
977        let layout = unsafe { Layout::array::<T>(value.capacity()).unwrap_unchecked() };
978        mem::forget(value);
979        Self {
980            data,
981            len,
982            layout,
983            #[cfg(feature = "pool")]
984            reservation: std::sync::Mutex::new(None),
985        }
986    }
987}
988
989impl MutableBuffer {
990    #[inline]
991    pub(super) fn extend_from_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
992        &mut self,
993        mut iterator: I,
994    ) {
995        let item_size = std::mem::size_of::<T>();
996        let (lower, _) = iterator.size_hint();
997        let additional = lower * item_size;
998        self.reserve(additional);
999
1000        // this is necessary because of https://github.com/rust-lang/rust/issues/32155
1001        let mut len = SetLenOnDrop::new(&mut self.len);
1002        let mut dst = unsafe { self.data.as_ptr().add(len.local_len) };
1003        let capacity = self.layout.size();
1004
1005        while len.local_len + item_size <= capacity {
1006            if let Some(item) = iterator.next() {
1007                unsafe {
1008                    let src = item.to_byte_slice().as_ptr();
1009                    std::ptr::copy_nonoverlapping(src, dst, item_size);
1010                    dst = dst.add(item_size);
1011                }
1012                len.local_len += item_size;
1013            } else {
1014                break;
1015            }
1016        }
1017        drop(len);
1018
1019        iterator.for_each(|item| self.push(item));
1020    }
1021
1022    /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length.
1023    /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.
1024    /// # Example
1025    /// ```
1026    /// # use arrow_buffer::buffer::MutableBuffer;
1027    /// let v = vec![1u32];
1028    /// let iter = v.iter().map(|x| x * 2);
1029    /// let buffer = unsafe { MutableBuffer::from_trusted_len_iter(iter) };
1030    /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
1031    /// ```
1032    ///
1033    /// # Panics
1034    ///
1035    /// Panics if the iterator does not report an upper bound via `size_hint`, or if the
1036    /// reported length does not match the number of items produced, or if allocating the
1037    /// required buffer fails for the same reasons as [`MutableBuffer::new`].
1038    ///
1039    /// # Safety
1040    /// This method assumes that the iterator's size is correct and is undefined behavior
1041    /// to use it on an iterator that reports an incorrect length.
1042    // This implementation is required for two reasons:
1043    // 1. there is no trait `TrustedLen` in stable rust and therefore
1044    //    we can't specialize `extend` for `TrustedLen` like `Vec` does.
1045    // 2. `from_trusted_len_iter` is faster.
1046    #[inline]
1047    pub unsafe fn from_trusted_len_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
1048        iterator: I,
1049    ) -> Self {
1050        let item_size = std::mem::size_of::<T>();
1051        let (_, upper) = iterator.size_hint();
1052        let upper = upper.expect("from_trusted_len_iter requires an upper limit");
1053        let len = upper * item_size;
1054
1055        let mut buffer = MutableBuffer::new(len);
1056
1057        let mut dst = buffer.data.as_ptr();
1058        for item in iterator {
1059            // note how there is no reserve here (compared with `extend_from_iter`)
1060            let src = item.to_byte_slice().as_ptr();
1061            unsafe { std::ptr::copy_nonoverlapping(src, dst, item_size) };
1062            dst = unsafe { dst.add(item_size) };
1063        }
1064        assert_eq!(
1065            unsafe { dst.offset_from(buffer.data.as_ptr()) } as usize,
1066            len,
1067            "Trusted iterator length was not accurately reported"
1068        );
1069        buffer.len = len;
1070        buffer
1071    }
1072
1073    /// Creates a [`MutableBuffer`] from a boolean [`Iterator`] with a trusted (upper) length.
1074    /// # use arrow_buffer::buffer::MutableBuffer;
1075    /// # Example
1076    /// ```
1077    /// # use arrow_buffer::buffer::MutableBuffer;
1078    /// let v = vec![false, true, false];
1079    /// let iter = v.iter().map(|x| *x || true);
1080    /// let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(iter) };
1081    /// assert_eq!(buffer.len(), 1) // 3 booleans have 1 byte
1082    /// ```
1083    ///
1084    /// # Panics
1085    ///
1086    /// Panics if the iterator does not report an upper bound via `size_hint`, or if it yields
1087    /// fewer items than reported.
1088    ///
1089    /// # Safety
1090    /// This method assumes that the iterator's size is correct and is undefined behavior
1091    /// to use it on an iterator that reports an incorrect length.
1092    // This implementation is required for two reasons:
1093    // 1. there is no trait `TrustedLen` in stable rust and therefore
1094    //    we can't specialize `extend` for `TrustedLen` like `Vec` does.
1095    // 2. `from_trusted_len_iter_bool` is faster.
1096    #[inline]
1097    pub unsafe fn from_trusted_len_iter_bool<I: Iterator<Item = bool>>(mut iterator: I) -> Self {
1098        let (_, upper) = iterator.size_hint();
1099        let len = upper.expect("from_trusted_len_iter requires an upper limit");
1100
1101        Self::collect_bool(len, |_| iterator.next().unwrap())
1102    }
1103
1104    /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length or errors
1105    /// if any of the items of the iterator is an error.
1106    /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.
1107    ///
1108    /// # Errors
1109    ///
1110    /// Returns the first error yielded by the iterator.
1111    ///
1112    /// # Panics
1113    ///
1114    /// Note that unlike the [`Err`] cases, these panics are violations of the safety contract
1115    /// below, and are only checks that happen to be cheap enough to keep:
1116    ///
1117    /// Panics if the iterator does not report an upper bound via `size_hint`, or if the
1118    /// reported length does not match the number of items produced before an error-free finish,
1119    /// or if allocating the required buffer fails for the same reasons as
1120    /// [`MutableBuffer::new`].
1121    ///
1122    /// # Safety
1123    /// This method assumes that the iterator's size is correct and is undefined behavior
1124    /// to use it on an iterator that reports an incorrect length.
1125    #[inline]
1126    pub unsafe fn try_from_trusted_len_iter<
1127        E,
1128        T: ArrowNativeType,
1129        I: Iterator<Item = Result<T, E>>,
1130    >(
1131        iterator: I,
1132    ) -> Result<Self, E> {
1133        let item_size = std::mem::size_of::<T>();
1134        let (_, upper) = iterator.size_hint();
1135        let upper = upper.expect("try_from_trusted_len_iter requires an upper limit");
1136        let len = upper * item_size;
1137
1138        let mut buffer = MutableBuffer::new(len);
1139
1140        let mut dst = buffer.data.as_ptr();
1141        for item in iterator {
1142            let item = item?;
1143            // note how there is no reserve here (compared with `extend_from_iter`)
1144            let src = item.to_byte_slice().as_ptr();
1145            unsafe { std::ptr::copy_nonoverlapping(src, dst, item_size) };
1146            dst = unsafe { dst.add(item_size) };
1147        }
1148        // try_from_trusted_len_iter is instantiated a lot, so we extract part of it into a less
1149        // generic method to reduce compile time
1150        unsafe fn finalize_buffer(dst: *mut u8, buffer: &mut MutableBuffer, len: usize) {
1151            unsafe {
1152                assert_eq!(
1153                    dst.offset_from(buffer.data.as_ptr()) as usize,
1154                    len,
1155                    "Trusted iterator length was not accurately reported"
1156                );
1157                buffer.len = len;
1158            }
1159        }
1160        unsafe { finalize_buffer(dst, &mut buffer, len) };
1161        Ok(buffer)
1162    }
1163}
1164
1165impl Default for MutableBuffer {
1166    fn default() -> Self {
1167        Self::with_capacity(0)
1168    }
1169}
1170
1171impl std::ops::Deref for MutableBuffer {
1172    type Target = [u8];
1173
1174    fn deref(&self) -> &[u8] {
1175        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.len) }
1176    }
1177}
1178
1179impl std::ops::DerefMut for MutableBuffer {
1180    fn deref_mut(&mut self) -> &mut [u8] {
1181        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
1182    }
1183}
1184
1185impl AsRef<[u8]> for &MutableBuffer {
1186    fn as_ref(&self) -> &[u8] {
1187        self.as_slice()
1188    }
1189}
1190
1191impl Drop for MutableBuffer {
1192    fn drop(&mut self) {
1193        if self.layout.size() != 0 {
1194            // Safety: data was allocated with standard allocator with given layout
1195            unsafe { std::alloc::dealloc(self.data.as_ptr().cast(), self.layout) };
1196        }
1197    }
1198}
1199
1200impl PartialEq for MutableBuffer {
1201    fn eq(&self, other: &MutableBuffer) -> bool {
1202        if self.len != other.len {
1203            return false;
1204        }
1205        if self.layout != other.layout {
1206            return false;
1207        }
1208        self.as_slice() == other.as_slice()
1209    }
1210}
1211
1212unsafe impl Sync for MutableBuffer {}
1213unsafe impl Send for MutableBuffer {}
1214
1215struct SetLenOnDrop<'a> {
1216    len: &'a mut usize,
1217    local_len: usize,
1218}
1219
1220impl<'a> SetLenOnDrop<'a> {
1221    #[inline]
1222    fn new(len: &'a mut usize) -> Self {
1223        SetLenOnDrop {
1224            local_len: *len,
1225            len,
1226        }
1227    }
1228}
1229
1230impl Drop for SetLenOnDrop<'_> {
1231    #[inline]
1232    fn drop(&mut self) {
1233        *self.len = self.local_len;
1234    }
1235}
1236
1237/// Creating a `MutableBuffer` instance by setting bits according to the boolean values
1238impl std::iter::FromIterator<bool> for MutableBuffer {
1239    fn from_iter<I>(iter: I) -> Self
1240    where
1241        I: IntoIterator<Item = bool>,
1242    {
1243        let mut iterator = iter.into_iter();
1244        let mut result = {
1245            let byte_capacity: usize = iterator.size_hint().0.saturating_add(7) / 8;
1246            MutableBuffer::new(byte_capacity)
1247        };
1248
1249        loop {
1250            let mut exhausted = false;
1251            let mut byte_accum: u8 = 0;
1252            let mut mask: u8 = 1;
1253
1254            //collect (up to) 8 bits into a byte
1255            while mask != 0 {
1256                if let Some(value) = iterator.next() {
1257                    byte_accum |= match value {
1258                        true => mask,
1259                        false => 0,
1260                    };
1261                    mask <<= 1;
1262                } else {
1263                    exhausted = true;
1264                    break;
1265                }
1266            }
1267
1268            // break if the iterator was exhausted before it provided a bool for this byte
1269            if exhausted && mask == 1 {
1270                break;
1271            }
1272
1273            //ensure we have capacity to write the byte
1274            if result.len() == result.capacity() {
1275                //no capacity for new byte, allocate 1 byte more (plus however many more the iterator advertises)
1276                let additional_byte_capacity = 1usize.saturating_add(
1277                    iterator.size_hint().0.saturating_add(7) / 8, //convert bit count to byte count, rounding up
1278                );
1279                result.reserve(additional_byte_capacity)
1280            }
1281
1282            // Soundness: capacity was allocated above
1283            unsafe { result.push_unchecked(byte_accum) };
1284            if exhausted {
1285                break;
1286            }
1287        }
1288        result
1289    }
1290}
1291
1292impl<T: ArrowNativeType> std::iter::FromIterator<T> for MutableBuffer {
1293    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
1294        let mut buffer = Self::default();
1295        buffer.extend_from_iter(iter.into_iter());
1296        buffer
1297    }
1298}
1299
1300#[cfg(test)]
1301mod tests {
1302    use super::*;
1303
1304    #[test]
1305    fn test_mutable_new() {
1306        let buf = MutableBuffer::new(63);
1307        assert_eq!(64, buf.capacity());
1308        assert_eq!(0, buf.len());
1309        assert!(buf.is_empty());
1310    }
1311
1312    #[test]
1313    fn test_mutable_default() {
1314        let buf = MutableBuffer::default();
1315        assert_eq!(0, buf.capacity());
1316        assert_eq!(0, buf.len());
1317        assert!(buf.is_empty());
1318
1319        let mut buf = MutableBuffer::default();
1320        buf.extend_from_slice(b"hello");
1321        assert_eq!(5, buf.len());
1322        assert_eq!(b"hello", buf.as_slice());
1323    }
1324
1325    #[test]
1326    fn test_mutable_extend_from_slice() {
1327        let mut buf = MutableBuffer::new(100);
1328        buf.extend_from_slice(b"hello");
1329        assert_eq!(5, buf.len());
1330        assert_eq!(b"hello", buf.as_slice());
1331
1332        buf.extend_from_slice(b" world");
1333        assert_eq!(11, buf.len());
1334        assert_eq!(b"hello world", buf.as_slice());
1335
1336        buf.clear();
1337        assert_eq!(0, buf.len());
1338        buf.extend_from_slice(b"hello arrow");
1339        assert_eq!(11, buf.len());
1340        assert_eq!(b"hello arrow", buf.as_slice());
1341    }
1342
1343    #[test]
1344    fn mutable_extend_from_iter() {
1345        let mut buf = MutableBuffer::new(0);
1346        buf.extend(vec![1u32, 2]);
1347        assert_eq!(8, buf.len());
1348        assert_eq!(&[1u8, 0, 0, 0, 2, 0, 0, 0], buf.as_slice());
1349
1350        buf.extend(vec![3u32, 4]);
1351        assert_eq!(16, buf.len());
1352        assert_eq!(
1353            &[1u8, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4, 0, 0, 0],
1354            buf.as_slice()
1355        );
1356    }
1357
1358    #[test]
1359    fn mutable_extend_from_iter_unaligned_u64() {
1360        let mut buf = MutableBuffer::new(16);
1361        buf.push(1_u8);
1362        buf.extend([1_u64]);
1363        assert_eq!(9, buf.len());
1364        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1365    }
1366
1367    #[test]
1368    fn mutable_extend_from_slice_unaligned_u64() {
1369        let mut buf = MutableBuffer::new(16);
1370        buf.extend_from_slice(&[1_u8]);
1371        buf.extend_from_slice(&[1_u64]);
1372        assert_eq!(9, buf.len());
1373        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1374    }
1375
1376    #[test]
1377    fn mutable_push_unaligned_u64() {
1378        let mut buf = MutableBuffer::new(16);
1379        buf.push(1_u8);
1380        buf.push(1_u64);
1381        assert_eq!(9, buf.len());
1382        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1383    }
1384
1385    #[test]
1386    fn mutable_push_unchecked_unaligned_u64() {
1387        let mut buf = MutableBuffer::new(16);
1388        unsafe {
1389            buf.push_unchecked(1_u8);
1390            buf.push_unchecked(1_u64);
1391        }
1392        assert_eq!(9, buf.len());
1393        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1394    }
1395
1396    #[test]
1397    fn test_from_trusted_len_iter() {
1398        let iter = vec![1u32, 2].into_iter();
1399        let buf = unsafe { MutableBuffer::from_trusted_len_iter(iter) };
1400        assert_eq!(8, buf.len());
1401        assert_eq!(&[1u8, 0, 0, 0, 2, 0, 0, 0], buf.as_slice());
1402    }
1403
1404    #[test]
1405    fn test_mutable_reserve() {
1406        let mut buf = MutableBuffer::new(1);
1407        assert_eq!(64, buf.capacity());
1408
1409        // Reserving a smaller capacity should have no effect.
1410        buf.reserve(10);
1411        assert_eq!(64, buf.capacity());
1412
1413        buf.reserve(80);
1414        assert_eq!(128, buf.capacity());
1415
1416        buf.reserve(129);
1417        assert_eq!(256, buf.capacity());
1418    }
1419
1420    #[test]
1421    fn test_mutable_resize() {
1422        let mut buf = MutableBuffer::new(1);
1423        assert_eq!(64, buf.capacity());
1424        assert_eq!(0, buf.len());
1425
1426        buf.resize(20, 0);
1427        assert_eq!(64, buf.capacity());
1428        assert_eq!(20, buf.len());
1429
1430        buf.resize(10, 0);
1431        assert_eq!(64, buf.capacity());
1432        assert_eq!(10, buf.len());
1433
1434        buf.resize(100, 0);
1435        assert_eq!(128, buf.capacity());
1436        assert_eq!(100, buf.len());
1437
1438        buf.resize(30, 0);
1439        assert_eq!(128, buf.capacity());
1440        assert_eq!(30, buf.len());
1441
1442        buf.resize(0, 0);
1443        assert_eq!(128, buf.capacity());
1444        assert_eq!(0, buf.len());
1445    }
1446
1447    #[test]
1448    fn test_mutable_into() {
1449        let mut buf = MutableBuffer::new(1);
1450        buf.extend_from_slice(b"aaaa bbbb cccc dddd");
1451        assert_eq!(19, buf.len());
1452        assert_eq!(64, buf.capacity());
1453        assert_eq!(b"aaaa bbbb cccc dddd", buf.as_slice());
1454
1455        let immutable_buf: Buffer = buf.into();
1456        assert_eq!(19, immutable_buf.len());
1457        assert_eq!(64, immutable_buf.capacity());
1458        assert_eq!(b"aaaa bbbb cccc dddd", immutable_buf.as_slice());
1459    }
1460
1461    #[test]
1462    fn test_mutable_equal() {
1463        let mut buf = MutableBuffer::new(1);
1464        let mut buf2 = MutableBuffer::new(1);
1465
1466        buf.extend_from_slice(&[0xaa]);
1467        buf2.extend_from_slice(&[0xaa, 0xbb]);
1468        assert!(buf != buf2);
1469
1470        buf.extend_from_slice(&[0xbb]);
1471        assert_eq!(buf, buf2);
1472
1473        buf2.reserve(65);
1474        assert!(buf != buf2);
1475    }
1476
1477    #[test]
1478    fn test_mutable_shrink_to_fit() {
1479        let mut buffer = MutableBuffer::new(128);
1480        assert_eq!(buffer.capacity(), 128);
1481        buffer.push(1);
1482        buffer.push(2);
1483
1484        buffer.shrink_to_fit();
1485        assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
1486    }
1487
1488    #[test]
1489    fn test_mutable_set_null_bits() {
1490        let mut buffer = MutableBuffer::new(8).with_bitset(8, true);
1491
1492        for i in 0..=buffer.capacity() {
1493            buffer.set_null_bits(i, 0);
1494            assert_eq!(buffer[..8], [255; 8][..]);
1495        }
1496
1497        buffer.set_null_bits(1, 4);
1498        assert_eq!(buffer[..8], [255, 0, 0, 0, 0, 255, 255, 255][..]);
1499    }
1500
1501    #[test]
1502    #[should_panic = "out of bounds for buffer of length"]
1503    fn test_mutable_set_null_bits_oob() {
1504        let mut buffer = MutableBuffer::new(64);
1505        buffer.set_null_bits(1, buffer.capacity());
1506    }
1507
1508    #[test]
1509    #[should_panic = "out of bounds for buffer of length"]
1510    fn test_mutable_set_null_bits_oob_by_overflow() {
1511        let mut buffer = MutableBuffer::new(0);
1512        buffer.set_null_bits(1, usize::MAX);
1513    }
1514
1515    #[test]
1516    fn from_iter() {
1517        let buffer = [1u16, 2, 3, 4].into_iter().collect::<MutableBuffer>();
1518        assert_eq!(buffer.len(), 4 * mem::size_of::<u16>());
1519        assert_eq!(buffer.as_slice(), &[1, 0, 2, 0, 3, 0, 4, 0]);
1520    }
1521
1522    #[test]
1523    #[should_panic(expected = "invalid allocation layout for requested capacity")]
1524    fn test_with_capacity_panics_above_max_capacity() {
1525        let max_capacity = isize::MAX as usize - (isize::MAX as usize % ALIGNMENT);
1526        let _ = MutableBuffer::with_capacity(max_capacity + 1);
1527    }
1528
1529    #[cfg(feature = "pool")]
1530    mod pool_tests {
1531        use super::*;
1532        use crate::pool::{MemoryPool, TrackingMemoryPool};
1533
1534        #[test]
1535        fn test_reallocate_with_pool() {
1536            let pool = TrackingMemoryPool::default();
1537            let mut buffer = MutableBuffer::with_capacity(100);
1538            buffer.claim(&pool);
1539
1540            // Initial capacity should be 128 (multiple of 64)
1541            assert_eq!(buffer.capacity(), 128);
1542            assert_eq!(pool.used(), 128);
1543
1544            // Reallocate to a larger size
1545            buffer.try_reallocate(200).unwrap();
1546
1547            // The capacity is exactly the requested size, not rounded up
1548            assert_eq!(buffer.capacity(), 200);
1549            assert_eq!(pool.used(), 200);
1550
1551            // Reallocate to a smaller size
1552            buffer.try_reallocate(50).unwrap();
1553
1554            // The capacity is exactly the requested size, not rounded up
1555            assert_eq!(buffer.capacity(), 50);
1556            assert_eq!(pool.used(), 50);
1557        }
1558
1559        #[test]
1560        fn test_truncate_with_pool() {
1561            let pool = TrackingMemoryPool::default();
1562            let mut buffer = MutableBuffer::with_capacity(100);
1563
1564            // Fill buffer with some data
1565            buffer.resize(80, 1);
1566            assert_eq!(buffer.len(), 80);
1567
1568            buffer.claim(&pool);
1569            assert_eq!(pool.used(), 128);
1570
1571            // Truncate buffer
1572            buffer.truncate(40);
1573            assert_eq!(buffer.len(), 40);
1574            assert_eq!(pool.used(), 40);
1575
1576            // Truncate to zero
1577            buffer.clear();
1578            assert_eq!(buffer.len(), 0);
1579            assert_eq!(pool.used(), 0);
1580        }
1581
1582        #[test]
1583        fn test_resize_with_pool() {
1584            let pool = TrackingMemoryPool::default();
1585            let mut buffer = MutableBuffer::with_capacity(100);
1586            buffer.claim(&pool);
1587
1588            // Initial state
1589            assert_eq!(buffer.len(), 0);
1590            assert_eq!(pool.used(), 128);
1591
1592            // Resize to increase length
1593            buffer.resize(50, 1);
1594            assert_eq!(buffer.len(), 50);
1595            assert_eq!(pool.used(), 50);
1596
1597            // Resize to increase length beyond capacity
1598            buffer.resize(150, 1);
1599            assert_eq!(buffer.len(), 150);
1600            assert_eq!(buffer.capacity(), 256);
1601            assert_eq!(pool.used(), 150);
1602
1603            // Resize to decrease length
1604            buffer.resize(30, 1);
1605            assert_eq!(buffer.len(), 30);
1606            assert_eq!(pool.used(), 30);
1607        }
1608
1609        #[test]
1610        fn test_buffer_lifecycle_with_pool() {
1611            let pool = TrackingMemoryPool::default();
1612
1613            // Create a buffer with memory reservation
1614            let mut mutable = MutableBuffer::with_capacity(100);
1615            mutable.resize(80, 1);
1616            mutable.claim(&pool);
1617
1618            // Memory reservation is based on capacity when using claim()
1619            assert_eq!(pool.used(), 128);
1620
1621            // Convert to immutable Buffer
1622            let buffer = mutable.into_buffer();
1623
1624            // Memory reservation should be preserved
1625            assert_eq!(pool.used(), 128);
1626
1627            // Drop the buffer and the reservation should be released
1628            drop(buffer);
1629            assert_eq!(pool.used(), 0);
1630        }
1631    }
1632
1633    fn create_expected_repeated_slice<T: ArrowNativeType>(
1634        slice_to_repeat: &[T],
1635        repeat_count: usize,
1636    ) -> Buffer {
1637        let mut expected = MutableBuffer::new(size_of_val(slice_to_repeat) * repeat_count);
1638        for _ in 0..repeat_count {
1639            // Not using push_slice_repeated as this is the function under test
1640            expected.extend_from_slice(slice_to_repeat);
1641        }
1642        expected.into()
1643    }
1644
1645    // Helper to test a specific repeat count with various slice sizes
1646    fn test_repeat_count<T: ArrowNativeType + PartialEq + std::fmt::Debug>(
1647        repeat_count: usize,
1648        test_data: &[T],
1649    ) {
1650        let mut buffer = MutableBuffer::new(0);
1651        buffer.repeat_slice_n_times(test_data, repeat_count);
1652
1653        let expected = create_expected_repeated_slice(test_data, repeat_count);
1654        let result: Buffer = buffer.into();
1655
1656        assert_eq!(
1657            result,
1658            expected,
1659            "Failed for repeat_count={}, slice_len={}",
1660            repeat_count,
1661            test_data.len()
1662        );
1663    }
1664
1665    #[test]
1666    fn test_repeat_slice_count_edge_cases() {
1667        // Empty slice
1668        test_repeat_count(100, &[] as &[i32]);
1669
1670        // Zero repeats
1671        test_repeat_count(0, &[1i32, 2, 3]);
1672    }
1673
1674    #[test]
1675    #[should_panic(expected = "buffer length overflow")]
1676    fn test_repeat_slice_count_multiply_overflow() {
1677        let mut buffer = MutableBuffer::new(0);
1678        buffer.repeat_slice_n_times(&[0_u64], usize::MAX / mem::size_of::<u64>() + 1);
1679    }
1680
1681    #[test]
1682    #[should_panic(expected = "buffer length overflow")]
1683    fn test_repeat_slice_count_len_overflow() {
1684        let mut buffer = MutableBuffer::new(0);
1685        buffer.push(0_u8);
1686        buffer.repeat_slice_n_times(&[0_u8], usize::MAX);
1687    }
1688
1689    #[test]
1690    fn test_small_repeats_counts() {
1691        // test any special implementation for small repeat counts
1692        let data = &[1u8, 2, 3, 4, 5];
1693
1694        for _ in 1..=10 {
1695            test_repeat_count(2, data);
1696        }
1697    }
1698
1699    #[test]
1700    fn test_different_size_of_i32_repeat_slice() {
1701        let data: &[i32] = &[1, 2, 3];
1702        let data_with_single_item: &[i32] = &[42];
1703
1704        for data in &[data, data_with_single_item] {
1705            for item in 1..=9 {
1706                let base_repeat_count = 2_usize.pow(item);
1707                test_repeat_count(base_repeat_count - 1, data);
1708                test_repeat_count(base_repeat_count, data);
1709                test_repeat_count(base_repeat_count + 1, data);
1710            }
1711        }
1712    }
1713
1714    #[test]
1715    fn test_different_size_of_u8_repeat_slice() {
1716        let data: &[u8] = &[1, 2, 3];
1717        let data_with_single_item: &[u8] = &[10];
1718
1719        for data in &[data, data_with_single_item] {
1720            for item in 1..=9 {
1721                let base_repeat_count = 2_usize.pow(item);
1722                test_repeat_count(base_repeat_count - 1, data);
1723                test_repeat_count(base_repeat_count, data);
1724                test_repeat_count(base_repeat_count + 1, data);
1725            }
1726        }
1727    }
1728
1729    #[test]
1730    fn test_different_size_of_u16_repeat_slice() {
1731        let data: &[u16] = &[1, 2, 3];
1732        let data_with_single_item: &[u16] = &[10];
1733
1734        for data in &[data, data_with_single_item] {
1735            for item in 1..=9 {
1736                let base_repeat_count = 2_usize.pow(item);
1737                test_repeat_count(base_repeat_count - 1, data);
1738                test_repeat_count(base_repeat_count, data);
1739                test_repeat_count(base_repeat_count + 1, data);
1740            }
1741        }
1742    }
1743
1744    #[test]
1745    fn test_various_slice_lengths() {
1746        // Test different slice lengths with same repeat pattern
1747        let repeat_count = 37; // Arbitrary non-power-of-2
1748
1749        // Single element
1750        test_repeat_count(repeat_count, &[42i32]);
1751
1752        // Small slices
1753        test_repeat_count(repeat_count, &[1i32, 2]);
1754        test_repeat_count(repeat_count, &[1i32, 2, 3]);
1755        test_repeat_count(repeat_count, &[1i32, 2, 3, 4]);
1756        test_repeat_count(repeat_count, &[1i32, 2, 3, 4, 5]);
1757
1758        // Larger slices
1759        let data_10: Vec<i32> = (0..10).collect();
1760        test_repeat_count(repeat_count, &data_10);
1761
1762        let data_100: Vec<i32> = (0..100).collect();
1763        test_repeat_count(repeat_count, &data_100);
1764
1765        let data_1000: Vec<i32> = (0..1000).collect();
1766        test_repeat_count(repeat_count, &data_1000);
1767    }
1768
1769    #[test]
1770    #[should_panic(expected = "invalid allocation layout for requested capacity")]
1771    fn test_mutable_new_capacity_overflow() {
1772        // Tests overflow during initial allocation
1773        let _ = MutableBuffer::new(usize::MAX - 10);
1774    }
1775
1776    #[test]
1777    #[should_panic(expected = "buffer length overflow")]
1778    fn test_mutable_reserve_overflow() {
1779        // Tests overflow during growth (checked_add)
1780        let mut buf = MutableBuffer::new(1);
1781        buf.push(1u8);
1782        buf.reserve(usize::MAX);
1783    }
1784}
arrow_buffer/buffer/mutable.rs

arrow_buffer/buffer/
mutable.rs
