lzma_rust2/enc/
lzma_writer.rs

1use super::{
2    encoder::{LzmaEncoder, LzmaEncoderModes},
3    range_enc::RangeEncoder,
4    LzmaOptions,
5};
6use crate::{error_invalid_input, error_unsupported, AutoFinish, AutoFinisher, Write};
7
8/// A single-threaded LZMA compressor.
9pub struct LzmaWriter<W: Write> {
10    rc: RangeEncoder<W>,
11    lzma: LzmaEncoder,
12    use_end_marker: bool,
13    current_uncompressed_size: u64,
14    expected_uncompressed_size: Option<u64>,
15    props: u8,
16    mode: LzmaEncoderModes,
17}
18
19impl<W: Write> LzmaWriter<W> {
20    /// Creates a new LZMA writer with full control over formatting options.
21    pub fn new(
22        mut out: W,
23        options: &LzmaOptions,
24        use_header: bool,
25        use_end_marker: bool,
26        expected_uncompressed_size: Option<u64>,
27    ) -> crate::Result<LzmaWriter<W>> {
28        let (mut lzma, mode) = LzmaEncoder::new(
29            options.mode,
30            options.lc,
31            options.lp,
32            options.pb,
33            options.mf,
34            options.depth_limit,
35            options.dict_size,
36            options.nice_len as usize,
37        );
38        if let Some(preset_dict) = &options.preset_dict {
39            if use_header {
40                return Err(error_unsupported(
41                    "header is not supported with preset dict",
42                ));
43            }
44            lzma.lz.set_preset_dict(options.dict_size, preset_dict);
45        }
46
47        let props = options.get_props();
48        if use_header {
49            out.write_all(&[props])?;
50            let dict_size = options.dict_size;
51            out.write_all(&dict_size.to_le_bytes())?;
52            let expected_compressed_size = expected_uncompressed_size.unwrap_or(u64::MAX);
53            out.write_all(&expected_compressed_size.to_le_bytes())?;
54        }
55
56        let rc = RangeEncoder::new(out);
57        Ok(LzmaWriter {
58            rc,
59            lzma,
60            use_end_marker,
61            current_uncompressed_size: 0,
62            expected_uncompressed_size,
63            props,
64            mode,
65        })
66    }
67
68    /// Creates a new LZMA writer that includes a .lzma header with the specified input size.
69    #[inline]
70    pub fn new_use_header(
71        out: W,
72        options: &LzmaOptions,
73        input_size: Option<u64>,
74    ) -> crate::Result<Self> {
75        Self::new(out, options, true, input_size.is_none(), input_size)
76    }
77
78    /// Creates a new LZMA writer without a .lzma header.
79    #[inline]
80    pub fn new_no_header(
81        out: W,
82        options: &LzmaOptions,
83        use_end_marker: bool,
84    ) -> crate::Result<Self> {
85        Self::new(out, options, false, use_end_marker, None)
86    }
87
88    /// Returns a wrapper around `self` that will finish the stream on drop.
89    pub fn auto_finish(self) -> AutoFinisher<Self> {
90        AutoFinisher(Some(self))
91    }
92
93    /// Returns the LZMA properties byte.
94    #[inline]
95    pub fn props(&self) -> u8 {
96        self.props
97    }
98
99    /// Returns the number of uncompressed bytes written so far.
100    #[inline]
101    pub fn get_uncompressed_size(&self) -> u64 {
102        self.current_uncompressed_size
103    }
104
105    /// Unwraps the writer, returning the underlying writer.
106    pub fn into_inner(self) -> W {
107        self.rc.into_inner()
108    }
109
110    /// Returns a reference to the inner writer.
111    pub fn inner(&self) -> &W {
112        self.rc.inner()
113    }
114
115    /// Returns a mutable reference to the inner writer.
116    pub fn inner_mut(&mut self) -> &mut W {
117        self.rc.inner_mut()
118    }
119
120    /// Finishes the compression and returns the underlying writer.
121    pub fn finish(mut self) -> crate::Result<W> {
122        if let Some(exp) = self.expected_uncompressed_size {
123            if exp != self.current_uncompressed_size {
124                return Err(error_invalid_input(
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
136        let Self { rc, .. } = self;
137
138        Ok(rc.into_inner())
139    }
140}
141
142impl<W: Write> Write for LzmaWriter<W> {
143    fn write(&mut self, buf: &[u8]) -> crate::Result<usize> {
144        if let Some(exp) = self.expected_uncompressed_size {
145            if exp < self.current_uncompressed_size + buf.len() as u64 {
146                return Err(error_invalid_input(
147                    "expected compressed size does not match actual compressed size",
148                ));
149            }
150        }
151        self.current_uncompressed_size += buf.len() as u64;
152        let mut len = buf.len();
153        let mut off = 0;
154        while len > 0 {
155            let used = self.lzma.lz.fill_window(&buf[off..]);
156            off += used;
157            len -= used;
158            self.lzma.encode_for_lzma1(&mut self.rc, &mut self.mode)?;
159        }
160
161        Ok(off)
162    }
163
164    fn flush(&mut self) -> crate::Result<()> {
165        Ok(())
166    }
167}
168
169impl<W: Write> AutoFinish for LzmaWriter<W> {
170    fn finish_ignore_error(self) {
171        let _ = self.finish();
172    }
173}
lzma_rust2/enc/lzma_writer.rs

lzma_rust2/enc/
lzma_writer.rs
