Write output files

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 idat_original_size = png.idat_data.len();
-    let file_current_size = filepath.metadata().unwrap().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!("    IDAT size = {} bytes", idat_original_size);
-        println!("    File size = {} bytes", file_current_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
+    // TODO: Force reencoding if interlacing was changed
-    let mut best: Option<(u8, u8, u8, u8)> = None;
     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 {
+            png.ihdr_data.filter = better.0;
-            println!("IDAT already optimized");
         }
     }
 
-    // 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(())
 }

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 <directory>")
                                .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() {
+    handle_optimization(matches.values_of("files")
-        opts.out_file = PathBuf::from(input);
+                               .unwrap()
-        // TODO: Handle wildcards
+                               .iter()
-        match optipng::optimize(Path::new(input), &opts) {
+                               .map(PathBuf::from)
+                               .collect(),
+                        &mut opts);
+}
+
+fn handle_optimization(inputs: Vec<PathBuf>, 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,

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<u8>,
     pub ihdr_data: IhdrData,
     pub raw_data: Vec<u8>,
     pub palette: Option<Vec<u8>>,
+    pub aux_headers: HashMap<String, Vec<u8>>,
 }
 
 #[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<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();
+        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::<BigEndian>(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::<BigEndian>(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::<BigEndian>(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::<BigEndian>(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::<BigEndian>(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::<BigEndian>(crc).ok();
+        }
+        // IDAT data
+        let mut idat_data = self.idat_data.clone();
+        output.reserve(idat_data.len() + 12);
+        output.write_u32::<BigEndian>(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::<BigEndian>(crc).ok();
+        // Stream end
+        let mut iend_data = vec![];
+        output.reserve(iend_data.len() + 12);
+        output.write_u32::<BigEndian>(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::<BigEndian>(crc).ok();
+
+        output
+    }
 }
 
 fn file_header_is_valid(bytes: &[u8]) -> bool {