Skip to main content

photon_protocol/compressor/
zstd.rs

1use std::fmt;
2
3#[cfg(not(target_arch = "wasm32"))]
4mod native {
5    use std::sync::Mutex;
6
7    use bytes::BytesMut;
8    use zstd::bulk::{Compressor as Encoder, Decompressor as Decoder};
9
10    use crate::ports::compress::{CompressionError, Compressor};
11
12    pub struct ZstdCompressor {
13        pub level: i32,
14        compressor: Mutex<Encoder<'static>>,
15        decompressor: Mutex<Decoder<'static>>,
16    }
17
18    impl ZstdCompressor {
19        pub const NAME: &str = "zstd";
20
21        pub fn new(level: i32) -> Self {
22            Self {
23                level,
24                compressor: Mutex::new(
25                    Encoder::new(level).expect("failed to create zstd compressor"),
26                ),
27                decompressor: Mutex::new(
28                    Decoder::new().expect("failed to create zstd decompressor"),
29                ),
30            }
31        }
32    }
33
34    impl Compressor for ZstdCompressor {
35        fn compress(&self, input: &[u8], output: &mut BytesMut) -> Result<(), CompressionError> {
36            let compressed = self
37                .compressor
38                .lock()
39                .unwrap()
40                .compress(input)
41                .map_err(|e| CompressionError::Internal(e.to_string()))?;
42
43            output.extend_from_slice(&compressed);
44            Ok(())
45        }
46
47        fn decompress(&self, input: &[u8], output: &mut BytesMut) -> Result<(), CompressionError> {
48            let capacity = zstd::zstd_safe::get_frame_content_size(input)
49                .ok()
50                .flatten()
51                .map_or(input.len() * 10, |s| s as usize);
52
53            let decompressed = self
54                .decompressor
55                .lock()
56                .unwrap()
57                .decompress(input, capacity)
58                .map_err(|_| CompressionError::CorruptPayload {
59                    compressor_name: self.name().to_owned(),
60                })?;
61
62            output.extend_from_slice(&decompressed);
63            Ok(())
64        }
65
66        fn name(&self) -> &'static str {
67            Self::NAME
68        }
69    }
70}
71
72#[cfg(target_arch = "wasm32")]
73mod wasm {
74    use std::io::Read;
75
76    use bytes::BytesMut;
77    use ruzstd::decoding::StreamingDecoder;
78    use ruzstd::encoding::{CompressionLevel, compress_to_vec};
79
80    use crate::ports::compress::{CompressionError, Compressor};
81
82    pub struct ZstdCompressor {
83        pub level: i32,
84    }
85
86    impl ZstdCompressor {
87        pub const NAME: &str = "zstd";
88
89        pub fn new(level: i32) -> Self {
90            Self { level }
91        }
92    }
93
94    impl Compressor for ZstdCompressor {
95        fn compress(&self, input: &[u8], output: &mut BytesMut) -> Result<(), CompressionError> {
96            let level = match self.level {
97                0 => CompressionLevel::Uncompressed,
98                1..=3 => CompressionLevel::Fastest,
99                4..=6 => CompressionLevel::Default,
100                7..=9 => CompressionLevel::Better,
101                _ => CompressionLevel::Best,
102            };
103
104            output.extend_from_slice(&compress_to_vec(input, level));
105            Ok(())
106        }
107
108        fn decompress(&self, input: &[u8], output: &mut BytesMut) -> Result<(), CompressionError> {
109            let mut decoder = StreamingDecoder::new(input)
110                .map_err(|e| CompressionError::Internal(format!("zstd frame init: {e}")))?;
111
112            let mut buf = Vec::new();
113            decoder
114                .read_to_end(&mut buf)
115                .map_err(|e| CompressionError::Internal(format!("zstd decompress: {e}")))?;
116
117            output.extend_from_slice(&buf);
118            Ok(())
119        }
120
121        fn name(&self) -> &'static str {
122            Self::NAME
123        }
124    }
125}
126
127#[cfg(not(target_arch = "wasm32"))]
128pub use native::ZstdCompressor;
129#[cfg(target_arch = "wasm32")]
130pub use wasm::ZstdCompressor;
131
132impl Default for ZstdCompressor {
133    fn default() -> Self {
134        Self::new(3)
135    }
136}
137
138impl Clone for ZstdCompressor {
139    fn clone(&self) -> Self {
140        Self::new(self.level)
141    }
142}
143
144#[allow(clippy::missing_fields_in_debug)]
145impl fmt::Debug for ZstdCompressor {
146    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
147        f.debug_struct("ZstdCompressor")
148            .field("level", &self.level)
149            .finish()
150    }
151}