arrow_buffer/util/
bit_chunk_iterator.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
18//! Types for iterating over bitmasks in 64-bit chunks
19
20use crate::util::bit_util::ceil;
21use std::fmt::Debug;
22
23/// Iterates over an arbitrarily aligned byte buffer
24///
25/// Yields an iterator of aligned u64, along with the leading and trailing
26/// u64 necessary to align the buffer to a 8-byte boundary
27///
28/// This is unlike [`BitChunkIterator`] which only exposes a trailing u64,
29/// and consequently has to perform more work for each read
30#[derive(Debug)]
31pub struct UnalignedBitChunk<'a> {
32    lead_padding: usize,
33    trailing_padding: usize,
34
35    prefix: Option<u64>,
36    chunks: &'a [u64],
37    suffix: Option<u64>,
38}
39
40impl<'a> UnalignedBitChunk<'a> {
41    /// Create from a byte array, offset and length in bits
42    pub fn new(buffer: &'a [u8], offset: usize, len: usize) -> Self {
43        if len == 0 {
44            return Self {
45                lead_padding: 0,
46                trailing_padding: 0,
47                prefix: None,
48                chunks: &[],
49                suffix: None,
50            };
51        }
52
53        let byte_offset = offset / 8;
54        let offset_padding = offset % 8;
55
56        let bytes_len = (len + offset_padding).div_ceil(8);
57        let buffer = &buffer[byte_offset..byte_offset + bytes_len];
58
59        let prefix_mask = compute_prefix_mask(offset_padding);
60
61        // If less than 8 bytes, read into prefix
62        if buffer.len() <= 8 {
63            let (suffix_mask, trailing_padding) = compute_suffix_mask(len, offset_padding);
64            let prefix = read_u64(buffer) & suffix_mask & prefix_mask;
65
66            return Self {
67                lead_padding: offset_padding,
68                trailing_padding,
69                prefix: Some(prefix),
70                chunks: &[],
71                suffix: None,
72            };
73        }
74
75        // If less than 16 bytes, read into prefix and suffix
76        if buffer.len() <= 16 {
77            let (suffix_mask, trailing_padding) = compute_suffix_mask(len, offset_padding);
78            let prefix = read_u64(&buffer[..8]) & prefix_mask;
79            let suffix = read_u64(&buffer[8..]) & suffix_mask;
80
81            return Self {
82                lead_padding: offset_padding,
83                trailing_padding,
84                prefix: Some(prefix),
85                chunks: &[],
86                suffix: Some(suffix),
87            };
88        }
89
90        // Read into prefix and suffix as needed
91        // Safety: u64 has no invalid bit patterns so reinterpreting initialized u8 bytes as u64 is sound.
92        let (prefix, mut chunks, suffix) = unsafe { buffer.align_to::<u64>() };
93        assert!(
94            prefix.len() < 8 && suffix.len() < 8,
95            "align_to did not return largest possible aligned slice"
96        );
97
98        let (alignment_padding, prefix) = match (offset_padding, prefix.is_empty()) {
99            (0, true) => (0, None),
100            (_, true) => {
101                let prefix = chunks[0] & prefix_mask;
102                chunks = &chunks[1..];
103                (0, Some(prefix))
104            }
105            (_, false) => {
106                let alignment_padding = (8 - prefix.len()) * 8;
107
108                let prefix = (read_u64(prefix) & prefix_mask) << alignment_padding;
109                (alignment_padding, Some(prefix))
110            }
111        };
112
113        let lead_padding = offset_padding + alignment_padding;
114        let (suffix_mask, trailing_padding) = compute_suffix_mask(len, lead_padding);
115
116        let suffix = match (trailing_padding, suffix.is_empty()) {
117            (0, _) => None,
118            (_, true) => {
119                let suffix = chunks[chunks.len() - 1] & suffix_mask;
120                chunks = &chunks[..chunks.len() - 1];
121                Some(suffix)
122            }
123            (_, false) => Some(read_u64(suffix) & suffix_mask),
124        };
125
126        Self {
127            lead_padding,
128            trailing_padding,
129            prefix,
130            chunks,
131            suffix,
132        }
133    }
134
135    /// Returns the number of leading padding bits
136    pub fn lead_padding(&self) -> usize {
137        self.lead_padding
138    }
139
140    /// Returns the number of trailing padding bits
141    pub fn trailing_padding(&self) -> usize {
142        self.trailing_padding
143    }
144
145    /// Returns the prefix, if any
146    pub fn prefix(&self) -> Option<u64> {
147        self.prefix
148    }
149
150    /// Returns the suffix, if any
151    pub fn suffix(&self) -> Option<u64> {
152        self.suffix
153    }
154
155    /// Returns reference to the chunks
156    pub fn chunks(&self) -> &'a [u64] {
157        self.chunks
158    }
159
160    /// Returns an iterator over the chunks
161    #[inline]
162    pub fn iter(&self) -> UnalignedBitChunkIterator<'a> {
163        self.prefix
164            .into_iter()
165            .chain(self.chunks.iter().copied())
166            .chain(self.suffix)
167    }
168
169    /// Counts the number of ones
170    pub fn count_ones(&self) -> usize {
171        self.iter().map(|x| x.count_ones() as usize).sum()
172    }
173}
174
175impl<'a> IntoIterator for &UnalignedBitChunk<'a> {
176    type Item = u64;
177    type IntoIter = UnalignedBitChunkIterator<'a>;
178
179    fn into_iter(self) -> Self::IntoIter {
180        self.iter()
181    }
182}
183
184/// Iterator over an [`UnalignedBitChunk`]
185pub type UnalignedBitChunkIterator<'a> = std::iter::Chain<
186    std::iter::Chain<std::option::IntoIter<u64>, std::iter::Copied<std::slice::Iter<'a, u64>>>,
187    std::option::IntoIter<u64>,
188>;
189
190#[inline]
191fn read_u64(input: &[u8]) -> u64 {
192    let len = input.len().min(8);
193    let mut buf = [0_u8; 8];
194    buf[..len].copy_from_slice(input);
195    u64::from_le_bytes(buf)
196}
197
198#[inline]
199fn compute_prefix_mask(lead_padding: usize) -> u64 {
200    !((1 << lead_padding) - 1)
201}
202
203#[inline]
204fn compute_suffix_mask(len: usize, lead_padding: usize) -> (u64, usize) {
205    let trailing_bits = (len + lead_padding) % 64;
206
207    if trailing_bits == 0 {
208        return (u64::MAX, 0);
209    }
210
211    let trailing_padding = 64 - trailing_bits;
212    let suffix_mask = (1 << trailing_bits) - 1;
213    (suffix_mask, trailing_padding)
214}
215
216/// Iterates over an arbitrarily aligned byte buffer 64 bits at a time
217///
218/// [`Self::iter`] yields iterator of `u64`, and a remainder. The first byte in the buffer
219/// will be the least significant byte in output u64
220#[derive(Debug)]
221pub struct BitChunks<'a> {
222    buffer: &'a [u8],
223    /// offset inside a byte, guaranteed to be between 0 and 7 (inclusive)
224    bit_offset: usize,
225    /// number of complete u64 chunks
226    chunk_len: usize,
227    /// number of remaining bits, guaranteed to be between 0 and 63 (inclusive)
228    remainder_len: usize,
229}
230
231impl<'a> BitChunks<'a> {
232    /// Create a new [`BitChunks`] from a byte array, and an offset and length in bits
233    pub fn new(buffer: &'a [u8], offset: usize, len: usize) -> Self {
234        let end = offset.checked_add(len).expect("offset + len out of bounds");
235        assert!(ceil(end, 8) <= buffer.len(), "offset + len out of bounds");
236
237        let byte_offset = offset / 8;
238        let bit_offset = offset % 8;
239
240        // number of complete u64 chunks
241        let chunk_len = len / 64;
242        // number of remaining bits
243        let remainder_len = len % 64;
244
245        BitChunks::<'a> {
246            buffer: &buffer[byte_offset..],
247            bit_offset,
248            chunk_len,
249            remainder_len,
250        }
251    }
252}
253
254/// Iterator over chunks of 64 bits represented as an u64
255#[derive(Debug)]
256pub struct BitChunkIterator<'a> {
257    buffer: &'a [u8],
258    bit_offset: usize,
259    chunk_len: usize,
260    index: usize,
261}
262
263impl<'a> BitChunks<'a> {
264    /// Returns the number of remaining bits, guaranteed to be between 0 and 63 (inclusive)
265    #[inline]
266    pub const fn remainder_len(&self) -> usize {
267        self.remainder_len
268    }
269
270    /// Returns the number of `u64` chunks
271    #[inline]
272    pub const fn chunk_len(&self) -> usize {
273        self.chunk_len
274    }
275
276    /// Returns the bitmask of remaining bits
277    #[inline]
278    pub fn remainder_bits(&self) -> u64 {
279        let bit_len = self.remainder_len;
280        if bit_len == 0 {
281            0
282        } else {
283            let bit_offset = self.bit_offset;
284            // number of bytes to read
285            // might be one more than sizeof(u64) if the offset is in the middle of a byte
286            let byte_len = ceil(bit_len + bit_offset, 8);
287            // pointer to remainder bytes after all complete chunks
288            // Safety: the buffer contains `chunk_len * 8 + ceil(remainder_len + bit_offset, 8)`
289            // bytes, so offsetting by `chunk_len * 8` and reading `byte_len` bytes is in-bounds.
290            let base = unsafe {
291                self.buffer
292                    .as_ptr()
293                    .add(self.chunk_len * std::mem::size_of::<u64>())
294            };
295
296            let mut bits = unsafe { std::ptr::read(base) } as u64 >> bit_offset;
297            for i in 1..byte_len {
298                let byte = unsafe { std::ptr::read(base.add(i)) };
299                bits |= (byte as u64) << (i * 8 - bit_offset);
300            }
301
302            bits & ((1 << bit_len) - 1)
303        }
304    }
305
306    /// Return the number of `u64` that are needed to represent all bits
307    /// (including remainder).
308    ///
309    /// This is equal to `chunk_len + 1` if there is a remainder,
310    /// otherwise it is equal to `chunk_len`.
311    #[inline]
312    pub fn num_u64s(&self) -> usize {
313        if self.remainder_len == 0 {
314            self.chunk_len
315        } else {
316            self.chunk_len + 1
317        }
318    }
319
320    /// Return the number of *bytes* that are needed to represent all bits
321    /// (including remainder).
322    #[inline]
323    pub fn num_bytes(&self) -> usize {
324        ceil(self.chunk_len * 64 + self.remainder_len, 8)
325    }
326
327    /// Returns the `index`th complete chunk of 64 bits, the value
328    /// [`Self::iter`] yields at that position
329    ///
330    /// # Panics
331    ///
332    /// Panics if `index >= self.chunk_len()`
333    ///
334    /// # Example
335    ///
336    /// ```
337    /// # use arrow_buffer::Buffer;
338    /// let buffer = Buffer::from((0..17_u8).collect::<Vec<_>>());
339    /// // Start at bit 8, so chunk 1 holds bytes 9 through 16
340    /// let chunks = buffer.bit_chunks(8, 128);
341    /// assert_eq!(chunks.chunk_len(), 2);
342    /// assert_eq!(chunks.chunk(1), 0x100f0e0d0c0b0a09);
343    /// assert_eq!(chunks.iter().nth(1), Some(chunks.chunk(1)));
344    /// ```
345    #[inline]
346    pub fn chunk(&self, index: usize) -> u64 {
347        assert!(index < self.chunk_len, "chunk index out of bounds");
348        // Safety: `index < chunk_len`, and the constructor checked the buffer
349        // covers every complete chunk plus the remainder byte
350        unsafe { read_chunk(self.buffer, self.bit_offset, index) }
351    }
352
353    /// Returns an iterator over chunks of 64 bits represented as an `u64`
354    #[inline]
355    pub const fn iter(&self) -> BitChunkIterator<'a> {
356        BitChunkIterator::<'a> {
357            buffer: self.buffer,
358            bit_offset: self.bit_offset,
359            chunk_len: self.chunk_len,
360            index: 0,
361        }
362    }
363
364    /// Returns an iterator over chunks of 64 bits, with the remaining bits zero padded to 64-bits
365    #[inline]
366    pub fn iter_padded(&self) -> impl Iterator<Item = u64> + 'a {
367        self.iter().chain(std::iter::once(self.remainder_bits()))
368    }
369}
370
371impl<'a> IntoIterator for BitChunks<'a> {
372    type Item = u64;
373    type IntoIter = BitChunkIterator<'a>;
374
375    fn into_iter(self) -> Self::IntoIter {
376        self.iter()
377    }
378}
379
380impl<'a> IntoIterator for &BitChunks<'a> {
381    type Item = u64;
382    type IntoIter = BitChunkIterator<'a>;
383
384    fn into_iter(self) -> Self::IntoIter {
385        self.iter()
386    }
387}
388
389/// Reads the `index`th complete 64-bit chunk of `buffer`, whose bits start
390/// at `bit_offset` (in `0..8`)
391///
392/// # Safety
393///
394/// `index` must be less than the number of complete chunks, so that the
395/// buffer holds at least `index * 8 + 8` bytes, plus one more byte when
396/// `bit_offset != 0` (the remainder byte the constructor guarantees)
397#[inline]
398unsafe fn read_chunk(buffer: &[u8], bit_offset: usize, index: usize) -> u64 {
399    debug_assert!(bit_offset < 8);
400    debug_assert!(buffer.len() >= index * 8 + 8 + usize::from(bit_offset != 0));
401    // cast to *const u64 should be fine since we are using read_unaligned below
402    #[expect(clippy::cast_ptr_alignment)]
403    let raw_data = buffer.as_ptr().cast::<u64>();
404
405    // bit-packed buffers are stored starting with the least-significant byte first
406    // so when reading as u64 on a big-endian machine, the bytes need to be swapped
407    // Safety: the caller guarantees `raw_data.add(index)` is in bounds;
408    // `read_unaligned` handles any pointer alignment.
409    let current = unsafe { std::ptr::read_unaligned(raw_data.add(index)).to_le() };
410
411    if bit_offset == 0 {
412        current
413    } else {
414        // the constructor ensures that bit_offset is in 0..8
415        // that means we need to read at most one additional byte to fill in the high bits
416        // Safety: the caller guarantees the byte at `index + 1` is in bounds
417        let next = unsafe { std::ptr::read_unaligned(raw_data.add(index + 1).cast::<u8>()) as u64 };
418
419        (current >> bit_offset) | (next << (64 - bit_offset))
420    }
421}
422
423impl Iterator for BitChunkIterator<'_> {
424    type Item = u64;
425
426    #[inline]
427    fn next(&mut self) -> Option<u64> {
428        let index = self.index;
429        if index >= self.chunk_len {
430            return None;
431        }
432
433        // Safety: `index < self.chunk_len`, and the constructor checked the
434        // buffer covers every complete chunk plus the remainder byte
435        let combined = unsafe { read_chunk(self.buffer, self.bit_offset, index) };
436
437        self.index = index + 1;
438
439        Some(combined)
440    }
441
442    #[inline]
443    fn size_hint(&self) -> (usize, Option<usize>) {
444        (
445            self.chunk_len - self.index,
446            Some(self.chunk_len - self.index),
447        )
448    }
449}
450
451impl ExactSizeIterator for BitChunkIterator<'_> {
452    #[inline]
453    fn len(&self) -> usize {
454        self.chunk_len - self.index
455    }
456}
457
458#[cfg(test)]
459mod tests {
460    use rand::distr::uniform::UniformSampler;
461    use rand::distr::uniform::UniformUsize;
462    use rand::prelude::*;
463    use rand::rng;
464
465    use crate::buffer::Buffer;
466    use crate::util::bit_chunk_iterator::UnalignedBitChunk;
467
468    #[test]
469    fn test_iter_aligned() {
470        let input: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7];
471        let buffer: Buffer = Buffer::from(input);
472
473        let bitchunks = buffer.bit_chunks(0, 64);
474        let result = bitchunks.into_iter().collect::<Vec<_>>();
475
476        assert_eq!(vec![0x0706050403020100], result);
477    }
478
479    #[test]
480    fn test_iter_unaligned() {
481        let input: &[u8] = &[
482            0b00000000, 0b00000001, 0b00000010, 0b00000100, 0b00001000, 0b00010000, 0b00100000,
483            0b01000000, 0b11111111,
484        ];
485        let buffer: Buffer = Buffer::from(input);
486
487        let bitchunks = buffer.bit_chunks(4, 64);
488
489        assert_eq!(0, bitchunks.remainder_len());
490        assert_eq!(0, bitchunks.remainder_bits());
491
492        let result = bitchunks.into_iter().collect::<Vec<_>>();
493
494        assert_eq!(
495            vec![0b1111010000000010000000010000000010000000010000000010000000010000],
496            result
497        );
498    }
499
500    #[test]
501    fn test_iter_unaligned_remainder_1_byte() {
502        let input: &[u8] = &[
503            0b00000000, 0b00000001, 0b00000010, 0b00000100, 0b00001000, 0b00010000, 0b00100000,
504            0b01000000, 0b11111111,
505        ];
506        let buffer: Buffer = Buffer::from(input);
507
508        let bitchunks = buffer.bit_chunks(4, 66);
509
510        assert_eq!(2, bitchunks.remainder_len());
511        assert_eq!(0b00000011, bitchunks.remainder_bits());
512
513        let result = bitchunks.into_iter().collect::<Vec<_>>();
514
515        assert_eq!(
516            vec![0b1111010000000010000000010000000010000000010000000010000000010000],
517            result
518        );
519    }
520
521    #[test]
522    fn test_iter_unaligned_remainder_bits_across_bytes() {
523        let input: &[u8] = &[0b00111111, 0b11111100];
524        let buffer: Buffer = Buffer::from(input);
525
526        // remainder contains bits from both bytes
527        // result should be the highest 2 bits from first byte followed by lowest 5 bits of second bytes
528        let bitchunks = buffer.bit_chunks(6, 7);
529
530        assert_eq!(7, bitchunks.remainder_len());
531        assert_eq!(0b1110000, bitchunks.remainder_bits());
532    }
533
534    #[test]
535    fn test_iter_unaligned_remainder_bits_large() {
536        let input: &[u8] = &[
537            0b11111111, 0b00000000, 0b11111111, 0b00000000, 0b11111111, 0b00000000, 0b11111111,
538            0b00000000, 0b11111111,
539        ];
540        let buffer: Buffer = Buffer::from(input);
541
542        let bitchunks = buffer.bit_chunks(2, 63);
543
544        assert_eq!(63, bitchunks.remainder_len());
545        assert_eq!(
546            0b100_0000_0011_1111_1100_0000_0011_1111_1100_0000_0011_1111_1100_0000_0011_1111,
547            bitchunks.remainder_bits()
548        );
549    }
550
551    #[test]
552    fn test_chunk_aligned() {
553        let input: Vec<u8> = (0..24).collect();
554        let buffer = Buffer::from(input);
555
556        let bitchunks = buffer.bit_chunks(0, 24 * 8);
557        assert_eq!(3, bitchunks.chunk_len());
558        assert_eq!(0x0706050403020100, bitchunks.chunk(0));
559        assert_eq!(0x0f0e0d0c0b0a0908, bitchunks.chunk(1));
560        assert_eq!(0x1716151413121110, bitchunks.chunk(2));
561    }
562
563    #[test]
564    fn test_chunk_matches_iter() {
565        // 26 bytes cover three complete chunks at every offset below, with
566        // the last one needing the byte after its eight when the bit offset
567        // is not zero
568        let input: Vec<u8> = (0..26_u8).map(|i| i.wrapping_mul(13)).collect();
569        let buffer = Buffer::from(input);
570
571        // Bit offsets within a byte and across one, lengths ending on and
572        // off a word boundary
573        for offset in (0..8).chain([8, 13]) {
574            for len in [64, 65, 128, 130, 191, 192] {
575                let bitchunks = buffer.bit_chunks(offset, len);
576                assert_eq!(len / 64, bitchunks.chunk_len());
577                for index in 0..bitchunks.chunk_len() {
578                    assert_eq!(
579                        bitchunks.iter().nth(index),
580                        Some(bitchunks.chunk(index)),
581                        "offset {offset} len {len} chunk {index}"
582                    );
583                }
584            }
585        }
586    }
587
588    #[test]
589    #[should_panic(expected = "chunk index out of bounds")]
590    fn test_chunk_out_of_bounds() {
591        let buffer = Buffer::from(vec![0xFF_u8; 16]);
592        let bitchunks = buffer.bit_chunks(0, 128);
593        assert_eq!(2, bitchunks.chunk_len());
594        bitchunks.chunk(2);
595    }
596
597    #[test]
598    #[should_panic(expected = "chunk index out of bounds")]
599    fn test_chunk_out_of_bounds_no_complete_chunk() {
600        let buffer = Buffer::from(vec![0xFF_u8; 8]);
601        let bitchunks = buffer.bit_chunks(1, 63);
602        assert_eq!(0, bitchunks.chunk_len());
603        bitchunks.chunk(0);
604    }
605
606    #[test]
607    fn test_iter_remainder_out_of_bounds() {
608        // allocating a full page should trigger a fault when reading out of bounds
609        const ALLOC_SIZE: usize = 4 * 1024;
610        let input = vec![0xFF_u8; ALLOC_SIZE];
611
612        let buffer: Buffer = Buffer::from_vec(input);
613
614        let bitchunks = buffer.bit_chunks(57, ALLOC_SIZE * 8 - 57);
615
616        assert_eq!(u64::MAX, bitchunks.iter().last().unwrap());
617        assert_eq!(0x7F, bitchunks.remainder_bits());
618    }
619
620    #[test]
621    #[should_panic(expected = "offset + len out of bounds")]
622    fn test_out_of_bound_should_panic_length_is_more_than_buffer_length() {
623        const ALLOC_SIZE: usize = 4 * 1024;
624        let input = vec![0xFF_u8; ALLOC_SIZE];
625
626        let buffer: Buffer = Buffer::from_vec(input);
627
628        // We are reading more than exists in the buffer
629        buffer.bit_chunks(0, (ALLOC_SIZE + 1) * 8);
630    }
631
632    #[test]
633    #[should_panic(expected = "offset + len out of bounds")]
634    fn test_out_of_bound_should_panic_length_is_more_than_buffer_length_but_not_when_not_using_ceil()
635     {
636        const ALLOC_SIZE: usize = 4 * 1024;
637        let input = vec![0xFF_u8; ALLOC_SIZE];
638
639        let buffer: Buffer = Buffer::from_vec(input);
640
641        // We are reading more than exists in the buffer
642        buffer.bit_chunks(0, (ALLOC_SIZE * 8) + 1);
643    }
644
645    #[test]
646    #[should_panic(expected = "offset + len out of bounds")]
647    fn test_out_of_bound_should_panic_when_offset_is_not_zero_and_length_is_the_entire_buffer_length()
648     {
649        const ALLOC_SIZE: usize = 4 * 1024;
650        let input = vec![0xFF_u8; ALLOC_SIZE];
651
652        let buffer: Buffer = Buffer::from_vec(input);
653
654        // We are reading more than exists in the buffer
655        buffer.bit_chunks(8, ALLOC_SIZE * 8);
656    }
657
658    #[test]
659    #[should_panic(expected = "offset + len out of bounds")]
660    fn test_out_of_bound_should_panic_when_offset_is_not_zero_and_length_is_the_entire_buffer_length_with_ceil()
661     {
662        const ALLOC_SIZE: usize = 4 * 1024;
663        let input = vec![0xFF_u8; ALLOC_SIZE];
664
665        let buffer: Buffer = Buffer::from_vec(input);
666
667        // We are reading more than exists in the buffer
668        buffer.bit_chunks(1, ALLOC_SIZE * 8);
669    }
670
671    #[test]
672    #[should_panic(expected = "offset + len out of bounds")]
673    fn test_out_of_bound_should_panic_when_offset_and_length_overflow() {
674        let buffer = Buffer::from(vec![0xFF_u8; 8]);
675        buffer.bit_chunks(1, usize::MAX);
676    }
677
678    #[test]
679    fn test_unaligned_bit_chunk_iterator() {
680        let buffer = Buffer::from(&[0xFF; 5]);
681        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 0, 40);
682
683        assert!(unaligned.chunks().is_empty()); // Less than 128 elements
684        assert_eq!(unaligned.lead_padding(), 0);
685        assert_eq!(unaligned.trailing_padding(), 24);
686        // 24x 1 bit then 40x 0 bits
687        assert_eq!(
688            unaligned.prefix(),
689            Some(0b0000000000000000000000001111111111111111111111111111111111111111)
690        );
691        assert_eq!(unaligned.suffix(), None);
692
693        let buffer = buffer.slice(1);
694        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 0, 32);
695
696        assert!(unaligned.chunks().is_empty()); // Less than 128 elements
697        assert_eq!(unaligned.lead_padding(), 0);
698        assert_eq!(unaligned.trailing_padding(), 32);
699        // 32x 1 bit then 32x 0 bits
700        assert_eq!(
701            unaligned.prefix(),
702            Some(0b0000000000000000000000000000000011111111111111111111111111111111)
703        );
704        assert_eq!(unaligned.suffix(), None);
705
706        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 5, 27);
707
708        assert!(unaligned.chunks().is_empty()); // Less than 128 elements
709        assert_eq!(unaligned.lead_padding(), 5); // 5 % 8 == 5
710        assert_eq!(unaligned.trailing_padding(), 32);
711        // 5x 0 bit, 27x 1 bit then 32x 0 bits
712        assert_eq!(
713            unaligned.prefix(),
714            Some(0b0000000000000000000000000000000011111111111111111111111111100000)
715        );
716        assert_eq!(unaligned.suffix(), None);
717
718        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 12, 20);
719
720        assert!(unaligned.chunks().is_empty()); // Less than 128 elements
721        assert_eq!(unaligned.lead_padding(), 4); // 12 % 8 == 4
722        assert_eq!(unaligned.trailing_padding(), 40);
723        // 4x 0 bit, 20x 1 bit then 40x 0 bits
724        assert_eq!(
725            unaligned.prefix(),
726            Some(0b0000000000000000000000000000000000000000111111111111111111110000)
727        );
728        assert_eq!(unaligned.suffix(), None);
729
730        let buffer = Buffer::from(&[0xFF; 14]);
731
732        // Verify buffer alignment
733        let (prefix, aligned, suffix) = unsafe { buffer.as_slice().align_to::<u64>() };
734        assert_eq!(prefix.len(), 0);
735        assert_eq!(aligned.len(), 1);
736        assert_eq!(suffix.len(), 6);
737
738        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 0, 112);
739
740        assert!(unaligned.chunks().is_empty()); // Less than 128 elements
741        assert_eq!(unaligned.lead_padding(), 0); // No offset and buffer aligned on 64-bit boundary
742        assert_eq!(unaligned.trailing_padding(), 16);
743        assert_eq!(unaligned.prefix(), Some(u64::MAX));
744        assert_eq!(unaligned.suffix(), Some((1 << 48) - 1));
745
746        let buffer = Buffer::from(&[0xFF; 16]);
747
748        // Verify buffer alignment
749        let (prefix, aligned, suffix) = unsafe { buffer.as_slice().align_to::<u64>() };
750        assert_eq!(prefix.len(), 0);
751        assert_eq!(aligned.len(), 2);
752        assert_eq!(suffix.len(), 0);
753
754        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 0, 128);
755
756        assert_eq!(unaligned.prefix(), Some(u64::MAX));
757        assert_eq!(unaligned.suffix(), Some(u64::MAX));
758        assert!(unaligned.chunks().is_empty()); // Exactly 128 elements
759
760        let buffer = Buffer::from(&[0xFF; 64]);
761
762        // Verify buffer alignment
763        let (prefix, aligned, suffix) = unsafe { buffer.as_slice().align_to::<u64>() };
764        assert_eq!(prefix.len(), 0);
765        assert_eq!(aligned.len(), 8);
766        assert_eq!(suffix.len(), 0);
767
768        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 0, 512);
769
770        // Buffer is completely aligned and larger than 128 elements -> all in chunks array
771        assert_eq!(unaligned.suffix(), None);
772        assert_eq!(unaligned.prefix(), None);
773        assert_eq!(unaligned.chunks(), [u64::MAX; 8].as_slice());
774        assert_eq!(unaligned.lead_padding(), 0);
775        assert_eq!(unaligned.trailing_padding(), 0);
776
777        let buffer = buffer.slice(1); // Offset buffer 1 byte off 64-bit alignment
778
779        // Verify buffer alignment
780        let (prefix, aligned, suffix) = unsafe { buffer.as_slice().align_to::<u64>() };
781        assert_eq!(prefix.len(), 7);
782        assert_eq!(aligned.len(), 7);
783        assert_eq!(suffix.len(), 0);
784
785        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 0, 504);
786
787        // Need a prefix with 1 byte of lead padding to bring the buffer into alignment
788        assert_eq!(unaligned.prefix(), Some(u64::MAX - 0xFF));
789        assert_eq!(unaligned.suffix(), None);
790        assert_eq!(unaligned.chunks(), [u64::MAX; 7].as_slice());
791        assert_eq!(unaligned.lead_padding(), 8);
792        assert_eq!(unaligned.trailing_padding(), 0);
793
794        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 17, 300);
795
796        // Out of 64-bit alignment by 8 bits from buffer, and 17 bits from provided offset
797        //   => need 8 + 17 = 25 bits of lead padding + 39 bits in prefix
798        //
799        // This leaves 300 - 17 = 261 bits remaining
800        //   => 4x 64-bit aligned 64-bit chunks + 5 remaining bits
801        //   => trailing padding of 59 bits
802        assert_eq!(unaligned.lead_padding(), 25);
803        assert_eq!(unaligned.trailing_padding(), 59);
804        assert_eq!(unaligned.prefix(), Some(u64::MAX - (1 << 25) + 1));
805        assert_eq!(unaligned.suffix(), Some(0b11111));
806        assert_eq!(unaligned.chunks(), [u64::MAX; 4].as_slice());
807
808        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 17, 0);
809
810        assert_eq!(unaligned.prefix(), None);
811        assert_eq!(unaligned.suffix(), None);
812        assert!(unaligned.chunks().is_empty());
813        assert_eq!(unaligned.lead_padding(), 0);
814        assert_eq!(unaligned.trailing_padding(), 0);
815
816        let unaligned = UnalignedBitChunk::new(buffer.as_slice(), 17, 1);
817
818        assert_eq!(unaligned.prefix(), Some(2));
819        assert_eq!(unaligned.suffix(), None);
820        assert!(unaligned.chunks().is_empty());
821        assert_eq!(unaligned.lead_padding(), 1);
822        assert_eq!(unaligned.trailing_padding(), 62);
823    }
824
825    #[test]
826    fn fuzz_unaligned_bit_chunk_iterator() {
827        let mut rng = rng();
828
829        let uusize = UniformUsize::new(usize::MIN, usize::MAX).unwrap();
830        for _ in 0..100 {
831            let mask_len = rng.random_range(0..1024);
832            let bools: Vec<_> = std::iter::from_fn(|| Some(rng.random()))
833                .take(mask_len)
834                .collect();
835
836            let buffer = Buffer::from_iter(bools.iter().copied());
837
838            let max_offset = 64.min(mask_len);
839            let offset = uusize.sample(&mut rng).checked_rem(max_offset).unwrap_or(0);
840
841            let max_truncate = 128.min(mask_len - offset);
842            let truncate = uusize
843                .sample(&mut rng)
844                .checked_rem(max_truncate)
845                .unwrap_or(0);
846
847            let unaligned =
848                UnalignedBitChunk::new(buffer.as_slice(), offset, mask_len - offset - truncate);
849
850            let bool_slice = &bools[offset..mask_len - truncate];
851
852            let count = unaligned.count_ones();
853            let expected_count = bool_slice.iter().filter(|x| **x).count();
854
855            assert_eq!(count, expected_count);
856
857            let collected: Vec<u64> = unaligned.iter().collect();
858
859            let get_bit = |idx: usize| -> bool {
860                let padded_index = idx + unaligned.lead_padding();
861                let byte_idx = padded_index / 64;
862                let bit_idx = padded_index % 64;
863                (collected[byte_idx] & (1 << bit_idx)) != 0
864            };
865
866            for (idx, b) in bool_slice.iter().enumerate() {
867                assert_eq!(*b, get_bit(idx))
868            }
869        }
870    }
871}
arrow_buffer/util/bit_chunk_iterator.rs

arrow_buffer/util/
bit_chunk_iterator.rs
