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), /// 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 { 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, } #[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, /// The headers stored in the IHDR chunk pub ihdr_data: IhdrData, /// The uncompressed, optionally filtered data from the IDAT chunk pub raw_data: Vec, /// The palette containing colors used in an Indexed image /// Contains 3 bytes per color (R+G+B), up to 768 pub palette: Option>, /// The pixel value that should be rendered as transparent pub transparency_pixel: Option>, /// A map of how transparent each color in the palette should be pub transparency_palette: Option>, /// All non-critical headers from the PNG are stored here pub aux_headers: HashMap>, } #[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 { 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 = 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 = 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> = HashMap::new(); let mut idat_headers: Vec = 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 { // 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::(self.ihdr_data.width).ok(); ihdr_data.write_u32::(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 { 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 = 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 { 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 = 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> = 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 = 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 = 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 { 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 { 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)> { 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> { 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> { 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, Vec, Vec)> { 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, Vec)> { 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, Vec)> { 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> { 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> { 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::>(); 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> { 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)>, String> { let mut rdr = Cursor::new(byte_data.iter() .skip(*byte_offset) .take(4) .cloned() .collect::>()); let length: u32 = match rdr.read_u32::() { Ok(x) => x, Err(_) => return Err("Invalid data found--unable to read PNG file".to_owned()), }; *byte_offset += 4; let mut header_bytes: Vec = 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 = 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::>()); let crc: u32 = match rdr.read_u32::() { 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 { 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::().unwrap(), height: rdr.read_u32::().unwrap(), compression: byte_data[10], filter: byte_data[11], interlaced: byte_data[12], }) } fn write_png_block(key: &[u8], header: &[u8], output: &mut Vec) { 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::(header_data.len() as u32 - 4).ok(); let crc = crc32::checksum_ieee(&header_data); output.append(&mut header_data); output.write_u32::(crc).ok(); }