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};
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            _ => return Err(bytes),
236        };
237
238        let len = bytes.len();
239        let data = bytes.ptr();
240        #[cfg(feature = "pool")]
241        let reservation = bytes.reservation.lock().unwrap().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) = self.reservation.lock().unwrap().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) = self.reservation.lock().unwrap().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) = self.reservation.lock().unwrap().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) = self.reservation.lock().unwrap().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 = self.reservation.lock().unwrap().take();
612            *bytes.reservation.lock().unwrap() = 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() as *const 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    #[inline]
807    pub unsafe fn extend_bool_trusted_len<I: Iterator<Item = bool>>(
808        &mut self,
809        mut iter: I,
810        offset: usize,
811    ) {
812        let (lower, upper) = iter.size_hint();
813        let len = upper.expect("Iterator must have exact size_hint");
814        assert_eq!(lower, len, "Iterator must have exact size_hint");
815        debug_assert!(
816            offset <= self.len * 8,
817            "offset must be <= buffer length in bits"
818        );
819
820        if len == 0 {
821            return;
822        }
823
824        let start_len = offset;
825        let end_bit = start_len + len;
826
827        // SAFETY: we will initialize all newly exposed bytes before they are read
828        let new_len_bytes = bit_util::ceil(end_bit, 8);
829        if new_len_bytes > self.len {
830            self.reserve(new_len_bytes - self.len);
831            // SAFETY: caller will initialize all newly exposed bytes before they are read
832            unsafe { self.set_len(new_len_bytes) };
833        }
834
835        let slice = self.as_slice_mut();
836
837        let mut bit_idx = start_len;
838
839        // ---- Unaligned prefix: advance to the next 64-bit boundary ----
840        let misalignment = bit_idx & 63;
841        let prefix_bits = if misalignment == 0 {
842            0
843        } else {
844            (64 - misalignment).min(end_bit - bit_idx)
845        };
846
847        if prefix_bits != 0 {
848            let byte_start = bit_idx / 8;
849            let byte_end = bit_util::ceil(bit_idx + prefix_bits, 8);
850            let bit_offset = bit_idx % 8;
851
852            // Clear any newly-visible bits in the existing partial byte
853            if bit_offset != 0 {
854                let keep_mask = (1u8 << bit_offset).wrapping_sub(1);
855                slice[byte_start] &= keep_mask;
856            }
857
858            // Zero any new bytes we will partially fill in this prefix
859            let zero_from = if bit_offset == 0 {
860                byte_start
861            } else {
862                byte_start + 1
863            };
864            if byte_end > zero_from {
865                slice[zero_from..byte_end].fill(0);
866            }
867
868            for _ in 0..prefix_bits {
869                let v = iter.next().unwrap();
870                if v {
871                    let byte_idx = bit_idx / 8;
872                    let bit = bit_idx % 8;
873                    slice[byte_idx] |= 1 << bit;
874                }
875                bit_idx += 1;
876            }
877        }
878
879        if bit_idx < end_bit {
880            // ---- Aligned middle: write u64 chunks ----
881            debug_assert_eq!(bit_idx & 63, 0);
882            let remaining_bits = end_bit - bit_idx;
883            let chunks = remaining_bits / 64;
884
885            let words_start = bit_idx / 8;
886            let words_end = words_start + chunks * 8;
887            for dst in slice[words_start..words_end].chunks_exact_mut(8) {
888                let mut packed: u64 = 0;
889                for i in 0..64 {
890                    packed |= (iter.next().unwrap() as u64) << i;
891                }
892                dst.copy_from_slice(&packed.to_le_bytes());
893                bit_idx += 64;
894            }
895
896            // ---- Unaligned suffix: remaining < 64 bits ----
897            let suffix_bits = end_bit - bit_idx;
898            if suffix_bits != 0 {
899                debug_assert_eq!(bit_idx % 8, 0);
900                let byte_start = bit_idx / 8;
901                let byte_end = bit_util::ceil(end_bit, 8);
902                slice[byte_start..byte_end].fill(0);
903
904                for _ in 0..suffix_bits {
905                    let v = iter.next().unwrap();
906                    if v {
907                        let byte_idx = bit_idx / 8;
908                        let bit = bit_idx % 8;
909                        slice[byte_idx] |= 1 << bit;
910                    }
911                    bit_idx += 1;
912                }
913            }
914        }
915
916        // Clear any unused bits in the last byte
917        let remainder = end_bit % 8;
918        if remainder != 0 {
919            let mask = (1u8 << remainder).wrapping_sub(1);
920            slice[bit_util::ceil(end_bit, 8) - 1] &= mask;
921        }
922
923        debug_assert_eq!(bit_idx, end_bit);
924    }
925
926    /// Register this [`MutableBuffer`] with the provided [`MemoryPool`]
927    ///
928    /// This claims the memory used by this buffer in the pool, allowing for
929    /// accurate accounting of memory usage. Any prior reservation will be
930    /// released so this works well when the buffer is being shared among
931    /// multiple arrays.
932    #[cfg(feature = "pool")]
933    pub fn claim(&self, pool: &dyn MemoryPool) {
934        *self.reservation.lock().unwrap() = Some(pool.reserve(self.capacity()));
935    }
936}
937
938/// Creates a non-null pointer with alignment of [`ALIGNMENT`]
939///
940/// This is similar to [`NonNull::dangling`]
941#[inline]
942pub(crate) fn dangling_ptr() -> NonNull<u8> {
943    // SAFETY: ALIGNMENT is a non-zero usize which is then cast
944    // to a *mut u8. Therefore, `ptr` is not null and the conditions for
945    // calling new_unchecked() are respected.
946    #[cfg(miri)]
947    {
948        // Since miri implies a nightly rust version we can use the unstable strict_provenance feature
949        unsafe { NonNull::new_unchecked(std::ptr::without_provenance_mut(ALIGNMENT)) }
950    }
951    #[cfg(not(miri))]
952    {
953        unsafe { NonNull::new_unchecked(ALIGNMENT as *mut u8) }
954    }
955}
956
957impl<A: ArrowNativeType> Extend<A> for MutableBuffer {
958    #[inline]
959    fn extend<T: IntoIterator<Item = A>>(&mut self, iter: T) {
960        let iterator = iter.into_iter();
961        self.extend_from_iter(iterator)
962    }
963}
964
965impl<T: ArrowNativeType> From<Vec<T>> for MutableBuffer {
966    fn from(mut value: Vec<T>) -> Self {
967        // Safety
968        // Vec::as_mut_ptr guaranteed to not be null and ArrowNativeType are trivially transmutable
969        let data = unsafe { NonNull::new_unchecked(value.as_mut_ptr().cast()) };
970        let len = value.len() * mem::size_of::<T>();
971        // Safety
972        // Vec guaranteed to have a valid layout matching that of `Layout::array`
973        // This is based on `RawVec::current_memory`
974        let layout = unsafe { Layout::array::<T>(value.capacity()).unwrap_unchecked() };
975        mem::forget(value);
976        Self {
977            data,
978            len,
979            layout,
980            #[cfg(feature = "pool")]
981            reservation: std::sync::Mutex::new(None),
982        }
983    }
984}
985
986impl MutableBuffer {
987    #[inline]
988    pub(super) fn extend_from_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
989        &mut self,
990        mut iterator: I,
991    ) {
992        let item_size = std::mem::size_of::<T>();
993        let (lower, _) = iterator.size_hint();
994        let additional = lower * item_size;
995        self.reserve(additional);
996
997        // this is necessary because of https://github.com/rust-lang/rust/issues/32155
998        let mut len = SetLenOnDrop::new(&mut self.len);
999        let mut dst = unsafe { self.data.as_ptr().add(len.local_len) };
1000        let capacity = self.layout.size();
1001
1002        while len.local_len + item_size <= capacity {
1003            if let Some(item) = iterator.next() {
1004                unsafe {
1005                    let src = item.to_byte_slice().as_ptr();
1006                    std::ptr::copy_nonoverlapping(src, dst, item_size);
1007                    dst = dst.add(item_size);
1008                }
1009                len.local_len += item_size;
1010            } else {
1011                break;
1012            }
1013        }
1014        drop(len);
1015
1016        iterator.for_each(|item| self.push(item));
1017    }
1018
1019    /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length.
1020    /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.
1021    /// # Example
1022    /// ```
1023    /// # use arrow_buffer::buffer::MutableBuffer;
1024    /// let v = vec![1u32];
1025    /// let iter = v.iter().map(|x| x * 2);
1026    /// let buffer = unsafe { MutableBuffer::from_trusted_len_iter(iter) };
1027    /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
1028    /// ```
1029    ///
1030    /// # Panics
1031    ///
1032    /// Panics if the iterator does not report an upper bound via `size_hint`, or if the
1033    /// reported length does not match the number of items produced, or if allocating the
1034    /// required buffer fails for the same reasons as [`MutableBuffer::new`].
1035    ///
1036    /// # Safety
1037    /// This method assumes that the iterator's size is correct and is undefined behavior
1038    /// to use it on an iterator that reports an incorrect length.
1039    // This implementation is required for two reasons:
1040    // 1. there is no trait `TrustedLen` in stable rust and therefore
1041    //    we can't specialize `extend` for `TrustedLen` like `Vec` does.
1042    // 2. `from_trusted_len_iter` is faster.
1043    #[inline]
1044    pub unsafe fn from_trusted_len_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
1045        iterator: I,
1046    ) -> Self {
1047        let item_size = std::mem::size_of::<T>();
1048        let (_, upper) = iterator.size_hint();
1049        let upper = upper.expect("from_trusted_len_iter requires an upper limit");
1050        let len = upper * item_size;
1051
1052        let mut buffer = MutableBuffer::new(len);
1053
1054        let mut dst = buffer.data.as_ptr();
1055        for item in iterator {
1056            // note how there is no reserve here (compared with `extend_from_iter`)
1057            let src = item.to_byte_slice().as_ptr();
1058            unsafe { std::ptr::copy_nonoverlapping(src, dst, item_size) };
1059            dst = unsafe { dst.add(item_size) };
1060        }
1061        assert_eq!(
1062            unsafe { dst.offset_from(buffer.data.as_ptr()) } as usize,
1063            len,
1064            "Trusted iterator length was not accurately reported"
1065        );
1066        buffer.len = len;
1067        buffer
1068    }
1069
1070    /// Creates a [`MutableBuffer`] from a boolean [`Iterator`] with a trusted (upper) length.
1071    /// # use arrow_buffer::buffer::MutableBuffer;
1072    /// # Example
1073    /// ```
1074    /// # use arrow_buffer::buffer::MutableBuffer;
1075    /// let v = vec![false, true, false];
1076    /// let iter = v.iter().map(|x| *x || true);
1077    /// let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(iter) };
1078    /// assert_eq!(buffer.len(), 1) // 3 booleans have 1 byte
1079    /// ```
1080    ///
1081    /// # Panics
1082    ///
1083    /// Panics if the iterator does not report an upper bound via `size_hint`, or if it yields
1084    /// fewer items than reported.
1085    ///
1086    /// # Safety
1087    /// This method assumes that the iterator's size is correct and is undefined behavior
1088    /// to use it on an iterator that reports an incorrect length.
1089    // This implementation is required for two reasons:
1090    // 1. there is no trait `TrustedLen` in stable rust and therefore
1091    //    we can't specialize `extend` for `TrustedLen` like `Vec` does.
1092    // 2. `from_trusted_len_iter_bool` is faster.
1093    #[inline]
1094    pub unsafe fn from_trusted_len_iter_bool<I: Iterator<Item = bool>>(mut iterator: I) -> Self {
1095        let (_, upper) = iterator.size_hint();
1096        let len = upper.expect("from_trusted_len_iter requires an upper limit");
1097
1098        Self::collect_bool(len, |_| iterator.next().unwrap())
1099    }
1100
1101    /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length or errors
1102    /// if any of the items of the iterator is an error.
1103    /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.
1104    ///
1105    /// # Panics
1106    ///
1107    /// Panics if the iterator does not report an upper bound via `size_hint`, or if the
1108    /// reported length does not match the number of items produced before an error-free finish,
1109    /// or if allocating the required buffer fails for the same reasons as
1110    /// [`MutableBuffer::new`].
1111    ///
1112    /// # Safety
1113    /// This method assumes that the iterator's size is correct and is undefined behavior
1114    /// to use it on an iterator that reports an incorrect length.
1115    #[inline]
1116    pub unsafe fn try_from_trusted_len_iter<
1117        E,
1118        T: ArrowNativeType,
1119        I: Iterator<Item = Result<T, E>>,
1120    >(
1121        iterator: I,
1122    ) -> Result<Self, E> {
1123        let item_size = std::mem::size_of::<T>();
1124        let (_, upper) = iterator.size_hint();
1125        let upper = upper.expect("try_from_trusted_len_iter requires an upper limit");
1126        let len = upper * item_size;
1127
1128        let mut buffer = MutableBuffer::new(len);
1129
1130        let mut dst = buffer.data.as_ptr();
1131        for item in iterator {
1132            let item = item?;
1133            // note how there is no reserve here (compared with `extend_from_iter`)
1134            let src = item.to_byte_slice().as_ptr();
1135            unsafe { std::ptr::copy_nonoverlapping(src, dst, item_size) };
1136            dst = unsafe { dst.add(item_size) };
1137        }
1138        // try_from_trusted_len_iter is instantiated a lot, so we extract part of it into a less
1139        // generic method to reduce compile time
1140        unsafe fn finalize_buffer(dst: *mut u8, buffer: &mut MutableBuffer, len: usize) {
1141            unsafe {
1142                assert_eq!(
1143                    dst.offset_from(buffer.data.as_ptr()) as usize,
1144                    len,
1145                    "Trusted iterator length was not accurately reported"
1146                );
1147                buffer.len = len;
1148            }
1149        }
1150        unsafe { finalize_buffer(dst, &mut buffer, len) };
1151        Ok(buffer)
1152    }
1153}
1154
1155impl Default for MutableBuffer {
1156    fn default() -> Self {
1157        Self::with_capacity(0)
1158    }
1159}
1160
1161impl std::ops::Deref for MutableBuffer {
1162    type Target = [u8];
1163
1164    fn deref(&self) -> &[u8] {
1165        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.len) }
1166    }
1167}
1168
1169impl std::ops::DerefMut for MutableBuffer {
1170    fn deref_mut(&mut self) -> &mut [u8] {
1171        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
1172    }
1173}
1174
1175impl AsRef<[u8]> for &MutableBuffer {
1176    fn as_ref(&self) -> &[u8] {
1177        self.as_slice()
1178    }
1179}
1180
1181impl Drop for MutableBuffer {
1182    fn drop(&mut self) {
1183        if self.layout.size() != 0 {
1184            // Safety: data was allocated with standard allocator with given layout
1185            unsafe { std::alloc::dealloc(self.data.as_ptr() as _, self.layout) };
1186        }
1187    }
1188}
1189
1190impl PartialEq for MutableBuffer {
1191    fn eq(&self, other: &MutableBuffer) -> bool {
1192        if self.len != other.len {
1193            return false;
1194        }
1195        if self.layout != other.layout {
1196            return false;
1197        }
1198        self.as_slice() == other.as_slice()
1199    }
1200}
1201
1202unsafe impl Sync for MutableBuffer {}
1203unsafe impl Send for MutableBuffer {}
1204
1205struct SetLenOnDrop<'a> {
1206    len: &'a mut usize,
1207    local_len: usize,
1208}
1209
1210impl<'a> SetLenOnDrop<'a> {
1211    #[inline]
1212    fn new(len: &'a mut usize) -> Self {
1213        SetLenOnDrop {
1214            local_len: *len,
1215            len,
1216        }
1217    }
1218}
1219
1220impl Drop for SetLenOnDrop<'_> {
1221    #[inline]
1222    fn drop(&mut self) {
1223        *self.len = self.local_len;
1224    }
1225}
1226
1227/// Creating a `MutableBuffer` instance by setting bits according to the boolean values
1228impl std::iter::FromIterator<bool> for MutableBuffer {
1229    fn from_iter<I>(iter: I) -> Self
1230    where
1231        I: IntoIterator<Item = bool>,
1232    {
1233        let mut iterator = iter.into_iter();
1234        let mut result = {
1235            let byte_capacity: usize = iterator.size_hint().0.saturating_add(7) / 8;
1236            MutableBuffer::new(byte_capacity)
1237        };
1238
1239        loop {
1240            let mut exhausted = false;
1241            let mut byte_accum: u8 = 0;
1242            let mut mask: u8 = 1;
1243
1244            //collect (up to) 8 bits into a byte
1245            while mask != 0 {
1246                if let Some(value) = iterator.next() {
1247                    byte_accum |= match value {
1248                        true => mask,
1249                        false => 0,
1250                    };
1251                    mask <<= 1;
1252                } else {
1253                    exhausted = true;
1254                    break;
1255                }
1256            }
1257
1258            // break if the iterator was exhausted before it provided a bool for this byte
1259            if exhausted && mask == 1 {
1260                break;
1261            }
1262
1263            //ensure we have capacity to write the byte
1264            if result.len() == result.capacity() {
1265                //no capacity for new byte, allocate 1 byte more (plus however many more the iterator advertises)
1266                let additional_byte_capacity = 1usize.saturating_add(
1267                    iterator.size_hint().0.saturating_add(7) / 8, //convert bit count to byte count, rounding up
1268                );
1269                result.reserve(additional_byte_capacity)
1270            }
1271
1272            // Soundness: capacity was allocated above
1273            unsafe { result.push_unchecked(byte_accum) };
1274            if exhausted {
1275                break;
1276            }
1277        }
1278        result
1279    }
1280}
1281
1282impl<T: ArrowNativeType> std::iter::FromIterator<T> for MutableBuffer {
1283    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
1284        let mut buffer = Self::default();
1285        buffer.extend_from_iter(iter.into_iter());
1286        buffer
1287    }
1288}
1289
1290#[cfg(test)]
1291mod tests {
1292    use super::*;
1293
1294    #[test]
1295    fn test_mutable_new() {
1296        let buf = MutableBuffer::new(63);
1297        assert_eq!(64, buf.capacity());
1298        assert_eq!(0, buf.len());
1299        assert!(buf.is_empty());
1300    }
1301
1302    #[test]
1303    fn test_mutable_default() {
1304        let buf = MutableBuffer::default();
1305        assert_eq!(0, buf.capacity());
1306        assert_eq!(0, buf.len());
1307        assert!(buf.is_empty());
1308
1309        let mut buf = MutableBuffer::default();
1310        buf.extend_from_slice(b"hello");
1311        assert_eq!(5, buf.len());
1312        assert_eq!(b"hello", buf.as_slice());
1313    }
1314
1315    #[test]
1316    fn test_mutable_extend_from_slice() {
1317        let mut buf = MutableBuffer::new(100);
1318        buf.extend_from_slice(b"hello");
1319        assert_eq!(5, buf.len());
1320        assert_eq!(b"hello", buf.as_slice());
1321
1322        buf.extend_from_slice(b" world");
1323        assert_eq!(11, buf.len());
1324        assert_eq!(b"hello world", buf.as_slice());
1325
1326        buf.clear();
1327        assert_eq!(0, buf.len());
1328        buf.extend_from_slice(b"hello arrow");
1329        assert_eq!(11, buf.len());
1330        assert_eq!(b"hello arrow", buf.as_slice());
1331    }
1332
1333    #[test]
1334    fn mutable_extend_from_iter() {
1335        let mut buf = MutableBuffer::new(0);
1336        buf.extend(vec![1u32, 2]);
1337        assert_eq!(8, buf.len());
1338        assert_eq!(&[1u8, 0, 0, 0, 2, 0, 0, 0], buf.as_slice());
1339
1340        buf.extend(vec![3u32, 4]);
1341        assert_eq!(16, buf.len());
1342        assert_eq!(
1343            &[1u8, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4, 0, 0, 0],
1344            buf.as_slice()
1345        );
1346    }
1347
1348    #[test]
1349    fn mutable_extend_from_iter_unaligned_u64() {
1350        let mut buf = MutableBuffer::new(16);
1351        buf.push(1_u8);
1352        buf.extend([1_u64]);
1353        assert_eq!(9, buf.len());
1354        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1355    }
1356
1357    #[test]
1358    fn mutable_extend_from_slice_unaligned_u64() {
1359        let mut buf = MutableBuffer::new(16);
1360        buf.extend_from_slice(&[1_u8]);
1361        buf.extend_from_slice(&[1_u64]);
1362        assert_eq!(9, buf.len());
1363        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1364    }
1365
1366    #[test]
1367    fn mutable_push_unaligned_u64() {
1368        let mut buf = MutableBuffer::new(16);
1369        buf.push(1_u8);
1370        buf.push(1_u64);
1371        assert_eq!(9, buf.len());
1372        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1373    }
1374
1375    #[test]
1376    fn mutable_push_unchecked_unaligned_u64() {
1377        let mut buf = MutableBuffer::new(16);
1378        unsafe {
1379            buf.push_unchecked(1_u8);
1380            buf.push_unchecked(1_u64);
1381        }
1382        assert_eq!(9, buf.len());
1383        assert_eq!(&[1u8, 1u8, 0, 0, 0, 0, 0, 0, 0], buf.as_slice());
1384    }
1385
1386    #[test]
1387    fn test_from_trusted_len_iter() {
1388        let iter = vec![1u32, 2].into_iter();
1389        let buf = unsafe { MutableBuffer::from_trusted_len_iter(iter) };
1390        assert_eq!(8, buf.len());
1391        assert_eq!(&[1u8, 0, 0, 0, 2, 0, 0, 0], buf.as_slice());
1392    }
1393
1394    #[test]
1395    fn test_mutable_reserve() {
1396        let mut buf = MutableBuffer::new(1);
1397        assert_eq!(64, buf.capacity());
1398
1399        // Reserving a smaller capacity should have no effect.
1400        buf.reserve(10);
1401        assert_eq!(64, buf.capacity());
1402
1403        buf.reserve(80);
1404        assert_eq!(128, buf.capacity());
1405
1406        buf.reserve(129);
1407        assert_eq!(256, buf.capacity());
1408    }
1409
1410    #[test]
1411    fn test_mutable_resize() {
1412        let mut buf = MutableBuffer::new(1);
1413        assert_eq!(64, buf.capacity());
1414        assert_eq!(0, buf.len());
1415
1416        buf.resize(20, 0);
1417        assert_eq!(64, buf.capacity());
1418        assert_eq!(20, buf.len());
1419
1420        buf.resize(10, 0);
1421        assert_eq!(64, buf.capacity());
1422        assert_eq!(10, buf.len());
1423
1424        buf.resize(100, 0);
1425        assert_eq!(128, buf.capacity());
1426        assert_eq!(100, buf.len());
1427
1428        buf.resize(30, 0);
1429        assert_eq!(128, buf.capacity());
1430        assert_eq!(30, buf.len());
1431
1432        buf.resize(0, 0);
1433        assert_eq!(128, buf.capacity());
1434        assert_eq!(0, buf.len());
1435    }
1436
1437    #[test]
1438    fn test_mutable_into() {
1439        let mut buf = MutableBuffer::new(1);
1440        buf.extend_from_slice(b"aaaa bbbb cccc dddd");
1441        assert_eq!(19, buf.len());
1442        assert_eq!(64, buf.capacity());
1443        assert_eq!(b"aaaa bbbb cccc dddd", buf.as_slice());
1444
1445        let immutable_buf: Buffer = buf.into();
1446        assert_eq!(19, immutable_buf.len());
1447        assert_eq!(64, immutable_buf.capacity());
1448        assert_eq!(b"aaaa bbbb cccc dddd", immutable_buf.as_slice());
1449    }
1450
1451    #[test]
1452    fn test_mutable_equal() {
1453        let mut buf = MutableBuffer::new(1);
1454        let mut buf2 = MutableBuffer::new(1);
1455
1456        buf.extend_from_slice(&[0xaa]);
1457        buf2.extend_from_slice(&[0xaa, 0xbb]);
1458        assert!(buf != buf2);
1459
1460        buf.extend_from_slice(&[0xbb]);
1461        assert_eq!(buf, buf2);
1462
1463        buf2.reserve(65);
1464        assert!(buf != buf2);
1465    }
1466
1467    #[test]
1468    fn test_mutable_shrink_to_fit() {
1469        let mut buffer = MutableBuffer::new(128);
1470        assert_eq!(buffer.capacity(), 128);
1471        buffer.push(1);
1472        buffer.push(2);
1473
1474        buffer.shrink_to_fit();
1475        assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
1476    }
1477
1478    #[test]
1479    fn test_mutable_set_null_bits() {
1480        let mut buffer = MutableBuffer::new(8).with_bitset(8, true);
1481
1482        for i in 0..=buffer.capacity() {
1483            buffer.set_null_bits(i, 0);
1484            assert_eq!(buffer[..8], [255; 8][..]);
1485        }
1486
1487        buffer.set_null_bits(1, 4);
1488        assert_eq!(buffer[..8], [255, 0, 0, 0, 0, 255, 255, 255][..]);
1489    }
1490
1491    #[test]
1492    #[should_panic = "out of bounds for buffer of length"]
1493    fn test_mutable_set_null_bits_oob() {
1494        let mut buffer = MutableBuffer::new(64);
1495        buffer.set_null_bits(1, buffer.capacity());
1496    }
1497
1498    #[test]
1499    #[should_panic = "out of bounds for buffer of length"]
1500    fn test_mutable_set_null_bits_oob_by_overflow() {
1501        let mut buffer = MutableBuffer::new(0);
1502        buffer.set_null_bits(1, usize::MAX);
1503    }
1504
1505    #[test]
1506    fn from_iter() {
1507        let buffer = [1u16, 2, 3, 4].into_iter().collect::<MutableBuffer>();
1508        assert_eq!(buffer.len(), 4 * mem::size_of::<u16>());
1509        assert_eq!(buffer.as_slice(), &[1, 0, 2, 0, 3, 0, 4, 0]);
1510    }
1511
1512    #[test]
1513    #[should_panic(expected = "invalid allocation layout for requested capacity")]
1514    fn test_with_capacity_panics_above_max_capacity() {
1515        let max_capacity = isize::MAX as usize - (isize::MAX as usize % ALIGNMENT);
1516        let _ = MutableBuffer::with_capacity(max_capacity + 1);
1517    }
1518
1519    #[cfg(feature = "pool")]
1520    mod pool_tests {
1521        use super::*;
1522        use crate::pool::{MemoryPool, TrackingMemoryPool};
1523
1524        #[test]
1525        fn test_reallocate_with_pool() {
1526            let pool = TrackingMemoryPool::default();
1527            let mut buffer = MutableBuffer::with_capacity(100);
1528            buffer.claim(&pool);
1529
1530            // Initial capacity should be 128 (multiple of 64)
1531            assert_eq!(buffer.capacity(), 128);
1532            assert_eq!(pool.used(), 128);
1533
1534            // Reallocate to a larger size
1535            buffer.try_reallocate(200).unwrap();
1536
1537            // The capacity is exactly the requested size, not rounded up
1538            assert_eq!(buffer.capacity(), 200);
1539            assert_eq!(pool.used(), 200);
1540
1541            // Reallocate to a smaller size
1542            buffer.try_reallocate(50).unwrap();
1543
1544            // The capacity is exactly the requested size, not rounded up
1545            assert_eq!(buffer.capacity(), 50);
1546            assert_eq!(pool.used(), 50);
1547        }
1548
1549        #[test]
1550        fn test_truncate_with_pool() {
1551            let pool = TrackingMemoryPool::default();
1552            let mut buffer = MutableBuffer::with_capacity(100);
1553
1554            // Fill buffer with some data
1555            buffer.resize(80, 1);
1556            assert_eq!(buffer.len(), 80);
1557
1558            buffer.claim(&pool);
1559            assert_eq!(pool.used(), 128);
1560
1561            // Truncate buffer
1562            buffer.truncate(40);
1563            assert_eq!(buffer.len(), 40);
1564            assert_eq!(pool.used(), 40);
1565
1566            // Truncate to zero
1567            buffer.clear();
1568            assert_eq!(buffer.len(), 0);
1569            assert_eq!(pool.used(), 0);
1570        }
1571
1572        #[test]
1573        fn test_resize_with_pool() {
1574            let pool = TrackingMemoryPool::default();
1575            let mut buffer = MutableBuffer::with_capacity(100);
1576            buffer.claim(&pool);
1577
1578            // Initial state
1579            assert_eq!(buffer.len(), 0);
1580            assert_eq!(pool.used(), 128);
1581
1582            // Resize to increase length
1583            buffer.resize(50, 1);
1584            assert_eq!(buffer.len(), 50);
1585            assert_eq!(pool.used(), 50);
1586
1587            // Resize to increase length beyond capacity
1588            buffer.resize(150, 1);
1589            assert_eq!(buffer.len(), 150);
1590            assert_eq!(buffer.capacity(), 256);
1591            assert_eq!(pool.used(), 150);
1592
1593            // Resize to decrease length
1594            buffer.resize(30, 1);
1595            assert_eq!(buffer.len(), 30);
1596            assert_eq!(pool.used(), 30);
1597        }
1598
1599        #[test]
1600        fn test_buffer_lifecycle_with_pool() {
1601            let pool = TrackingMemoryPool::default();
1602
1603            // Create a buffer with memory reservation
1604            let mut mutable = MutableBuffer::with_capacity(100);
1605            mutable.resize(80, 1);
1606            mutable.claim(&pool);
1607
1608            // Memory reservation is based on capacity when using claim()
1609            assert_eq!(pool.used(), 128);
1610
1611            // Convert to immutable Buffer
1612            let buffer = mutable.into_buffer();
1613
1614            // Memory reservation should be preserved
1615            assert_eq!(pool.used(), 128);
1616
1617            // Drop the buffer and the reservation should be released
1618            drop(buffer);
1619            assert_eq!(pool.used(), 0);
1620        }
1621    }
1622
1623    fn create_expected_repeated_slice<T: ArrowNativeType>(
1624        slice_to_repeat: &[T],
1625        repeat_count: usize,
1626    ) -> Buffer {
1627        let mut expected = MutableBuffer::new(size_of_val(slice_to_repeat) * repeat_count);
1628        for _ in 0..repeat_count {
1629            // Not using push_slice_repeated as this is the function under test
1630            expected.extend_from_slice(slice_to_repeat);
1631        }
1632        expected.into()
1633    }
1634
1635    // Helper to test a specific repeat count with various slice sizes
1636    fn test_repeat_count<T: ArrowNativeType + PartialEq + std::fmt::Debug>(
1637        repeat_count: usize,
1638        test_data: &[T],
1639    ) {
1640        let mut buffer = MutableBuffer::new(0);
1641        buffer.repeat_slice_n_times(test_data, repeat_count);
1642
1643        let expected = create_expected_repeated_slice(test_data, repeat_count);
1644        let result: Buffer = buffer.into();
1645
1646        assert_eq!(
1647            result,
1648            expected,
1649            "Failed for repeat_count={}, slice_len={}",
1650            repeat_count,
1651            test_data.len()
1652        );
1653    }
1654
1655    #[test]
1656    fn test_repeat_slice_count_edge_cases() {
1657        // Empty slice
1658        test_repeat_count(100, &[] as &[i32]);
1659
1660        // Zero repeats
1661        test_repeat_count(0, &[1i32, 2, 3]);
1662    }
1663
1664    #[test]
1665    #[should_panic(expected = "buffer length overflow")]
1666    fn test_repeat_slice_count_multiply_overflow() {
1667        let mut buffer = MutableBuffer::new(0);
1668        buffer.repeat_slice_n_times(&[0_u64], usize::MAX / mem::size_of::<u64>() + 1);
1669    }
1670
1671    #[test]
1672    #[should_panic(expected = "buffer length overflow")]
1673    fn test_repeat_slice_count_len_overflow() {
1674        let mut buffer = MutableBuffer::new(0);
1675        buffer.push(0_u8);
1676        buffer.repeat_slice_n_times(&[0_u8], usize::MAX);
1677    }
1678
1679    #[test]
1680    fn test_small_repeats_counts() {
1681        // test any special implementation for small repeat counts
1682        let data = &[1u8, 2, 3, 4, 5];
1683
1684        for _ in 1..=10 {
1685            test_repeat_count(2, data);
1686        }
1687    }
1688
1689    #[test]
1690    fn test_different_size_of_i32_repeat_slice() {
1691        let data: &[i32] = &[1, 2, 3];
1692        let data_with_single_item: &[i32] = &[42];
1693
1694        for data in &[data, data_with_single_item] {
1695            for item in 1..=9 {
1696                let base_repeat_count = 2_usize.pow(item);
1697                test_repeat_count(base_repeat_count - 1, data);
1698                test_repeat_count(base_repeat_count, data);
1699                test_repeat_count(base_repeat_count + 1, data);
1700            }
1701        }
1702    }
1703
1704    #[test]
1705    fn test_different_size_of_u8_repeat_slice() {
1706        let data: &[u8] = &[1, 2, 3];
1707        let data_with_single_item: &[u8] = &[10];
1708
1709        for data in &[data, data_with_single_item] {
1710            for item in 1..=9 {
1711                let base_repeat_count = 2_usize.pow(item);
1712                test_repeat_count(base_repeat_count - 1, data);
1713                test_repeat_count(base_repeat_count, data);
1714                test_repeat_count(base_repeat_count + 1, data);
1715            }
1716        }
1717    }
1718
1719    #[test]
1720    fn test_different_size_of_u16_repeat_slice() {
1721        let data: &[u16] = &[1, 2, 3];
1722        let data_with_single_item: &[u16] = &[10];
1723
1724        for data in &[data, data_with_single_item] {
1725            for item in 1..=9 {
1726                let base_repeat_count = 2_usize.pow(item);
1727                test_repeat_count(base_repeat_count - 1, data);
1728                test_repeat_count(base_repeat_count, data);
1729                test_repeat_count(base_repeat_count + 1, data);
1730            }
1731        }
1732    }
1733
1734    #[test]
1735    fn test_various_slice_lengths() {
1736        // Test different slice lengths with same repeat pattern
1737        let repeat_count = 37; // Arbitrary non-power-of-2
1738
1739        // Single element
1740        test_repeat_count(repeat_count, &[42i32]);
1741
1742        // Small slices
1743        test_repeat_count(repeat_count, &[1i32, 2]);
1744        test_repeat_count(repeat_count, &[1i32, 2, 3]);
1745        test_repeat_count(repeat_count, &[1i32, 2, 3, 4]);
1746        test_repeat_count(repeat_count, &[1i32, 2, 3, 4, 5]);
1747
1748        // Larger slices
1749        let data_10: Vec<i32> = (0..10).collect();
1750        test_repeat_count(repeat_count, &data_10);
1751
1752        let data_100: Vec<i32> = (0..100).collect();
1753        test_repeat_count(repeat_count, &data_100);
1754
1755        let data_1000: Vec<i32> = (0..1000).collect();
1756        test_repeat_count(repeat_count, &data_1000);
1757    }
1758
1759    #[test]
1760    #[should_panic(expected = "invalid allocation layout for requested capacity")]
1761    fn test_mutable_new_capacity_overflow() {
1762        // Tests overflow during initial allocation
1763        let _ = MutableBuffer::new(usize::MAX - 10);
1764    }
1765
1766    #[test]
1767    #[should_panic(expected = "buffer length overflow")]
1768    fn test_mutable_reserve_overflow() {
1769        // Tests overflow during growth (checked_add)
1770        let mut buf = MutableBuffer::new(1);
1771        buf.push(1u8);
1772        buf.reserve(usize::MAX);
1773    }
1774}
arrow_buffer/buffer/mutable.rs

arrow_buffer/buffer/
mutable.rs
