fib_quant/
bitpack.rs

1use crate::{FibQuantError, Result};
2
3/// Pack fixed-width indices into little-endian bit order.
4pub fn pack_indices(indices: &[u32], width: u8) -> Result<Vec<u8>> {
5    if width == 0 || width > 32 {
6        return Err(FibQuantError::CorruptPayload(format!(
7            "invalid bit width {width}"
8        )));
9    }
10    let total_bits = indices
11        .len()
12        .checked_mul(width as usize)
13        .ok_or_else(|| FibQuantError::ResourceLimitExceeded("index bit count overflow".into()))?;
14    let expected_bytes = total_bits
15        .checked_add(7)
16        .ok_or_else(|| FibQuantError::ResourceLimitExceeded("index byte count overflow".into()))?
17        / 8;
18    let mut out = vec![0u8; expected_bytes];
19    let max = if width == 32 {
20        u32::MAX
21    } else {
22        (1u32 << width) - 1
23    };
24    for (idx, &value) in indices.iter().enumerate() {
25        if value > max {
26            return Err(FibQuantError::IndexOutOfRange {
27                index: value,
28                codebook_size: max,
29            });
30        }
31        let start = idx.checked_mul(width as usize).ok_or_else(|| {
32            FibQuantError::ResourceLimitExceeded("index bit offset overflow".into())
33        })?;
34        for bit in 0..width as usize {
35            if ((value >> bit) & 1) == 1 {
36                let pos = start.checked_add(bit).ok_or_else(|| {
37                    FibQuantError::ResourceLimitExceeded("index bit position overflow".into())
38                })?;
39                out[pos / 8] |= 1 << (pos % 8);
40            }
41        }
42    }
43    Ok(out)
44}
45
46/// Unpack fixed-width indices and reject nonzero padding bits.
47pub fn unpack_indices(bytes: &[u8], count: usize, width: u8) -> Result<Vec<u32>> {
48    if width == 0 || width > 32 {
49        return Err(FibQuantError::CorruptPayload(format!(
50            "invalid bit width {width}"
51        )));
52    }
53    let total_bits = count
54        .checked_mul(width as usize)
55        .ok_or_else(|| FibQuantError::ResourceLimitExceeded("index bit count overflow".into()))?;
56    let expected_bytes = total_bits
57        .checked_add(7)
58        .ok_or_else(|| FibQuantError::ResourceLimitExceeded("index byte count overflow".into()))?
59        / 8;
60    if bytes.len() != expected_bytes {
61        return Err(FibQuantError::CorruptPayload(format!(
62            "payload has {} bytes, expected {expected_bytes}",
63            bytes.len()
64        )));
65    }
66    for pos in total_bits..(expected_bytes * 8) {
67        if ((bytes[pos / 8] >> (pos % 8)) & 1) == 1 {
68            return Err(FibQuantError::CorruptPayload(
69                "nonzero fixed-rate padding bits".into(),
70            ));
71        }
72    }
73    let mut indices = Vec::with_capacity(count);
74    for idx in 0..count {
75        let start = idx.checked_mul(width as usize).ok_or_else(|| {
76            FibQuantError::ResourceLimitExceeded("index bit offset overflow".into())
77        })?;
78        let mut value = 0u32;
79        for bit in 0..width as usize {
80            let pos = start.checked_add(bit).ok_or_else(|| {
81                FibQuantError::ResourceLimitExceeded("index bit position overflow".into())
82            })?;
83            if ((bytes[pos / 8] >> (pos % 8)) & 1) == 1 {
84                value |= 1 << bit;
85            }
86        }
87        indices.push(value);
88    }
89    Ok(indices)
90}
fib_quant/bitpack.rs

fib_quant/
bitpack.rs
