/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 | | } |