vector/sinks/util/
zstd.rs

1use std::{fmt::Display, io};
2
3use super::buffer::compression::CompressionLevel;
4
5#[derive(Debug)]
6pub struct ZstdCompressionLevel(i32);
7
8impl From<CompressionLevel> for ZstdCompressionLevel {
9    fn from(value: CompressionLevel) -> Self {
10        let val: i32 = match value {
11            CompressionLevel::None => 0,
12            CompressionLevel::Default => zstd::DEFAULT_COMPRESSION_LEVEL,
13            CompressionLevel::Best => 21,
14            CompressionLevel::Fast => 1,
15            CompressionLevel::Val(v) => v.clamp(1, 21) as i32,
16        };
17        ZstdCompressionLevel(val)
18    }
19}
20
21impl Display for ZstdCompressionLevel {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        write!(f, "{}", self.0)
24    }
25}
26
27pub struct ZstdEncoder<W: io::Write> {
28    inner: zstd::Encoder<'static, W>,
29}
30
31impl<W: io::Write> ZstdEncoder<W> {
32    pub fn new(writer: W, level: ZstdCompressionLevel) -> io::Result<Self> {
33        let encoder = zstd::Encoder::new(writer, level.0)?;
34        Ok(Self { inner: encoder })
35    }
36
37    pub fn finish(self) -> io::Result<W> {
38        self.inner.finish()
39    }
40
41    pub fn get_ref(&self) -> &W {
42        self.inner.get_ref()
43    }
44}
45
46impl<W: io::Write> io::Write for ZstdEncoder<W> {
47    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
48        #[allow(clippy::disallowed_methods)] // Caller handles the result of `write`.
49        self.inner.write(buf)
50    }
51
52    fn flush(&mut self) -> io::Result<()> {
53        self.inner.flush()
54    }
55}
56
57impl<W: io::Write + std::fmt::Debug> std::fmt::Debug for ZstdEncoder<W> {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        f.debug_struct("ZstdEncoder")
60            .field("inner", &self.get_ref())
61            .finish()
62    }
63}
64
65/// Safety:
66/// 1. There is no sharing references to zstd encoder. `Write` requires unique reference, and `finish` moves the instance itself.
67/// 2. Sharing only internal writer, which implements `Sync`
68unsafe impl<W: io::Write + Sync> Sync for ZstdEncoder<W> {}