blackhole_compress/
stream.rs

1use std::io::{Read, Write};
2use crate::error::{ReadError, WriteError};
3
4pub struct StreamReader<'a,R> where R: Read +'a {
5    reader:&'a mut R,
6    buf: [u8; 256],
7    buf_size: usize,
8    current_index:usize,
9    current_bits:usize
10}
11impl<'a,R> StreamReader<'a,R> where R: Read +'a {
12    pub fn new(reader:&'a mut R) -> StreamReader<'a,R> {
13        StreamReader {
14            reader: reader,
15            buf: [0;256],
16            buf_size: 0,
17            current_index:0,
18            current_bits:0
19        }
20    }
21
22    fn read_next(&mut self) -> Result<usize,ReadError> {
23        self.buf_size = self.reader.read(&mut self.buf)?;
24
25        self.current_index = 0;
26
27        Ok(self.buf_size)
28    }
29    
30    pub fn get_bit_from_lsb(&mut self) -> Result<u8,ReadError> {
31        if self.buf_size == 0 || self.current_index >= self.buf.len() {
32            if self.read_next()? == 0 {
33                Err(ReadError::UnexpectedEofError)
34            } else {
35                self.get_bit_from_lsb()
36            }
37        } else {
38            let bit = if self.buf[self.current_index] & (1u8 << self.current_bits as u8) == 0 {
39                0
40            } else {
41                1
42            };
43
44            self.current_bits += 1;
45
46            if self.current_bits >= 8 {
47                self.current_index += 1;
48                self.current_bits = 0;
49            }
50
51            Ok(bit)
52        }
53    }
54
55    pub fn get_bits_from_lsb(&mut self, size:usize) -> Result<u8,ReadError> {
56        if size > 8 {
57            Err(ReadError::InvalidArgumentError(String::from("The specified read size is too long.")))
58        } else {
59            let mut bits = 0;
60
61            for i in 0..size {
62                bits |= self.get_bit_from_lsb()? << i;
63            }
64
65            Ok(bits)
66        }
67    }
68
69    pub fn read_once(&mut self) -> Result<Option<u8>,ReadError> {
70        if self.buf_size == 0 || self.current_index >= self.buf_size {
71            if self.read_next()? == 0 {
72                return Ok(None);
73            }
74        }
75
76        if self.current_bits == 0 {
77            let r = self.buf[self.current_index];
78
79            self.current_index += 1;
80
81            Ok(Some(r))
82        } else {
83            self.get_bits_from_lsb(8).map(|b| Some(b))
84        }
85    }
86
87    pub fn read_u8(&mut self) -> Result<u8,ReadError> {
88        Ok(self.read_once()?.ok_or(ReadError::UnexpectedEofError)?)
89    }
90
91    pub fn read_u16(&mut self) -> Result<u16,ReadError> {
92        Ok(self.read_once()?.ok_or(ReadError::UnexpectedEofError)? as u16 |
93               (self.read_once()?.ok_or(ReadError::UnexpectedEofError)? as u16) << 8
94        )
95    }
96
97    pub fn read_u32(&mut self) -> Result<u32,ReadError> {
98        Ok(self.read_once()?.ok_or(ReadError::UnexpectedEofError)? as u32 |
99           (self.read_once()?.ok_or(ReadError::UnexpectedEofError)? as u32) <<  8 |
100           (self.read_once()?.ok_or(ReadError::UnexpectedEofError)? as u32) << 16 |
101           (self.read_once()?.ok_or(ReadError::UnexpectedEofError)? as u32) << 24
102        )
103    }
104
105    pub fn read_u64(&mut self) -> Result<u64,ReadError> {
106        Ok((self.read_u32()? as u64) | (self.read_u32()? as u64) << 32)
107    }
108
109    pub fn read_until(&mut self,size:usize) -> Result<Vec<u8>,ReadError> {
110        let mut r = Vec::with_capacity(size);
111
112        for _ in 0..size {
113            r.push(self.read_once()?.ok_or(ReadError::UnexpectedEofError)?);
114        }
115
116        Ok(r)
117    }
118
119    pub fn skip_bits(&mut self,count:usize) -> Result<(),ReadError> {
120        self.current_index += count / 8;
121        self.current_bits += count % 8;
122
123        Ok(())
124    }
125
126    pub fn skip_surplus_bits(&mut self) {
127        self.current_index += 1;
128        self.current_bits = 0;
129    }
130}
131pub struct StreamWriter<'a,W> where W: Write +'a {
132    writer:&'a mut W,
133    buf: [u8; 256],
134    current_index:usize,
135    current_bits:usize,
136    written_size:usize
137}
138impl<'a,W> StreamWriter<'a,W> where W: Write +'a {
139    pub fn new(writer:&'a mut W) -> StreamWriter<'a,W> {
140        StreamWriter {
141            writer:writer,
142            buf: [0;256],
143            current_index:0,
144            current_bits:0,
145            written_size:0
146        }
147    }
148
149    pub fn flush(&mut self) -> Result<(),WriteError> {
150        let mut slice = &mut self.buf[0..self.current_index];
151        let mut start_index = 0;
152
153        while start_index < self.current_index {
154            let size = self.writer.write(&mut slice)?;
155
156            if size == 0 {
157                return Err(WriteError::InvalidState(String::from("An error occurred in writing data.")))
158            }
159
160            self.written_size += size;
161
162            start_index += size;
163
164            slice = &mut slice[size..];
165        }
166
167        if self.current_bits == 0 {
168            self.current_index = 0;
169            self.current_bits = 0;
170        } else {
171            self.buf[0] = self.buf[self.current_index];
172            self.current_index = 0;
173        };
174
175        Ok(())
176    }
177
178    pub fn write(&mut self, b:u8) -> Result<(),WriteError> {
179        if self.current_index >= 256 {
180            self.flush()?;
181        }
182
183        if self.current_bits == 0 {
184            self.buf[self.current_index] = b;
185            self.current_index += 1;
186
187            Ok(())
188        } else {
189            self.write_u8_to_bits(b)
190        }
191    }
192
193    pub fn write_bit(&mut self,b:bool) -> Result<(),WriteError> {
194        if self.current_index >= 256 {
195            self.flush()?;
196        }
197
198        if self.current_bits == 0 {
199            self.buf[self.current_index] = 0;
200        }
201
202        if b {
203            self.buf[self.current_index] |= 1u8 << self.current_bits;
204        }
205
206        self.current_bits += 1;
207
208        if self.current_bits >= 8 {
209            self.current_index += 1;
210            self.current_bits = 0;
211        }
212
213        Ok(())
214    }
215
216    pub fn write_u16(&mut self,value:u16) -> Result<(),WriteError> {
217        self.write((value & 0xFF) as u8)?;
218        self.write(((value >> 8) & 0xFF) as u8)?;
219
220        Ok(())
221    }
222
223    pub fn write_u32(&mut self,value:u32) -> Result<(),WriteError> {
224        self.write((value & 0xFF) as u8)?;
225        self.write(((value >> 8) & 0xFF) as u8)?;
226        self.write(((value >> 16) & 0xFF) as u8)?;
227        self.write(((value >> 24) & 0xFF) as u8)?;
228
229        Ok(())
230    }
231
232    pub fn write_u64(&mut self,value:u64) -> Result<(),WriteError> {
233        self.write_u32((value & 0xFFFFFFFF) as u32)?;
234        self.write_u32(((value >> 32) & 0xFFFFFFFF) as u32)?;
235
236        Ok(())
237    }
238
239    fn write_u8_to_bits(&mut self,value:u8) -> Result<(),WriteError> {
240        for i in 0..8 {
241            if value & (1u8 << i) != 0 {
242                self.write_bit(true)?;
243            } else {
244                self.write_bit(false)?;
245            }
246        }
247
248        Ok(())
249    }
250
251    pub fn write_bits(&mut self,value:u64,len:usize) -> Result<(),WriteError> {
252        for i in 0..len {
253            if value & (1u64 << i) != 0 {
254                self.write_bit(true)?;
255            } else {
256                self.write_bit(false)?;
257            }
258        }
259
260        Ok(())
261    }
262
263    pub fn pad_zeros(&mut self) -> Result<(),WriteError> {
264        if self.current_bits > 0 {
265            self.current_bits = 0;
266            self.current_index += 1;
267        }
268
269        Ok(())
270    }
271
272    pub fn write_bytes(&mut self,bytes:&[u8]) -> Result<(),WriteError> {
273        for &b in bytes.into_iter() {
274            self.write(b)?;
275        }
276
277        Ok(())
278    }
279
280    pub fn written_size(&self) -> usize {
281        self.written_size
282    }
283}