use zfp_sys;
use anyhow::{Result, Error};
use bytes::ByteOrder;
use ndarray::{Array1, Array, Array2};
use std::ops::Deref;
pub trait Zfp<T> {
fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>>;
fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>>;
fn from_zfp_compressed_bytes(compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>>;
fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>>;
}
fn zfp_compress_raw(field: *mut zfp_sys::zfp_field, tolerance: f64) -> Result<Vec<u8>> {
let zfp = unsafe {
zfp_sys::zfp_stream_open(std::ptr::null_mut() as *mut zfp_sys::bitstream)
};
unsafe {
zfp_sys::zfp_stream_set_accuracy(zfp, tolerance);
}
let bufsize = unsafe { zfp_sys::zfp_stream_maximum_size(zfp, field) };
let mut buffer: Vec<u8> = vec![0; bufsize];
let stream = unsafe { zfp_sys::stream_open(buffer.as_mut_ptr() as *mut std::ffi::c_void, bufsize) };
unsafe { zfp_sys::zfp_stream_set_bit_stream(zfp, stream); }
let zfpsize = unsafe { zfp_sys::zfp_compress(zfp, field) };
unsafe {
zfp_sys::zfp_field_free(field);
zfp_sys::zfp_stream_close(zfp);
zfp_sys::stream_close(stream);
}
if zfpsize == 0 {
return Err(anyhow::anyhow!("compress failed"));
}
return Ok(buffer[0..zfpsize].to_vec());
}
fn zfp_decompress_raw(mut compressed_bytes: Vec<u8>, shape: &[usize], data_type: zfp_sys::zfp_type, tolerance: f64) -> Result<Vec<u8>> {
let mut buffer: Vec<u8> = Vec::new();
buffer.resize(shape.iter().fold(1, |acc, e| acc * e) * match data_type {
zfp_sys::zfp_type_zfp_type_double => std::mem::size_of::<f64>(),
zfp_sys::zfp_type_zfp_type_float => std::mem::size_of::<f32>(),
zfp_sys::zfp_type_zfp_type_int32 => std::mem::size_of::<i32>(),
zfp_sys::zfp_type_zfp_type_int64 => std::mem::size_of::<i64>(),
_ => { return Err(anyhow::anyhow!("no such type supported")); }
}, 0);
let field = match shape.len() {
1 => {
unsafe {
zfp_sys::zfp_field_1d(
buffer.as_mut_ptr() as *mut std::ffi::c_void,
data_type,
*(shape.get(0).unwrap()) as u32,
)
}
}
2 => {
unsafe {
zfp_sys::zfp_field_2d(
buffer.as_mut_ptr() as *mut std::ffi::c_void,
data_type,
*(shape.get(0).unwrap()) as u32,
*(shape.get(1).unwrap()) as u32,
)
}
}
_ => { unreachable!() }
};
let zfp = unsafe {
zfp_sys::zfp_stream_open(std::ptr::null_mut() as *mut zfp_sys::bitstream)
};
unsafe {
zfp_sys::zfp_stream_set_accuracy(zfp, tolerance);
}
let stream = unsafe { zfp_sys::stream_open(compressed_bytes.as_mut_ptr() as *mut std::ffi::c_void, compressed_bytes.len()) };
unsafe {
zfp_sys::zfp_stream_set_bit_stream(zfp, stream);
}
let ret = unsafe {
zfp_sys::zfp_decompress(zfp, field)
};
unsafe {
zfp_sys::zfp_field_free(field);
zfp_sys::zfp_stream_close(zfp);
zfp_sys::stream_close(stream);
}
if ret == 0 {
return Err(anyhow::anyhow!("cannot decompress zfp stream"));
} else {
return Ok(buffer);
}
}
impl Zfp<f64> for Array1<f64> {
fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>> {
let data_type = zfp_sys::zfp_type_zfp_type_double;
let field = unsafe {
zfp_sys::zfp_field_1d(
self.as_mut_ptr() as *mut std::ffi::c_void,
data_type,
self.len() as u32,
)
};
return zfp_compress_raw(field, tolerance);
}
fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>, Error> {
let mut payload = Vec::new();
let mut buf = [0; 8];
bytes::LittleEndian::write_f64(&mut buf, tolerance);
payload.append(&mut buf.to_vec());
let mut buf = [0; 8];
bytes::LittleEndian::write_u64(&mut buf, self.len() as u64);
payload.append(&mut buf.to_vec());
payload.append(&mut self.zfp_compress(tolerance)?);
return Ok(payload);
}
fn from_zfp_compressed_bytes(mut compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>> {
let data_type = zfp_sys::zfp_type_zfp_type_double;
let original_length = shape.iter().fold(1, |acc, n| acc * (*n));
let decompressed_bytes = zfp_decompress_raw(compressed_bytes, shape, data_type, tolerance)?;
if shape.len() != 1 {
return Err(anyhow::anyhow!("invalid shape"));
}
let result = Array1::from_shape_vec(
(shape[0], ),
unsafe { std::slice::from_raw_parts(decompressed_bytes.as_ptr() as *const f64, original_length).to_vec() },
)?;
return Ok(Box::new(result));
}
fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>> {
let tolerance = bytes::LittleEndian::read_f64(compressed_bytes[..std::mem::size_of::<f64>()].as_ref());
let compressed_bytes = &compressed_bytes[std::mem::size_of::<f64>()..];
let length = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
return Array1::<f64>::from_zfp_compressed_bytes(compressed_bytes.to_vec(), &[length as usize], tolerance);
}
}
impl Zfp<f32> for Array1<f32> {
fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>> {
let data_type = zfp_sys::zfp_type_zfp_type_float;
let field = unsafe {
zfp_sys::zfp_field_1d(
self.as_mut_ptr() as *mut std::ffi::c_void,
data_type,
self.len() as u32,
)
};
return zfp_compress_raw(field, tolerance);
}
fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>, Error> {
let mut payload = Vec::new();
let mut buf = [0; 8];
bytes::LittleEndian::write_f64(&mut buf, tolerance);
payload.append(&mut buf.to_vec());
let mut buf = [0; 8];
bytes::LittleEndian::write_u64(&mut buf, self.len() as u64);
payload.append(&mut buf.to_vec());
payload.append(&mut self.zfp_compress(tolerance)?);
return Ok(payload);
}
fn from_zfp_compressed_bytes(mut compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>> {
let data_type = zfp_sys::zfp_type_zfp_type_float;
let original_length = shape.iter().fold(1, |acc, n| acc * (*n));
let decompressed_bytes = zfp_decompress_raw(compressed_bytes, shape, data_type, tolerance)?;
if shape.len() != 1 {
return Err(anyhow::anyhow!("invalid shape"));
}
let result = Array1::from_shape_vec(
(shape[0], ),
unsafe { std::slice::from_raw_parts(decompressed_bytes.as_ptr() as *const f32, original_length).to_vec() },
)?;
return Ok(Box::new(result));
}
fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>> {
let tolerance = bytes::LittleEndian::read_f64(compressed_bytes[..std::mem::size_of::<f64>()].as_ref());
let compressed_bytes = &compressed_bytes[std::mem::size_of::<f64>()..];
let length = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
return Array1::<f32>::from_zfp_compressed_bytes(compressed_bytes.to_vec(), &[length as usize], tolerance);
}
}
impl Zfp<f32> for Array2<f32> {
fn zfp_compress(mut self, tolerance: f64) -> Result<Vec<u8>> {
let data_type = zfp_sys::zfp_type_zfp_type_float;
let field = unsafe {
zfp_sys::zfp_field_2d(
self.as_mut_ptr() as *mut std::ffi::c_void,
data_type,
*(self.shape().get(0).unwrap()) as u32,
*(self.shape().get(1).unwrap()) as u32,
)
};
return zfp_compress_raw(field, tolerance);
}
fn zfp_compress_with_header(mut self, tolerance: f64) -> Result<Vec<u8>, Error> {
let mut payload = Vec::new();
let mut buf = [0; 8];
bytes::LittleEndian::write_f64(&mut buf, tolerance);
payload.append(&mut buf.to_vec());
let mut buf = [0; 8];
bytes::LittleEndian::write_u64(&mut buf, *(self.shape().get(0).unwrap()) as u64);
payload.append(&mut buf.to_vec());
let mut buf = [0; 8];
bytes::LittleEndian::write_u64(&mut buf, *(self.shape().get(1).unwrap()) as u64);
payload.append(&mut buf.to_vec());
payload.append(&mut self.zfp_compress(tolerance)?);
return Ok(payload);
}
fn from_zfp_compressed_bytes(mut compressed_bytes: Vec<u8>, shape: &[usize], tolerance: f64) -> Result<Box<Self>> {
let data_type = zfp_sys::zfp_type_zfp_type_float;
let original_length = shape.iter().fold(1, |acc, n| acc * (*n));
let decompressed_bytes = zfp_decompress_raw(compressed_bytes, shape, data_type, tolerance)?;
if shape.len() != 2 {
return Err(anyhow::anyhow!("invalid shape"));
}
let result = Array2::from_shape_vec(
(shape[0], shape[1]),
unsafe { std::slice::from_raw_parts(decompressed_bytes.as_ptr() as *const f32, original_length).to_vec() },
)?;
return Ok(Box::new(result));
}
fn from_zfp_compressed_bytes_with_header(compressed_bytes: Vec<u8>) -> Result<Box<Self>> {
let tolerance = bytes::LittleEndian::read_f64(compressed_bytes[..std::mem::size_of::<f64>()].as_ref());
let compressed_bytes = &compressed_bytes[std::mem::size_of::<f64>()..];
let shape_0 = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
let shape_1 = bytes::LittleEndian::read_u64(compressed_bytes[..std::mem::size_of::<u64>()].as_ref());
let compressed_bytes = &compressed_bytes[std::mem::size_of::<u64>()..];
return Array2::<f32>::from_zfp_compressed_bytes(compressed_bytes.to_vec(), &[shape_0 as usize, shape_1 as usize], tolerance);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_zfp_compress_1d() -> Result<()> {
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])?;
let compressed = dbg!(input_vec.zfp_compress(0.001)?);
dbg!(Array1::<f64>::from_zfp_compressed_bytes(compressed, &[10], 0.001));
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])?;
let compressed = dbg!(input_vec.zfp_compress(0.001)?);
dbg!(Array1::<f32>::from_zfp_compressed_bytes(compressed, &[10], 0.001));
Ok(())
}
#[test]
fn test_zfp_compress_1d_with_header() -> Result<()> {
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])?;
let compressed = input_vec.zfp_compress_with_header(0.001)?;
dbg!(Array1::<f64>::from_zfp_compressed_bytes_with_header(compressed));
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])?;
let compressed = input_vec.zfp_compress_with_header(0.001)?;
dbg!(Array1::<f32>::from_zfp_compressed_bytes_with_header(compressed));
Ok(())
}
#[test]
fn test_zfp_compress_2d_with_header() -> Result<()> {
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])?;
let compressed = input_vec.zfp_compress_with_header(0.001)?;
dbg!(Array2::<f32>::from_zfp_compressed_bytes_with_header(compressed));
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])?;
let compressed = input_vec.zfp_compress_with_header(0.001)?;
dbg!(Array2::<f32>::from_zfp_compressed_bytes_with_header(compressed));
Ok(())
}
}