lzma_rust2/enc/
lzma2_writer.rs

1use std::io::Write;
2
3use byteorder::WriteBytesExt;
4
5use super::{
6    encoder::{EncodeMode, LZMAEncoder, LZMAEncoderModes},
7    lz::MFType,
8    range_enc::{RangeEncoder, RangeEncoderBuffer},
9};
10
11#[derive(Debug, Clone)]
12pub struct LZMA2Options {
13    pub dict_size: u32,
14    pub lc: u32,
15    pub lp: u32,
16    pub pb: u32,
17    pub mode: EncodeMode,
18    pub nice_len: u32,
19    pub mf: MFType,
20    pub depth_limit: i32,
21    pub preset_dict: Option<Vec<u8>>,
22}
23
24impl Default for LZMA2Options {
25    fn default() -> Self {
26        Self::with_preset(6)
27    }
28}
29
30impl LZMA2Options {
31    pub const LC_DEFAULT: u32 = 3;
32    pub const LP_DEFAULT: u32 = 0;
33    pub const PB_DEFAULT: u32 = 2;
34    pub const NICE_LEN_MAX: u32 = 273;
35    pub const NICE_LEN_MIN: u32 = 8;
36    pub const DICT_SIZE_DEFAULT: u32 = 8 << 20;
37    const PRESET_TO_DICT_SIZE: &'static [u32] = &[
38        1 << 18,
39        1 << 20,
40        1 << 21,
41        1 << 22,
42        1 << 22,
43        1 << 23,
44        1 << 23,
45        1 << 24,
46        1 << 25,
47        1 << 26,
48    ];
49    const PRESET_TO_DEPTH_LIMIT: &'static [i32] = &[4, 8, 24, 48];
50
51    #[allow(clippy::too_many_arguments)]
52    pub fn new(
53        dict_size: u32,
54        lc: u32,
55        lp: u32,
56        pb: u32,
57        mode: EncodeMode,
58        nice_len: u32,
59        mf: MFType,
60        depth_limit: i32,
61    ) -> Self {
62        Self {
63            dict_size,
64            lc,
65            lp,
66            pb,
67            mode,
68            nice_len,
69            mf,
70            depth_limit,
71            preset_dict: None,
72        }
73    }
74
75    /// preset: [0..9]
76    #[inline]
77    pub fn with_preset(preset: u32) -> Self {
78        let mut opt = Self {
79            dict_size: Default::default(),
80            lc: Default::default(),
81            lp: Default::default(),
82            pb: Default::default(),
83            mode: EncodeMode::Normal,
84            nice_len: Default::default(),
85            mf: Default::default(),
86            depth_limit: Default::default(),
87            preset_dict: Default::default(),
88        };
89        opt.set_preset(preset);
90        opt
91    }
92
93    /// preset: [0..9]
94    pub fn set_preset(&mut self, preset: u32) {
95        if preset > 9 {
96            return;
97        }
98        self.lc = Self::LC_DEFAULT;
99        self.lp = Self::LP_DEFAULT;
100        self.pb = Self::PB_DEFAULT;
101        self.dict_size = Self::PRESET_TO_DICT_SIZE[preset as usize];
102        if preset <= 3 {
103            self.mode = EncodeMode::Fast;
104            self.mf = MFType::HC4;
105            self.nice_len = if preset <= 1 { 128 } else { Self::NICE_LEN_MAX };
106            self.depth_limit = Self::PRESET_TO_DEPTH_LIMIT[preset as usize];
107        } else {
108            self.mode = EncodeMode::Normal;
109            self.mf = MFType::BT4;
110            self.nice_len = if preset == 4 {
111                16
112            } else if preset == 5 {
113                32
114            } else {
115                64
116            };
117            self.depth_limit = 0;
118        }
119    }
120
121    pub fn get_memory_usage(&self) -> u32 {
122        let dict_size = self.dict_size;
123        let extra_size_before = get_extra_size_before(dict_size);
124        70 + LZMAEncoder::get_mem_usage(self.mode, dict_size, extra_size_before, self.mf)
125    }
126
127    #[inline(always)]
128    pub fn get_props(&self) -> u8 {
129        ((self.pb * 5 + self.lp) * 9 + self.lc) as u8
130    }
131}
132
133const COMPRESSED_SIZE_MAX: u32 = 64 << 10;
134
135pub fn get_extra_size_before(dict_size: u32) -> u32 {
136    COMPRESSED_SIZE_MAX.saturating_sub(dict_size)
137}
138
139/// LZMA2 format writer
140/// # Examples
141/// ```
142/// use std::io::Write;
143///
144/// use lzma_rust2::{LZMA2Options, LZMA2Writer};
145///
146/// let mut writer = LZMA2Writer::new(Vec::new(), &LZMA2Options::default());
147/// writer.write_all(b"hello world").unwrap();
148/// writer.finish().unwrap();
149/// ```
150pub struct LZMA2Writer<W: Write> {
151    inner: W,
152    rc: RangeEncoder<RangeEncoderBuffer>,
153    lzma: LZMAEncoder,
154    mode: LZMAEncoderModes,
155    props: u8,
156    dict_reset_needed: bool,
157    state_reset_needed: bool,
158    props_needed: bool,
159    pending_size: u32,
160}
161
162impl<W: Write> LZMA2Writer<W> {
163    pub fn new(inner: W, options: &LZMA2Options) -> Self {
164        let dict_size = options.dict_size;
165        let rc = RangeEncoder::new_buffer(COMPRESSED_SIZE_MAX as usize);
166        let (mut lzma, mode) = LZMAEncoder::new(
167            options.mode,
168            options.lc,
169            options.lp,
170            options.pb,
171            options.mf,
172            options.depth_limit,
173            options.dict_size,
174            options.nice_len as usize,
175        );
176
177        let props = options.get_props();
178        let mut dict_reset_needed = true;
179        if let Some(preset_dict) = &options.preset_dict {
180            lzma.lz.set_preset_dict(dict_size, preset_dict);
181            dict_reset_needed = false;
182        }
183        Self {
184            inner,
185            rc,
186            lzma,
187            mode,
188            props,
189            dict_reset_needed,
190            state_reset_needed: true,
191            props_needed: true,
192            pending_size: 0,
193        }
194    }
195
196    fn write_lzma(&mut self, uncompressed_size: u32, compressed_size: u32) -> std::io::Result<()> {
197        let mut control = if self.props_needed {
198            if self.dict_reset_needed {
199                0x80 + (3 << 5)
200            } else {
201                0x80 + (2 << 5)
202            }
203        } else if self.state_reset_needed {
204            0x80 + (1 << 5)
205        } else {
206            0x80
207        };
208        control |= (uncompressed_size - 1) >> 16;
209        let mut chunk_header = [0u8; 6];
210        chunk_header[0] = control as u8;
211        chunk_header[1] = ((uncompressed_size - 1) >> 8) as u8;
212        chunk_header[2] = (uncompressed_size - 1) as u8;
213        chunk_header[3] = ((compressed_size - 1) >> 8) as u8;
214        chunk_header[4] = (compressed_size - 1) as u8;
215        if self.props_needed {
216            chunk_header[5] = self.props;
217            self.inner.write_all(&chunk_header)?;
218        } else {
219            self.inner.write_all(&chunk_header[..5])?;
220        }
221
222        self.rc.write_to(&mut self.inner)?;
223        self.props_needed = false;
224        self.state_reset_needed = false;
225        self.dict_reset_needed = false;
226        Ok(())
227    }
228
229    fn write_uncompressed(&mut self, mut uncompressed_size: u32) -> std::io::Result<()> {
230        while uncompressed_size > 0 {
231            let chunk_size = uncompressed_size.min(COMPRESSED_SIZE_MAX);
232            let mut chunk_header = [0u8; 3];
233            chunk_header[0] = if self.dict_reset_needed { 0x01 } else { 0x02 };
234            chunk_header[1] = ((chunk_size - 1) >> 8) as u8;
235            chunk_header[2] = (chunk_size - 1) as u8;
236            self.inner.write_all(&chunk_header)?;
237            self.lzma.lz.copy_uncompressed(
238                &mut self.inner,
239                uncompressed_size as i32,
240                chunk_size as usize,
241            )?;
242            uncompressed_size -= chunk_size;
243            self.dict_reset_needed = false;
244        }
245        self.state_reset_needed = true;
246        Ok(())
247    }
248
249    fn write_chunk(&mut self) -> std::io::Result<()> {
250        let compressed_size = self.rc.finish_buffer()?.unwrap_or_default() as u32;
251        let mut uncompressed_size = self.lzma.data.uncompressed_size;
252        assert!(compressed_size > 0);
253        assert!(
254            uncompressed_size > 0,
255            "uncompressed_size is 0, read_pos={}",
256            self.lzma.lz.read_pos
257        );
258        if compressed_size + 2 < uncompressed_size {
259            self.write_lzma(uncompressed_size, compressed_size)?;
260        } else {
261            self.lzma.reset(&mut self.mode);
262            uncompressed_size = self.lzma.data.uncompressed_size;
263            assert!(uncompressed_size > 0);
264            self.write_uncompressed(uncompressed_size)?;
265        }
266        self.pending_size -= uncompressed_size;
267        self.lzma.reset_uncompressed_size();
268        self.rc.reset_buffer();
269        Ok(())
270    }
271
272    pub fn inner(&mut self) -> &mut W {
273        &mut self.inner
274    }
275
276    pub fn finish(mut self) -> std::io::Result<W> {
277        self.lzma.lz.set_finishing();
278
279        while self.pending_size > 0 {
280            self.lzma.encode_for_lzma2(&mut self.rc, &mut self.mode)?;
281            self.write_chunk()?;
282        }
283
284        self.inner.write_u8(0x00)?;
285
286        Ok(self.inner)
287    }
288}
289
290impl<W: Write> Write for LZMA2Writer<W> {
291    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
292        let mut len = buf.len();
293
294        let mut off = 0;
295        while len > 0 {
296            let used = self.lzma.lz.fill_window(&buf[off..(off + len)]);
297            off += used;
298            len -= used;
299            self.pending_size += used as u32;
300            if self.lzma.encode_for_lzma2(&mut self.rc, &mut self.mode)? {
301                self.write_chunk()?;
302            }
303        }
304        Ok(off)
305    }
306
307    fn flush(&mut self) -> std::io::Result<()> {
308        self.lzma.lz.set_flushing();
309        while self.pending_size > 0 {
310            self.lzma.encode_for_lzma2(&mut self.rc, &mut self.mode)?;
311            self.write_chunk()?;
312        }
313        self.inner.flush()
314    }
315}
lzma_rust2/enc/lzma2_writer.rs

lzma_rust2/enc/
lzma2_writer.rs
