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

numcodecs_zfp/
ffi.rs
