irox_bits/
bitstream.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2025 IROX Contributors
3//
4
5use crate::{Bits, BitsError, BitsErrorKind, BitsWrapper, MutBits};
6use core::cmp::Ordering;
7
8pub struct BitStreamEncoder<'a, T: MutBits> {
9    delegate: BitsWrapper<'a, T>,
10    buf: u32,
11    remaining: u8,
12}
13impl<T: MutBits> Drop for BitStreamEncoder<'_, T> {
14    fn drop(&mut self) {
15        let [a, b, c, d] = self.buf.to_be_bytes();
16        if self.remaining < 8 {
17            // write 4
18            let _ = self.delegate.write_all_bytes(&[a, b, c, d]);
19        } else if self.remaining < 16 {
20            // write 3
21            let _ = self.delegate.write_all_bytes(&[a, b, c]);
22        } else if self.remaining < 24 {
23            // write 2
24            let _ = self.delegate.write_all_bytes(&[a, b]);
25        } else if self.remaining < 32 {
26            // write 1
27            let _ = self.delegate.write_u8(a);
28        }
29    }
30}
31impl<'a, T: MutBits> BitStreamEncoder<'a, T> {
32    pub fn new(delegate: BitsWrapper<'a, T>) -> Self {
33        Self {
34            delegate,
35            buf: 0u32,
36            remaining: 32u8,
37        }
38    }
39    pub fn write_u8_bits(&mut self, val: u8, num_bits: u8) -> Result<(), BitsError> {
40        self.write_u32_bits(val as u32, num_bits)
41    }
42    pub fn write_u16_bits(&mut self, val: u16, num_bits: u8) -> Result<(), BitsError> {
43        self.write_u32_bits(val as u32, num_bits)
44    }
45    pub fn write_u32_bits(&mut self, val: u32, mut num_bits: u8) -> Result<(), BitsError> {
46        if num_bits > 32 {
47            return Err(BitsErrorKind::InvalidInput.into());
48        }
49        while num_bits > 0 {
50            match num_bits.cmp(&self.remaining) {
51                Ordering::Less => {
52                    let shift = self.remaining - num_bits;
53                    let mask = (1u32 << num_bits) - 1;
54                    self.buf |= (val & mask) << shift;
55                    self.remaining -= num_bits;
56                    num_bits = 0;
57                }
58                Ordering::Equal => {
59                    let mask = (1u32 << num_bits) - 1;
60                    self.buf |= val & mask;
61                    num_bits = 0;
62                    self.delegate.write_be_u32(self.buf)?;
63                    self.remaining = 32;
64                    self.buf = 0;
65                }
66                Ordering::Greater => {
67                    let touse = self.remaining;
68                    let shift = num_bits - self.remaining;
69                    let mask = (1u32 << touse) - 1;
70                    self.buf |= (val >> shift) & mask;
71                    self.delegate.write_be_u32(self.buf)?;
72                    self.remaining = 32;
73                    self.buf = 0;
74                    num_bits -= touse;
75                }
76            }
77        }
78        Ok(())
79    }
80}
81
82pub struct BitStreamDecoder<'a, T: Bits> {
83    delegate: BitsWrapper<'a, T>,
84    buf: u32,
85    used: u8,
86}
87impl<'a, T: Bits> BitStreamDecoder<'a, T> {
88    pub fn new(delegate: BitsWrapper<'a, T>) -> Self {
89        Self {
90            delegate,
91            buf: 0,
92            used: 0,
93        }
94    }
95    pub fn read_u32_bits(&mut self, num_bits: u8) -> Result<u32, BitsError> {
96        if num_bits > 32 {
97            return Err(BitsErrorKind::InvalidInput.into());
98        }
99        loop {
100            match self.used.cmp(&num_bits) {
101                Ordering::Less => {
102                    // used < numbits - add more.
103                    let v = self.delegate.read_u8()?;
104                    self.buf = (self.buf << 8) | v as u32;
105                    self.used += 8;
106                }
107                Ordering::Equal => {
108                    let mask = (1u32 << num_bits) - 1;
109                    self.used = 0;
110                    let b = self.buf & mask;
111                    self.buf = 0;
112                    return Ok(b);
113                }
114                Ordering::Greater => {
115                    let rem = self.used - num_bits;
116                    let mask = (1u32 << num_bits) - 1;
117                    let b = (self.buf >> rem) & mask;
118                    self.used -= num_bits;
119                    return Ok(b);
120                }
121            }
122        }
123    }
124}
125#[cfg(all(test, feature = "std"))]
126mod test {
127    use crate::{BitStreamDecoder, BitStreamEncoder, BitsError, BitsWrapper};
128
129    #[test]
130    pub fn test_dec() -> Result<(), BitsError> {
131        let buf = vec![0xAB, 0xCD, 0xAB, 0xCD];
132        let mut dec = BitStreamDecoder::new(BitsWrapper::Owned(buf));
133        assert_eq!(0xA, dec.read_u32_bits(4)?);
134        assert_eq!(0xB, dec.read_u32_bits(4)?);
135        assert_eq!(0xC, dec.read_u32_bits(4)?);
136        assert_eq!(0xD, dec.read_u32_bits(4)?);
137        assert_eq!(0xABCD, dec.read_u32_bits(16)?);
138        Ok(())
139    }
140
141    #[test]
142    pub fn test_dec2() -> Result<(), BitsError> {
143        let buf = vec![0x03, 0xC0, 0x81, 0x00, 0x88, 0x10, 0x1A, 0x02];
144        let mut dec = BitStreamDecoder::new(BitsWrapper::Owned(buf));
145        assert_eq!(7, dec.read_u32_bits(9)?);
146        assert_eq!(258, dec.read_u32_bits(9)?);
147        assert_eq!(8, dec.read_u32_bits(9)?);
148        assert_eq!(8, dec.read_u32_bits(9)?);
149        assert_eq!(258, dec.read_u32_bits(9)?);
150        assert_eq!(6, dec.read_u32_bits(9)?);
151        assert_eq!(257, dec.read_u32_bits(9)?);
152
153        Ok(())
154    }
155
156    #[test]
157    pub fn test_enc() -> Result<(), BitsError> {
158        let mut buf = Vec::<u8>::new();
159        {
160            let wrap = BitsWrapper::Borrowed(&mut buf);
161            let mut enc = BitStreamEncoder::new(wrap);
162
163            enc.write_u16_bits(0xAAAA, 4)?;
164            enc.write_u16_bits(0xBBBB, 4)?;
165            enc.write_u16_bits(0xCCCC, 4)?;
166            enc.write_u16_bits(0xDDDD, 4)?;
167            enc.write_u16_bits(0xABCD, 16)?;
168        }
169        // println!("{:?}", buf);
170        assert_eq!(buf, [0xAB, 0xCD, 0xAB, 0xCD]);
171        Ok(())
172    }
173
174    #[test]
175    pub fn test_enc2() -> Result<(), BitsError> {
176        let mut buf = Vec::<u8>::new();
177        {
178            let wrap = BitsWrapper::Borrowed(&mut buf);
179            let mut enc = BitStreamEncoder::new(wrap);
180
181            enc.write_u16_bits(7, 9)?;
182            enc.write_u16_bits(258, 9)?;
183            enc.write_u16_bits(8, 9)?;
184            enc.write_u16_bits(8, 9)?;
185            enc.write_u16_bits(258, 9)?;
186            enc.write_u16_bits(6, 9)?;
187            enc.write_u16_bits(257, 9)?;
188        }
189        assert_eq!(buf, [0x03, 0xC0, 0x81, 0x00, 0x88, 0x10, 0x1A, 0x02]);
190        // 0x007  0b000000111
191        // 0x102             100000010
192        // 0x008                      000001000
193        // 0x008                               000001000
194        // 0x102                                        100000010
195        // 0x006                                                 000000011
196        // 0x03 = 0b00000011
197        // 0xC0 =           11000000
198        // 0x81 =                   10000001
199        // 0x00 =                           00000000
200        // 0x88 =                                   10001000
201        // 0x10 =                                           00010000
202
203        Ok(())
204    }
205}
irox_bits/bitstream.rs

irox_bits/
bitstream.rs
