sit_algos/
huffman.rs

1use std::io::{self, Read, Seek};
2
3use bitstream_io::{BigEndian, BitRead, BitReader};
4
5use super::huffman_decoder::HuffmanDecoder;
6
7#[derive(thiserror::Error, Debug)]
8pub enum Error {
9    #[error("Invalid huffman tree")]
10    InvalidTree,
11
12    #[error(transparent)]
13    Io(#[from] io::Error),
14}
15
16pub struct HuffmanReader<R: Read + Seek> {
17    inner: bitstream_io::BitReader<R, BigEndian>,
18    uncompressed_size: u64,
19    offset: u64,
20    tree: HuffmanDecoder,
21}
22
23impl<R: Read + Seek> HuffmanReader<R> {
24    pub fn try_from(inner: R, uncompressed_size: u64) -> Result<Self, Error> {
25        let mut inner = bitstream_io::BitReader::new(inner);
26        let tree = Self::read_hufftree(&mut inner)?;
27
28        Ok(Self {
29            inner,
30            uncompressed_size,
31            offset: 0,
32            tree,
33        })
34    }
35
36    fn read_hufftree(
37        reader: &mut bitstream_io::BitReader<R, BigEndian>,
38    ) -> Result<HuffmanDecoder, Error> {
39        let mut tree = HuffmanDecoder::new();
40        Self::read_hufftree_node(reader, &mut tree, 0)?;
41        tree.make_table(false);
42        Ok(tree)
43    }
44
45    #[inline]
46    fn read_hufftree_node(
47        reader: &mut BitReader<R, BigEndian>,
48        tree: &mut HuffmanDecoder,
49        node: i32,
50    ) -> Result<(), Error> {
51        match reader.read_bit()? {
52            true => {
53                let value: u8 = reader.read_var(8)?;
54                tree.set_leaf_value(node, value as i32);
55            }
56            false => {
57                let left = tree.new_node();
58                Self::read_hufftree_node(reader, tree, left)?;
59
60                let right = tree.new_node();
61                Self::read_hufftree_node(reader, tree, right)?;
62
63                tree.set_branch(node, false, left);
64                tree.set_branch(node, true, right);
65            }
66        }
67
68        Ok(())
69    }
70
71    pub fn into_inner(self) -> R {
72        self.inner.into_reader()
73    }
74}
75
76impl<R: io::Read + io::Seek> io::Read for HuffmanReader<R> {
77    #[inline]
78    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
79        for (idx, byte) in buf.iter_mut().enumerate() {
80            if self.stream_position()? >= self.stream_len()? {
81                return Ok(idx);
82            }
83
84            match self.tree.next_symbol(&mut self.inner) {
85                Ok(value) => {
86                    *byte = value as u8;
87                    self.offset += 1;
88                }
89                Err(e) => return Err(e),
90            }
91        }
92
93        Ok(buf.len())
94    }
95}
96
97impl<R: io::Read + io::Seek> io::Seek for HuffmanReader<R> {
98    fn seek(&mut self, _: io::SeekFrom) -> io::Result<u64> {
99        todo!()
100    }
101
102    #[inline]
103    fn stream_len(&mut self) -> io::Result<u64> {
104        Ok(self.uncompressed_size)
105    }
106
107    #[inline]
108    fn stream_position(&mut self) -> io::Result<u64> {
109        Ok(self.offset)
110    }
111}
sit_algos/huffman.rs

sit_algos/
huffman.rs
