#[must_use]
pub fn unpack_scalar(input: &[u8], bits: u8, count: usize, output: &mut [u32]) -> Option<usize> {
if bits > 32 {
return None;
}
if bits == 0 {
for slot in output.iter_mut().take(count) {
*slot = 0;
}
return Some(0);
}
let width = u32::from(bits);
let needed = (count * bits as usize).div_ceil(8);
if input.len() < needed || output.len() < count {
return None;
}
let mask = if bits == 32 {
u32::MAX
} else {
(1_u32 << width) - 1
};
let mut accumulator = 0_u64;
let mut filled = 0_u32;
let mut cursor = 0_usize;
for slot in output.iter_mut().take(count) {
while filled < width {
let byte = input.get(cursor).copied().unwrap_or(0);
cursor += 1;
accumulator |= u64::from(byte) << filled;
filled += 8;
}
*slot = (accumulator as u32) & mask;
accumulator >>= width;
filled -= width;
}
Some(needed)
}
pub fn prefix_sum_scalar(values: &mut [u32], base: u32) {
let mut running = base;
for slot in values.iter_mut() {
running = running.saturating_add(*slot);
*slot = running;
}
}
pub const NARROW_MAX_BITS: u8 = 8;
#[must_use]
pub fn unpack(input: &[u8], bits: u8, count: usize, output: &mut [u32]) -> Option<usize> {
if bits > 0 && bits <= NARROW_MAX_BITS {
return unpack_narrow(input, bits, count, output);
}
#[cfg(target_arch = "aarch64")]
{
if bits > 0
&& bits <= 32
&& count >= NEON_MIN_COUNT
&& std::arch::is_aarch64_feature_detected!("neon")
{
return unsafe { neon::unpack_neon(input, bits, count, output) };
}
}
unpack_scalar(input, bits, count, output)
}
#[must_use]
pub fn unpack_narrow(input: &[u8], bits: u8, count: usize, output: &mut [u32]) -> Option<usize> {
if bits == 0 || bits > NARROW_MAX_BITS {
return None;
}
let width = usize::from(bits);
let needed = count.checked_mul(width)?.div_ceil(8);
if input.len() < needed || output.len() < count {
return None;
}
let mask = (1_u32 << width) - 1;
let mut index = 0_usize;
let mut byte = 0_usize;
#[cfg(target_arch = "aarch64")]
if count >= 8 && std::arch::is_aarch64_feature_detected!("neon") {
let (consumed_values, consumed_bytes) =
unsafe { neon::unpack_narrow_neon(input, width, count, output) };
index = consumed_values;
byte = consumed_bytes;
}
while index + 8 <= count {
let window = match input.get(byte..byte + 8) {
Some(bytes) => {
let mut buffer = [0_u8; 8];
buffer.copy_from_slice(bytes);
u64::from_le_bytes(buffer)
}
None => {
let mut assembled = 0_u64;
for step in 0..width {
let value = input.get(byte + step).copied().unwrap_or(0);
assembled |= u64::from(value) << (8 * step);
}
assembled
}
};
for lane in 0..8 {
let slot = output.get_mut(index + lane)?;
*slot = ((window >> (lane * width)) as u32) & mask;
}
index += 8;
byte += width;
}
let remaining = count - index;
if remaining > 0 {
let tail = input.get(byte..)?;
let target = output.get_mut(index..)?;
unpack_scalar(tail, bits, remaining, target)?;
}
Some(needed)
}
pub fn prefix_sum(values: &mut [u32], base: u32) {
prefix_sum_wide(values, base);
}
#[cfg(target_arch = "aarch64")]
const WIDE_GROUP: usize = 16;
pub fn prefix_sum_wide(values: &mut [u32], base: u32) {
#[cfg(target_arch = "aarch64")]
{
if values.len() >= WIDE_GROUP && std::arch::is_aarch64_feature_detected!("neon") {
let (consumed, carry) = unsafe { neon::prefix_sum_wide_neon(values, base) };
if let Some(tail) = values.get_mut(consumed..) {
prefix_sum_narrow(tail, carry);
}
return;
}
}
prefix_sum_narrow(values, base);
}
fn prefix_sum_narrow(values: &mut [u32], base: u32) {
#[cfg(target_arch = "aarch64")]
{
if values.len() >= NEON_MIN_COUNT && std::arch::is_aarch64_feature_detected!("neon") {
unsafe { neon::prefix_sum_neon(values, base) };
return;
}
}
prefix_sum_scalar(values, base);
}
#[cfg(target_arch = "aarch64")]
const NEON_MIN_COUNT: usize = 8;
#[cfg(target_arch = "aarch64")]
mod neon {
use std::arch::aarch64::{
uint8x16_t, uint32x4_t, vandq_u16, vandq_u32, vdupq_n_u16, vdupq_n_u32, vextq_u32,
vget_low_u16, vgetq_lane_u32, vld1q_s16, vld1q_u8, vld1q_u32, vmovl_high_u16, vmovl_u16,
vqaddq_u32, vqtbl1q_u8, vreinterpretq_u16_u8, vshlq_u16, vshlq_u32, vst1q_u32,
};
#[target_feature(enable = "neon")]
pub(super) unsafe fn unpack_narrow_neon(
input: &[u8],
width: usize,
count: usize,
output: &mut [u32],
) -> (usize, usize) {
let mut selectors = [0_u8; 16];
let mut shifts = [0_i16; 8];
for lane in 0..8_usize {
let offset = lane * width;
let byte = u8::try_from(offset / 8).unwrap_or(0);
let Some(low) = selectors.get_mut(lane * 2) else {
return (0, 0);
};
*low = byte;
let Some(high) = selectors.get_mut(lane * 2 + 1) else {
return (0, 0);
};
*high = byte.saturating_add(1);
let Some(shift) = shifts.get_mut(lane) else {
return (0, 0);
};
*shift = -i16::try_from(offset % 8).unwrap_or(0);
}
let mask_value = ((1_u32 << width) - 1) as u16;
let mut index = 0_usize;
let mut byte = 0_usize;
let (indices, counts, mask) = unsafe {
(
vld1q_u8(selectors.as_ptr()),
vld1q_s16(shifts.as_ptr()),
vdupq_n_u16(mask_value),
)
};
while index + 8 <= count {
let Some(window) = input.get(byte..byte + 16) else {
break;
};
let Some(target) = output.get_mut(index..index + 8) else {
break;
};
unsafe {
let table: uint8x16_t = vld1q_u8(window.as_ptr());
let gathered = vreinterpretq_u16_u8(vqtbl1q_u8(table, indices));
let values = vandq_u16(vshlq_u16(gathered, counts), mask);
vst1q_u32(target.as_mut_ptr(), vmovl_u16(vget_low_u16(values)));
vst1q_u32(target.as_mut_ptr().add(4), vmovl_high_u16(values));
}
index += 8;
byte += width;
}
(index, byte)
}
#[target_feature(enable = "neon")]
pub(super) unsafe fn unpack_neon(
input: &[u8],
bits: u8,
count: usize,
output: &mut [u32],
) -> Option<usize> {
let width = u32::from(bits);
let needed = (count * bits as usize).div_ceil(8);
if input.len() < needed || output.len() < count {
return None;
}
if width > 25 {
return super::unpack_scalar(input, bits, count, output);
}
let mask_value = if width == 32 {
u32::MAX
} else {
(1_u32 << width) - 1
};
let mask = vdupq_n_u32(mask_value);
let mut index = 0_usize;
while index + 4 <= count {
let mut words = [0_u32; 4];
let mut shifts = [0_i32; 4];
for lane in 0..4 {
let bit_offset = (index + lane) * bits as usize;
let byte = bit_offset / 8;
let shift = bit_offset % 8;
let mut word = 0_u32;
for step in 0..4 {
let value = input.get(byte + step).copied().unwrap_or(0);
word |= u32::from(value) << (8 * step);
}
let slot = words.get_mut(lane)?;
*slot = word;
let slot = shifts.get_mut(lane)?;
*slot = -(i32::try_from(shift).unwrap_or(0));
}
unsafe {
let loaded: uint32x4_t = vld1q_u32(words.as_ptr());
let counts: std::arch::aarch64::int32x4_t =
std::arch::aarch64::vld1q_s32(shifts.as_ptr());
let shifted = vshlq_u32(loaded, counts);
let masked = vandq_u32(shifted, mask);
let target = output.get_mut(index..index + 4)?;
vst1q_u32(target.as_mut_ptr(), masked);
}
index += 4;
}
if index < count {
let tail = output.get_mut(index..count)?;
let bit_offset = index * bits as usize;
let byte = bit_offset / 8;
let shift = bit_offset % 8;
if shift == 0 {
let rest = input.get(byte..)?;
super::unpack_scalar(rest, bits, count - index, tail)?;
} else {
return super::unpack_scalar(input, bits, count, output).map(|_| needed);
}
}
Some(needed)
}
#[target_feature(enable = "neon")]
pub(super) unsafe fn prefix_sum_wide_neon(values: &mut [u32], base: u32) -> (usize, u32) {
let mut carry = base;
let mut index = 0_usize;
let zero = vdupq_n_u32(0);
while index + super::WIDE_GROUP <= values.len() {
let Some(window) = values.get_mut(index..index + super::WIDE_GROUP) else {
break;
};
let pointer = window.as_mut_ptr();
unsafe {
let mut v0: uint32x4_t = vld1q_u32(pointer);
let mut v1: uint32x4_t = vld1q_u32(pointer.add(4));
let mut v2: uint32x4_t = vld1q_u32(pointer.add(8));
let mut v3: uint32x4_t = vld1q_u32(pointer.add(12));
v0 = vqaddq_u32(v0, vextq_u32(zero, v0, 3));
v1 = vqaddq_u32(v1, vextq_u32(zero, v1, 3));
v2 = vqaddq_u32(v2, vextq_u32(zero, v2, 3));
v3 = vqaddq_u32(v3, vextq_u32(zero, v3, 3));
v0 = vqaddq_u32(v0, vextq_u32(zero, v0, 2));
v1 = vqaddq_u32(v1, vextq_u32(zero, v1, 2));
v2 = vqaddq_u32(v2, vextq_u32(zero, v2, 2));
v3 = vqaddq_u32(v3, vextq_u32(zero, v3, 2));
v0 = vqaddq_u32(v0, vdupq_n_u32(carry));
v1 = vqaddq_u32(v1, vdupq_n_u32(vgetq_lane_u32::<3>(v0)));
v2 = vqaddq_u32(v2, vdupq_n_u32(vgetq_lane_u32::<3>(v1)));
v3 = vqaddq_u32(v3, vdupq_n_u32(vgetq_lane_u32::<3>(v2)));
carry = vgetq_lane_u32::<3>(v3);
vst1q_u32(pointer, v0);
vst1q_u32(pointer.add(4), v1);
vst1q_u32(pointer.add(8), v2);
vst1q_u32(pointer.add(12), v3);
}
index += super::WIDE_GROUP;
}
(index, carry)
}
#[target_feature(enable = "neon")]
pub(super) unsafe fn prefix_sum_neon(values: &mut [u32], base: u32) {
let mut carry = base;
let mut index = 0_usize;
while index + 4 <= values.len() {
let Some(window) = values.get_mut(index..index + 4) else {
return;
};
unsafe {
let mut vector: uint32x4_t = vld1q_u32(window.as_ptr());
let zero = vdupq_n_u32(0);
vector = vqaddq_u32(vector, vextq_u32(zero, vector, 3));
vector = vqaddq_u32(vector, vextq_u32(zero, vector, 2));
vector = vqaddq_u32(vector, vdupq_n_u32(carry));
vst1q_u32(window.as_mut_ptr(), vector);
}
carry = window.get(3).copied().unwrap_or(carry);
index += 4;
}
if index < values.len()
&& let Some(tail) = values.get_mut(index..)
{
super::prefix_sum_scalar(tail, carry);
}
}
}
#[cfg(test)]
#[allow(
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic,
clippy::unwrap_used
)]
mod tests {
use super::*;
fn pack(values: &[u32], bits: u8) -> Vec<u8> {
let mut out = Vec::new();
if bits == 0 {
return out;
}
let width = u32::from(bits);
let mut accumulator = 0_u64;
let mut filled = 0_u32;
for value in values {
accumulator |= u64::from(*value) << filled;
filled += width;
while filled >= 8 {
out.push((accumulator & 0xFF) as u8);
accumulator >>= 8;
filled -= 8;
}
}
if filled > 0 {
out.push((accumulator & 0xFF) as u8);
}
out
}
fn sample(bits: u8, count: usize) -> Vec<u32> {
let mask = if bits == 0 {
0
} else if bits >= 32 {
u32::MAX
} else {
(1_u32 << bits) - 1
};
(0..count)
.map(|index| {
let mixed = (index as u64)
.wrapping_mul(2_654_435_761)
.wrapping_add(0x9E37_79B9);
(mixed as u32) & mask
})
.collect()
}
#[test]
fn the_narrow_path_equals_the_oracle_at_every_count_and_exact_buffer() {
for bits in 1..=NARROW_MAX_BITS {
for count in 1..=200_usize {
let values = sample(bits, count);
let mut packed = pack(&values, bits);
let needed = (count * usize::from(bits)).div_ceil(8);
packed.truncate(needed);
assert_eq!(packed.len(), needed, "the fixture must be exact");
let mut narrow = vec![0_u32; count];
let consumed = unpack_narrow(&packed, bits, count, &mut narrow)
.expect("an exact buffer is enough");
assert_eq!(consumed, needed, "bits {bits} count {count}");
assert_eq!(
narrow, values,
"the narrow path diverged at {bits} bits, {count} values"
);
let mut dispatched = vec![0_u32; count];
unpack(&packed, bits, count, &mut dispatched).expect("dispatches");
assert_eq!(dispatched, values);
}
}
}
#[test]
fn the_narrow_path_refuses_what_it_cannot_decode() {
let mut output = [0_u32; 8];
assert!(unpack_narrow(&[0_u8; 8], 0, 8, &mut output).is_none());
assert!(unpack_narrow(&[0_u8; 8], 9, 8, &mut output).is_none());
assert!(unpack_narrow(&[], 1, 8, &mut output).is_none());
assert!(unpack_narrow(&[0_u8; 8], 1, 9, &mut output).is_none());
}
#[test]
fn prop_neon_unpack_equals_scalar() {
for bits in 0..=32_u8 {
for count in [1_usize, 3, 4, 7, 8, 15, 16, 63, 64, 65, 128] {
let values = sample(bits, count);
let packed = pack(&values, bits);
let mut oracle = vec![0_u32; count];
let mut dispatched = vec![0_u32; count];
let expected = unpack_scalar(&packed, bits, count, &mut oracle);
let actual = unpack(&packed, bits, count, &mut dispatched);
assert_eq!(
expected.is_some(),
actual.is_some(),
"bits {bits} count {count}"
);
assert_eq!(oracle, values, "the oracle itself is wrong at {bits} bits");
assert_eq!(
dispatched, oracle,
"dispatched arm diverged at {bits} bits, {count} values"
);
}
}
}
#[test]
fn neon_prefix_sum_equals_scalar_for_arbitrary_blocks() {
for count in [0_usize, 1, 3, 4, 5, 8, 63, 64, 65, 200] {
for base in [0_u32, 1, 1_000, u32::MAX - 10] {
let gaps: Vec<u32> = (0..count).map(|index| (index as u32 % 17) + 1).collect();
let mut oracle = gaps.clone();
let mut dispatched = gaps;
prefix_sum_scalar(&mut oracle, base);
prefix_sum(&mut dispatched, base);
assert_eq!(
dispatched, oracle,
"prefix sum diverged at count {count}, base {base}"
);
}
}
}
#[test]
fn the_wide_prefix_sum_carries_across_every_group_boundary() {
for count in 0..=80_usize {
for base in [0_u32, 1, 1_000, u32::MAX - 10, u32::MAX] {
for span in [1_u32, 17, 1_000_000_007] {
let gaps: Vec<u32> = (0..count)
.map(|index| (index as u32).wrapping_mul(span) % span.max(2) + 1)
.collect();
let mut oracle = gaps.clone();
let mut wide = gaps;
prefix_sum_scalar(&mut oracle, base);
prefix_sum_wide(&mut wide, base);
assert_eq!(
wide, oracle,
"the wide prefix sum diverged at count {count}, \
base {base}, span {span}"
);
}
}
}
}
#[test]
fn the_prefix_sum_is_a_running_total() {
let mut values = vec![1_u32, 2, 3, 4, 5];
prefix_sum_scalar(&mut values, 10);
assert_eq!(values, vec![11, 13, 16, 20, 25]);
}
#[test]
fn the_prefix_sum_saturates_rather_than_wrapping() {
let mut values = vec![u32::MAX, 5];
prefix_sum(&mut values, 1);
assert_eq!(values, vec![u32::MAX, u32::MAX]);
}
#[test]
fn an_oversized_width_or_short_buffer_is_refused_by_both_arms() {
let mut output = [0_u32; 8];
assert_eq!(unpack_scalar(&[0; 64], 33, 8, &mut output), None);
assert_eq!(unpack(&[0; 64], 33, 8, &mut output), None);
assert_eq!(unpack_scalar(&[0; 1], 8, 8, &mut output), None);
assert_eq!(unpack(&[0; 1], 8, 8, &mut output), None);
let mut small = [0_u32; 2];
assert_eq!(unpack(&[0; 64], 8, 8, &mut small), None);
}
#[test]
fn a_zero_width_run_decodes_to_zeros_and_consumes_nothing() {
let mut output = vec![9_u32; 16];
assert_eq!(unpack(&[], 0, 16, &mut output), Some(0));
assert_eq!(output, vec![0_u32; 16]);
}
#[test]
fn the_dispatched_arm_matches_the_oracle_on_maximal_values() {
for bits in 1..=32_u8 {
let maximum = if bits == 32 {
u32::MAX
} else {
(1_u32 << bits) - 1
};
let values = vec![maximum; 64];
let packed = pack(&values, bits);
let mut oracle = vec![0_u32; 64];
let mut dispatched = vec![0_u32; 64];
unpack_scalar(&packed, bits, 64, &mut oracle).expect("scalar unpacks");
unpack(&packed, bits, 64, &mut dispatched).expect("dispatch unpacks");
assert_eq!(oracle, values, "oracle wrong at {bits}");
assert_eq!(dispatched, oracle, "dispatch wrong at {bits}");
}
}
}