lzma_rust2/enc/
lzma_writer.rs

1use std::io::Write;
2
3use byteorder::WriteBytesExt;
4
5use super::{range_enc::RangeEncoder, CountingWriter, LZMA2Options};
6
7use super::encoder::{LZMAEncoder, LZMAEncoderModes};
8
9/// Compresses into the legacy .lzma file format or into a raw LZMA stream
10///
11/// # Examples
12/// ```
13/// use std::io::Write;
14/// use lzma_rust2::{CountingWriter, LZMA2Options, LZMAWriter};
15///
16/// let s = b"Hello, world!";
17/// let mut out = Vec::new();
18/// let mut options = LZMA2Options::with_preset(6);
19/// options.dict_size = LZMA2Options::DICT_SIZE_DEFAULT;
20///
21/// let mut w = LZMAWriter::new_no_header(CountingWriter::new(&mut out), &options, false).unwrap();
22/// w.write_all(s).unwrap();
23/// w.write(&[]).unwrap();
24///
25/// ```
26///
27pub struct LZMAWriter<W: Write> {
28    rc: RangeEncoder<CountingWriter<W>>,
29    lzma: LZMAEncoder,
30    use_end_marker: bool,
31    finished: bool,
32    current_uncompressed_size: u64,
33    expected_uncompressed_size: Option<u64>,
34    props: u8,
35    mode: LZMAEncoderModes,
36}
37
38impl<W: Write> LZMAWriter<W> {
39    pub fn new(
40        mut out: CountingWriter<W>,
41        options: &LZMA2Options,
42        use_header: bool,
43        use_end_marker: bool,
44        expected_uncompressed_size: Option<u64>,
45    ) -> Result<LZMAWriter<W>, std::io::Error> {
46        let (mut lzma, mode) = LZMAEncoder::new(
47            options.mode,
48            options.lc,
49            options.lp,
50            options.pb,
51            options.mf,
52            options.depth_limit,
53            options.dict_size,
54            options.nice_len as usize,
55        );
56        if let Some(preset_dict) = &options.preset_dict {
57            if use_header {
58                return Err(std::io::Error::new(
59                    std::io::ErrorKind::Unsupported,
60                    "Header is not supported with preset dict",
61                ));
62            }
63            lzma.lz.set_preset_dict(options.dict_size, preset_dict);
64        }
65
66        let props = options.get_props();
67        if use_header {
68            out.write_u8(props as _)?;
69            let mut dict_size = options.dict_size;
70            for _i in 0..4 {
71                out.write_u8((dict_size & 0xFF) as u8)?;
72                dict_size >>= 8;
73            }
74            let expected_compressed_size = expected_uncompressed_size.unwrap_or(u64::MAX);
75            for i in 0..8 {
76                out.write_u8(((expected_compressed_size >> (i * 8)) & 0xFF) as u8)?;
77            }
78        }
79
80        let rc = RangeEncoder::new(out);
81        Ok(LZMAWriter {
82            rc,
83            lzma,
84            use_end_marker,
85            finished: false,
86            current_uncompressed_size: 0,
87            expected_uncompressed_size,
88            props,
89            mode,
90        })
91    }
92
93    #[inline]
94    pub fn new_use_header(
95        out: CountingWriter<W>,
96        options: &LZMA2Options,
97        input_size: Option<u64>,
98    ) -> Result<Self, std::io::Error> {
99        Self::new(out, options, true, input_size.is_none(), input_size)
100    }
101
102    #[inline]
103    pub fn new_no_header(
104        out: CountingWriter<W>,
105        options: &LZMA2Options,
106        use_end_marker: bool,
107    ) -> Result<Self, std::io::Error> {
108        Self::new(out, options, false, use_end_marker, None)
109    }
110
111    #[inline]
112    pub fn props(&self) -> u8 {
113        self.props
114    }
115
116    #[inline]
117    pub fn get_uncompressed_size(&self) -> u64 {
118        self.current_uncompressed_size
119    }
120
121    pub fn finish(&mut self) -> std::io::Result<()> {
122        if !self.finished {
123            if let Some(exp) = self.expected_uncompressed_size {
124                if exp != self.current_uncompressed_size {
125                    return Err(std::io::Error::new(
126                        std::io::ErrorKind::InvalidInput,
127                        "Expected compressed size does not match actual compressed size",
128                    ));
129                }
130            }
131            self.lzma.lz.set_finishing();
132            self.lzma.encode_for_lzma1(&mut self.rc, &mut self.mode)?;
133            if self.use_end_marker {
134                self.lzma.encode_lzma1_end_marker(&mut self.rc)?;
135            }
136            self.rc.finish()?;
137            self.finished = true;
138        }
139        Ok(())
140    }
141}
142
143impl<W: Write> Write for LZMAWriter<W> {
144    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
145        if self.finished {
146            return Err(std::io::Error::new(
147                std::io::ErrorKind::InvalidInput,
148                "Already finished",
149            ));
150        }
151        if buf.is_empty() {
152            self.finish()?;
153            self.rc.inner().write(buf)?;
154            return Ok(0);
155        }
156        if let Some(exp) = self.expected_uncompressed_size {
157            if exp < self.current_uncompressed_size + buf.len() as u64 {
158                return Err(std::io::Error::new(
159                    std::io::ErrorKind::InvalidInput,
160                    "Expected compressed size does not match actual compressed size",
161                ));
162            }
163        }
164        self.current_uncompressed_size += buf.len() as u64;
165        let mut len = buf.len();
166        let mut off = 0;
167        while len > 0 {
168            let used = self.lzma.lz.fill_window(&buf[off..]);
169            off += used;
170            len -= used;
171            self.lzma.encode_for_lzma1(&mut self.rc, &mut self.mode)?;
172        }
173
174        Ok(off)
175    }
176
177    fn flush(&mut self) -> std::io::Result<()> {
178        Ok(())
179    }
180}
lzma_rust2/enc/lzma_writer.rs

lzma_rust2/enc/
lzma_writer.rs
