## Coverage Report

#### Created: 2025-11-17 14:14

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/arrow-rs/arrow-buffer/src/buffer/mutable.rs

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