compression/bitio/
reader.rs

1//! rust-compression
2//!
3//! # Licensing
4//! This Source Code is subject to the terms of the Mozilla Public License
5//! version 2.0 (the "License"). You can obtain a copy of the License at
6//! <http://mozilla.org/MPL/2.0/>.
7
8use crate::bitio::direction::Direction;
9use crate::bitio::small_bit_vec::SmallBitVec;
10use crate::cbuffer::CircularBuffer;
11use crate::core::cmp;
12use crate::core::iter::Iterator;
13use crate::core::marker::PhantomData;
14use crate::core::mem::size_of;
15use crate::core::ops::{BitOrAssign, Shl, Shr};
16#[cfg(not(feature = "std"))]
17use alloc::borrow::ToOwned;
18#[cfg(not(feature = "std"))]
19use alloc::string::String;
20#[cfg(not(feature = "std"))]
21#[allow(unused_imports)]
22use alloc::{vec, vec::Vec};
23use num_traits::sign::Unsigned;
24
25pub(crate) trait BitRead
26where
27    Self::Direction: Direction,
28{
29    type Direction;
30
31    fn peek_bits<T: Unsigned, R: Iterator<Item = u8>>(
32        &mut self,
33        len: usize,
34        iter: &mut R,
35    ) -> Result<SmallBitVec<T>, String>
36    where
37        T: BitOrAssign
38            + Shl<usize, Output = T>
39            + Shr<usize, Output = T>
40            + From<u8>;
41
42    fn skip_bits<R: Iterator<Item = u8>>(
43        &mut self,
44        len: usize,
45        iter: &mut R,
46    ) -> Result<usize, String>;
47    fn read_bits<T: Unsigned, R: Iterator<Item = u8>>(
48        &mut self,
49        len: usize,
50        iter: &mut R,
51    ) -> Result<SmallBitVec<T>, String>
52    where
53        T: BitOrAssign
54            + Shl<usize, Output = T>
55            + Shr<usize, Output = T>
56            + From<u8>;
57
58    fn skip_to_next_byte(&mut self) -> usize;
59}
60
61#[derive(Clone, Debug)]
62pub(crate) struct BitReader<D: Direction> {
63    buf: u8,
64    counter: usize,
65    cbuf: CircularBuffer<u8>,
66    pos: usize,
67    phantom: PhantomData<fn() -> D>,
68}
69
70impl<D: Direction> BitRead for BitReader<D> {
71    type Direction = D;
72
73    fn peek_bits<T: Unsigned, R: Iterator<Item = u8>>(
74        &mut self,
75        len: usize,
76        iter: &mut R,
77    ) -> Result<SmallBitVec<T>, String>
78    where
79        T: BitOrAssign
80            + Shl<usize, Output = T>
81            + Shr<usize, Output = T>
82            + From<u8>,
83    {
84        let firstlen = cmp::min(len, self.counter);
85        let needlen = (len - firstlen + 7) >> 3;
86
87        if needlen > 0 {
88            // バッファに読み込む
89            if needlen > self.pos {
90                if needlen + self.pos > self.buffer_cap() {
91                    return Err("len is too long".to_owned());
92                }
93                let rbuf = iter.take(needlen - self.pos).collect::<Vec<u8>>();
94                self.cbuf.append(&rbuf);
95                self.pos += rbuf.len();
96            }
97            let mut ret = Self::conv_u8_to_t(self.buf);
98            let mut count = self.counter;
99            for i in (0..cmp::min(self.pos, needlen))
100                .map(|x| Self::conv_u8_to_t(self.cbuf[self.pos - x - 1]))
101            {
102                ret |= D::backward(i, count);
103                count += size_of::<u8>() << 3;
104            }
105            let retlen = cmp::min(count, len);
106            Ok(if retlen != 0 {
107                SmallBitVec::new(
108                    D::convert(ret, size_of::<T>() << 3, retlen),
109                    retlen,
110                )
111            } else {
112                SmallBitVec::new(T::zero(), 0)
113            })
114        } else {
115            Ok(SmallBitVec::new(
116                D::convert(T::from(self.buf), size_of::<u8>() << 3, firstlen),
117                firstlen,
118            ))
119        }
120    }
121
122    fn skip_bits<R: Iterator<Item = u8>>(
123        &mut self,
124        len: usize,
125        iter: &mut R,
126    ) -> Result<usize, String> {
127        let firstlen = cmp::min(len, self.counter);
128        let midlen = (len - firstlen) >> 3;
129        let mut lastlen = (len - firstlen) & 0x07;
130
131        if lastlen > 0 || midlen > 0 {
132            // バッファに読み込む
133            let needlen = midlen + if lastlen > 0 { 1 } else { 0 };
134            if needlen > self.pos {
135                if needlen + self.pos > self.buffer_cap() {
136                    return Err("len is too long".to_owned());
137                }
138                let rbuf = iter.take(needlen - self.pos).collect::<Vec<u8>>();
139                self.cbuf.append(&rbuf);
140                self.pos += rbuf.len();
141            }
142            if midlen == self.pos {
143                lastlen = 0;
144            }
145            let readlen = cmp::min(self.pos, needlen);
146            self.pos -= readlen;
147            if lastlen > 0 {
148                self.buf = D::forward(self.cbuf[self.pos], lastlen);
149                self.counter = (size_of::<u8>() << 3) - lastlen;
150            } else {
151                self.buf = 0;
152                self.counter = 0;
153            }
154            Ok(firstlen + cmp::min((midlen << 3) + lastlen, readlen << 3))
155        } else {
156            self.buf = D::forward(self.buf, firstlen);
157            self.counter -= firstlen;
158            Ok(firstlen)
159        }
160    }
161
162    fn read_bits<T: Unsigned, R: Iterator<Item = u8>>(
163        &mut self,
164        len: usize,
165        iter: &mut R,
166    ) -> Result<SmallBitVec<T>, String>
167    where
168        T: BitOrAssign
169            + Shl<usize, Output = T>
170            + Shr<usize, Output = T>
171            + From<u8>,
172    {
173        let r = self.peek_bits::<T, R>(len, iter);
174        if let Ok(ref l) = r {
175            let _ = self.skip_bits::<_>(l.len(), iter)?;
176        }
177        r
178    }
179
180    fn skip_to_next_byte(&mut self) -> usize {
181        let len = self.counter;
182        self.buf = 0;
183        self.counter = 0;
184        len
185    }
186}
187
188const DEFAULT_BUF_SIZE: usize = 8; // u64まで対応可能
189
190impl<D: Direction> BitReader<D> {
191    #[inline]
192    pub(crate) fn new() -> Self {
193        Self::with_capacity(DEFAULT_BUF_SIZE)
194    }
195
196    #[inline]
197    pub(crate) fn with_capacity(cap: usize) -> Self {
198        Self {
199            buf: 0,
200            counter: 0,
201            cbuf: CircularBuffer::<u8>::new(cap),
202            pos: 0,
203            phantom: PhantomData,
204        }
205    }
206
207    #[inline]
208    pub(crate) fn buffer_cap(&self) -> usize {
209        self.cbuf.cap() - self.pos
210    }
211
212    #[inline]
213    fn conv_u8_to_t<T: Unsigned>(value: u8) -> T
214    where
215        T: Shl<usize, Output = T> + Shr<usize, Output = T> + From<u8>,
216    {
217        D::convert(T::from(value), size_of::<u8>() << 3, size_of::<T>() << 3)
218    }
219}
220
221impl<D: Direction> Default for BitReader<D> {
222    #[inline]
223    fn default() -> Self {
224        Self::new()
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use crate::action::Action;
232    use crate::bitio::direction::left::Left;
233    use crate::bitio::direction::right::Right;
234    use crate::bitio::writer::{BitWriteExt, BitWriter};
235
236    #[test]
237    fn leftbitreader_read() {
238        let cursor = vec![0b1100_1100];
239        let mut iter = cursor.into_iter();
240
241        let mut reader = BitReader::<Left>::new();
242
243        assert_eq!(
244            reader.read_bits::<u32, _>(1, &mut iter).ok(),
245            Some(SmallBitVec::new(0b1, 1))
246        );
247        assert_eq!(
248            reader.read_bits::<u32, _>(2, &mut iter).ok(),
249            Some(SmallBitVec::new(0b10, 2))
250        );
251        assert_eq!(
252            reader.read_bits::<u32, _>(3, &mut iter).ok(),
253            Some(SmallBitVec::new(0b011, 3))
254        );
255        assert_eq!(
256            reader.read_bits::<u32, _>(2, &mut iter).ok(),
257            Some(SmallBitVec::new(0b00, 2))
258        );
259    }
260
261    #[test]
262    fn leftbitreader_readmulti() {
263        let cursor = vec![243, 221, 190, 200];
264        let mut iter = cursor.into_iter();
265
266        let mut reader = BitReader::<Left>::new();
267
268        assert_eq!(
269            reader.read_bits::<u32, _>(10, &mut iter).ok(),
270            Some(SmallBitVec::new(975, 10))
271        );
272        assert_eq!(
273            reader.read_bits::<u32, _>(10, &mut iter).ok(),
274            Some(SmallBitVec::new(475, 10))
275        );
276        assert_eq!(
277            reader.read_bits::<u32, _>(12, &mut iter).ok(),
278            Some(SmallBitVec::new(3784, 12))
279        );
280    }
281
282    #[test]
283    fn rightbitreader_read() {
284        let cursor = vec![0b1100_1100];
285        let mut iter = cursor.into_iter();
286
287        let mut reader = BitReader::<Right>::new();
288
289        assert_eq!(
290            reader.read_bits::<u32, _>(1, &mut iter).ok(),
291            Some(SmallBitVec::new(0b0, 1))
292        );
293        assert_eq!(
294            reader.read_bits::<u32, _>(2, &mut iter).ok(),
295            Some(SmallBitVec::new(0b10, 2))
296        );
297        assert_eq!(
298            reader.read_bits::<u32, _>(3, &mut iter).ok(),
299            Some(SmallBitVec::new(0b0001, 3))
300        );
301        assert_eq!(
302            reader.read_bits::<u32, _>(2, &mut iter).ok(),
303            Some(SmallBitVec::new(0b11, 2))
304        );
305    }
306
307    #[test]
308    fn rightbitreader_multi() {
309        let cursor = vec![0xCF, 0x6F, 0x87, 0xEC];
310
311        let mut reader = BitReader::<Right>::new();
312        let mut iter = cursor.into_iter();
313
314        assert_eq!(
315            reader.read_bits::<u32, _>(10, &mut iter).ok(),
316            Some(SmallBitVec::new(975, 10))
317        );
318        assert_eq!(
319            reader.read_bits::<u32, _>(10, &mut iter).ok(),
320            Some(SmallBitVec::new(475, 10))
321        );
322        assert_eq!(
323            reader.read_bits::<u32, _>(12, &mut iter).ok(),
324            Some(SmallBitVec::new(3784, 12))
325        );
326    }
327
328    #[test]
329    fn leftbitreader_peek() {
330        let mut writer = BitWriter::<Left>::new();
331        let mut ret = vec![
332            SmallBitVec::new(0b1_u32, 1),
333            SmallBitVec::new(0b10, 2),
334            SmallBitVec::new(0b011, 3),
335            SmallBitVec::new(0b00, 2),
336        ]
337        .to_bytes(&mut writer, Action::Flush);
338
339        let mut reader = BitReader::<Left>::new();
340
341        assert_eq!(
342            reader.peek_bits::<u32, _>(1, &mut ret).ok(),
343            Some(SmallBitVec::new(0b1, 1))
344        );
345        assert_eq!(
346            reader.read_bits::<u32, _>(1, &mut ret).ok(),
347            Some(SmallBitVec::new(0b1, 1))
348        );
349        assert_eq!(
350            reader.peek_bits::<u32, _>(2, &mut ret).ok(),
351            Some(SmallBitVec::new(0b10, 2))
352        );
353        assert_eq!(
354            reader.read_bits::<u32, _>(2, &mut ret).ok(),
355            Some(SmallBitVec::new(0b10, 2))
356        );
357        assert_eq!(
358            reader.peek_bits::<u32, _>(3, &mut ret).ok(),
359            Some(SmallBitVec::new(0b011, 3))
360        );
361        assert_eq!(
362            reader.read_bits::<u32, _>(3, &mut ret).ok(),
363            Some(SmallBitVec::new(0b011, 3))
364        );
365        assert_eq!(
366            reader.peek_bits::<u32, _>(2, &mut ret).ok(),
367            Some(SmallBitVec::new(0b00, 2))
368        );
369        assert_eq!(
370            reader.read_bits::<u32, _>(2, &mut ret).ok(),
371            Some(SmallBitVec::new(0b00, 2))
372        );
373    }
374
375    #[test]
376    fn leftbitreader_peek_big() {
377        let mut writer = BitWriter::<Left>::new();
378        let mut ret = vec![
379            SmallBitVec::new(975_u32, 10),
380            SmallBitVec::new(475, 10),
381            SmallBitVec::new(3784, 12),
382        ]
383        .to_bytes(&mut writer, Action::Flush);
384
385        let mut reader = BitReader::<Left>::new();
386
387        assert_eq!(
388            reader.peek_bits::<u32, _>(10, &mut ret).ok(),
389            Some(SmallBitVec::new(975, 10))
390        );
391        assert_eq!(
392            reader.read_bits::<u32, _>(10, &mut ret).ok(),
393            Some(SmallBitVec::new(975, 10))
394        );
395        assert_eq!(
396            reader.peek_bits::<u32, _>(10, &mut ret).ok(),
397            Some(SmallBitVec::new(475, 10))
398        );
399        assert_eq!(
400            reader.read_bits::<u32, _>(10, &mut ret).ok(),
401            Some(SmallBitVec::new(475, 10))
402        );
403        assert_eq!(
404            reader.peek_bits::<u32, _>(15, &mut ret).ok(),
405            Some(SmallBitVec::new(3784, 12))
406        );
407        assert_eq!(
408            reader.read_bits::<u32, _>(15, &mut ret).ok(),
409            Some(SmallBitVec::new(3784, 12))
410        );
411    }
412
413    #[test]
414    fn leftbitreader_zeros() {
415        let mut writer = BitWriter::<Left>::new();
416        let mut ret = vec![
417            SmallBitVec::new(32_u32, 16),
418            SmallBitVec::new(8, 5),
419            SmallBitVec::new(0, 3),
420            SmallBitVec::new(1, 3),
421            SmallBitVec::new(0, 3),
422            SmallBitVec::new(3, 2),
423            SmallBitVec::new(0, 3),
424        ]
425        .to_bytes(&mut writer, Action::Flush);
426
427        let mut reader = BitReader::<Left>::new();
428
429        assert_eq!(
430            reader.read_bits::<u32, _>(16, &mut ret).ok(),
431            Some(SmallBitVec::new(32, 16))
432        );
433        assert_eq!(
434            reader.read_bits::<u32, _>(5, &mut ret).ok(),
435            Some(SmallBitVec::new(8, 5))
436        );
437        assert_eq!(
438            reader.read_bits::<u32, _>(3, &mut ret).ok(),
439            Some(SmallBitVec::new(0, 3))
440        );
441        assert_eq!(
442            reader.read_bits::<u32, _>(3, &mut ret).ok(),
443            Some(SmallBitVec::new(1, 3))
444        );
445        assert_eq!(
446            reader.read_bits::<u32, _>(3, &mut ret).ok(),
447            Some(SmallBitVec::new(0, 3))
448        );
449        assert_eq!(
450            reader.read_bits::<u32, _>(2, &mut ret).ok(),
451            Some(SmallBitVec::new(3, 2))
452        );
453        assert_eq!(
454            reader.read_bits::<u32, _>(3, &mut ret).ok(),
455            Some(SmallBitVec::new(0, 3))
456        );
457    }
458
459    #[test]
460    fn leftbitreader_skip() {
461        let mut writer = BitWriter::<Left>::new();
462        let mut ret = vec![
463            SmallBitVec::new(0b1_u32, 1),
464            SmallBitVec::new(0b10, 2),
465            SmallBitVec::new(0b011, 3),
466            SmallBitVec::new(0b00, 2),
467        ]
468        .to_bytes(&mut writer, Action::Flush);
469
470        let mut reader = BitReader::<Left>::new();
471
472        assert_eq!(
473            reader.peek_bits::<u32, _>(1, &mut ret).ok(),
474            Some(SmallBitVec::new(0b1, 1))
475        );
476        assert_eq!(reader.skip_bits::<_>(1, &mut ret).ok(), Some(1));
477        assert_eq!(
478            reader.peek_bits::<u32, _>(2, &mut ret).ok(),
479            Some(SmallBitVec::new(0b10, 2))
480        );
481        assert_eq!(reader.skip_bits::<_>(2, &mut ret).ok(), Some(2));
482        assert_eq!(
483            reader.peek_bits::<u32, _>(3, &mut ret).ok(),
484            Some(SmallBitVec::new(0b011, 3))
485        );
486        assert_eq!(reader.skip_bits::<_>(3, &mut ret).ok(), Some(3));
487        assert_eq!(
488            reader.peek_bits::<u32, _>(2, &mut ret).ok(),
489            Some(SmallBitVec::new(0b00, 2))
490        );
491        assert_eq!(reader.skip_to_next_byte(), 2);
492    }
493
494    #[test]
495    fn leftbitreader_skip_big() {
496        let mut writer = BitWriter::<Left>::new();
497        let mut ret = vec![
498            SmallBitVec::new(975_u32, 10),
499            SmallBitVec::new(475, 10),
500            SmallBitVec::new(3784, 12),
501        ]
502        .to_bytes(&mut writer, Action::Flush);
503
504        let mut reader = BitReader::<Left>::new();
505
506        assert_eq!(
507            reader.peek_bits::<u32, _>(10, &mut ret).ok(),
508            Some(SmallBitVec::new(975, 10))
509        );
510        assert_eq!(reader.skip_bits::<_>(20, &mut ret).ok(), Some(20));
511        assert_eq!(
512            reader.peek_bits::<u32, _>(15, &mut ret).ok(),
513            Some(SmallBitVec::new(3784, 12))
514        );
515        assert_eq!(reader.skip_to_next_byte(), 4);
516        assert_eq!(
517            reader.peek_bits::<u32, _>(15, &mut ret).ok(),
518            Some(SmallBitVec::new(200, 8))
519        );
520    }
521}