oxipng/src/png.rs
2016-04-05 10:52:02 -04:00

1445 lines
52 KiB
Rust

use bit_vec::BitVec;
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use crc::crc32;
use std::collections::HashMap;
use std::fmt;
use std::fs::File;
use std::io::Cursor;
use std::io::prelude::*;
use std::iter::Iterator;
use std::path::Path;
#[derive(Debug,PartialEq,Clone,Copy)]
/// The color type used to represent this image
pub enum ColorType {
/// Grayscale, with one color channel
Grayscale,
/// RGB, with three color channels
RGB,
/// Indexed, with one byte per pixel representing one of up to 256 colors in the image
Indexed,
/// Grayscale + Alpha, with two color channels
GrayscaleAlpha,
/// RGBA, with four color channels
RGBA,
}
impl fmt::Display for ColorType {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f,
"{}",
match *self {
ColorType::Grayscale => "Grayscale",
ColorType::RGB => "RGB",
ColorType::Indexed => "Indexed",
ColorType::GrayscaleAlpha => "Grayscale + Alpha",
ColorType::RGBA => "RGB + Alpha",
})
}
}
impl ColorType {
fn png_header_code(&self) -> u8 {
match *self {
ColorType::Grayscale => 0,
ColorType::RGB => 2,
ColorType::Indexed => 3,
ColorType::GrayscaleAlpha => 4,
ColorType::RGBA => 6,
}
}
}
#[derive(Debug,PartialEq,Clone,Copy)]
/// The number of bits to be used per channel per pixel
pub enum BitDepth {
/// One bit per channel per pixel
One,
/// Two bits per channel per pixel
Two,
/// Four bits per channel per pixel
Four,
/// Eight bits per channel per pixel
Eight,
/// Sixteen bits per channel per pixel
Sixteen,
}
impl fmt::Display for BitDepth {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f,
"{}",
match *self {
BitDepth::One => "1",
BitDepth::Two => "2",
BitDepth::Four => "4",
BitDepth::Eight => "8",
BitDepth::Sixteen => "16",
})
}
}
impl BitDepth {
/// Retrieve the number of bits per channel per pixel as a `u8`
pub fn as_u8(&self) -> u8 {
match *self {
BitDepth::One => 1,
BitDepth::Two => 2,
BitDepth::Four => 4,
BitDepth::Eight => 8,
BitDepth::Sixteen => 16,
}
}
/// Parse a number of bits per channel per pixel into a `BitDepth`
pub fn from_u8(depth: u8) -> BitDepth {
match depth {
1 => BitDepth::One,
2 => BitDepth::Two,
4 => BitDepth::Four,
8 => BitDepth::Eight,
16 => BitDepth::Sixteen,
_ => panic!("Unsupported bit depth"),
}
}
}
#[derive(Debug,PartialEq,Clone)]
/// Options to use for performing operations on headers (such as stripping)
pub enum Headers {
/// None
None,
/// Some, with a list of 4-character chunk codes
Some(Vec<String>),
/// Headers that won't affect rendering (all but cHRM, gAMA, iCCP, sBIT, sRGB, bKGD, hIST, pHYs, sPLT)
Safe,
/// All non-critical headers
All,
}
#[derive(Debug,Clone)]
/// An iterator over the scan lines of a PNG image
pub struct ScanLines<'a> {
/// A reference to the PNG image being iterated upon
pub png: &'a PngData,
start: usize,
end: usize,
/// Current pass number, and 0-indexed row within the pass
pass: Option<(u8, u32)>,
}
impl<'a> Iterator for ScanLines<'a> {
type Item = ScanLine;
fn next(&mut self) -> Option<Self::Item> {
if self.end == self.png.raw_data.len() {
None
} else if self.png.ihdr_data.interlaced == 1 {
// Scanlines for interlaced PNG files
if self.pass.is_none() {
self.pass = Some((1, 0));
}
// Handle edge cases for images smaller than 5 pixels in either direction
if self.png.ihdr_data.width < 5 && self.pass.unwrap().0 == 2 {
if let Some(pass) = self.pass.as_mut() {
pass.0 = 3;
pass.1 = 4;
}
}
// Intentionally keep these separate so that they can be applied one after another
if self.png.ihdr_data.height < 5 && self.pass.unwrap().0 == 3 {
if let Some(pass) = self.pass.as_mut() {
pass.0 = 4;
pass.1 = 0;
}
}
let bits_per_pixel = self.png.ihdr_data.bit_depth.as_u8() as usize *
self.png.channels_per_pixel() as usize;
let mut bits_per_line = self.png.ihdr_data.width as usize * bits_per_pixel;
let y_steps;
match self.pass {
Some((1, _)) | Some((2, _)) => {
bits_per_line = (bits_per_line as f32 / 8f32).ceil() as usize;
y_steps = 8;
}
Some((3, _)) => {
bits_per_line = (bits_per_line as f32 / 4f32).ceil() as usize;
y_steps = 8;
}
Some((4, _)) => {
bits_per_line = (bits_per_line as f32 / 4f32).ceil() as usize;
y_steps = 4;
}
Some((5, _)) => {
bits_per_line = (bits_per_line as f32 / 2f32).ceil() as usize;
y_steps = 4;
}
Some((6, _)) => {
bits_per_line = (bits_per_line as f32 / 2f32).ceil() as usize;
y_steps = 2;
}
Some((7, _)) => {
y_steps = 2;
}
_ => unreachable!(),
}
// Determine whether to trim the last (overflow) pixel for rows on this pass
let gap = bits_per_line % bits_per_pixel;
if gap != 0 {
let x_start = bits_per_pixel *
match self.pass.unwrap().0 {
2 => 4,
4 => 2,
6 => 1,
_ => 0,
};
if gap >= x_start {
bits_per_line += bits_per_pixel - gap;
} else {
bits_per_line -= gap;
}
}
let bytes_per_line = (bits_per_line as f32 / 8f32).ceil() as usize;
self.start = self.end;
self.end = self.start + bytes_per_line + 1;
if let Some(pass) = self.pass.as_mut() {
if pass.1 + y_steps >= self.png.ihdr_data.height {
pass.0 += 1;
pass.1 = match pass.0 {
3 => 4,
5 => 2,
7 => 1,
_ => 0,
};
} else {
pass.1 += y_steps;
}
}
Some(ScanLine {
filter: self.png.raw_data[self.start],
data: self.png.raw_data[(self.start + 1)..self.end].to_owned(),
})
} else {
// Standard, non-interlaced PNG scanlines
let bits_per_line = self.png.ihdr_data.width as usize *
self.png.ihdr_data.bit_depth.as_u8() as usize *
self.png.channels_per_pixel() as usize;
let bytes_per_line = (bits_per_line as f32 / 8f32).ceil() as usize;
self.start = self.end;
self.end = self.start + bytes_per_line + 1;
Some(ScanLine {
filter: self.png.raw_data[self.start],
data: self.png.raw_data[(self.start + 1)..self.end].to_owned(),
})
}
}
}
#[derive(Debug,Clone)]
/// A scan line in a PNG image
pub struct ScanLine {
/// The filter type used to encode the current scan line (0-4)
pub filter: u8,
/// The byte data for the current scan line, encoded with the filter specified in the `filter` field
pub data: Vec<u8>,
}
#[derive(Debug,Clone)]
/// Contains all data relevant to a PNG image
pub struct PngData {
/// The filtered and compressed data of the IDAT chunk
pub idat_data: Vec<u8>,
/// The headers stored in the IHDR chunk
pub ihdr_data: IhdrData,
/// The uncompressed, optionally filtered data from the IDAT chunk
pub raw_data: Vec<u8>,
/// The palette containing colors used in an Indexed image
/// Contains 3 bytes per color (R+G+B), up to 768
pub palette: Option<Vec<u8>>,
/// The pixel value that should be rendered as transparent
pub transparency_pixel: Option<Vec<u8>>,
/// A map of how transparent each color in the palette should be
pub transparency_palette: Option<Vec<u8>>,
/// All non-critical headers from the PNG are stored here
pub aux_headers: HashMap<String, Vec<u8>>,
}
#[derive(Debug,Clone,Copy)]
/// Headers from the IHDR chunk of the image
pub struct IhdrData {
/// The width of the image in pixels
pub width: u32,
/// The height of the image in pixels
pub height: u32,
/// The color type of the image
pub color_type: ColorType,
/// The bit depth of the image
pub bit_depth: BitDepth,
/// The compression method used for this image (0 for DEFLATE)
pub compression: u8,
/// The filter mode used for this image (currently only 0 is valid)
pub filter: u8,
/// The interlacing mode of the image (0 = None, 1 = Adam7)
pub interlaced: u8,
}
impl PngData {
/// Create a new `PngData` struct by opening a file
pub fn new(filepath: &Path) -> Result<PngData, String> {
let mut file = match File::open(filepath) {
Ok(f) => f,
Err(_) => return Err("Failed to open file for reading".to_owned()),
};
let mut byte_data: Vec<u8> = Vec::new();
// Read raw png data into memory
match file.read_to_end(&mut byte_data) {
Ok(_) => (),
Err(_) => return Err("Failed to read from file".to_owned()),
}
let mut byte_offset: usize = 0;
// Test that png header is valid
let header: Vec<u8> = byte_data.iter().take(8).cloned().collect();
if !file_header_is_valid(header.as_ref()) {
return Err("Invalid PNG header detected".to_owned());
}
byte_offset += 8;
// Read the data headers
let mut aux_headers: HashMap<String, Vec<u8>> = HashMap::new();
let mut idat_headers: Vec<u8> = Vec::new();
loop {
let header = parse_next_header(byte_data.as_ref(), &mut byte_offset);
let header = match header {
Ok(x) => x,
Err(x) => return Err(x),
};
let header = match header {
Some(x) => x,
None => break,
};
if header.0 == "IDAT" {
idat_headers.extend(header.1);
} else {
aux_headers.insert(header.0, header.1);
}
}
// Parse the headers into our PngData
if idat_headers.is_empty() {
return Err("Image data was empty, skipping".to_owned());
}
if aux_headers.get("IHDR").is_none() {
return Err("Image header data was missing, skipping".to_owned());
}
let ihdr_header = match parse_ihdr_header(aux_headers.remove("IHDR").unwrap().as_ref()) {
Ok(x) => x,
Err(x) => return Err(x),
};
let raw_data = match super::deflate::deflate::inflate(idat_headers.as_ref()) {
Ok(x) => x,
Err(x) => return Err(x),
};
// Handle transparency header
let mut has_transparency_pixel = false;
let mut has_transparency_palette = false;
if aux_headers.contains_key("tRNS") {
if ihdr_header.color_type == ColorType::Indexed {
has_transparency_palette = true;
} else {
has_transparency_pixel = true;
}
}
let mut png_data = PngData {
idat_data: idat_headers.clone(),
ihdr_data: ihdr_header,
raw_data: raw_data,
palette: aux_headers.remove("PLTE"),
transparency_pixel: if has_transparency_pixel {
aux_headers.remove("tRNS")
} else {
None
},
transparency_palette: if has_transparency_palette {
aux_headers.remove("tRNS")
} else {
None
},
aux_headers: aux_headers,
};
png_data.raw_data = png_data.unfilter_image();
// Return the PngData
Ok(png_data)
}
/// Return the number of channels in the image, based on color type
pub fn channels_per_pixel(&self) -> u8 {
match self.ihdr_data.color_type {
ColorType::Grayscale | ColorType::Indexed => 1,
ColorType::GrayscaleAlpha => 2,
ColorType::RGB => 3,
ColorType::RGBA => 4,
}
}
/// Format the `PngData` struct into a valid PNG bytestream
pub fn output(&self) -> Vec<u8> {
// PNG header
let mut output = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
// IHDR
let mut ihdr_data = Vec::with_capacity(13);
ihdr_data.write_u32::<BigEndian>(self.ihdr_data.width).ok();
ihdr_data.write_u32::<BigEndian>(self.ihdr_data.height).ok();
ihdr_data.write_u8(self.ihdr_data.bit_depth.as_u8()).ok();
ihdr_data.write_u8(self.ihdr_data.color_type.png_header_code()).ok();
ihdr_data.write_u8(0).ok(); // Compression -- deflate
ihdr_data.write_u8(0).ok(); // Filter method -- 5-way adaptive filtering
ihdr_data.write_u8(self.ihdr_data.interlaced).ok();
write_png_block(b"IHDR", &ihdr_data, &mut output);
// Ancillary headers
for (key, header) in self.aux_headers.iter().filter(|&(ref key, _)| {
!(**key == "bKGD" || **key == "hIST" || **key == "tRNS")
}) {
write_png_block(&key.as_bytes(), &header, &mut output);
}
// Palette
if let Some(palette) = self.palette.clone() {
write_png_block(b"PLTE", &palette, &mut output);
if let Some(transparency_palette) = self.transparency_palette.clone() {
// Transparency pixel
write_png_block(b"tRNS", &transparency_palette, &mut output);
}
} else if let Some(transparency_pixel) = self.transparency_pixel.clone() {
// Transparency pixel
write_png_block(b"tRNS", &transparency_pixel, &mut output);
}
// Special ancillary headers that need to come after PLTE but before IDAT
for (key, header) in self.aux_headers.iter().filter(|&(ref key, _)| {
**key == "bKGD" || **key == "hIST" || **key == "tRNS"
}) {
write_png_block(&key.as_bytes(), &header, &mut output);
}
// IDAT data
write_png_block(b"IDAT", &self.idat_data, &mut output);
// Stream end
write_png_block(b"IEND", &[], &mut output);
output
}
/// Return an iterator over the scanlines of the image
pub fn scan_lines(&self) -> ScanLines {
ScanLines {
png: &self,
start: 0,
end: 0,
pass: None,
}
}
/// Reverse all filters applied on the image, returning an unfiltered IDAT bytestream
pub fn unfilter_image(&self) -> Vec<u8> {
let mut unfiltered = Vec::with_capacity(self.raw_data.len());
let bpp = (((self.ihdr_data.bit_depth.as_u8() * self.channels_per_pixel()) as f32) /
8f32)
.ceil() as usize;
let mut last_line: Vec<u8> = Vec::new();
for line in self.scan_lines() {
let unfiltered_line = unfilter_line(line.filter, bpp, &line.data, &last_line);
unfiltered.push(0);
unfiltered.extend_from_slice(&unfiltered_line);
last_line = unfiltered_line;
}
unfiltered
}
/// Apply the specified filter type to all rows in the image
/// 0: None
/// 1: Sub
/// 2: Up
/// 3: Average
/// 4: Paeth
/// 5: All (heuristically pick the best filter for each line)
pub fn filter_image(&self, filter: u8) -> Vec<u8> {
let mut filtered = Vec::with_capacity(self.raw_data.len());
let bpp = (((self.ihdr_data.bit_depth.as_u8() * self.channels_per_pixel()) as f32) /
8f32)
.ceil() as usize;
let mut last_line: Vec<u8> = Vec::new();
for line in self.scan_lines() {
match filter {
0 | 1 | 2 | 3 | 4 => {
filtered.push(filter);
filtered.extend(filter_line(filter, bpp, &line.data, &last_line));
}
5 => {
// Heuristically guess best filter per line
// Uses MSAD algorithm mentioned in libpng reference docs
// http://www.libpng.org/pub/png/book/chapter09.html
let mut trials: HashMap<u8, Vec<u8>> = HashMap::with_capacity(5);
for filter in 0..5 {
trials.insert(filter, filter_line(filter, bpp, &line.data, &last_line));
}
let (best_filter, best_line) = trials.iter()
.min_by_key(|x| {
x.1.iter().fold(0u64, |acc, &x| {
let signed = x as i8;
acc + (signed as i16).abs() as u64
})
})
.unwrap();
filtered.push(*best_filter);
filtered.extend_from_slice(best_line);
}
_ => unreachable!(),
}
last_line = line.data.clone();
}
filtered
}
/// Attempt to reduce the bit depth of the image
/// Returns true if the bit depth was reduced, false otherwise
pub fn reduce_bit_depth(&mut self) -> bool {
if self.ihdr_data.bit_depth != BitDepth::Sixteen {
if self.ihdr_data.color_type == ColorType::Indexed ||
self.ihdr_data.color_type == ColorType::Grayscale {
return match reduce_bit_depth_8_or_less(self) {
Some((data, depth)) => {
self.raw_data = data;
self.ihdr_data.bit_depth = BitDepth::from_u8(depth);
true
}
None => false,
};
}
return false;
}
// Reduce from 16 to 8 bits per channel per pixel
let mut reduced =
Vec::with_capacity((self.ihdr_data.width * self.ihdr_data.height *
self.channels_per_pixel() as u32 +
self.ihdr_data.height) as usize);
let mut high_byte = 0;
for line in self.scan_lines() {
reduced.push(line.filter);
for (i, byte) in line.data.iter().enumerate() {
if i % 2 == 0 {
// High byte
high_byte = *byte;
} else {
// Low byte
if high_byte != *byte {
// Can't reduce, exit early
return false;
}
reduced.push(*byte);
}
}
}
self.ihdr_data.bit_depth = BitDepth::Eight;
self.raw_data = reduced;
true
}
/// Attempt to reduce the number of colors in the palette
/// Returns true if the palette was reduced, false otherwise
pub fn reduce_palette(&mut self) -> bool {
// TODO: Implement
false
}
/// Attempt to reduce the color type of the image
/// Returns true if the color type was reduced, false otherwise
pub fn reduce_color_type(&mut self) -> bool {
let mut changed = false;
let mut should_reduce_bit_depth = false;
// Go down one step at a time
// Maybe not the most efficient, but it's safe
if self.ihdr_data.color_type == ColorType::RGBA {
if let Some(data) = reduce_rgba_to_grayscale_alpha(self) {
self.raw_data = data;
self.ihdr_data.color_type = ColorType::GrayscaleAlpha;
changed = true;
} else if let Some(data) = reduce_rgba_to_rgb(self) {
self.raw_data = data;
self.ihdr_data.color_type = ColorType::RGB;
changed = true;
} else if let Some((data, palette, trans)) = reduce_rgba_to_palette(self) {
self.raw_data = data;
self.palette = Some(palette);
if trans.iter().any(|x| *x != 255) {
self.transparency_palette = Some(trans);
} else {
self.transparency_palette = None;
}
self.ihdr_data.color_type = ColorType::Indexed;
changed = true;
should_reduce_bit_depth = true;
}
}
if self.ihdr_data.color_type == ColorType::GrayscaleAlpha {
if let Some(data) = reduce_grayscale_alpha_to_grayscale(self) {
self.raw_data = data;
self.ihdr_data.color_type = ColorType::Grayscale;
changed = true;
should_reduce_bit_depth = true;
}
}
if self.ihdr_data.color_type == ColorType::RGB {
if let Some(data) = reduce_rgb_to_grayscale(self) {
self.raw_data = data;
self.ihdr_data.color_type = ColorType::Grayscale;
changed = true;
should_reduce_bit_depth = true;
} else if let Some((data, palette)) = reduce_rgb_to_palette(self) {
self.raw_data = data;
self.palette = Some(palette);
self.ihdr_data.color_type = ColorType::Indexed;
changed = true;
should_reduce_bit_depth = true;
}
}
if self.ihdr_data.color_type == ColorType::Indexed && self.transparency_palette.is_none() &&
self.palette.as_ref().map(|x| x.len()).unwrap() > 128 {
if let Some(data) = reduce_palette_to_grayscale(self) {
self.raw_data = data;
self.palette = None;
self.ihdr_data.color_type = ColorType::Grayscale;
changed = true;
should_reduce_bit_depth = false;
}
} else if self.ihdr_data.color_type == ColorType::Grayscale {
if let Some((data, palette)) = reduce_grayscale_to_palette(self) {
self.raw_data = data;
self.palette = Some(palette);
self.ihdr_data.color_type = ColorType::Indexed;
changed = true;
should_reduce_bit_depth = true;
}
}
if should_reduce_bit_depth {
// Some conversions will allow us to perform bit depth reduction that
// wasn't possible before
if let Some((data, depth)) = reduce_bit_depth_8_or_less(self) {
self.raw_data = data;
self.ihdr_data.bit_depth = BitDepth::from_u8(depth);
}
}
changed
}
/// Convert the image to the specified interlacing type
/// Returns true if the interlacing was changed, false otherwise
/// The `interlace` parameter specifies the *new* interlacing mode
/// Assumes that the data has already been de-filtered
pub fn change_interlacing(&mut self, interlace: u8) -> bool {
if interlace == self.ihdr_data.interlaced {
return false;
}
if interlace == 1 {
// Convert progressive to interlaced data
interlace_image(self);
} else {
// Convert interlaced to progressive data
deinterlace_image(self);
}
true
}
}
fn interlace_image(png: &mut PngData) {
let mut passes: Vec<BitVec> = Vec::with_capacity(7);
for _ in 0..7 {
passes.push(BitVec::new());
}
let bits_per_pixel = png.ihdr_data.bit_depth.as_u8() * png.channels_per_pixel();
for (index, line) in png.scan_lines().enumerate() {
match index % 8 {
// Add filter bytes to appropriate lines
0 => {
passes[0].extend(BitVec::from_elem(8, false));
passes[3].extend(BitVec::from_elem(8, false));
passes[5].extend(BitVec::from_elem(8, false));
if png.ihdr_data.width > 4 {
passes[1].extend(BitVec::from_elem(8, false));
}
}
4 => {
passes[3].extend(BitVec::from_elem(8, false));
passes[5].extend(BitVec::from_elem(8, false));
passes[2].extend(BitVec::from_elem(8, false));
}
2 | 6 => {
passes[4].extend(BitVec::from_elem(8, false));
passes[5].extend(BitVec::from_elem(8, false));
}
_ => {
passes[6].extend(BitVec::from_elem(8, false));
}
}
let bit_vec = BitVec::from_bytes(&line.data);
for (i, bit) in bit_vec.iter().enumerate() {
// Avoid moving padded 0's into new image
if i >= (png.ihdr_data.width * bits_per_pixel as u32) as usize {
break;
}
// Copy pixels into interlaced passes
let pix_modulo = (((i / bits_per_pixel as usize) as f32).floor() as usize) % 8;
match index % 8 {
0 => {
match pix_modulo {
0 => passes[0].push(bit),
4 => passes[1].push(bit),
2 | 6 => passes[3].push(bit),
_ => passes[5].push(bit),
}
}
4 => {
match pix_modulo {
0 | 4 => passes[2].push(bit),
2 | 6 => passes[3].push(bit),
_ => passes[5].push(bit),
}
}
2 | 6 => {
match pix_modulo % 2 {
0 => passes[4].push(bit),
_ => passes[5].push(bit),
}
}
_ => {
passes[6].push(bit);
}
}
}
// Pad end of line on each pass to get 8 bits per byte
for pass in &mut passes {
while pass.len() % 8 != 0 {
pass.push(false);
}
}
}
let mut output = Vec::new();
for pass in &passes {
output.extend(pass.to_bytes());
}
png.raw_data = output;
}
fn deinterlace_image(png: &mut PngData) {
let bits_per_pixel = png.ihdr_data.bit_depth.as_u8() * png.channels_per_pixel();
let mut lines: Vec<BitVec> = Vec::with_capacity(png.ihdr_data.height as usize);
for _ in 0..png.ihdr_data.height {
// Initialize each output line with a starting filter byte of 0
// as well as some blank data
lines.push(BitVec::from_elem(8 + bits_per_pixel as usize * png.ihdr_data.width as usize,
false));
}
let mut current_pass = 1;
let mut pass_constants = interlaced_constants(current_pass);
let mut current_y: usize = pass_constants.y_shift as usize;
for line in png.scan_lines() {
let bit_vec = BitVec::from_bytes(&line.data);
let bits_in_line = ((png.ihdr_data.width - pass_constants.x_shift as u32) as f32 /
pass_constants.x_step as f32)
.ceil() as usize *
bits_per_pixel as usize;
for (i, bit) in bit_vec.iter().enumerate() {
// Avoid moving padded 0's into new image
if i >= bits_in_line {
break;
}
let current_x: usize = pass_constants.x_shift as usize +
(i / bits_per_pixel as usize) * pass_constants.x_step as usize;
// Copy this bit into the output line, offset by 8 because of filter byte
let index = 8 + (i % bits_per_pixel as usize) + current_x * bits_per_pixel as usize;
lines[current_y].set(index, bit);
}
// Calculate the next line and move to next pass if necessary
current_y += pass_constants.y_step as usize;
if current_y >= png.ihdr_data.height as usize {
if current_pass == 7 {
break;
}
current_pass += 1;
if current_pass == 2 && png.ihdr_data.width <= 4 {
current_pass += 1;
}
if current_pass == 3 && png.ihdr_data.height <= 4 {
current_pass += 1;
}
pass_constants = interlaced_constants(current_pass);
current_y = pass_constants.y_shift as usize;
}
}
let mut output = Vec::new();
for line in &mut lines {
while line.len() % 8 != 0 {
line.push(false);
}
output.extend(line.to_bytes());
}
png.raw_data = output;
}
struct InterlacedConstants {
x_shift: u8,
y_shift: u8,
x_step: u8,
y_step: u8,
}
fn interlaced_constants(pass: u8) -> InterlacedConstants {
match pass {
1 => {
InterlacedConstants {
x_shift: 0,
y_shift: 0,
x_step: 8,
y_step: 8,
}
}
2 => {
InterlacedConstants {
x_shift: 4,
y_shift: 0,
x_step: 8,
y_step: 8,
}
}
3 => {
InterlacedConstants {
x_shift: 0,
y_shift: 4,
x_step: 4,
y_step: 8,
}
}
4 => {
InterlacedConstants {
x_shift: 2,
y_shift: 0,
x_step: 4,
y_step: 4,
}
}
5 => {
InterlacedConstants {
x_shift: 0,
y_shift: 2,
x_step: 2,
y_step: 4,
}
}
6 => {
InterlacedConstants {
x_shift: 1,
y_shift: 0,
x_step: 2,
y_step: 2,
}
}
7 => {
InterlacedConstants {
x_shift: 0,
y_shift: 1,
x_step: 1,
y_step: 2,
}
}
_ => unreachable!(),
}
}
fn filter_line(filter: u8, bpp: usize, data: &[u8], last_line: &[u8]) -> Vec<u8> {
let mut filtered = Vec::with_capacity(data.len());
match filter {
0 => {
filtered.extend_from_slice(data);
}
1 => {
for (i, byte) in data.iter().enumerate() {
filtered.push(match i.checked_sub(bpp) {
Some(x) => byte.wrapping_sub(data[x]),
None => *byte,
});
}
}
2 => {
for (i, byte) in data.iter().enumerate() {
if last_line.is_empty() {
filtered.push(*byte);
} else {
filtered.push(byte.wrapping_sub(last_line[i]));
};
}
}
3 => {
for (i, byte) in data.iter().enumerate() {
if last_line.is_empty() {
filtered.push(match i.checked_sub(bpp) {
Some(x) => byte.wrapping_sub(data[x] >> 1),
None => *byte,
});
} else {
filtered.push(match i.checked_sub(bpp) {
Some(x) => {
byte.wrapping_sub(((data[x] as u16 + last_line[i] as u16) >> 1) as u8)
}
None => byte.wrapping_sub(last_line[i] >> 1),
});
};
}
}
4 => {
for (i, byte) in data.iter().enumerate() {
if last_line.is_empty() {
filtered.push(match i.checked_sub(bpp) {
Some(x) => byte.wrapping_sub(data[x]),
None => *byte,
});
} else {
filtered.push(match i.checked_sub(bpp) {
Some(x) => {
byte.wrapping_sub(paeth_predictor(data[x], last_line[i], last_line[x]))
}
None => byte.wrapping_sub(last_line[i]),
});
};
}
}
_ => unreachable!(),
}
filtered
}
fn unfilter_line(filter: u8, bpp: usize, data: &[u8], last_line: &[u8]) -> Vec<u8> {
let mut unfiltered = Vec::with_capacity(data.len());
match filter {
0 => {
unfiltered.extend_from_slice(data);
}
1 => {
for (i, byte) in data.iter().enumerate() {
match i.checked_sub(bpp) {
Some(x) => {
let b = unfiltered[x];
unfiltered.push(byte.wrapping_add(b));
}
None => {
unfiltered.push(*byte);
}
};
}
}
2 => {
for (i, byte) in data.iter().enumerate() {
if last_line.is_empty() {
unfiltered.push(*byte);
} else {
unfiltered.push(byte.wrapping_add(last_line[i]));
};
}
}
3 => {
for (i, byte) in data.iter().enumerate() {
if last_line.is_empty() {
match i.checked_sub(bpp) {
Some(x) => {
let b = unfiltered[x];
unfiltered.push(byte.wrapping_add(b >> 1));
}
None => {
unfiltered.push(*byte);
}
};
} else {
match i.checked_sub(bpp) {
Some(x) => {
let b = unfiltered[x];
unfiltered.push(byte.wrapping_add(((b as u16 + last_line[i] as u16) >> 1) as u8));
}
None => {
unfiltered.push(byte.wrapping_add(last_line[i] >> 1));
}
};
};
}
}
4 => {
for (i, byte) in data.iter().enumerate() {
if last_line.is_empty() {
match i.checked_sub(bpp) {
Some(x) => {
let b = unfiltered[x];
unfiltered.push(byte.wrapping_add(b));
}
None => {
unfiltered.push(*byte);
}
};
} else {
match i.checked_sub(bpp) {
Some(x) => {
let b = unfiltered[x];
unfiltered.push(byte.wrapping_add(paeth_predictor(b,
last_line[i],
last_line[x])));
}
None => {
unfiltered.push(byte.wrapping_add(last_line[i]));
}
};
};
}
}
_ => unreachable!(),
}
unfiltered
}
fn reduce_bit_depth_8_or_less(png: &PngData) -> Option<(Vec<u8>, u8)> {
let mut reduced = BitVec::with_capacity(png.raw_data.len() * 8);
let bit_depth: usize = png.ihdr_data.bit_depth.as_u8() as usize;
let mut allowed_bits = 1;
for line in png.scan_lines() {
let bit_vec = BitVec::from_bytes(&line.data);
for (i, bit) in bit_vec.iter().enumerate() {
let bit_index = bit_depth - (i % bit_depth);
if bit && bit_index > allowed_bits {
allowed_bits = bit_index.next_power_of_two();
if allowed_bits == bit_depth {
// Not reducable
return None;
}
}
}
}
for line in png.scan_lines() {
reduced.extend(BitVec::from_bytes(&[line.filter]));
let bit_vec = BitVec::from_bytes(&line.data);
for (i, bit) in bit_vec.iter().enumerate() {
let bit_index = bit_depth - (i % bit_depth);
if bit_index <= allowed_bits {
reduced.push(bit);
}
}
// Pad end of line to get 8 bits per byte
while reduced.len() % 8 != 0 {
reduced.push(false);
}
}
Some((reduced.to_bytes(), allowed_bits as u8))
}
fn reduce_rgba_to_rgb(png: &PngData) -> Option<Vec<u8>> {
let mut reduced = Vec::with_capacity(png.raw_data.len());
let byte_depth = png.ihdr_data.bit_depth.as_u8() >> 3;
let bpp: usize = 4 * byte_depth as usize;
for line in png.scan_lines() {
reduced.push(line.filter);
for (i, byte) in line.data.iter().enumerate() {
if i % bpp >= (bpp - byte_depth as usize) {
if *byte != 255 {
return None;
}
} else {
reduced.push(*byte);
}
}
}
Some(reduced)
}
fn reduce_rgba_to_grayscale_alpha(png: &PngData) -> Option<Vec<u8>> {
let mut reduced = Vec::with_capacity(png.raw_data.len());
let byte_depth = png.ihdr_data.bit_depth.as_u8() >> 3;
let bpp: usize = 4 * byte_depth as usize;
for line in png.scan_lines() {
reduced.push(line.filter);
let mut low_bytes = Vec::with_capacity(4);
let mut high_bytes = Vec::with_capacity(4);
let mut trans_bytes = Vec::with_capacity(byte_depth as usize);
for (i, byte) in line.data.iter().enumerate() {
if i % bpp < (bpp - byte_depth as usize) {
if byte_depth == 1 || i % 2 == 1 {
low_bytes.push(*byte);
} else {
high_bytes.push(*byte);
}
} else {
trans_bytes.push(*byte);
}
if i % bpp == bpp - 1 {
low_bytes.sort();
low_bytes.dedup();
if low_bytes.len() > 1 {
return None;
}
if byte_depth == 2 {
high_bytes.sort();
high_bytes.dedup();
if high_bytes.len() > 1 {
return None;
}
reduced.push(high_bytes[0]);
high_bytes.clear();
}
reduced.push(low_bytes[0]);
low_bytes.clear();
reduced.extend_from_slice(&trans_bytes);
trans_bytes.clear();
}
}
}
Some(reduced)
}
fn reduce_rgba_to_palette(png: &PngData) -> Option<(Vec<u8>, Vec<u8>, Vec<u8>)> {
if png.ihdr_data.bit_depth != BitDepth::Eight {
return None;
}
let mut reduced = Vec::with_capacity(png.raw_data.len());
let mut palette = Vec::with_capacity(256);
let bpp: usize = (4 * png.ihdr_data.bit_depth.as_u8() as usize) >> 3;
for line in png.scan_lines() {
reduced.push(line.filter);
let mut cur_pixel = Vec::with_capacity(bpp);
for (i, byte) in line.data.iter().enumerate() {
cur_pixel.push(*byte);
if i % bpp == bpp - 1 {
if palette.contains(&cur_pixel) {
let idx = palette.iter().enumerate().find(|&x| x.1 == &cur_pixel).unwrap().0;
reduced.push(idx as u8);
} else {
let len = palette.len();
if len == 256 {
return None;
}
palette.push(cur_pixel.clone());
reduced.push(len as u8);
}
cur_pixel.clear();
}
}
}
let mut color_palette = Vec::with_capacity(palette.len() * 3);
let mut trans_palette = Vec::with_capacity(palette.len());
for color in &palette {
for (i, byte) in color.iter().enumerate() {
if i < 3 {
color_palette.push(*byte);
} else {
trans_palette.push(*byte);
}
}
}
Some((reduced, color_palette, trans_palette))
}
fn reduce_rgb_to_palette(png: &PngData) -> Option<(Vec<u8>, Vec<u8>)> {
if png.ihdr_data.bit_depth != BitDepth::Eight {
return None;
}
let mut reduced = Vec::with_capacity(png.raw_data.len());
let mut palette = Vec::with_capacity(256);
let bpp: usize = (3 * png.ihdr_data.bit_depth.as_u8() as usize) >> 3;
for line in png.scan_lines() {
reduced.push(line.filter);
let mut cur_pixel = Vec::with_capacity(bpp);
for (i, byte) in line.data.iter().enumerate() {
cur_pixel.push(*byte);
if i % bpp == bpp - 1 {
if palette.contains(&cur_pixel) {
let idx = palette.iter().enumerate().find(|&x| x.1 == &cur_pixel).unwrap().0;
reduced.push(idx as u8);
} else {
let len = palette.len();
if len == 256 {
return None;
}
palette.push(cur_pixel.clone());
reduced.push(len as u8);
}
cur_pixel.clear();
}
}
}
let mut color_palette = Vec::with_capacity(palette.len() * 3);
for color in &palette {
color_palette.extend_from_slice(&color);
}
Some((reduced, color_palette))
}
fn reduce_grayscale_to_palette(png: &PngData) -> Option<(Vec<u8>, Vec<u8>)> {
if png.ihdr_data.bit_depth == BitDepth::Sixteen {
return None;
}
let mut reduced = BitVec::with_capacity(png.raw_data.len() * 8);
// Only perform reduction if we can get to 4-bits or less
let mut palette = Vec::with_capacity(16);
let bpp: usize = png.ihdr_data.bit_depth.as_u8() as usize;
for line in png.scan_lines() {
reduced.extend(BitVec::from_bytes(&[line.filter]));
let bit_vec = BitVec::from_bytes(&line.data);
let mut cur_pixel = BitVec::with_capacity(bpp);
for (i, bit) in bit_vec.iter().enumerate() {
cur_pixel.push(bit);
if i % bpp == bpp - 1 {
let pix_value = cur_pixel.to_bytes()[0] >> (8 - bpp);
let pix_slice = vec![pix_value, pix_value, pix_value];
if palette.contains(&pix_slice) {
let index = palette.iter().enumerate().find(|&x| x.1 == &pix_slice).unwrap().0;
let idx = BitVec::from_bytes(&[(index as u8) << (8 - bpp)]);
for b in idx.iter().take(bpp) {
reduced.push(b);
}
} else {
let len = palette.len();
if len == 16 {
return None;
}
palette.push(pix_slice.clone());
let idx = BitVec::from_bytes(&[(len as u8) << (8 - bpp)]);
for b in idx.iter().take(bpp) {
reduced.push(b);
}
}
cur_pixel = BitVec::with_capacity(bpp);
}
}
// Pad end of line to get 8 bits per byte
while reduced.len() % 8 != 0 {
reduced.push(false);
}
}
let mut color_palette = Vec::with_capacity(palette.len() * 3);
for color in &palette {
color_palette.extend_from_slice(&color);
}
Some((reduced.to_bytes(), color_palette))
}
fn reduce_palette_to_grayscale(png: &PngData) -> Option<Vec<u8>> {
let mut reduced = BitVec::with_capacity(png.raw_data.len() * 8);
let mut cur_pixel = Vec::with_capacity(3);
let palette = png.palette.clone().unwrap();
// Iterate through palette and determine if all colors are grayscale
for byte in &palette {
cur_pixel.push(*byte);
if cur_pixel.len() == 3 {
cur_pixel.sort();
cur_pixel.dedup();
if cur_pixel.len() > 1 {
return None;
}
cur_pixel.clear();
}
}
// Iterate through scanlines and assign grayscale value to each pixel
let bit_depth: usize = png.ihdr_data.bit_depth.as_u8() as usize;
for line in png.scan_lines() {
reduced.extend(BitVec::from_bytes(&[line.filter]));
let bit_vec = BitVec::from_bytes(&line.data);
let mut cur_pixel = BitVec::with_capacity(bit_depth);
for bit in bit_vec {
// Handle bit depths less than 8-bits
// At the end of each pixel, push its grayscale value onto the reduced image
cur_pixel.push(bit);
if cur_pixel.len() == bit_depth {
// `to_bytes` gives us e.g. 10000000 for a 1-bit pixel, when we would want 00000001
let padded_pixel = cur_pixel.to_bytes()[0] >> (8 - bit_depth);
let palette_idx: usize = padded_pixel as usize * 3;
reduced.extend(BitVec::from_bytes(&[palette[palette_idx]]));
// BitVec's clear function doesn't set len to 0
cur_pixel = BitVec::with_capacity(bit_depth);
}
}
// Pad end of line to get 8 bits per byte
while reduced.len() % 8 != 0 {
reduced.push(false);
}
}
Some(reduced.to_bytes())
}
fn reduce_rgb_to_grayscale(png: &PngData) -> Option<Vec<u8>> {
let mut reduced = Vec::with_capacity(png.raw_data.len());
let byte_depth = png.ihdr_data.bit_depth.as_u8() >> 3;
let bpp: usize = 3 * byte_depth as usize;
let mut cur_pixel = Vec::with_capacity(bpp);
for line in png.scan_lines() {
reduced.push(line.filter);
for (i, byte) in line.data.iter().enumerate() {
cur_pixel.push(*byte);
if i % bpp == bpp - 1 {
if bpp == 3 {
cur_pixel.sort();
cur_pixel.dedup();
if cur_pixel.len() > 1 {
return None;
}
reduced.push(cur_pixel[0]);
} else {
let mut pixel_zip = cur_pixel.iter()
.enumerate()
.filter(|&(i, _)| i % 2 == 0)
.map(|(_, x)| *x)
.zip(cur_pixel.iter()
.enumerate()
.filter(|&(i, _)| i % 2 == 1)
.map(|(_, x)| *x))
.collect::<Vec<(u8, u8)>>();
pixel_zip.sort();
pixel_zip.dedup();
if pixel_zip.len() > 1 {
return None;
}
reduced.push(pixel_zip[0].0);
reduced.push(pixel_zip[0].1);
}
cur_pixel.clear();
}
}
}
Some(reduced)
}
fn reduce_grayscale_alpha_to_grayscale(png: &PngData) -> Option<Vec<u8>> {
let mut reduced = Vec::with_capacity(png.raw_data.len());
let byte_depth = png.ihdr_data.bit_depth.as_u8() >> 3;
let bpp: usize = 2 * byte_depth as usize;
for line in png.scan_lines() {
reduced.push(line.filter);
for (i, byte) in line.data.iter().enumerate() {
if i % bpp >= (bpp - byte_depth as usize) {
if *byte != 255 {
return None;
}
} else {
reduced.push(*byte);
}
}
}
Some(reduced)
}
fn paeth_predictor(a: u8, b: u8, c: u8) -> u8 {
let p = a as i32 + b as i32 - c as i32;
let pa = (p - a as i32).abs();
let pb = (p - b as i32).abs();
let pc = (p - c as i32).abs();
if pa <= pb && pa <= pc {
a
} else if pb <= pc {
b
} else {
c
}
}
fn file_header_is_valid(bytes: &[u8]) -> bool {
let expected_header: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
bytes.iter().zip(expected_header.iter()).all(|x| x.0 == x.1)
}
fn parse_next_header(byte_data: &[u8],
byte_offset: &mut usize)
-> Result<Option<(String, Vec<u8>)>, String> {
let mut rdr = Cursor::new(byte_data.iter()
.skip(*byte_offset)
.take(4)
.cloned()
.collect::<Vec<u8>>());
let length: u32 = match rdr.read_u32::<BigEndian>() {
Ok(x) => x,
Err(_) => return Err("Invalid data found--unable to read PNG file".to_owned()),
};
*byte_offset += 4;
let mut header_bytes: Vec<u8> = byte_data.iter().skip(*byte_offset).take(4).cloned().collect();
let header = match String::from_utf8(header_bytes.clone()) {
Ok(x) => x,
Err(_) => return Err("Invalid data found--unable to read PNG file".to_owned()),
};
if header == "IEND" {
// End of data
return Ok(None);
}
*byte_offset += 4;
let data: Vec<u8> = byte_data.iter()
.skip(*byte_offset)
.take(length as usize)
.cloned()
.collect();
*byte_offset += length as usize;
let mut rdr = Cursor::new(byte_data.iter()
.skip(*byte_offset)
.take(4)
.cloned()
.collect::<Vec<u8>>());
let crc: u32 = match rdr.read_u32::<BigEndian>() {
Ok(x) => x,
Err(_) => return Err("Invalid data found--unable to read PNG file".to_owned()),
};
*byte_offset += 4;
header_bytes.extend(data.clone());
if crc32::checksum_ieee(header_bytes.as_ref()) != crc {
return Err(format!("Corrupt data chunk found--CRC Mismatch in {}", header));
}
Ok(Some((header, data)))
}
fn parse_ihdr_header(byte_data: &[u8]) -> Result<IhdrData, String> {
let mut rdr = Cursor::new(&byte_data[0..8]);
Ok(IhdrData {
color_type: match byte_data[9] {
0 => ColorType::Grayscale,
2 => ColorType::RGB,
3 => ColorType::Indexed,
4 => ColorType::GrayscaleAlpha,
6 => ColorType::RGBA,
_ => return Err("Unexpected color type in header".to_owned()),
},
bit_depth: match byte_data[8] {
1 => BitDepth::One,
2 => BitDepth::Two,
4 => BitDepth::Four,
8 => BitDepth::Eight,
16 => BitDepth::Sixteen,
_ => return Err("Unexpected bit depth in header".to_owned()),
},
width: rdr.read_u32::<BigEndian>().unwrap(),
height: rdr.read_u32::<BigEndian>().unwrap(),
compression: byte_data[10],
filter: byte_data[11],
interlaced: byte_data[12],
})
}
fn write_png_block(key: &[u8], header: &[u8], output: &mut Vec<u8>) {
let mut header_data = Vec::with_capacity(header.len() + 4);
header_data.extend_from_slice(key);
header_data.extend_from_slice(header);
output.reserve(header_data.len() + 8);
output.write_u32::<BigEndian>(header_data.len() as u32 - 4).ok();
let crc = crc32::checksum_ieee(&header_data);
output.append(&mut header_data);
output.write_u32::<BigEndian>(crc).ok();
}