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