use alloc::collections::BTreeMap;
use core::convert::TryFrom;
use core::marker::Copy;
use std::io::Read;
use flate2::read::ZlibDecoder;
use crate::parsers::common::Skip;
use crate::parsers::{extract, Endian, FromSlice};
use crate::record::{StateMetadata, Value};
use crate::EtError;
use crate::{impl_reader, impl_record};
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum PngColorType {
Indexed,
Grayscale,
AlphaGrayscale,
Color,
AlphaColor,
}
impl Default for PngColorType {
fn default() -> Self {
PngColorType::Indexed
}
}
impl PngColorType {
fn from_byte(byte: u8) -> Result<Self, EtError> {
match byte {
0 => Ok(PngColorType::Grayscale),
2 => Ok(PngColorType::Color),
3 => Ok(PngColorType::Indexed),
4 => Ok(PngColorType::AlphaGrayscale),
6 => Ok(PngColorType::AlphaColor),
_ => Err("Unknown PNG color type".into()),
}
}
fn pixel_size(self) -> usize {
match self {
PngColorType::Indexed | PngColorType::Grayscale => 1,
PngColorType::AlphaGrayscale => 2,
PngColorType::Color => 3,
PngColorType::AlphaColor => 4,
}
}
}
#[derive(Clone, Debug, Default)]
pub struct PngState {
color_type: PngColorType,
bit_depth: u8,
width: usize,
height: usize,
cur_x: usize,
cur_y: usize,
image_data: Vec<u8>,
palette: Option<Vec<(u16, u16, u16)>>,
}
impl PngState {
fn line_len(&self) -> usize {
1 + (self.width * self.color_type.pixel_size() * usize::from(self.bit_depth) + 7) / 8
}
fn unfilter_line(&mut self, line_num: usize) -> Result<(), EtError> {
let bytes_per_pixel = (self.color_type.pixel_size() * usize::from(self.bit_depth) + 7) / 8;
let line_len = self.line_len();
for pos in line_num * line_len + 1..(line_num + 1) * line_len {
let left = if pos < line_num * line_len + 1 + bytes_per_pixel {
0
} else {
self.image_data[pos - bytes_per_pixel]
};
let above = if line_num == 0 {
0
} else {
self.image_data[pos - line_len]
};
self.image_data[pos] = match self.image_data.get(line_num * line_len) {
Some(0) => self.image_data[pos],
Some(1) => self.image_data[pos].wrapping_add(left),
Some(2) => self.image_data[pos].wrapping_add(above),
Some(3) => {
let mut average = (left >> 1) + (above >> 1);
if left & above & 1 == 1 {
average += 1;
}
self.image_data[pos].wrapping_add(average)
}
Some(4) => {
let immediate_left = if pos == line_num * line_len + 1 {
0
} else {
self.image_data[pos - 1]
};
let above_left = if pos == line_num * line_len + 1 || line_num == 0 {
0
} else {
self.image_data[pos - 1 - line_len]
};
let estimate =
i16::from(immediate_left) + i16::from(above) - i16::from(above_left);
let pred_left = (estimate - i16::from(immediate_left)).abs();
let pred_above = (estimate - i16::from(above)).abs();
let pred_above_left = (estimate - i16::from(above_left)).abs();
let paeth = if pred_left <= pred_above && pred_left <= pred_above_left {
immediate_left
} else if pred_above <= pred_above_left {
above
} else {
above_left
};
self.image_data[pos].wrapping_add(paeth)
}
_ => return Err("Unknown line filter".into()),
}
}
self.image_data[line_num * line_len] = 0;
Ok(())
}
}
impl StateMetadata for PngState {
fn metadata(&self) -> BTreeMap<String, Value> {
let mut metadata = BTreeMap::new();
drop(metadata.insert("height".to_string(), (self.height as u64).into()));
drop(metadata.insert("width".to_string(), (self.width as u64).into()));
metadata
}
fn header(&self) -> Vec<&str> {
vec!["x", "y", "red", "green", "blue", "alpha"]
}
}
impl<'b: 's, 's> FromSlice<'b, 's> for PngState {
type State = ();
fn parse(
rb: &[u8],
_eof: bool,
consumed: &mut usize,
_state: &mut Self::State,
) -> Result<bool, EtError> {
let con = &mut 0;
if extract::<&[u8]>(rb, con, &mut 8)? != b"\x89PNG\r\n\x1A\n" {
return Err("Invalid PNG magic".into());
}
if extract::<&[u8]>(rb, con, &mut 8)? != b"\x00\x00\x00\x0DIHDR" {
return Err("Invalid PNG header".into());
}
let _ = extract::<Skip>(rb, con, &mut 10)?;
if extract::<u8>(rb, con, &mut Endian::Big)? != 0 {
return Err("PNG compression must be type 0".into());
}
if extract::<u8>(rb, con, &mut Endian::Big)? != 0 {
return Err("PNG filtering must be type 0".into());
}
if extract::<u8>(rb, con, &mut Endian::Big)? != 0 {
return Err("PNG interlacing not supported yet".into());
}
loop {
let _ = extract::<&[u8]>(rb, con, &mut 4)?;
let mut chunk_size = extract::<u32>(rb, con, &mut Endian::Big)? as usize;
let chunk_header = extract::<&[u8]>(rb, con, &mut 4)?;
if &chunk_header == b"IEND" {
break;
}
let _ = extract::<&[u8]>(rb, con, &mut chunk_size)?;
}
*consumed += *con;
Ok(true)
}
fn get(&mut self, rb: &'b [u8], _state: &'s Self::State) -> Result<(), EtError> {
let con = &mut 16;
self.width = extract::<u32>(rb, con, &mut Endian::Big)? as usize;
self.height = extract::<u32>(rb, con, &mut Endian::Big)? as usize;
self.bit_depth = extract(rb, con, &mut Endian::Big)?;
self.color_type = PngColorType::from_byte(extract(rb, con, &mut Endian::Big)?)?;
*con += 3;
let mut compressed_data = Vec::new();
loop {
let _ = extract::<&[u8]>(rb, con, &mut 4)?;
let mut chunk_size = extract::<u32>(rb, con, &mut Endian::Big)? as usize;
let chunk_header = extract::<&[u8]>(rb, con, &mut 4)?;
match chunk_header {
b"PLTE" => {
let mut raw_palette = Vec::new();
for _ in 0..chunk_size / 3 {
let r: u8 = extract(rb, con, &mut Endian::Big)?;
let g: u8 = extract(rb, con, &mut Endian::Big)?;
let b: u8 = extract(rb, con, &mut Endian::Big)?;
raw_palette.push((
257 * u16::from(r),
257 * u16::from(g),
257 * u16::from(b),
));
}
self.palette = Some(raw_palette);
}
b"IDAT" => {
compressed_data.extend_from_slice(extract(rb, con, &mut chunk_size)?);
}
b"IEND" => {
break;
}
_ => {
let _ = extract::<&[u8]>(rb, con, &mut chunk_size)?;
}
}
}
let _ = ZlibDecoder::new(&compressed_data[..]).read_to_end(&mut self.image_data)?;
self.cur_x = usize::MAX;
self.cur_y = 0;
Ok(())
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct PngRecord {
x: u32,
y: u32,
red: u16,
green: u16,
blue: u16,
alpha: u16,
}
impl_record!(PngRecord: x, y, red, green, blue, alpha);
fn get_bits(data: &[u8], pos: usize, n_bits: usize, rescale: bool) -> Result<u16, EtError> {
if n_bits == 16 {
u16::extract(&data[pos * 2..], &Endian::Big)
} else {
let shift = n_bits * (pos % (8 / n_bits));
let mask = u8::try_from(2u16.pow(u32::try_from(n_bits)?) - 1)?;
let d = data[n_bits * pos / 8];
let value = mask & (d >> shift);
if rescale {
Ok(u16::try_from(
u32::from(value) * 65535 / (2u32.pow(u32::try_from(n_bits)?) - 1),
)?)
} else {
Ok(u16::from(value))
}
}
}
impl<'b: 's, 's> FromSlice<'b, 's> for PngRecord {
type State = PngState;
fn parse(
_rb: &[u8],
_eof: bool,
_consumed: &mut usize,
state: &mut Self::State,
) -> Result<bool, EtError> {
if state.cur_x == usize::MAX {
state.cur_x = 0;
} else {
state.cur_x += 1;
}
if state.cur_x == state.width {
state.cur_x = 0;
state.cur_y += 1;
}
if state.cur_y >= state.height {
return Ok(false);
}
if state.cur_x == 0 {
state.unfilter_line(state.cur_y)?;
}
Ok(true)
}
fn get(&mut self, _rb: &'b [u8], state: &'s Self::State) -> Result<(), EtError> {
let bd = usize::from(state.bit_depth);
let line = &state.image_data
[state.cur_y * state.line_len() + 1..(state.cur_y + 1) * state.line_len()];
let pos = state.cur_x * state.color_type.pixel_size();
let (red, green, blue, alpha) = match state.color_type {
PngColorType::Indexed => {
let palette_pos = get_bits(line, pos, bd, false)? as usize;
if let Some(palette) = &state.palette {
if palette_pos >= palette.len() {
return Err("Color index was outside palette dimensions".into());
}
let (red, green, blue) = palette[palette_pos];
(red, green, blue, u16::MAX)
} else {
return Err("No palette was provided".into());
}
}
PngColorType::Grayscale => {
let gray = get_bits(line, pos, bd, true)?;
(gray, gray, gray, u16::MAX)
}
PngColorType::AlphaGrayscale => {
let gray = get_bits(line, pos, bd, true)?;
let alpha = get_bits(line, pos + 1, bd, true)?;
(gray, gray, gray, alpha)
}
PngColorType::Color => {
let red = get_bits(line, pos, bd, true)?;
let green = get_bits(line, pos + 1, bd, true)?;
let blue = get_bits(line, pos + 2, bd, true)?;
(red, green, blue, u16::MAX)
}
PngColorType::AlphaColor => {
let red = get_bits(line, pos, bd, true)?;
let green = get_bits(line, pos + 1, bd, true)?;
let blue = get_bits(line, pos + 2, bd, true)?;
let alpha = get_bits(line, pos + 3, bd, true)?;
(red, green, blue, alpha)
}
};
self.red = red;
self.green = green;
self.blue = blue;
self.alpha = alpha;
self.x = u32::try_from(state.cur_x)?;
self.y = u32::try_from(state.cur_y)?;
Ok(())
}
}
impl_reader!(PngReader, PngRecord, PngRecord, PngState, ());
#[cfg(test)]
mod tests {
use super::*;
use crate::readers::RecordReader;
#[test]
fn test_png_reader() -> Result<(), EtError> {
let rb: &[u8] = &include_bytes!("../../tests/data/bmp_24.png")[..];
let mut reader = PngReader::new(rb, None)?;
let _ = reader.metadata();
let mut n_recs = 0;
while reader.next()?.is_some() {
n_recs += 1;
}
assert_eq!(n_recs, 40000);
Ok(())
}
#[test]
fn test_indexed_png() -> Result<(), EtError> {
let rb: &[u8] = &include_bytes!("../../tests/data/bmp_indexed.png")[..];
let mut reader = PngReader::new(rb, None)?;
let _ = reader.metadata();
let mut n_recs = 0;
while reader.next()?.is_some() {
n_recs += 1;
}
assert_eq!(n_recs, 40000);
Ok(())
}
#[test]
fn test_minimal_png() -> Result<(), EtError> {
const TEST_IMAGE: &[u8] = &[
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48,
0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00,
0x00, 0x90, 0x77, 0x53, 0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, 0x54, 0x08,
0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00, 0x03, 0x01, 0x01, 0x00, 0x18, 0xDD, 0x8D,
0xB0, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82,
];
let mut reader = PngReader::new(TEST_IMAGE, None)?;
let _ = reader.metadata();
let pixel = reader.next()?.expect("first pixel exists");
assert_eq!(pixel.x, 0);
assert_eq!(pixel.y, 0);
assert_eq!(pixel.red, 65535);
assert_eq!(pixel.green, 0);
assert_eq!(pixel.blue, 0);
assert_eq!(pixel.alpha, 65535);
assert!(reader.next()?.is_none());
Ok(())
}
}