use crate::engine::block::{VolumeShape, VoxelBlock};
use crate::engine::codec::{EndianCodec, decode_slice, encode_slice};
use crate::engine::endian::FileEndian;
use crate::iter::{RegionIter, SlabStepper, SliceStepper, TileStepper};
use crate::mode::{ComplexToRealStrategy, Float32Complex, Int16Complex, M0Interpretation, Voxel};
use crate::{Error, Header, Mode};
use std::borrow::Cow;
type VoxelIter<'a, T> = Box<dyn Iterator<Item = Result<VoxelBlock<T>, Error>> + 'a>;
mod private {
pub trait Sealed {}
}
#[doc(hidden)]
pub trait VoxelSource: private::Sealed {
fn vs_read_block_bytes<'a>(
&'a self,
offset: [usize; 3],
shape: [usize; 3],
) -> Result<Cow<'a, [u8]>, Error>;
fn vs_decode_block<T: Voxel>(&self, bytes: &[u8]) -> Result<Vec<T>, Error>;
}
#[doc(hidden)]
#[allow(dead_code)]
pub trait ReaderCore: VoxelSource {
fn shape(&self) -> VolumeShape;
fn mode(&self) -> Mode;
fn endian(&self) -> FileEndian;
fn header(&self) -> &Header;
fn ext_header_bytes(&self) -> &[u8];
}
macro_rules! impl_inherent_reader_methods {
($ty:ty) => {
impl $ty {
pub fn slices<T: Voxel>(&self) -> RegionIter<'_, T, $ty, SliceStepper> {
RegionIter::with_stepper(self, self.shape(), SliceStepper::new())
}
pub fn slabs<T: Voxel>(&self, k: usize) -> RegionIter<'_, T, $ty, SlabStepper> {
RegionIter::with_stepper(self, self.shape(), SlabStepper::new(k))
}
pub fn tiles<T: Voxel>(&self, tile_shape: [usize; 3]) -> RegionIter<'_, T, $ty, TileStepper> {
RegionIter::with_stepper(self, self.shape(), TileStepper::new(tile_shape))
}
pub fn images<T: Voxel>(&self) -> RegionIter<'_, T, $ty, SliceStepper> {
self.slices()
}
pub fn image_stack<T: Voxel>(&self, k: usize) -> RegionIter<'_, T, $ty, SlabStepper> {
self.slabs(k)
}
pub fn planes<T: Voxel>(&self) -> RegionIter<'_, T, $ty, SliceStepper> {
self.slices()
}
pub fn plane_stack<T: Voxel>(&self, k: usize) -> RegionIter<'_, T, $ty, SlabStepper> {
self.slabs(k)
}
pub fn volumes<T: Voxel>(&self) -> Result<RegionIter<'_, T, $ty, SlabStepper>, Error> {
let mz = self.header().mz.max(0) as usize;
if !self.header().is_volume_stack() || mz == 0 {
return Err(Error::NotAVolumeStack {
ispg: self.header().ispg,
mz: self.header().mz,
});
}
Ok(self.slabs(mz))
}
pub fn subregion<T: Voxel>(&self, offset: [usize; 3], shape: [usize; 3]) -> Result<VoxelBlock<T>, Error> {
let bytes = self.vs_read_block_bytes(offset, shape)?;
let data = self.vs_decode_block::<T>(&bytes)?;
Ok(VoxelBlock { offset, shape, data })
}
pub fn read_volume<T: Voxel>(&self) -> Result<VoxelBlock<T>, Error> {
let s = self.shape();
self.subregion([0, 0, 0], [s.nx, s.ny, s.nz])
}
pub fn read_volume_f32(&self) -> Result<VoxelBlock<f32>, Error> {
let shape = self.shape();
let offset = [0, 0, 0];
let block_shape = [shape.nx, shape.ny, shape.nz];
match self.mode() {
Mode::Float32 => self.subregion::<f32>(offset, block_shape),
Mode::Int16 => {
let block = self.subregion::<i16>(offset, block_shape)?;
let data = crate::engine::convert::convert_i16_slice_to_f32(&block.data);
Ok(VoxelBlock { offset, shape: block_shape, data })
}
Mode::Uint16 => {
let block = self.subregion::<u16>(offset, block_shape)?;
let data = crate::engine::convert::convert_u16_slice_to_f32(&block.data);
Ok(VoxelBlock { offset, shape: block_shape, data })
}
Mode::Int8 => {
let block = self.subregion::<i8>(offset, block_shape)?;
let data = crate::engine::convert::convert_i8_slice_to_f32(&block.data);
Ok(VoxelBlock { offset, shape: block_shape, data })
}
#[cfg(feature = "f16")]
Mode::Float16 => {
let block = self.subregion::<crate::f16>(offset, block_shape)?;
let data = block.data.iter().map(|&v| f32::from(v)).collect();
Ok(VoxelBlock { offset, shape: block_shape, data })
}
#[cfg(not(feature = "f16"))]
Mode::Float16 => Err(Error::UnsupportedMode),
Mode::Float32Complex => {
let block = self.subregion::<Float32Complex>(offset, block_shape)?;
let data = block.data.iter()
.map(|c| c.to_real(ComplexToRealStrategy::Magnitude))
.collect();
Ok(VoxelBlock { offset, shape: block_shape, data })
}
Mode::Int16Complex => {
let block = self.subregion::<Int16Complex>(offset, block_shape)?;
let data = block.data.iter()
.map(|c| c.to_real(ComplexToRealStrategy::Magnitude))
.collect();
Ok(VoxelBlock { offset, shape: block_shape, data })
}
_ => Err(Error::UnsupportedMode),
}
}
pub fn slices_f32(&self) -> VoxelIter<'_, f32> {
iter_f32_helper(
self.mode(),
self.slices::<f32>(),
self.slices::<i16>(),
self.slices::<u16>(),
self.slices::<i8>(),
#[cfg(feature = "f16")]
self.slices::<crate::f16>(),
self.slices::<Float32Complex>(),
self.slices::<Int16Complex>(),
)
}
pub fn slabs_f32(&self, k: usize) -> VoxelIter<'_, f32> {
iter_f32_helper(
self.mode(),
self.slabs::<f32>(k),
self.slabs::<i16>(k),
self.slabs::<u16>(k),
self.slabs::<i8>(k),
#[cfg(feature = "f16")]
self.slabs::<crate::f16>(k),
self.slabs::<Float32Complex>(k),
self.slabs::<Int16Complex>(k),
)
}
pub fn slices_u8(&self) -> VoxelIter<'_, u8> {
if self.mode() != Mode::Uint16 {
return Box::new(std::iter::once(Err(Error::ModeMismatch {
file_mode: self.mode(),
requested_mode: Mode::Uint16,
})));
}
Box::new(self.slices::<u16>().map(|b| {
let b = b?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: crate::engine::convert::convert_u16_slice_to_u8(&b.data)?,
})
}))
}
pub fn slices_mode0(&self, interp: M0Interpretation) -> VoxelIter<'_, f32> {
if self.mode() != Mode::Int8 {
return Box::new(std::iter::once(Err(Error::ModeMismatch {
file_mode: self.mode(),
requested_mode: Mode::Int8,
})));
}
let volume_shape = self.shape();
Box::new((0..volume_shape.nz).map(move |z| {
let bytes = self.vs_read_block_bytes([0, 0, z], [volume_shape.nx, volume_shape.ny, 1])?;
let data = crate::engine::convert::reinterpret_m0(&bytes, interp);
Ok(VoxelBlock {
offset: [0, 0, z],
shape: [volume_shape.nx, volume_shape.ny, 1],
data,
})
}))
}
pub fn slabs_mode0(&self, k: usize, interp: M0Interpretation) -> VoxelIter<'_, f32> {
if self.mode() != Mode::Int8 {
return Box::new(std::iter::once(Err(Error::ModeMismatch {
file_mode: self.mode(),
requested_mode: Mode::Int8,
})));
}
let volume_shape = self.shape();
let k = k.max(1);
let mut z = 0usize;
Box::new(std::iter::from_fn(move || {
if z >= volume_shape.nz { return None; }
let start = z;
let sz = k.min(volume_shape.nz - z);
z += sz;
let bytes = match self.vs_read_block_bytes([0, 0, start], [volume_shape.nx, volume_shape.ny, sz]) {
Ok(b) => b,
Err(e) => return Some(Err(e)),
};
let data = crate::engine::convert::reinterpret_m0(&bytes, interp);
Some(Ok(VoxelBlock {
offset: [0, 0, start],
shape: [volume_shape.nx, volume_shape.ny, sz],
data,
}))
}))
}
}
};
}
impl_inherent_reader_methods!(crate::Reader);
#[cfg(feature = "mmap")]
impl_inherent_reader_methods!(crate::MmapReader);
#[cfg(feature = "f16")]
#[allow(clippy::too_many_arguments)]
fn iter_f32_helper<'a, I, I16, I16E, U16, U16E, I8, I8E, IF16, IF16E, IC32, IC32E, IC16, IC16E>(
mode: Mode,
iter_f32: I,
iter_i16: I16,
iter_u16: U16,
iter_i8: I8,
iter_f16: IF16,
iter_c32: IC32,
iter_c16: IC16,
) -> VoxelIter<'a, f32>
where
I: Iterator<Item = Result<VoxelBlock<f32>, Error>> + 'a,
I16: Iterator<Item = Result<VoxelBlock<i16>, I16E>> + 'a,
I16E: Into<Error>,
U16: Iterator<Item = Result<VoxelBlock<u16>, U16E>> + 'a,
U16E: Into<Error>,
I8: Iterator<Item = Result<VoxelBlock<i8>, I8E>> + 'a,
I8E: Into<Error>,
IF16: Iterator<Item = Result<VoxelBlock<crate::f16>, IF16E>> + 'a,
IF16E: Into<Error>,
IC32: Iterator<Item = Result<VoxelBlock<Float32Complex>, IC32E>> + 'a,
IC32E: Into<Error>,
IC16: Iterator<Item = Result<VoxelBlock<Int16Complex>, IC16E>> + 'a,
IC16E: Into<Error>,
{
match mode {
Mode::Float32 => Box::new(iter_f32) as Box<dyn Iterator<Item = _>>,
Mode::Int16 => Box::new(iter_i16.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: crate::engine::convert::convert_i16_slice_to_f32(&b.data),
})
})),
Mode::Uint16 => Box::new(iter_u16.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: crate::engine::convert::convert_u16_slice_to_f32(&b.data),
})
})),
Mode::Int8 => Box::new(iter_i8.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: crate::engine::convert::convert_i8_slice_to_f32(&b.data),
})
})),
#[cfg(feature = "f16")]
Mode::Float16 => Box::new(iter_f16.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: b.data.iter().map(|&v| f32::from(v)).collect(),
})
})),
Mode::Float32Complex => Box::new(iter_c32.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: b
.data
.iter()
.map(|c| c.to_real(crate::mode::ComplexToRealStrategy::Magnitude))
.collect(),
})
})),
Mode::Int16Complex => Box::new(iter_c16.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: b
.data
.iter()
.map(|c| c.to_real(crate::mode::ComplexToRealStrategy::Magnitude))
.collect(),
})
})),
_ => Box::new(std::iter::once(Err(Error::UnsupportedMode))),
}
}
#[cfg(not(feature = "f16"))]
#[allow(clippy::too_many_arguments)]
fn iter_f32_helper<'a, I, I16, I16E, U16, U16E, I8, I8E, IC32, IC32E, IC16, IC16E>(
mode: Mode,
iter_f32: I,
iter_i16: I16,
iter_u16: U16,
iter_i8: I8,
iter_c32: IC32,
iter_c16: IC16,
) -> VoxelIter<'a, f32>
where
I: Iterator<Item = Result<VoxelBlock<f32>, Error>> + 'a,
I16: Iterator<Item = Result<VoxelBlock<i16>, I16E>> + 'a,
I16E: Into<Error>,
U16: Iterator<Item = Result<VoxelBlock<u16>, U16E>> + 'a,
U16E: Into<Error>,
I8: Iterator<Item = Result<VoxelBlock<i8>, I8E>> + 'a,
I8E: Into<Error>,
IC32: Iterator<Item = Result<VoxelBlock<Float32Complex>, IC32E>> + 'a,
IC32E: Into<Error>,
IC16: Iterator<Item = Result<VoxelBlock<Int16Complex>, IC16E>> + 'a,
IC16E: Into<Error>,
{
match mode {
Mode::Float32 => Box::new(iter_f32) as Box<dyn Iterator<Item = _>>,
Mode::Int16 => Box::new(iter_i16.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: crate::engine::convert::convert_i16_slice_to_f32(&b.data),
})
})),
Mode::Uint16 => Box::new(iter_u16.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: crate::engine::convert::convert_u16_slice_to_f32(&b.data),
})
})),
Mode::Int8 => Box::new(iter_i8.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: crate::engine::convert::convert_i8_slice_to_f32(&b.data),
})
})),
Mode::Float16 => Box::new(std::iter::once(Err(Error::UnsupportedMode))),
Mode::Float32Complex => Box::new(iter_c32.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: b
.data
.iter()
.map(|c| c.to_real(crate::mode::ComplexToRealStrategy::Magnitude))
.collect(),
})
})),
Mode::Int16Complex => Box::new(iter_c16.map(|b| {
let b = b.map_err(Into::into)?;
Ok(VoxelBlock {
offset: b.offset,
shape: b.shape,
data: b
.data
.iter()
.map(|c| c.to_real(crate::mode::ComplexToRealStrategy::Magnitude))
.collect(),
})
})),
_ => Box::new(std::iter::once(Err(Error::UnsupportedMode))),
}
}
pub(crate) fn validate_block_bounds(
volume_shape: VolumeShape,
mode: Mode,
data_len: usize,
offset: [usize; 3],
block_shape: [usize; 3],
) -> Result<usize, Error> {
let [nx, ny, nz] = [volume_shape.nx, volume_shape.ny, volume_shape.nz];
let [ox, oy, oz] = offset;
let [sx, sy, sz] = block_shape;
if ox.checked_add(sx).is_none_or(|end| end > nx)
|| oy.checked_add(sy).is_none_or(|end| end > ny)
|| oz.checked_add(sz).is_none_or(|end| end > nz)
{
return Err(Error::BoundsError);
}
if mode == Mode::Packed4Bit {
return Err(Error::UnsupportedMode);
}
let count = sx
.checked_mul(sy)
.and_then(|v| v.checked_mul(sz))
.ok_or(Error::BoundsError)?;
let byte_len = mode.byte_size_for_count(count);
if count == 0 {
return Ok(0);
}
let last_row_start = volume_shape
.checked_linear_index([ox, oy + sy - 1, oz + sz - 1])
.ok_or(Error::BoundsError)?;
let last_byte = last_row_start
.checked_mul(mode.byte_size())
.and_then(|s| s.checked_add(sx * mode.byte_size()))
.ok_or(Error::BoundsError)?;
if last_byte > data_len {
return Err(Error::BoundsError);
}
Ok(byte_len)
}
pub(crate) fn gather_block_bytes(
data: &[u8],
volume_shape: VolumeShape,
mode: Mode,
offset: [usize; 3],
block_shape: [usize; 3],
) -> Vec<u8> {
let [nx, ny, _nz] = [volume_shape.nx, volume_shape.ny, volume_shape.nz];
let [ox, oy, oz] = offset;
let [sx, sy, sz] = block_shape;
let b = mode.byte_size();
if ox == 0 && sx == nx && oy == 0 && sy == ny {
let linear = oz * nx * ny;
let start = linear * b;
let byte_len = sx * sy * sz * b;
return data[start..start + byte_len].to_vec();
}
let mut dst = vec![0u8; sx * sy * sz * b];
for z in 0..sz {
for y in 0..sy {
let src_linear = ox + (oy + y) * nx + (oz + z) * nx * ny;
let src_start = src_linear * b;
let dst_linear = y * sx + z * sx * sy;
let dst_start = dst_linear * b;
dst[dst_start..dst_start + sx * b]
.copy_from_slice(&data[src_start..src_start + sx * b]);
}
}
dst
}
pub(crate) fn encode_block_to_buf<T: EndianCodec + Sync>(
block: &VoxelBlock<T>,
volume_shape: VolumeShape,
bytes_per_voxel: usize,
file_endian: FileEndian,
data_offset: usize,
buf: &mut [u8],
) -> Result<(), Error> {
let [nx, ny, _nz] = [volume_shape.nx, volume_shape.ny, volume_shape.nz];
let [ox, oy, oz] = block.offset;
let [sx, sy, sz] = block.shape;
let b = bytes_per_voxel;
if ox == 0 && sx == nx && oy == 0 && sy == ny {
let linear = oz * nx * ny;
let start_byte = data_offset + linear * b;
let byte_len = sx * sy * sz * b;
let end_byte = start_byte + byte_len;
if end_byte > buf.len() {
return Err(Error::BoundsError);
}
encode_slice(&block.data, &mut buf[start_byte..end_byte], file_endian)?;
return Ok(());
}
for z in 0..sz {
for y in 0..sy {
let file_linear = ox + (oy + y) * nx + (oz + z) * nx * ny;
let file_start = data_offset + file_linear * b;
let block_idx = y * sx + z * sx * sy;
if block_idx + sx > block.data.len() {
return Err(Error::BoundsError);
}
let row_values = &block.data[block_idx..block_idx + sx];
let row_end = file_start + sx * b;
if row_end > buf.len() {
return Err(Error::BoundsError);
}
encode_slice(row_values, &mut buf[file_start..row_end], file_endian)?;
}
}
Ok(())
}
pub(crate) fn decode_block<T: Voxel>(
bytes: &[u8],
file_mode: Mode,
endian: FileEndian,
) -> Result<Vec<T>, Error> {
if T::MODE != file_mode {
return Err(Error::ModeMismatch {
file_mode,
requested_mode: T::MODE,
});
}
if endian == FileEndian::native() {
decode_native_endian(bytes)
} else {
decode_slice(bytes, endian)
}
}
pub(crate) fn decode_native_endian<T: EndianCodec + Copy>(bytes: &[u8]) -> Result<Vec<T>, Error> {
let n = bytes.len() / T::BYTE_SIZE;
debug_assert_eq!(
bytes.len() % T::BYTE_SIZE, 0,
"decode_native_endian: bytes.len() ({}) must be a multiple of T::BYTE_SIZE ({})",
bytes.len(), T::BYTE_SIZE
);
let mut result = Vec::with_capacity(n);
unsafe {
core::ptr::copy_nonoverlapping(bytes.as_ptr(), result.as_mut_ptr() as *mut u8, bytes.len());
result.set_len(n);
}
Ok(result)
}
pub(crate) fn parse_header(
header_bytes: &[u8; 1024],
permissive: bool,
) -> Result<(crate::Header, Vec<String>, crate::FileEndian, usize), crate::Error> {
let (header, endian_warning) = crate::Header::decode_from_bytes_with_info(header_bytes);
let mut warnings = if permissive {
header
.validate_permissive()
.map_err(crate::Error::InvalidHeaderDetailed)?
} else {
header
.validate_detailed()
.map_err(crate::Error::InvalidHeaderDetailed)?;
Vec::new()
};
if let Some(msg) = endian_warning {
warnings.push(msg.to_string());
}
let data_size = header.data_size().ok_or(crate::Error::InvalidHeader)?;
let endian = header.detect_endian();
Ok((header, warnings, endian, data_size))
}
pub(crate) struct DecompressedMrc {
pub header: crate::Header,
pub ext_header: Vec<u8>,
pub data: Vec<u8>,
pub warnings: Vec<String>,
pub endian: crate::FileEndian,
pub shape: VolumeShape,
}
pub(crate) fn open_compressed<D: std::io::Read>(
mut decoder: D,
permissive: bool,
) -> Result<DecompressedMrc, crate::Error> {
let mut buf = Vec::new();
decoder.read_to_end(&mut buf)?;
if buf.len() < 1024 {
return Err(crate::Error::InvalidHeader);
}
let mut header_bytes = [0u8; 1024];
header_bytes.copy_from_slice(&buf[..1024]);
let (header, mut warnings, endian, data_size) = parse_header(&header_bytes, permissive)?;
let ext_size = header.nsymbt as usize;
if !permissive {
if buf.len() != 1024 + ext_size + data_size {
return Err(crate::Error::FileSizeMismatch {
expected: 1024 + ext_size + data_size,
actual: buf.len(),
});
}
} else if buf.len() != 1024 + ext_size + data_size {
warnings.push(format!(
"File size mismatch: expected {} bytes, got {}",
1024 + ext_size + data_size,
buf.len()
));
}
let ext_end = (1024 + ext_size).min(buf.len());
let ext_header = buf[1024..ext_end].to_vec();
let data = if ext_end < buf.len() { buf[ext_end..].to_vec() } else { Vec::new() };
let shape = VolumeShape::new(header.nx as usize, header.ny as usize, header.nz as usize);
Ok(DecompressedMrc {
header,
ext_header,
data,
warnings,
endian,
shape,
})
}
impl private::Sealed for crate::Reader {}
impl VoxelSource for crate::Reader {
fn vs_read_block_bytes<'a>(
&'a self,
offset: [usize; 3],
shape: [usize; 3],
) -> Result<Cow<'a, [u8]>, Error> {
self.read_block_bytes(offset, shape).map(Cow::Owned)
}
fn vs_decode_block<T: Voxel>(&self, bytes: &[u8]) -> Result<Vec<T>, Error> {
self.decode_block(bytes)
}
}
impl ReaderCore for crate::Reader {
fn shape(&self) -> VolumeShape {
self.shape()
}
fn mode(&self) -> Mode {
self.mode()
}
fn endian(&self) -> FileEndian {
self.endian
}
fn header(&self) -> &Header {
&self.header
}
fn ext_header_bytes(&self) -> &[u8] {
&self.ext_header
}
}
#[cfg(feature = "mmap")]
impl private::Sealed for crate::MmapReader {}
#[cfg(feature = "mmap")]
impl VoxelSource for crate::MmapReader {
fn vs_read_block_bytes<'a>(
&'a self,
offset: [usize; 3],
shape: [usize; 3],
) -> Result<Cow<'a, [u8]>, Error> {
self.read_block_bytes_cow(offset, shape)
}
fn vs_decode_block<T: Voxel>(&self, bytes: &[u8]) -> Result<Vec<T>, Error> {
self.decode_block(bytes)
}
}
#[cfg(feature = "mmap")]
impl ReaderCore for crate::MmapReader {
fn shape(&self) -> VolumeShape {
self.shape()
}
fn mode(&self) -> Mode {
self.mode()
}
fn endian(&self) -> FileEndian {
self.endian()
}
fn header(&self) -> &Header {
self.header()
}
fn ext_header_bytes(&self) -> &[u8] {
self.ext_header_bytes()
}
}