numcodecs_zfp_classic/
ffi.rs

1#![expect(unsafe_code)] // FFI
2
3use std::{marker::PhantomData, mem::ManuallyDrop};
4
5use ndarray::{ArrayView, ArrayViewMut, Dimension, IxDyn};
6use numcodecs::ArrayDType;
7
8use crate::{ZfpClassicCodecError, ZfpCompressionMode, ZfpDType};
9
10const ZFP_HEADER_NO_META: u32 = zfp_sys::ZFP_HEADER_FULL & !zfp_sys::ZFP_HEADER_META;
11
12pub struct ZfpField<'a, T: ZfpCompressible> {
13    field: *mut zfp_sys::zfp_field,
14    dims: u32,
15    _marker: PhantomData<&'a T>,
16}
17
18impl<'a, T: ZfpCompressible> ZfpField<'a, T> {
19    #[expect(clippy::needless_pass_by_value)]
20    pub fn new<D: Dimension>(data: ArrayView<'a, T, D>) -> Result<Self, ZfpClassicCodecError> {
21        let pointer: *mut std::ffi::c_void = data.as_ptr().cast::<std::ffi::c_void>().cast_mut();
22
23        let (field, dims) = match (data.shape(), data.strides()) {
24            ([nx], [sx]) => unsafe {
25                let field = zfp_sys::zfp_field_1d(pointer, T::Z_TYPE, *nx);
26                zfp_sys::zfp_field_set_stride_1d(field, *sx);
27                (field, 1)
28            },
29            ([ny, nx], [sy, sx]) => unsafe {
30                let field = zfp_sys::zfp_field_2d(pointer, T::Z_TYPE, *nx, *ny);
31                zfp_sys::zfp_field_set_stride_2d(field, *sx, *sy);
32                (field, 2)
33            },
34            ([nz, ny, nx], [sz, sy, sx]) => unsafe {
35                let field = zfp_sys::zfp_field_3d(pointer, T::Z_TYPE, *nx, *ny, *nz);
36                zfp_sys::zfp_field_set_stride_3d(field, *sx, *sy, *sz);
37                (field, 3)
38            },
39            ([nw, nz, ny, nx], [sw, sz, sy, sx]) => unsafe {
40                let field = zfp_sys::zfp_field_4d(pointer, T::Z_TYPE, *nx, *ny, *nz, *nw);
41                zfp_sys::zfp_field_set_stride_4d(field, *sx, *sy, *sz, *sw);
42                (field, 4)
43            },
44            (shape, _strides) => {
45                return Err(ZfpClassicCodecError::ExcessiveDimensionality {
46                    shape: shape.to_vec(),
47                })
48            }
49        };
50
51        Ok(Self {
52            field,
53            dims,
54            _marker: PhantomData::<&'a T>,
55        })
56    }
57}
58
59impl<T: ZfpCompressible> Drop for ZfpField<'_, T> {
60    fn drop(&mut self) {
61        unsafe { zfp_sys::zfp_field_free(self.field) };
62    }
63}
64
65pub struct ZfpCompressionStream<T: ZfpCompressible> {
66    stream: *mut zfp_sys::zfp_stream,
67    _marker: PhantomData<T>,
68}
69
70impl<T: ZfpCompressible> ZfpCompressionStream<T> {
71    pub fn new(
72        field: &ZfpField<T>,
73        mode: &ZfpCompressionMode,
74    ) -> Result<Self, ZfpClassicCodecError> {
75        let stream = unsafe { zfp_sys::zfp_stream_open(std::ptr::null_mut()) };
76        let stream = Self {
77            stream,
78            _marker: PhantomData::<T>,
79        };
80
81        match mode {
82            ZfpCompressionMode::Expert {
83                min_bits,
84                max_bits,
85                max_prec,
86                min_exp,
87            } => {
88                #[expect(clippy::cast_possible_wrap)]
89                const ZFP_TRUE: zfp_sys::zfp_bool = zfp_sys::zfp_true as zfp_sys::zfp_bool;
90
91                if unsafe {
92                    zfp_sys::zfp_stream_set_params(
93                        stream.stream,
94                        *min_bits,
95                        *max_bits,
96                        *max_prec,
97                        *min_exp,
98                    )
99                } != ZFP_TRUE
100                {
101                    return Err(ZfpClassicCodecError::InvalidExpertMode { mode: mode.clone() });
102                }
103            }
104            ZfpCompressionMode::FixedRate { rate } => {
105                let _actual_rate: f64 = unsafe {
106                    zfp_sys::zfp_stream_set_rate(stream.stream, *rate, T::Z_TYPE, field.dims, 0)
107                };
108            }
109            ZfpCompressionMode::FixedPrecision { precision } => {
110                let _actual_precision: u32 =
111                    unsafe { zfp_sys::zfp_stream_set_precision(stream.stream, *precision) };
112            }
113            ZfpCompressionMode::FixedAccuracy { tolerance } => {
114                let _actual_tolerance: f64 =
115                    unsafe { zfp_sys::zfp_stream_set_accuracy(stream.stream, *tolerance) };
116            }
117            ZfpCompressionMode::Reversible => {
118                let () = unsafe { zfp_sys::zfp_stream_set_reversible(stream.stream) };
119            }
120        }
121
122        Ok(stream)
123    }
124
125    #[must_use]
126    pub fn with_bitstream<'a, 'b>(
127        self,
128        field: ZfpField<'a, T>,
129        buffer: &'b mut Vec<u8>,
130    ) -> ZfpCompressionStreamWithBitstream<'a, 'b, T> {
131        let this = ManuallyDrop::new(self);
132        let field = ManuallyDrop::new(field);
133
134        let capacity = unsafe { zfp_sys::zfp_stream_maximum_size(this.stream, field.field) };
135        buffer.reserve(capacity);
136
137        let bitstream = unsafe {
138            zfp_sys::stream_open(buffer.spare_capacity_mut().as_mut_ptr().cast(), capacity)
139        };
140
141        unsafe { zfp_sys::zfp_stream_set_bit_stream(this.stream, bitstream) };
142        unsafe { zfp_sys::zfp_stream_rewind(this.stream) };
143
144        ZfpCompressionStreamWithBitstream {
145            stream: this.stream,
146            bitstream,
147            field: field.field,
148            buffer,
149            _marker: PhantomData::<&'a T>,
150        }
151    }
152}
153
154impl<T: ZfpCompressible> Drop for ZfpCompressionStream<T> {
155    fn drop(&mut self) {
156        unsafe { zfp_sys::zfp_stream_close(self.stream) };
157    }
158}
159
160pub struct ZfpCompressionStreamWithBitstream<'a, 'b, T: ZfpCompressible> {
161    stream: *mut zfp_sys::zfp_stream,
162    bitstream: *mut zfp_sys::bitstream,
163    field: *mut zfp_sys::zfp_field,
164    buffer: &'b mut Vec<u8>,
165    _marker: PhantomData<&'a T>,
166}
167
168impl<'a, 'b, T: ZfpCompressible> ZfpCompressionStreamWithBitstream<'a, 'b, T> {
169    pub fn write_header(
170        self,
171    ) -> Result<ZfpCompressionStreamWithBitstreamWithHeader<'a, 'b, T>, ZfpClassicCodecError> {
172        if unsafe { zfp_sys::zfp_write_header(self.stream, self.field, ZFP_HEADER_NO_META) } == 0 {
173            return Err(ZfpClassicCodecError::HeaderEncodeFailed);
174        }
175
176        let mut this = ManuallyDrop::new(self);
177
178        Ok(ZfpCompressionStreamWithBitstreamWithHeader {
179            stream: this.stream,
180            bitstream: this.bitstream,
181            field: this.field,
182            // Safety: self is consumed, buffer is not read inside drop,
183            //         the lifetime is carried on
184            buffer: unsafe { &mut *std::ptr::from_mut(this.buffer) },
185            _marker: PhantomData::<&'a T>,
186        })
187    }
188}
189
190impl<T: ZfpCompressible> Drop for ZfpCompressionStreamWithBitstream<'_, '_, T> {
191    fn drop(&mut self) {
192        unsafe { zfp_sys::zfp_field_free(self.field) };
193        unsafe { zfp_sys::zfp_stream_close(self.stream) };
194        unsafe { zfp_sys::stream_close(self.bitstream) };
195    }
196}
197
198pub struct ZfpCompressionStreamWithBitstreamWithHeader<'a, 'b, T: ZfpCompressible> {
199    stream: *mut zfp_sys::zfp_stream,
200    bitstream: *mut zfp_sys::bitstream,
201    field: *mut zfp_sys::zfp_field,
202    buffer: &'b mut Vec<u8>,
203    _marker: PhantomData<&'a T>,
204}
205
206impl<T: ZfpCompressible> ZfpCompressionStreamWithBitstreamWithHeader<'_, '_, T> {
207    pub fn compress(self) -> Result<(), ZfpClassicCodecError> {
208        let compressed_size = unsafe { zfp_sys::zfp_compress(self.stream, self.field) };
209
210        if compressed_size == 0 {
211            return Err(ZfpClassicCodecError::ZfpEncodeFailed);
212        }
213
214        // Safety: compressed_size bytes of the spare capacity have now been
215        //         written to and initialized
216        unsafe {
217            self.buffer.set_len(self.buffer.len() + compressed_size);
218        }
219
220        Ok(())
221    }
222}
223
224impl<T: ZfpCompressible> Drop for ZfpCompressionStreamWithBitstreamWithHeader<'_, '_, T> {
225    fn drop(&mut self) {
226        unsafe { zfp_sys::zfp_field_free(self.field) };
227        unsafe { zfp_sys::zfp_stream_close(self.stream) };
228        unsafe { zfp_sys::stream_close(self.bitstream) };
229    }
230}
231
232pub struct ZfpDecompressionStream<'a> {
233    stream: *mut zfp_sys::zfp_stream,
234    bitstream: *mut zfp_sys::bitstream,
235    data: &'a [u8],
236}
237
238impl<'a> ZfpDecompressionStream<'a> {
239    #[must_use]
240    pub fn new(data: &'a [u8]) -> Self {
241        let bitstream = unsafe {
242            zfp_sys::stream_open(
243                data.as_ptr().cast::<std::ffi::c_void>().cast_mut(),
244                data.len(),
245            )
246        };
247
248        let stream = unsafe { zfp_sys::zfp_stream_open(bitstream) };
249
250        Self {
251            stream,
252            bitstream,
253            data,
254        }
255    }
256
257    pub fn read_header(self) -> Result<ZfpDecompressionStreamWithHeader<'a>, ZfpClassicCodecError> {
258        let this = ManuallyDrop::new(self);
259
260        let field = unsafe { zfp_sys::zfp_field_alloc() };
261
262        let stream = ZfpDecompressionStreamWithHeader {
263            stream: this.stream,
264            bitstream: this.bitstream,
265            field,
266            _data: this.data,
267        };
268
269        if unsafe { zfp_sys::zfp_read_header(this.stream, field, ZFP_HEADER_NO_META) } == 0 {
270            return Err(ZfpClassicCodecError::HeaderDecodeFailed);
271        }
272
273        Ok(stream)
274    }
275}
276
277impl Drop for ZfpDecompressionStream<'_> {
278    fn drop(&mut self) {
279        unsafe { zfp_sys::zfp_stream_close(self.stream) };
280        unsafe { zfp_sys::stream_close(self.bitstream) };
281    }
282}
283
284pub struct ZfpDecompressionStreamWithHeader<'a> {
285    stream: *mut zfp_sys::zfp_stream,
286    bitstream: *mut zfp_sys::bitstream,
287    field: *mut zfp_sys::zfp_field,
288    _data: &'a [u8],
289}
290
291impl ZfpDecompressionStreamWithHeader<'_> {
292    pub fn decompress_into<T: ZfpCompressible>(
293        self,
294        mut decompressed: ArrayViewMut<T, IxDyn>,
295    ) -> Result<(), ZfpClassicCodecError> {
296        unsafe { zfp_sys::zfp_field_set_type(self.field, T::Z_TYPE) };
297
298        match (decompressed.shape(), decompressed.strides()) {
299            ([nx], [sx]) => unsafe {
300                zfp_sys::zfp_field_set_size_1d(self.field, *nx);
301                zfp_sys::zfp_field_set_stride_1d(self.field, *sx);
302            },
303            ([ny, nx], [sy, sx]) => unsafe {
304                zfp_sys::zfp_field_set_size_2d(self.field, *nx, *ny);
305                zfp_sys::zfp_field_set_stride_2d(self.field, *sx, *sy);
306            },
307            ([nz, ny, nx], [sz, sy, sx]) => unsafe {
308                zfp_sys::zfp_field_set_size_3d(self.field, *nx, *ny, *nz);
309                zfp_sys::zfp_field_set_stride_3d(self.field, *sx, *sy, *sz);
310            },
311            ([nw, nz, ny, nx], [sw, sz, sy, sx]) => unsafe {
312                zfp_sys::zfp_field_set_size_4d(self.field, *nx, *ny, *nz, *nw);
313                zfp_sys::zfp_field_set_stride_4d(self.field, *sx, *sy, *sz, *sw);
314            },
315            (shape, _strides) => {
316                return Err(ZfpClassicCodecError::ExcessiveDimensionality {
317                    shape: shape.to_vec(),
318                })
319            }
320        };
321
322        unsafe {
323            zfp_sys::zfp_field_set_pointer(
324                self.field,
325                decompressed.as_mut_ptr().cast::<std::ffi::c_void>(),
326            );
327        }
328
329        if unsafe { zfp_sys::zfp_decompress(self.stream, self.field) } == 0 {
330            Err(ZfpClassicCodecError::ZfpDecodeFailed)
331        } else {
332            Ok(())
333        }
334    }
335}
336
337impl Drop for ZfpDecompressionStreamWithHeader<'_> {
338    fn drop(&mut self) {
339        unsafe { zfp_sys::zfp_field_free(self.field) };
340        unsafe { zfp_sys::zfp_stream_close(self.stream) };
341        unsafe { zfp_sys::stream_close(self.bitstream) };
342    }
343}
344
345pub trait ZfpCompressible: Copy + ArrayDType {
346    const D_TYPE: ZfpDType;
347    const Z_TYPE: zfp_sys::zfp_type;
348
349    fn is_finite(self) -> bool;
350}
351
352impl ZfpCompressible for i32 {
353    const D_TYPE: ZfpDType = ZfpDType::I32;
354    const Z_TYPE: zfp_sys::zfp_type = zfp_sys::zfp_type_zfp_type_int32;
355
356    fn is_finite(self) -> bool {
357        true
358    }
359}
360
361impl ZfpCompressible for i64 {
362    const D_TYPE: ZfpDType = ZfpDType::I64;
363    const Z_TYPE: zfp_sys::zfp_type = zfp_sys::zfp_type_zfp_type_int64;
364
365    fn is_finite(self) -> bool {
366        true
367    }
368}
369
370impl ZfpCompressible for f32 {
371    const D_TYPE: ZfpDType = ZfpDType::F32;
372    const Z_TYPE: zfp_sys::zfp_type = zfp_sys::zfp_type_zfp_type_float;
373
374    fn is_finite(self) -> bool {
375        Self::is_finite(self)
376    }
377}
378
379impl ZfpCompressible for f64 {
380    const D_TYPE: ZfpDType = ZfpDType::F64;
381    const Z_TYPE: zfp_sys::zfp_type = zfp_sys::zfp_type_zfp_type_double;
382
383    fn is_finite(self) -> bool {
384        Self::is_finite(self)
385    }
386}
numcodecs_zfp_classic/ffi.rs

numcodecs_zfp_classic/
ffi.rs
