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

lzma_rust/enc/
lzma_writer.rs
