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

arrow_buffer/util/
bit_chunk_iterator.rs
