lz4/
encoder.rs

1use super::liblz4::*;
2use super::size_t;
3use std::cmp;
4use std::io::Result;
5use std::io::Write;
6use std::ptr;
7
8#[derive(Debug)]
9struct EncoderContext {
10    c: LZ4FCompressionContext,
11}
12
13#[derive(Clone, Debug)]
14pub struct EncoderBuilder {
15    block_size: BlockSize,
16    block_mode: BlockMode,
17    // 1: each block followed by a checksum of block's compressed data; 0: disabled (default)
18    block_checksum: BlockChecksum,
19    checksum: ContentChecksum,
20    // 0 == default (fast mode); values above 16 count as 16; values below 0 count as 0
21    level: u32,
22    // 1 == always flush (reduce need for tmp buffer)
23    auto_flush: bool,
24    favor_dec_speed: bool,
25    content_size: u64,
26}
27
28#[derive(Debug)]
29pub struct Encoder<W> {
30    c: EncoderContext,
31    w: W,
32    limit: usize,
33    buffer: Vec<u8>,
34}
35
36impl EncoderBuilder {
37    pub fn new() -> Self {
38        EncoderBuilder {
39            block_size: BlockSize::Default,
40            block_mode: BlockMode::Linked,
41            checksum: ContentChecksum::ChecksumEnabled,
42            block_checksum: BlockChecksum::BlockChecksumEnabled,
43            level: 0,
44            auto_flush: false,
45            favor_dec_speed: false,
46            content_size: 0,
47        }
48    }
49
50    pub fn block_size(&mut self, block_size: BlockSize) -> &mut Self {
51        self.block_size = block_size;
52        self
53    }
54
55    pub fn block_mode(&mut self, block_mode: BlockMode) -> &mut Self {
56        self.block_mode = block_mode;
57        self
58    }
59
60    pub fn block_checksum(&mut self, block_checksum: BlockChecksum) -> &mut Self {
61        self.block_checksum = block_checksum;
62        self
63    }
64
65    pub fn checksum(&mut self, checksum: ContentChecksum) -> &mut Self {
66        self.checksum = checksum;
67        self
68    }
69
70    pub fn level(&mut self, level: u32) -> &mut Self {
71        self.level = level;
72        self
73    }
74
75    pub fn auto_flush(&mut self, auto_flush: bool) -> &mut Self {
76        self.auto_flush = auto_flush;
77        self
78    }
79
80    /// Favor decompression speed over compression ratio. Requires compression
81    /// level >=10.
82    pub fn favor_dec_speed(&mut self, favor_dec_speed: bool) -> &mut Self {
83        self.favor_dec_speed = favor_dec_speed;
84        self
85    }
86
87    pub fn content_size(&mut self, content_size: u64) -> &mut Self {
88        self.content_size = content_size;
89        self
90    }
91
92    pub fn build<W: Write>(&self, w: W) -> Result<Encoder<W>> {
93        let block_size = self.block_size.get_size();
94        let preferences = LZ4FPreferences {
95            frame_info: LZ4FFrameInfo {
96                block_size_id: self.block_size.clone(),
97                block_mode: self.block_mode.clone(),
98                content_checksum_flag: self.checksum.clone(),
99                content_size: self.content_size.clone(),
100                frame_type: FrameType::Frame,
101                dict_id: 0,
102                block_checksum_flag: self.block_checksum.clone(),
103            },
104            compression_level: self.level,
105            auto_flush: if self.auto_flush { 1 } else { 0 },
106            favor_dec_speed: if self.favor_dec_speed { 1 } else { 0 },
107            reserved: [0; 3],
108        };
109        let mut encoder = Encoder {
110            w,
111            c: EncoderContext::new()?,
112            limit: block_size,
113            buffer: Vec::with_capacity(check_error(unsafe {
114                LZ4F_compressBound(block_size as size_t, &preferences)
115            })?),
116        };
117        encoder.write_header(&preferences)?;
118        Ok(encoder)
119    }
120}
121
122impl<W: Write> Encoder<W> {
123    fn write_header(&mut self, preferences: &LZ4FPreferences) -> Result<()> {
124        unsafe {
125            let len = check_error(LZ4F_compressBegin(
126                self.c.c,
127                self.buffer.as_mut_ptr(),
128                self.buffer.capacity() as size_t,
129                preferences,
130            ))?;
131            self.buffer.set_len(len);
132        }
133        self.w.write_all(&self.buffer)
134    }
135
136    fn write_end(&mut self) -> Result<()> {
137        unsafe {
138            let len = check_error(LZ4F_compressEnd(
139                self.c.c,
140                self.buffer.as_mut_ptr(),
141                self.buffer.capacity() as size_t,
142                ptr::null(),
143            ))?;
144            self.buffer.set_len(len);
145        };
146        self.w.write_all(&self.buffer)
147    }
148
149    /// Immutable writer reference.
150    pub fn writer(&self) -> &W {
151        &self.w
152    }
153
154    /// This function is used to flag that this session of compression is done
155    /// with. The stream is finished up (final bytes are written), and then the
156    /// wrapped writer is returned.
157    pub fn finish(mut self) -> (W, Result<()>) {
158        let result = self.write_end();
159        (self.w, result)
160    }
161}
162
163impl<W: Write> Write for Encoder<W> {
164    fn write(&mut self, buffer: &[u8]) -> Result<usize> {
165        let mut offset = 0;
166        while offset < buffer.len() {
167            let size = cmp::min(buffer.len() - offset, self.limit);
168            unsafe {
169                let len = check_error(LZ4F_compressUpdate(
170                    self.c.c,
171                    self.buffer.as_mut_ptr(),
172                    self.buffer.capacity() as size_t,
173                    buffer[offset..].as_ptr(),
174                    size as size_t,
175                    ptr::null(),
176                ))?;
177                self.buffer.set_len(len);
178                self.w.write_all(&self.buffer)?;
179            }
180            offset += size;
181        }
182        Ok(buffer.len())
183    }
184
185    fn flush(&mut self) -> Result<()> {
186        loop {
187            unsafe {
188                let len = check_error(LZ4F_flush(
189                    self.c.c,
190                    self.buffer.as_mut_ptr(),
191                    self.buffer.capacity() as size_t,
192                    ptr::null(),
193                ))?;
194                if len == 0 {
195                    break;
196                }
197                self.buffer.set_len(len);
198            };
199            self.w.write_all(&self.buffer)?;
200        }
201        self.w.flush()
202    }
203}
204
205impl EncoderContext {
206    fn new() -> Result<EncoderContext> {
207        let mut context = LZ4FCompressionContext(ptr::null_mut());
208        check_error(unsafe { LZ4F_createCompressionContext(&mut context, LZ4F_VERSION) })?;
209        Ok(EncoderContext { c: context })
210    }
211}
212
213impl Drop for EncoderContext {
214    fn drop(&mut self) {
215        unsafe { LZ4F_freeCompressionContext(self.c) };
216    }
217}
218
219#[cfg(test)]
220mod test {
221    use super::EncoderBuilder;
222    use std::io::{Read, Write};
223
224    #[test]
225    fn test_encoder_smoke() {
226        let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
227        encoder.write(b"Some ").unwrap();
228        encoder.write(b"data").unwrap();
229        let (_, result) = encoder.finish();
230        result.unwrap();
231    }
232
233    #[test]
234    fn test_encoder_random() {
235        let mut encoder = EncoderBuilder::new().level(1).build(Vec::new()).unwrap();
236        let mut input = Vec::new();
237        let mut rnd: u32 = 42;
238        for _ in 0..1024 * 1024 {
239            input.push((rnd & 0xFF) as u8);
240            rnd = ((1664525 as u64) * (rnd as u64) + (1013904223 as u64)) as u32;
241        }
242        encoder.write(&input).unwrap();
243        let (compressed, result) = encoder.finish();
244        result.unwrap();
245
246        let mut dec = crate::decoder::Decoder::new(&compressed[..]).unwrap();
247        let mut output = Vec::new();
248        dec.read_to_end(&mut output).unwrap();
249        assert_eq!(input, output);
250    }
251
252    #[test]
253    fn test_encoder_content_size() {
254        let mut encoder = EncoderBuilder::new()
255            .level(1)
256            .content_size(1024 * 1024)
257            .build(Vec::new())
258            .unwrap();
259        let mut input = Vec::new();
260        let mut rnd: u32 = 42;
261        for _ in 0..1024 * 1024 {
262            input.push((rnd & 0xFF) as u8);
263            rnd = ((1664525 as u64) * (rnd as u64) + (1013904223 as u64)) as u32;
264        }
265        encoder.write(&input).unwrap();
266        let (compressed, result) = encoder.finish();
267        result.unwrap();
268
269        let mut dec = crate::decoder::Decoder::new(&compressed[..]).unwrap();
270        let mut output = Vec::new();
271        dec.read_to_end(&mut output).unwrap();
272        assert_eq!(input, output);
273    }
274
275    #[test]
276    fn test_encoder_send() {
277        fn check_send<S: Send>(_: &S) {}
278        let enc = EncoderBuilder::new().build(Vec::new());
279        check_send(&enc);
280    }
281
282    #[test]
283    fn test_favor_dec_speed() {
284        let mut encoder = EncoderBuilder::new()
285            .level(11)
286            .favor_dec_speed(true)
287            .build(Vec::new())
288            .unwrap();
289        let mut input = Vec::new();
290        let mut rnd: u32 = 42;
291        for _ in 0..1024 * 1024 {
292            input.push((rnd & 0xFF) as u8);
293            rnd = ((1664525 as u64) * (rnd as u64) + (1013904223 as u64)) as u32;
294        }
295        encoder.write(&input).unwrap();
296        let (compressed, result) = encoder.finish();
297        result.unwrap();
298
299        let mut dec = crate::decoder::Decoder::new(&compressed[..]).unwrap();
300
301        let mut output = Vec::new();
302        dec.read_to_end(&mut output).unwrap();
303        assert_eq!(input, output);
304    }
305}