|
| 1 | +//! Generic scalar bit-packing for 64-bit values. |
| 2 | +//! |
| 3 | +//! Packs and unpacks groups of 32 values at any bit width `0..=64` using the same |
| 4 | +//! little-endian bitstream layout as the hand-unrolled 32-bit kernels in |
| 5 | +//! [`bitpacking`](super::bitpacking): value `j` occupies bits `[j*bit, (j+1)*bit)` |
| 6 | +//! of the concatenated stream. Each call moves exactly `bit` `u32` words. |
| 7 | +
|
| 8 | +const fn low_mask(bit: u8) -> u64 { |
| 9 | + if bit >= 64 { u64::MAX } else { (1u64 << bit) - 1 } |
| 10 | +} |
| 11 | + |
| 12 | +/// Packs 32 values from `input[inpos..]` into `output[outpos..]` at `bit` bits each. |
| 13 | +pub fn pack_wide(input: &[u64], inpos: usize, output: &mut [u32], outpos: usize, bit: u8) { |
| 14 | + if bit == 0 { |
| 15 | + return; |
| 16 | + } |
| 17 | + let mask = u128::from(low_mask(bit)); |
| 18 | + let mut acc: u128 = 0; |
| 19 | + let mut filled: u32 = 0; |
| 20 | + let mut out = outpos; |
| 21 | + for j in 0..32 { |
| 22 | + acc |= (u128::from(input[inpos + j]) & mask) << filled; |
| 23 | + filled += u32::from(bit); |
| 24 | + while filled >= 32 { |
| 25 | + output[out] = acc as u32; |
| 26 | + out += 1; |
| 27 | + acc >>= 32; |
| 28 | + filled -= 32; |
| 29 | + } |
| 30 | + } |
| 31 | +} |
| 32 | + |
| 33 | +/// Unpacks 32 values from `input[inpos..]` into `output[outpos..]` at `bit` bits each. |
| 34 | +pub fn unpack_wide(input: &[u32], inpos: usize, output: &mut [u64], outpos: usize, bit: u8) { |
| 35 | + if bit == 0 { |
| 36 | + output[outpos..outpos + 32].fill(0); |
| 37 | + return; |
| 38 | + } |
| 39 | + let mask = u128::from(low_mask(bit)); |
| 40 | + let mut acc: u128 = 0; |
| 41 | + let mut avail: u32 = 0; |
| 42 | + let mut inp = inpos; |
| 43 | + for j in 0..32 { |
| 44 | + while avail < u32::from(bit) { |
| 45 | + acc |= u128::from(input[inp]) << avail; |
| 46 | + inp += 1; |
| 47 | + avail += 32; |
| 48 | + } |
| 49 | + output[outpos + j] = (acc & mask) as u64; |
| 50 | + acc >>= u32::from(bit); |
| 51 | + avail -= u32::from(bit); |
| 52 | + } |
| 53 | +} |
| 54 | + |
| 55 | +#[cfg(test)] |
| 56 | +mod tests { |
| 57 | + use super::*; |
| 58 | + use crate::rust::integer_compression::{bitpacking, bitunpacking}; |
| 59 | + |
| 60 | + /// The wide packer must produce byte-identical output to the proven u32 kernels |
| 61 | + /// for every width `1..=32`, validating its bit ordering without needing C++. |
| 62 | + #[test] |
| 63 | + fn wide_matches_u32_kernels() { |
| 64 | + let values32: [u32; 32] = std::array::from_fn(|i| (i as u32).wrapping_mul(2_654_435_761)); |
| 65 | + |
| 66 | + for bit in 1..=32u8 { |
| 67 | + let mask = if bit == 32 { u32::MAX } else { (1u32 << bit) - 1 }; |
| 68 | + let masked32: [u32; 32] = std::array::from_fn(|i| values32[i] & mask); |
| 69 | + let masked64: [u64; 32] = std::array::from_fn(|i| u64::from(masked32[i])); |
| 70 | + |
| 71 | + let mut out_ref = vec![0u32; bit as usize]; |
| 72 | + bitpacking::fast_pack(&masked32, 0, &mut out_ref, 0, bit); |
| 73 | + |
| 74 | + let mut out_wide = vec![0u32; bit as usize]; |
| 75 | + pack_wide(&masked64, 0, &mut out_wide, 0, bit); |
| 76 | + |
| 77 | + assert_eq!(out_ref, out_wide, "pack mismatch at bit={bit}"); |
| 78 | + |
| 79 | + let mut back_ref = vec![0u32; 32]; |
| 80 | + bitunpacking::fast_unpack(&out_ref, 0, &mut back_ref, 0, bit); |
| 81 | + let mut back_wide = vec![0u64; 32]; |
| 82 | + unpack_wide(&out_wide, 0, &mut back_wide, 0, bit); |
| 83 | + |
| 84 | + for i in 0..32 { |
| 85 | + assert_eq!(u64::from(back_ref[i]), back_wide[i], "unpack mismatch at bit={bit}"); |
| 86 | + } |
| 87 | + } |
| 88 | + } |
| 89 | + |
| 90 | + #[test] |
| 91 | + fn wide_roundtrip_all_widths() { |
| 92 | + for bit in 0..=64u8 { |
| 93 | + let mask = low_mask(bit); |
| 94 | + let values: [u64; 32] = |
| 95 | + std::array::from_fn(|i| (i as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) & mask); |
| 96 | + |
| 97 | + let mut packed = vec![0u32; bit as usize]; |
| 98 | + pack_wide(&values, 0, &mut packed, 0, bit); |
| 99 | + |
| 100 | + let mut back = vec![0u64; 32]; |
| 101 | + unpack_wide(&packed, 0, &mut back, 0, bit); |
| 102 | + |
| 103 | + assert_eq!(values.to_vec(), back, "roundtrip mismatch at bit={bit}"); |
| 104 | + } |
| 105 | + } |
| 106 | +} |
0 commit comments