compression_codecs/zstd/
encoder.rs

1use crate::zstd::params::CParameter;
2use crate::EncodeV2;
3use compression_core::{
4    unshared::Unshared,
5    util::{PartialBuffer, WriteBuffer},
6};
7use libzstd::stream::raw::{Encoder, Operation};
8use std::io;
9use std::io::Result;
10
11#[derive(Debug)]
12pub struct ZstdEncoder {
13    encoder: Unshared<Encoder<'static>>,
14}
15
16impl ZstdEncoder {
17    pub fn new(level: i32) -> Self {
18        Self {
19            encoder: Unshared::new(Encoder::new(level).unwrap()),
20        }
21    }
22
23    pub fn new_with_params(level: i32, params: &[CParameter]) -> Self {
24        let mut encoder = Encoder::new(level).unwrap();
25        for param in params {
26            encoder.set_parameter(param.as_zstd()).unwrap();
27        }
28        Self {
29            encoder: Unshared::new(encoder),
30        }
31    }
32
33    pub fn new_with_dict(level: i32, dictionary: &[u8]) -> io::Result<Self> {
34        let encoder = Encoder::with_dictionary(level, dictionary)?;
35        Ok(Self {
36            encoder: Unshared::new(encoder),
37        })
38    }
39
40    fn call_fn_on_out_buffer(
41        &mut self,
42        output: &mut WriteBuffer<'_>,
43        f: fn(&mut Encoder<'static>, &mut zstd_safe::OutBuffer<'_, [u8]>) -> io::Result<usize>,
44    ) -> io::Result<bool> {
45        let mut out_buf = zstd_safe::OutBuffer::around(output.initialize_unwritten());
46        let res = f(self.encoder.get_mut(), &mut out_buf);
47        let len = out_buf.as_slice().len();
48        output.advance(len);
49
50        res.map(|bytes_left| bytes_left == 0)
51    }
52}
53
54impl EncodeV2 for ZstdEncoder {
55    fn encode(
56        &mut self,
57        input: &mut PartialBuffer<&[u8]>,
58        output: &mut WriteBuffer<'_>,
59    ) -> Result<()> {
60        let status = self
61            .encoder
62            .get_mut()
63            .run_on_buffers(input.unwritten(), output.initialize_unwritten())?;
64        input.advance(status.bytes_read);
65        output.advance(status.bytes_written);
66        Ok(())
67    }
68
69    fn flush(&mut self, output: &mut WriteBuffer<'_>) -> Result<bool> {
70        self.call_fn_on_out_buffer(output, |encoder, out_buf| encoder.flush(out_buf))
71    }
72
73    fn finish(&mut self, output: &mut WriteBuffer<'_>) -> Result<bool> {
74        self.call_fn_on_out_buffer(output, |encoder, out_buf| encoder.finish(out_buf, true))
75    }
76}
compression_codecs/zstd/encoder.rs

compression_codecs/zstd/
encoder.rs
