Skip to main content

sit_algos/installer/
mod.rs

1use std::{io, marker::PhantomData};
2
3use bitstream_io::{BitRead as _, BitReader, LE};
4use cfor::cfor;
5
6pub(crate) mod hufftree;
7mod window;
8
9use hufftree::HuffTree;
10
11use crate::installer::window::Window;
12
13pub struct InstallerReader<R> {
14    inner: io::Cursor<Vec<u8>>,
15    marker: PhantomData<R>,
16}
17
18impl<R: io::Read + io::Seek> InstallerReader<R> {
19    pub fn try_new(inner: R) -> Result<Self, io::Error> {
20        let data = decompress(inner).map_err(io::Error::other)?;
21        let inner = io::Cursor::new(data);
22
23        Ok(Self {
24            inner,
25            marker: PhantomData,
26        })
27    }
28}
29
30// TODO: Improve error handling
31#[derive(thiserror::Error, Debug)]
32pub enum Error {
33    #[error(transparent)]
34    BinRw(#[from] binrw::Error),
35
36    #[error(transparent)]
37    Io(#[from] io::Error),
38
39    #[error("invalid")]
40    InvalidTree,
41
42    #[error("invalid")]
43    TreeEncodingUnknown,
44
45    #[error("invalid")]
46    DictionaryOutOfBounds,
47
48    #[error("invalid")]
49    TooMuchOutput,
50}
51
52pub const fn init_lengths<const N: usize>() -> ([u16; N], [u8; N]) {
53    let mut values = [0u16; N];
54    let mut paths = [0u8; N];
55
56    let mut k = 0;
57    cfor! {let mut i = 0; i < N; i+=1; {
58        let value = (i.saturating_sub(4) >> 2) as u8;
59
60        values[i] = k;
61        paths[i] = value;
62
63        k = k.wrapping_add(1u16.wrapping_shl(value as u32));
64    }}
65
66    (values, paths)
67}
68
69pub const fn init_offsets<const N: usize>() -> ([u32; N], [u8; N]) {
70    let mut paths = [0u32; N];
71    let mut values = [0u8; N];
72
73    let mut k = 1;
74    cfor! {let mut i = 0; i < N; i+=1; {
75        let value = (i.saturating_sub(3) >> 2) as u8;
76
77        paths[i] = k;
78        values[i] = value;
79
80        k = k.wrapping_add(1u32.wrapping_shl(value as u32));
81    }};
82
83    (paths, values)
84}
85
86pub fn decompress(mut reader: impl io::Read + io::Seek) -> Result<Vec<u8>, Error> {
87    // TODO: Decompress byte- or at least block-wise to avoid allocating too much memory
88    let (offsets, offsets_bits) = init_offsets::<0x4b>();
89    let (lengths, length_bits) = init_lengths::<0x34>();
90
91    let mut total_output = Vec::new();
92    let mut win = Window::<0x40000>::new();
93
94    let mut reader = BitReader::<_, LE>::new(&mut reader);
95    let block_count = reader.read_var::<u32>(16)? as usize;
96
97    for _ in 0..block_count {
98        let _compressed_block_size = reader.read_var::<u32>(32)? as usize;
99        let uncompressed_block_size = reader.read_var::<u32>(32)? as usize;
100        let syms_code = HuffTree::read_from(&mut reader, 0x134)?;
101        let offsets_codes = HuffTree::read_from(&mut reader, 0x4b)?;
102
103        let mut output = vec![0u8; uncompressed_block_size];
104        let mut output_ptr = 0;
105        while output_ptr < output.len() {
106            let mut sym = syms_code.read_code(0, &mut reader)? as usize;
107            if sym < 0x100 {
108                // Emit literal
109                output[output_ptr] = win.put((sym & 0xFF) as u8);
110                output_ptr += 1;
111                continue;
112            }
113            sym -= 0x100;
114
115            let mut length = lengths[sym] as usize + 4;
116            let bits = length_bits[sym] as u32;
117            length += reader.read_var::<u32>(bits)? as usize;
118
119            let sym = offsets_codes.read_code(0, &mut reader)? as usize;
120            let mut offset = offsets[sym] as usize;
121            let bits = offsets_bits[sym] as u32;
122            offset += reader.read_var::<u32>(bits)? as usize;
123
124            if output_ptr + length > output.len() {
125                return Err(Error::TooMuchOutput);
126            }
127
128            // Copy from window
129            for i in 0..length {
130                let v = win.get(offset as u32);
131                output[output_ptr + i] = win.put(v);
132            }
133
134            output_ptr += length;
135        }
136
137        reader.byte_align();
138        total_output.append(&mut output);
139    }
140
141    Ok(total_output)
142}
143
144impl<R: io::Read + io::Seek> io::Read for InstallerReader<R> {
145    #[inline]
146    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
147        self.inner.read(buf)
148    }
149}
150
151impl<R: io::Seek> io::Seek for InstallerReader<R> {
152    #[inline]
153    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
154        self.inner.seek(pos)
155    }
156
157    #[inline]
158    fn stream_position(&mut self) -> io::Result<u64> {
159        self.inner.stream_position()
160    }
161
162    #[inline]
163    fn stream_len(&mut self) -> io::Result<u64> {
164        self.inner.stream_len()
165    }
166}