msg_tool\utils/
bit_stream.rs1use crate::ext::io::*;
3use anyhow::Result;
4use std::io::{Read, Write};
5
6pub struct MsbBitStream<T: Read> {
8 pub m_input: T,
10 m_bits: u32,
11 pub m_cached_bits: u32,
13}
14
15impl<T: Read> MsbBitStream<T> {
16 pub fn new(input: T) -> Self {
18 MsbBitStream {
19 m_input: input,
20 m_bits: 0,
21 m_cached_bits: 0,
22 }
23 }
24
25 pub fn get_bits(&mut self, count: u32) -> Result<u32> {
27 while self.m_cached_bits < count {
28 let byte = self.m_input.read_u8()?;
29 self.m_bits = (self.m_bits << 8) | byte as u32;
30 self.m_cached_bits += 8;
31 }
32 let mask = (1 << count) - 1;
33 self.m_cached_bits -= count;
34 let result = (self.m_bits >> self.m_cached_bits) & mask;
35 Ok(result)
36 }
37
38 pub fn get_next_bit(&mut self) -> Result<bool> {
40 if self.m_cached_bits == 0 {
41 let byte = self.m_input.read_u8()?;
42 self.m_bits = (self.m_bits << 8) | byte as u32;
43 self.m_cached_bits += 8;
44 }
45 self.m_cached_bits -= 1;
46 let bit = (self.m_bits >> self.m_cached_bits) & 1 != 0;
47 Ok(bit)
48 }
49}
50
51pub struct MsbBitWriter<'a, T: Write> {
53 pub writer: &'a mut T,
55 buffer: u32,
56 buffer_size: u32,
57}
58
59impl<'a, T: Write> MsbBitWriter<'a, T> {
60 pub fn new(writer: &'a mut T) -> Self {
62 MsbBitWriter {
63 writer,
64 buffer: 0,
65 buffer_size: 0,
66 }
67 }
68
69 pub fn flush(&mut self) -> Result<()> {
72 if self.buffer_size > 0 {
73 self.writer
74 .write_u8(((self.buffer << (8 - self.buffer_size)) & 0xFF) as u8)?;
75 self.buffer = 0;
76 self.buffer_size = 0;
77 }
78 Ok(())
79 }
80
81 pub fn put_bits(&mut self, byte: u32, token_width: u8) -> Result<()> {
83 for i in 0..token_width {
84 self.put_bit((byte & (1 << (token_width - 1 - i))) != 0)?;
85 }
86 Ok(())
87 }
88
89 pub fn put_bit(&mut self, bit: bool) -> Result<()> {
91 self.buffer <<= 1;
92 if bit {
93 self.buffer |= 1;
94 }
95 self.buffer_size += 1;
96 if self.buffer_size == 8 {
97 self.writer.write_u8((self.buffer & 0xFF) as u8)?;
98 self.buffer_size -= 8;
99 }
100 Ok(())
101 }
102}
103
104pub struct LsbBitStream<T: Read> {
106 pub m_input: T,
108 m_bits: u32,
109 pub m_cached_bits: u32,
111}
112
113impl<T: Read> LsbBitStream<T> {
114 pub fn new(input: T) -> Self {
116 LsbBitStream {
117 m_input: input,
118 m_bits: 0,
119 m_cached_bits: 0,
120 }
121 }
122
123 pub fn get_bits(&mut self, mut count: u32) -> Result<u32> {
125 if self.m_cached_bits >= count {
126 let mask = (1 << count) - 1;
127 let value = self.m_bits & mask;
128 self.m_bits >>= count;
129 self.m_cached_bits -= count;
130 Ok(value)
131 } else {
132 let mut value = self.m_bits & ((1 << self.m_cached_bits) - 1);
133 count -= self.m_cached_bits;
134 let mut shift = self.m_cached_bits;
135 self.m_cached_bits = 0;
136 while count >= 8 {
137 let b = self.m_input.read_u8()?;
138 value |= (b as u32) << shift;
139 shift += 8;
140 count -= 8;
141 }
142 if count > 0 {
143 let b = self.m_input.read_u8()?;
144 value |= ((b as u32) & ((1 << count) - 1)) << shift;
145 self.m_bits = b as u32 >> count;
146 self.m_cached_bits = 8 - count;
147 }
148 Ok(value)
149 }
150 }
151
152 pub fn get_next_bit(&mut self) -> Result<bool> {
154 Ok(self.get_bits(1)? == 1)
155 }
156}