use crate::error::DecodeError;
use crate::filter;
use crate::format::{self, ChunkType, Header};
use crate::simd;
#[derive(Debug, Clone)]
pub struct Image {
pub width: u32,
pub height: u32,
pub pixels: Vec<u8>,
}
impl Image {
#[must_use]
pub fn new(width: u32, height: u32) -> Self {
Self {
width,
height,
pixels: vec![0; (width as usize) * (height as usize) * 4],
}
}
#[must_use]
pub const fn len(&self) -> usize {
(self.width as usize) * (self.height as usize)
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.width == 0 || self.height == 0
}
#[must_use]
pub const fn byte_len(&self) -> usize {
self.pixels.len()
}
#[must_use]
pub const fn contains(&self, x: u32, y: u32) -> bool {
x < self.width && y < self.height
}
#[must_use]
pub fn get_pixel(&self, x: u32, y: u32) -> [u8; 4] {
self.try_get_pixel(x, y)
.expect("pixel coordinate out of bounds")
}
#[must_use]
pub fn try_get_pixel(&self, x: u32, y: u32) -> Option<[u8; 4]> {
if !self.contains(x, y) {
return None;
}
let idx = ((y * self.width + x) as usize) * 4;
Some([
self.pixels[idx],
self.pixels[idx + 1],
self.pixels[idx + 2],
self.pixels[idx + 3],
])
}
pub fn set_pixel(&mut self, x: u32, y: u32, rgba: [u8; 4]) {
self.try_set_pixel(x, y, rgba)
.expect("pixel coordinate out of bounds");
}
pub fn try_set_pixel(&mut self, x: u32, y: u32, rgba: [u8; 4]) -> Option<()> {
if !self.contains(x, y) {
return None;
}
let idx = ((y * self.width + x) as usize) * 4;
self.pixels[idx] = rgba[0];
self.pixels[idx + 1] = rgba[1];
self.pixels[idx + 2] = rgba[2];
self.pixels[idx + 3] = rgba[3];
Some(())
}
}
pub fn decode(data: &[u8]) -> Result<Image, DecodeError> {
let header = Header::parse(data)?;
let expected_bytes = header.expected_pixel_bytes();
let mut offset = format::HEADER_SIZE;
let mut pixels: Option<Vec<u8>> = None;
while offset < data.len() {
let chunk_header = format::ChunkHeader::parse(&data[offset..])?;
offset += format::CHUNK_HEADER_SIZE;
let chunk_data_end = offset
.checked_add(chunk_header.compressed_size as usize)
.filter(|&end| end <= data.len())
.ok_or(DecodeError::TruncatedChunk)?;
let compressed = &data[offset..chunk_data_end];
offset = chunk_data_end;
match chunk_header.chunk_type {
ChunkType::PixelData => {
if pixels.is_some() {
return Err(DecodeError::MultiplePixelData);
}
let decompressed = zstd::decode_all(std::io::Cursor::new(compressed))
.map_err(DecodeError::ZstdError)?;
let planar = if header.flags & format::FLAG_FILTERED != 0 {
let width = header.width as usize;
let height = header.height as usize;
let expected_filtered = expected_bytes + filter::overhead_bytes(height, 4);
if decompressed.len() != expected_filtered {
return Err(DecodeError::SizeMismatch {
expected: expected_filtered,
actual: decompressed.len(),
});
}
filter::unfilter_planes(&decompressed, width, height, 4).ok_or(
DecodeError::SizeMismatch {
expected: expected_bytes,
actual: decompressed.len(),
},
)?
} else {
if decompressed.len() != expected_bytes {
return Err(DecodeError::SizeMismatch {
expected: expected_bytes,
actual: decompressed.len(),
});
}
decompressed
};
let rgba = if header.flags & format::FLAG_PLANAR != 0 {
let n_pixels = header.width as usize * header.height as usize;
simd::reinterleave_rgba(&planar, n_pixels)
} else {
planar
};
pixels = Some(rgba);
}
ChunkType::Metadata | ChunkType::IccProfile => {
}
}
}
let pixels = pixels.ok_or(DecodeError::MissingPixelData)?;
Ok(Image {
width: header.width,
height: header.height,
pixels,
})
}