lzma_rs/decode/
lzma2.rs

1use crate::decode::lzbuffer::LzBuffer;
2use crate::decode::lzma::{DecoderState, LzmaProperties};
3use crate::decode::{lzbuffer, rangecoder};
4use crate::error;
5use byteorder::{BigEndian, ReadBytesExt};
6use std::io;
7use std::io::Read;
8
9#[derive(Debug)]
10/// Raw decoder for LZMA2.
11pub struct Lzma2Decoder {
12    lzma_state: DecoderState,
13}
14
15impl Lzma2Decoder {
16    /// Creates a new object ready for decompressing data that it's given.
17    pub fn new() -> Lzma2Decoder {
18        Lzma2Decoder {
19            lzma_state: DecoderState::new(
20                LzmaProperties {
21                    lc: 0,
22                    lp: 0,
23                    pb: 0,
24                },
25                None,
26            ),
27        }
28    }
29
30    /// Performs the equivalent of replacing this decompression state with a freshly allocated copy.
31    ///
32    /// This function may not allocate memory and will attempt to reuse any previously allocated resources.
33    #[cfg(feature = "raw_decoder")]
34    pub fn reset(&mut self) {
35        self.lzma_state.reset_state(LzmaProperties {
36            lc: 0,
37            lp: 0,
38            pb: 0,
39        });
40    }
41
42    /// Decompresses the input data into the output, consuming only as much input as needed and writing as much output as possible.
43    pub fn decompress<W: io::Write, R: io::BufRead>(
44        &mut self,
45        input: &mut R,
46        output: &mut W,
47    ) -> error::Result<()> {
48        let mut accum = lzbuffer::LzAccumBuffer::from_stream(output, usize::MAX);
49
50        loop {
51            let status = input.read_u8().map_err(|e| {
52                error::Error::LzmaError(format!("LZMA2 expected new status: {}", e))
53            })?;
54
55            lzma_info!("LZMA2 status: {}", status);
56
57            if status == 0 {
58                lzma_info!("LZMA2 end of input");
59                break;
60            } else if status == 1 {
61                // uncompressed reset dict
62                Self::parse_uncompressed(&mut accum, input, true)?;
63            } else if status == 2 {
64                // uncompressed no reset
65                Self::parse_uncompressed(&mut accum, input, false)?;
66            } else {
67                self.parse_lzma(&mut accum, input, status)?;
68            }
69        }
70
71        accum.finish()?;
72        Ok(())
73    }
74
75    fn parse_lzma<R, W>(
76        &mut self,
77        accum: &mut lzbuffer::LzAccumBuffer<W>,
78        input: &mut R,
79        status: u8,
80    ) -> error::Result<()>
81    where
82        R: io::BufRead,
83        W: io::Write,
84    {
85        if status & 0x80 == 0 {
86            return Err(error::Error::LzmaError(format!(
87                "LZMA2 invalid status {}, must be 0, 1, 2 or >= 128",
88                status
89            )));
90        }
91
92        let reset_dict: bool;
93        let reset_state: bool;
94        let reset_props: bool;
95        match (status >> 5) & 0x3 {
96            0 => {
97                reset_dict = false;
98                reset_state = false;
99                reset_props = false;
100            }
101            1 => {
102                reset_dict = false;
103                reset_state = true;
104                reset_props = false;
105            }
106            2 => {
107                reset_dict = false;
108                reset_state = true;
109                reset_props = true;
110            }
111            3 => {
112                reset_dict = true;
113                reset_state = true;
114                reset_props = true;
115            }
116            _ => unreachable!(),
117        }
118
119        let unpacked_size = input
120            .read_u16::<BigEndian>()
121            .map_err(|e| error::Error::LzmaError(format!("LZMA2 expected unpacked size: {}", e)))?;
122        let unpacked_size = ((((status & 0x1F) as u64) << 16) | (unpacked_size as u64)) + 1;
123
124        let packed_size = input
125            .read_u16::<BigEndian>()
126            .map_err(|e| error::Error::LzmaError(format!("LZMA2 expected packed size: {}", e)))?;
127        let packed_size = (packed_size as u64) + 1;
128
129        lzma_info!(
130            "LZMA2 compressed block {{ unpacked_size: {}, packed_size: {}, reset_dict: {}, reset_state: {}, reset_props: {} }}",
131            unpacked_size,
132            packed_size,
133            reset_dict,
134            reset_state,
135            reset_props
136        );
137
138        if reset_dict {
139            accum.reset()?;
140        }
141
142        if reset_state {
143            let new_props = if reset_props {
144                let props = input.read_u8().map_err(|e| {
145                    error::Error::LzmaError(format!("LZMA2 expected new properties: {}", e))
146                })?;
147
148                let mut pb = props as u32;
149                if pb >= 225 {
150                    return Err(error::Error::LzmaError(format!(
151                        "LZMA2 invalid properties: {} must be < 225",
152                        pb
153                    )));
154                }
155
156                let lc = pb % 9;
157                pb /= 9;
158                let lp = pb % 5;
159                pb /= 5;
160
161                if lc + lp > 4 {
162                    return Err(error::Error::LzmaError(format!(
163                        "LZMA2 invalid properties: lc + lp ({} + {}) must be <= 4",
164                        lc, lp
165                    )));
166                }
167
168                lzma_info!("Properties {{ lc: {}, lp: {}, pb: {} }}", lc, lp, pb);
169                LzmaProperties { lc, lp, pb }
170            } else {
171                self.lzma_state.lzma_props
172            };
173
174            self.lzma_state.reset_state(new_props);
175        }
176
177        self.lzma_state
178            .set_unpacked_size(Some(unpacked_size + accum.len() as u64));
179
180        let mut taken = input.take(packed_size);
181        let mut rangecoder = rangecoder::RangeDecoder::new(&mut taken)
182            .map_err(|e| error::Error::LzmaError(format!("LZMA input too short: {}", e)))?;
183        self.lzma_state.process(accum, &mut rangecoder)
184    }
185
186    fn parse_uncompressed<R, W>(
187        accum: &mut lzbuffer::LzAccumBuffer<W>,
188        input: &mut R,
189        reset_dict: bool,
190    ) -> error::Result<()>
191    where
192        R: io::BufRead,
193        W: io::Write,
194    {
195        let unpacked_size = input
196            .read_u16::<BigEndian>()
197            .map_err(|e| error::Error::LzmaError(format!("LZMA2 expected unpacked size: {}", e)))?;
198        let unpacked_size = (unpacked_size as usize) + 1;
199
200        lzma_info!(
201            "LZMA2 uncompressed block {{ unpacked_size: {}, reset_dict: {} }}",
202            unpacked_size,
203            reset_dict
204        );
205
206        if reset_dict {
207            accum.reset()?;
208        }
209
210        let mut buf = vec![0; unpacked_size];
211        input.read_exact(buf.as_mut_slice()).map_err(|e| {
212            error::Error::LzmaError(format!(
213                "LZMA2 expected {} uncompressed bytes: {}",
214                unpacked_size, e
215            ))
216        })?;
217        accum.append_bytes(buf.as_slice());
218
219        Ok(())
220    }
221}