diff --git a/src/lib.rs b/src/lib.rs index d15ed1f9..ea79f7c9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,11 +1,16 @@ extern crate byteorder; extern crate crc; -extern crate libz_sys; extern crate libc; +extern crate libz_sys; +use std::collections::HashSet; +use std::fs; +use std::fs::File; +use std::io; +use std::io::BufWriter; +use std::io::Write; use std::path::Path; use std::path::PathBuf; -use std::collections::HashSet; pub mod deflate { pub mod deflate; @@ -20,6 +25,7 @@ pub struct Options { pub stdout: bool, pub fix_errors: bool, pub pretend: bool, + pub recursive: bool, pub clobber: bool, pub create: bool, pub preserve_attrs: bool, @@ -47,8 +53,8 @@ pub fn optimize(filepath: &Path, opts: &Options) -> Result<(), String> { }; // Print png info - let idat_current_size = png.idat_data.len(); - let file_current_size = filepath.metadata().unwrap().len(); + let idat_original_size = png.idat_data.len(); + let file_original_size = filepath.metadata().unwrap().len() as usize; if opts.verbosity.is_some() { println!(" {}x{} pixels, PNG format", png.ihdr_data.width, @@ -63,22 +69,24 @@ pub fn optimize(filepath: &Path, opts: &Options) -> Result<(), String> { png.ihdr_data.bit_depth, png.ihdr_data.color_type); } - println!(" IDAT size = {} bytes", idat_current_size); - println!(" File size = {} bytes", file_current_size); + println!(" IDAT size = {} bytes", idat_original_size); + println!(" File size = {} bytes", file_original_size); } // TODO: Bit depth/palette reduction // // TODO: Apply interlacing changes // - // Go through selected permutations and determine the best - let mut best: Option<(u8, u8, u8, u8)> = None; + // TODO: Force reencoding if interlacing was changed if opts.idat_recoding { + // Go through selected permutations and determine the best + let mut best: Option<(u8, u8, u8, u8)> = None; let combinations = opts.filter.len() * opts.compression.len() * opts.memory.len() * opts.strategies.len(); println!("Trying: {} combinations", combinations); // TODO: Force reencoding if interlacing was changed // TODO: Multithreading for f in &opts.filter { + // TODO: Apply filter for zc in &opts.compression { for zm in &opts.memory { for zs in &opts.strategies { @@ -116,14 +124,48 @@ pub fn optimize(filepath: &Path, opts: &Options) -> Result<(), String> { better.3, better.0, png.idat_data.len()); - } else { - println!("IDAT already optimized"); + png.ihdr_data.filter = better.0; } } - // TODO: Backup before writing? + // TODO: Perform stripping - // TODO: Write output file + let output_data = png.output(); + if file_original_size <= output_data.len() { + println!("File already optimized"); + return Ok(()) + } + + if opts.pretend { + println!("Running in pretend mode, no output"); + } else { + if opts.backup { + match fs::copy(in_file, in_file.with_extension(format!("bak.{}", in_file.extension().unwrap().to_str().unwrap()))) { + Ok(x) => x, + Err(_) => return Err(format!("Unable to write to backup file at {}", opts.out_file.display())) + }; + } + + if opts.stdout { + let mut buffer = BufWriter::new(io::stdout()); + match buffer.write_all(&output_data) { + Ok(_) => (), + Err(_) => return Err(format!("Unable to write to stdout")) + } + } else { + let out_file = match File::create(&opts.out_file) { + Ok(x) => x, + Err(_) => return Err(format!("Unable to write to file {}", opts.out_file.display())) + }; + let mut buffer = BufWriter::new(out_file); + match buffer.write_all(&output_data) { + Ok(_) => println!("Output: {}", opts.out_file.display()), + Err(_) => return Err(format!("Unable to write to file {}", opts.out_file.display())) + } + } + } + println!(" IDAT size = {} bytes ({} bytes decrease)", png.idat_data.len(), idat_original_size - png.idat_data.len()); + println!(" file size = {} bytes ({} bytes = {:.2}% decrease)", output_data.len(), file_original_size - output_data.len(), (file_original_size - output_data.len()) as f64 / file_original_size as f64); Ok(()) } diff --git a/src/main.rs b/src/main.rs index bf33f2eb..561ef08b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,7 +5,6 @@ extern crate regex; use clap::{App, Arg, ArgMatches}; use regex::Regex; use std::collections::HashSet; -use std::path::Path; use std::path::PathBuf; fn main() { @@ -27,6 +26,7 @@ fn main() { out_dir: None, stdout: false, pretend: false, + recursive: false, fix_errors: false, clobber: true, create: true, @@ -68,7 +68,12 @@ fn main() { .possible_value("6")) .arg(Arg::with_name("backup") .help("Back up modified files") + .short("b") .long("backup")) + .arg(Arg::with_name("recursive") + .help("Recurse into subdirectories") + .short("r") + .long("recursive")) .arg(Arg::with_name("output_dir") .help("Write output file(s) to ") .long("dir") @@ -207,19 +212,46 @@ fn main() { let mut opts = default_opts; - parse_opts_into_struct(&matches, &mut opts); + match parse_opts_into_struct(&matches, &mut opts) { + Ok(_) => (), + Err(x) => { + println!("{}", x); + return (); + } + } - for input in matches.values_of("files").unwrap() { - opts.out_file = PathBuf::from(input); - // TODO: Handle wildcards - match optipng::optimize(Path::new(input), &opts) { + handle_optimization(matches.values_of("files") + .unwrap() + .iter() + .map(PathBuf::from) + .collect(), + &mut opts); +} + +fn handle_optimization(inputs: Vec, opts: &mut optipng::Options) { + for input in inputs { + if input.is_dir() { + if opts.recursive { + handle_optimization(input.read_dir().unwrap().map(|x| x.unwrap().path()).collect(), + opts) + } else { + println!("{} is a directory, skipping", input.display()); + } + continue; + } + if let Some(out_dir) = opts.out_dir.clone() { + opts.out_file = out_dir.join(input.file_name().unwrap()); + } else { + opts.out_file = input.clone(); + } + match optipng::optimize(&input, opts) { Ok(_) => (), Err(x) => println!("{}", x), }; } } -fn parse_opts_into_struct(matches: &ArgMatches, opts: &mut optipng::Options) { +fn parse_opts_into_struct(matches: &ArgMatches, opts: &mut optipng::Options) -> Result<(), String> { match matches.value_of("optimization") { Some("0") => { opts.idat_recoding = false; @@ -355,7 +387,13 @@ fn parse_opts_into_struct(matches: &ArgMatches, opts: &mut optipng::Options) { } if let Some(x) = matches.value_of("output_dir") { - opts.out_dir = Some(PathBuf::from(x)); + let path = PathBuf::from(x); + if !path.exists() { + + } else if !path.is_dir() { + return Err(format!("{} is an existing file (not a directory), cannot create directory", x)); + } + opts.out_dir = Some(path); } if let Some(x) = matches.value_of("output_file") { @@ -370,6 +408,10 @@ fn parse_opts_into_struct(matches: &ArgMatches, opts: &mut optipng::Options) { opts.backup = true; } + if matches.is_present("recursive") { + opts.recursive = true; + } + if matches.is_present("fix") { opts.fix_errors = true; } @@ -419,6 +461,8 @@ fn parse_opts_into_struct(matches: &ArgMatches, opts: &mut optipng::Options) { if matches.is_present("strip") { opts.strip = true; } + + Ok(()) } fn parse_numeric_range_opts(input: &str, diff --git a/src/png.rs b/src/png.rs index e82c6745..f5d72359 100644 --- a/src/png.rs +++ b/src/png.rs @@ -1,5 +1,5 @@ use std::io::Cursor; -use byteorder::{BigEndian, ReadBytesExt}; +use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; use crc::crc32; use std::collections::HashMap; use std::fmt; @@ -30,6 +30,18 @@ impl fmt::Display for ColorType { } } +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)] pub enum BitDepth { One, @@ -53,12 +65,25 @@ impl fmt::Display for BitDepth { } } +impl BitDepth { + fn as_u8(&self) -> u8 { + match *self { + BitDepth::One => 1, + BitDepth::Two => 2, + BitDepth::Four => 4, + BitDepth::Eight => 8, + BitDepth::Sixteen => 16, + } + } +} + #[derive(Debug)] pub struct PngData { pub idat_data: Vec, pub ihdr_data: IhdrData, pub raw_data: Vec, pub palette: Option>, + pub aux_headers: HashMap>, } #[derive(Debug)] @@ -131,7 +156,8 @@ impl PngData { Err(x) => return Err(x), }, raw_data: raw_data, - palette: aux_headers.get("PLTE").cloned(), + palette: aux_headers.remove("PLTE"), + aux_headers: aux_headers, }) } pub fn bits_per_pixel(&self) -> u8 { @@ -143,6 +169,68 @@ impl PngData { ColorType::RGBA => 4, } } + 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(); + ihdr_data.write_u8(self.ihdr_data.filter).ok(); + ihdr_data.write_u8(self.ihdr_data.interlaced).ok(); + output.reserve(ihdr_data.len() + 12); + output.write_u32::(ihdr_data.len() as u32).ok(); + let mut type_head = "IHDR".as_bytes().to_owned(); + let crc = crc32::checksum_ieee(&ihdr_data); + output.append(&mut type_head); + output.append(&mut ihdr_data); + output.write_u32::(crc).ok(); + // Ancillary headers + for (key, header) in self.aux_headers.iter() { + let mut header_data = header.clone(); + output.reserve(header_data.len() + 12); + output.write_u32::(header_data.len() as u32).ok(); + let mut type_head = key.as_bytes().to_owned(); + let crc = crc32::checksum_ieee(&header_data); + output.append(&mut type_head); + output.append(&mut header_data); + output.write_u32::(crc).ok(); + } + // Palette + if let Some(palette) = self.palette.clone() { + let mut palette_data = palette.clone(); + output.reserve(palette_data.len() + 12); + output.write_u32::(palette_data.len() as u32).ok(); + let mut type_head = "PLTE".as_bytes().to_owned(); + let crc =crc32::checksum_ieee(&palette_data); + output.append(&mut type_head); + output.append(&mut palette_data); + output.write_u32::(crc).ok(); + } + // IDAT data + let mut idat_data = self.idat_data.clone(); + output.reserve(idat_data.len() + 12); + output.write_u32::(idat_data.len() as u32).ok(); + let mut type_head = "IDAT".as_bytes().to_owned(); + let crc = crc32::checksum_ieee(&idat_data); + output.append(&mut type_head); + output.append(&mut idat_data); + output.write_u32::(crc).ok(); + // Stream end + let mut iend_data = vec![]; + output.reserve(iend_data.len() + 12); + output.write_u32::(iend_data.len() as u32).ok(); + let mut type_head = "IEND".as_bytes().to_owned(); + let crc = crc32::checksum_ieee(&iend_data); + output.append(&mut type_head); + output.append(&mut iend_data); + output.write_u32::(crc).ok(); + + output + } } fn file_header_is_valid(bytes: &[u8]) -> bool {