Skip to main content

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