mid_compression/zstd/
compressor.rs

1use std::{
2    num::NonZeroUsize,
3    ops::RangeInclusive,
4    ptr,
5};
6
7use zstd_sys::{
8    ZSTD_CCtx,
9    ZSTD_compressCCtx,
10    ZSTD_createCCtx,
11    ZSTD_freeCCtx,
12    ZSTD_isError,
13    ZSTD_maxCLevel,
14};
15
16use crate::{
17    error,
18    interface::ICompressor,
19};
20
21/// ZStandard compression context
22pub struct ZStdCctx {
23    ptr: ptr::NonNull<ZSTD_CCtx>,
24    level: u8,
25}
26
27impl ICompressor for ZStdCctx {
28    fn level(&self) -> usize {
29        self.level as _
30    }
31
32    fn set_level(&mut self, level: usize) {
33        self.level = level as _;
34    }
35
36    fn try_compress(
37        &mut self,
38        buf: &[u8],
39        preallocated: &mut Vec<u8>,
40    ) -> Result<std::num::NonZeroUsize, error::CompressError> {
41        if !Self::levels_range().contains(&(self.level as usize)) {
42            crate::cold();
43            return Err(error::CompressError::InvalidLevel);
44        }
45
46        let result = unsafe {
47            ZSTD_compressCCtx(
48                self.ptr.as_ptr(),
49                preallocated.as_ptr() as *mut _,
50                preallocated.capacity(),
51                buf.as_ptr() as *const _,
52                buf.len() as _,
53                self.level as _,
54            )
55        };
56
57        if unsafe { ZSTD_isError(result) } == 0 {
58            unsafe { preallocated.set_len(result) };
59            Ok(unsafe { NonZeroUsize::new_unchecked(result) })
60        } else {
61            crate::cold();
62            Err(error::CompressError::TooShortBuffer)
63        }
64    }
65
66    fn supported_levels(&self) -> std::ops::RangeInclusive<usize> {
67        Self::levels_range()
68    }
69}
70
71impl ZStdCctx {
72    /// Create ZStd compression context.
73    ///
74    /// # Panics
75    ///
76    /// - if underlying call to the `ZSTD_createCCtx`
77    /// fails (returns NULL)
78    /// - if level is invalid
79    pub fn new(level: u8) -> Self {
80        if !Self::levels_range().contains(&(level as usize)) {
81            panic!(
82                "Invalid level, supported levels are: {:?}",
83                Self::levels_range()
84            );
85        }
86
87        let cctx = unsafe { ZSTD_createCCtx() };
88        let cctx = ptr::NonNull::new(cctx)
89            .expect("Failed to create ZStdandard compression context");
90        Self { ptr: cctx, level }
91    }
92
93    fn levels_range() -> RangeInclusive<usize> {
94        1..=unsafe { ZSTD_maxCLevel() as usize }
95    }
96}
97
98unsafe impl Send for ZStdCctx {}
99
100impl Drop for ZStdCctx {
101    fn drop(&mut self) {
102        if unsafe { ZSTD_isError(ZSTD_freeCCtx(self.ptr.as_ptr())) } == 1 {
103            panic!("Failed to deallocate ZStd compression context");
104        }
105    }
106}