Skip to main content

ruvector_temporal_tensor/
bitpack.rs

1//! Bitstream packer/unpacker for arbitrary bit widths (1-8).
2//!
3//! Uses a 64-bit accumulator for sub-byte codes with no alignment padding.
4
5/// Pack unsigned codes of `bits` width into a byte stream.
6///
7/// Each code occupies exactly `bits` bits in the output with no alignment
8/// padding between codes. A trailing partial byte is emitted if needed.
9///
10/// For 8-bit codes, writes bytes directly without bit accumulation.
11#[inline]
12pub fn pack(codes: &[u32], bits: u32, out: &mut Vec<u8>) {
13    // Fast path: 8-bit codes map 1:1 to bytes.
14    if bits == 8 {
15        out.extend(codes.iter().map(|&c| c as u8));
16        return;
17    }
18
19    let mut acc: u64 = 0;
20    let mut acc_bits: u32 = 0;
21
22    for &code in codes {
23        acc |= (code as u64) << acc_bits;
24        acc_bits += bits;
25        while acc_bits >= 8 {
26            out.push((acc & 0xFF) as u8);
27            acc >>= 8;
28            acc_bits -= 8;
29        }
30    }
31
32    if acc_bits > 0 {
33        out.push((acc & 0xFF) as u8);
34    }
35}
36
37/// Unpack `count` unsigned codes of `bits` width from a byte stream.
38///
39/// Stops early if the data is exhausted before `count` codes are extracted.
40///
41/// For 8-bit codes, reads bytes directly without bit accumulation.
42#[inline]
43pub fn unpack(data: &[u8], bits: u32, count: usize, out: &mut Vec<u32>) {
44    // Fast path: 8-bit codes map 1:1 from bytes.
45    if bits == 8 {
46        let n = count.min(data.len());
47        out.extend(data[..n].iter().map(|&b| b as u32));
48        return;
49    }
50
51    let mask = (1u64 << bits) - 1;
52    let mut acc: u64 = 0;
53    let mut acc_bits: u32 = 0;
54    let mut byte_idx = 0usize;
55    let mut decoded = 0usize;
56
57    while decoded < count {
58        while acc_bits < bits && byte_idx < data.len() {
59            acc |= (data[byte_idx] as u64) << acc_bits;
60            acc_bits += 8;
61            byte_idx += 1;
62        }
63        if acc_bits < bits {
64            break;
65        }
66
67        out.push((acc & mask) as u32);
68        acc >>= bits;
69        acc_bits -= bits;
70        decoded += 1;
71    }
72}
73
74/// Compute qmax for a given bit width: `2^(bits-1) - 1`.
75///
76/// Returns 0 for invalid bit widths (0 or >8).
77///
78/// | bits | qmax |
79/// |------|------|
80/// | 8    | 127  |
81/// | 7    | 63   |
82/// | 5    | 15   |
83/// | 3    | 3    |
84#[inline]
85pub fn qmax_from_bits(bits: u8) -> i32 {
86    if bits == 0 || bits > 8 {
87        return 0;
88    }
89    (1i32 << (bits - 1)) - 1
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_roundtrip_8bit() {
98        let codes: Vec<u32> = (0..256).collect();
99        let mut packed = Vec::new();
100        pack(&codes, 8, &mut packed);
101        assert_eq!(packed.len(), 256);
102
103        let mut unpacked = Vec::new();
104        unpack(&packed, 8, 256, &mut unpacked);
105        assert_eq!(codes, unpacked);
106    }
107
108    #[test]
109    fn test_roundtrip_3bit() {
110        let codes: Vec<u32> = (0..7).collect();
111        let mut packed = Vec::new();
112        pack(&codes, 3, &mut packed);
113
114        let mut unpacked = Vec::new();
115        unpack(&packed, 3, 7, &mut unpacked);
116        assert_eq!(codes, unpacked);
117    }
118
119    #[test]
120    fn test_roundtrip_5bit() {
121        let codes: Vec<u32> = (0..31).collect();
122        let mut packed = Vec::new();
123        pack(&codes, 5, &mut packed);
124
125        let mut unpacked = Vec::new();
126        unpack(&packed, 5, 31, &mut unpacked);
127        assert_eq!(codes, unpacked);
128    }
129
130    #[test]
131    fn test_roundtrip_7bit() {
132        let codes: Vec<u32> = (0..127).collect();
133        let mut packed = Vec::new();
134        pack(&codes, 7, &mut packed);
135
136        let mut unpacked = Vec::new();
137        unpack(&packed, 7, 127, &mut unpacked);
138        assert_eq!(codes, unpacked);
139    }
140
141    #[test]
142    fn test_packing_density() {
143        let codes = vec![5u32; 100];
144        let mut packed = Vec::new();
145        pack(&codes, 3, &mut packed);
146        assert_eq!(packed.len(), 38); // ceil(300/8) = 38
147    }
148
149    #[test]
150    fn test_qmax() {
151        assert_eq!(qmax_from_bits(8), 127);
152        assert_eq!(qmax_from_bits(7), 63);
153        assert_eq!(qmax_from_bits(5), 15);
154        assert_eq!(qmax_from_bits(3), 3);
155        assert_eq!(qmax_from_bits(1), 0);
156        assert_eq!(qmax_from_bits(0), 0);
157    }
158}