1use zfp_sys;
2use anyhow::{Result, Error};
3use bytes::ByteOrder;
4use ndarray::{Array1, Array, Array2};
5use std::ops::Deref;
6
7pub trait Zfp<T> {
8 fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>>;
9 fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>>;
10 fn from_zfp_compressed_bytes(compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>>;
11 fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>>;
12}
13
14fn zfp_compress_raw(field: *mut zfp_sys::zfp_field, tolerance: f64) -> Result<Vec<u8>> {
15 let zfp = unsafe {
16 zfp_sys::zfp_stream_open(std::ptr::null_mut() as *mut zfp_sys::bitstream)
17 };
18 unsafe {
19 zfp_sys::zfp_stream_set_accuracy(zfp, tolerance);
20 }
21
22 let bufsize = unsafe { zfp_sys::zfp_stream_maximum_size(zfp, field) };
23 let mut buffer: Vec<u8> = vec![0; bufsize];
24
25 let stream = unsafe { zfp_sys::stream_open(buffer.as_mut_ptr() as *mut std::ffi::c_void, bufsize) };
26 unsafe { zfp_sys::zfp_stream_set_bit_stream(zfp, stream); }
27
28 let zfpsize = unsafe { zfp_sys::zfp_compress(zfp, field) };
29
30 unsafe {
31 zfp_sys::zfp_field_free(field);
32 zfp_sys::zfp_stream_close(zfp);
33 zfp_sys::stream_close(stream);
34 }
35
36 if zfpsize == 0 {
37 return Err(anyhow::anyhow!("compress failed"));
38 }
39 return Ok(buffer[0..zfpsize].to_vec());
40}
41
42fn zfp_decompress_raw(mut compressed_bytes: Vec<u8>, shape: &[usize], data_type: zfp_sys::zfp_type, tolerance: f64) -> Result<Vec<u8>> {
43 let mut buffer: Vec<u8> = Vec::new();
44 buffer.resize(shape.iter().fold(1, |acc, e| acc * e) * match data_type {
45 zfp_sys::zfp_type_zfp_type_double => std::mem::size_of::<f64>(),
46 zfp_sys::zfp_type_zfp_type_float => std::mem::size_of::<f32>(),
47 zfp_sys::zfp_type_zfp_type_int32 => std::mem::size_of::<i32>(),
48 zfp_sys::zfp_type_zfp_type_int64 => std::mem::size_of::<i64>(),
49 _ => { return Err(anyhow::anyhow!("no such type supported")); }
50 }, 0);
51
52 let field = match shape.len() {
53 1 => {
54 unsafe {
55 zfp_sys::zfp_field_1d(
56 buffer.as_mut_ptr() as *mut std::ffi::c_void,
57 data_type,
58 *(shape.get(0).unwrap()) as u32,
59 )
60 }
61 }
62 2 => {
63 unsafe {
64 zfp_sys::zfp_field_2d(
65 buffer.as_mut_ptr() as *mut std::ffi::c_void,
66 data_type,
67 *(shape.get(0).unwrap()) as u32,
68 *(shape.get(1).unwrap()) as u32,
69 )
70 }
71 }
72 _ => { unreachable!() }
73 };
74
75 let zfp = unsafe {
76 zfp_sys::zfp_stream_open(std::ptr::null_mut() as *mut zfp_sys::bitstream)
77 };
78 unsafe {
79 zfp_sys::zfp_stream_set_accuracy(zfp, tolerance);
80 }
81
82 let stream = unsafe { zfp_sys::stream_open(compressed_bytes.as_mut_ptr() as *mut std::ffi::c_void, compressed_bytes.len()) };
83 unsafe {
84 zfp_sys::zfp_stream_set_bit_stream(zfp, stream);
85 }
86
87 let ret = unsafe {
88 zfp_sys::zfp_decompress(zfp, field)
89 };
90
91 unsafe {
92 zfp_sys::zfp_field_free(field);
93 zfp_sys::zfp_stream_close(zfp);
94 zfp_sys::stream_close(stream);
95 }
96 if ret == 0 {
97 return Err(anyhow::anyhow!("cannot decompress zfp stream"));
98 } else {
99 return Ok(buffer);
100 }
101}
102
103impl Zfp<f64> for Array1<f64> {
104 fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>> {
105 let data_type = zfp_sys::zfp_type_zfp_type_double;
106 let field = unsafe {
107 zfp_sys::zfp_field_1d(
108 self.as_mut_ptr() as *mut std::ffi::c_void,
109 data_type,
110 self.len() as u32,
111 )
112 };
113 return zfp_compress_raw(field, tolerance);
114 }
115
116 fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>, Error> {
117 let mut payload = Vec::new();
118 let mut buf = [0; 8];
119 bytes::LittleEndian::write_f64(&mut buf, tolerance);
120 payload.append(&mut buf.to_vec());
121 let mut buf = [0; 8];
122 bytes::LittleEndian::write_u64(&mut buf, self.len() as u64);
123 payload.append(&mut buf.to_vec());
124 payload.append(&mut self.zfp_compress(tolerance)?);
125 return Ok(payload);
126 }
127
128 fn from_zfp_compressed_bytes(mut compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>> {
129 let data_type = zfp_sys::zfp_type_zfp_type_double;
130 let original_length = shape.iter().fold(1, |acc, n| acc * (*n));
131 let decompressed_bytes = zfp_decompress_raw(compressed_bytes, shape, data_type, tolerance)?;
132 if shape.len() != 1 {
133 return Err(anyhow::anyhow!("invalid shape"));
134 }
135 let result = Array1::from_shape_vec(
136 (shape[0], ),
137 unsafe { std::slice::from_raw_parts(decompressed_bytes.as_ptr() as *const f64, original_length).to_vec() },
138 )?;
139 return Ok(Box::new(result));
140 }
141
142 fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>> {
143 let tolerance = bytes::LittleEndian::read_f64(compressed_bytes[..std::mem::size_of::<f64>()].as_ref());
144 let compressed_bytes = &compressed_bytes[std::mem::size_of::<f64>()..];
145 let length = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
146 let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
147 return Array1::<f64>::from_zfp_compressed_bytes(compressed_bytes.to_vec(), &[length as usize], tolerance);
148 }
149}
150
151
152impl Zfp<f32> for Array1<f32> {
153 fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>> {
154 let data_type = zfp_sys::zfp_type_zfp_type_float;
155 let field = unsafe {
156 zfp_sys::zfp_field_1d(
157 self.as_mut_ptr() as *mut std::ffi::c_void,
158 data_type,
159 self.len() as u32,
160 )
161 };
162 return zfp_compress_raw(field, tolerance);
163 }
164
165 fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>, Error> {
166 let mut payload = Vec::new();
167 let mut buf = [0; 8];
168 bytes::LittleEndian::write_f64(&mut buf, tolerance);
169 payload.append(&mut buf.to_vec());
170 let mut buf = [0; 8];
171 bytes::LittleEndian::write_u64(&mut buf, self.len() as u64);
172 payload.append(&mut buf.to_vec());
173 payload.append(&mut self.zfp_compress(tolerance)?);
174 return Ok(payload);
175 }
176
177 fn from_zfp_compressed_bytes(mut compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>> {
178 let data_type = zfp_sys::zfp_type_zfp_type_float;
179 let original_length = shape.iter().fold(1, |acc, n| acc * (*n));
180 let decompressed_bytes = zfp_decompress_raw(compressed_bytes, shape, data_type, tolerance)?;
181 if shape.len() != 1 {
182 return Err(anyhow::anyhow!("invalid shape"));
183 }
184 let result = Array1::from_shape_vec(
185 (shape[0], ),
186 unsafe { std::slice::from_raw_parts(decompressed_bytes.as_ptr() as *const f32, original_length).to_vec() },
187 )?;
188 return Ok(Box::new(result));
189 }
190
191 fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>> {
192 let tolerance = bytes::LittleEndian::read_f64(compressed_bytes[..std::mem::size_of::<f64>()].as_ref());
193 let compressed_bytes = &compressed_bytes[std::mem::size_of::<f64>()..];
194 let length = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
195 let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
196 return Array1::<f32>::from_zfp_compressed_bytes(compressed_bytes.to_vec(), &[length as usize], tolerance);
197 }
198}
199
200
201impl Zfp<f32> for Array2<f32> {
202 fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>> {
203 let data_type = zfp_sys::zfp_type_zfp_type_float;
204 let field = unsafe {
205 zfp_sys::zfp_field_2d(
206 self.as_mut_ptr() as *mut std::ffi::c_void,
207 data_type,
208 *(self.shape().get(0).unwrap()) as u32,
209 *(self.shape().get(1).unwrap()) as u32,
210 )
211 };
212 return zfp_compress_raw(field, tolerance);
213 }
214
215 fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>, Error> {
216 let mut payload = Vec::new();
217 let mut buf = [0; 8];
218 bytes::LittleEndian::write_f64(&mut buf, tolerance);
219 payload.append(&mut buf.to_vec());
220 let mut buf = [0; 8];
221 bytes::LittleEndian::write_u64(&mut buf, *(self.shape().get(0).unwrap()) as u64);
222 payload.append(&mut buf.to_vec());
223 let mut buf = [0; 8];
224 bytes::LittleEndian::write_u64(&mut buf, *(self.shape().get(1).unwrap()) as u64);
225 payload.append(&mut buf.to_vec());
226 payload.append(&mut self.zfp_compress(tolerance)?);
227 return Ok(payload);
228 }
229
230 fn from_zfp_compressed_bytes(mut compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>> {
231 let data_type = zfp_sys::zfp_type_zfp_type_float;
232 let original_length = shape.iter().fold(1, |acc, n| acc * (*n));
233 let decompressed_bytes = zfp_decompress_raw(compressed_bytes, shape, data_type, tolerance)?;
234 if shape.len() != 2 {
235 return Err(anyhow::anyhow!("invalid shape"));
236 }
237 let result = Array2::from_shape_vec(
238 (shape[0], shape[1]),
239 unsafe { std::slice::from_raw_parts(decompressed_bytes.as_ptr() as *const f32, original_length).to_vec() },
240 )?;
241 return Ok(Box::new(result));
242 }
243
244 fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>> {
245 let tolerance = bytes::LittleEndian::read_f64(compressed_bytes[..std::mem::size_of::<f64>()].as_ref());
246 let compressed_bytes = &compressed_bytes[std::mem::size_of::<f64>()..];
247 let shape_0 = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
248 let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
249 let shape_1 = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
250 let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
251 return Array2::<f32>::from_zfp_compressed_bytes(compressed_bytes.to_vec(), &[shape_0 as usize, shape_1 as usize], tolerance);
252 }
253}
254
255#[cfg(test)]
256mod tests {
257 use super::*;
258
259 #[test]
260 fn test_zfp_compress_1d() -> Result<()> {
261 let input_vec = Array1::from_shape_vec((10, ), vec![1.0, 2.0, 3.0, 4.0, 5.0, 5.1, 5.2, 5.3, 5.4, 5.5])?;
262 let compressed = dbg!(input_vec.zfp_compress(0.001)?);
263 dbg!(Array1::<f64>::from_zfp_compressed_bytes(compressed, &[10], 0.001));
264 let input_vec = Array1::from_shape_vec((10, ), vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 5.1, 5.2, 5.3, 5.4, 5.5])?;
265 let compressed = dbg!(input_vec.zfp_compress(0.001)?);
266 dbg!(Array1::<f32>::from_zfp_compressed_bytes(compressed, &[10], 0.001));
267 Ok(())
268 }
269
270 #[test]
271 fn test_zfp_compress_1d_with_header() -> Result<()> {
272 let input_vec = Array1::from_shape_vec((10, ), vec![1.0, 2.0, 3.0, 4.0, 5.0, 5.1, 5.2, 5.3, 5.4, 5.5])?;
273 let compressed = input_vec.zfp_compress_with_header(0.001)?;
274 dbg!(Array1::<f64>::from_zfp_compressed_bytes_with_header(compressed));
275 let input_vec = Array1::from_shape_vec((10, ), vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 5.1, 5.2, 5.3, 5.4, 5.5])?;
276 let compressed = input_vec.zfp_compress_with_header(0.001)?;
277 dbg!(Array1::<f32>::from_zfp_compressed_bytes_with_header(compressed));
278 Ok(())
279 }
280
281 #[test]
282 fn test_zfp_compress_2d_with_header() -> Result<()> {
283 let input_vec = Array2::from_shape_vec((5, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 5.1, 5.2, 5.3, 5.4, 5.5])?;
284 let compressed = input_vec.zfp_compress_with_header(0.001)?;
285 dbg!(Array2::<f32>::from_zfp_compressed_bytes_with_header(compressed));
286 let input_vec = Array2::from_shape_vec((5, 2), vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 5.1, 5.2, 5.3, 5.4, 5.5])?;
287 let compressed = input_vec.zfp_compress_with_header(0.001)?;
288 dbg!(Array2::<f32>::from_zfp_compressed_bytes_with_header(compressed));
289 Ok(())
290 }
291}