Merge pull request google#34 from shepmaster/writer

Output to a generic that implements Write

Cargo.toml

@@ -13,6 +13,7 @@ readme = "README"
 libc = "0.2"
 crc = "1.2"
 adler32 = "0.2"
+byteorder = "0.5.3"
 
 [profile.release]
 debug = true

src/gzip.rs

@@ -1,35 +1,33 @@
+use std::io::{self, Write};
 use crc::crc32;
+use byteorder::{LittleEndian, WriteBytesExt};
 
 use deflate::{deflate, BlockType};
 use Options;
 
-/// Compresses the data according to the gzip specification, RFC 1952.
-pub fn gzip_compress(options: &Options, in_data: &[u8], out: &mut Vec<u8>) {
-    out.push(31);  /* ID1 */
-    out.push(139);  /* ID2 */
-    out.push(8);  /* CM */
-    out.push(0);  /* FLG */
-    /* MTIME */
-    out.push(0);
-    out.push(0);
-    out.push(0);
-    out.push(0);
+static HEADER: &'static [u8] = &[
+    31,  // ID1
+    139, // ID2
+    8,   // CM
+    0,   // FLG
 
-    out.push(2);  /* XFL, 2 indicates best compression. */
-    out.push(3);  /* OS follows Unix conventions. */
+    0,   // MTIME
+    0,
+    0,
+    0,
 
-    deflate(options, BlockType::Dynamic, in_data, out);
+    2,   // XFL, 2 indicates best compression.
+    3,   // OS follows Unix conventions.
+];
 
-    let crc = crc32::checksum_ieee(in_data);
+/// Compresses the data according to the gzip specification, RFC 1952.
+pub fn gzip_compress<W>(options: &Options, in_data: &[u8], mut out: W) -> io::Result<()>
+    where W: Write
+{
+    try!(out.by_ref().write_all(HEADER));
 
-    out.push(crc as u8);
-    out.push((crc >> 8) as u8);
-    out.push((crc >> 16) as u8);
-    out.push((crc >> 24) as u8);
+    try!(deflate(options, BlockType::Dynamic, in_data, out.by_ref()));
 
-    let input_size = in_data.len() as u32;
-    out.push(input_size as u8);
-    out.push((input_size >> 8) as u8);
-    out.push((input_size >> 16) as u8);
-    out.push((input_size >> 24) as u8);
+    try!(out.by_ref().write_u32::<LittleEndian>(crc32::checksum_ieee(in_data)));
+    out.write_u32::<LittleEndian>(in_data.len() as u32)
 }

src/lib.rs

@@ -3,6 +3,7 @@
 extern crate libc;
 extern crate crc;
 extern crate adler32;
+extern crate byteorder;
 
 mod iter;
 mod blocksplitter;
@@ -18,6 +19,8 @@ mod tree;
 mod util;
 mod zlib;
 
+use std::io::{self, Write};
+
 use deflate::{deflate, BlockType};
 use gzip::gzip_compress;
 use zlib::zlib_compress;
@@ -58,7 +61,9 @@ pub enum Format {
     Deflate,
 }
 
-pub fn compress(options: &Options, output_type: &Format, in_data: &[u8], out: &mut Vec<u8>) {
+pub fn compress<W>(options: &Options, output_type: &Format, in_data: &[u8], out: W) -> io::Result<()>
+    where W: Write
+{
     match *output_type {
         Format::Gzip => gzip_compress(options, in_data, out),
         Format::Zlib => zlib_compress(options, in_data, out),

src/main.rs

@@ -1,5 +1,6 @@
 use std::env;
 use std::io::prelude::*;
+use std::io::{self, BufWriter};
 use std::fs::File;
 
 extern crate zopfli;
@@ -27,22 +28,49 @@ fn main() {
         file.read_to_end(&mut data)
             .unwrap_or_else(|why| panic!("couldn't read {}: {}", filename, why));
 
-        let mut out_data = vec![];
         let out_filename = format!("{}{}", filename, extension);
 
         // Attempt to create the output file, panic if the output file could not be opened
-        let mut out_file = File::create(&out_filename)
+        let out_file = File::create(&out_filename)
             .unwrap_or_else(|why| panic!("couldn't create output file {}: {}", out_filename, why));
+        let mut out_file = WriteStatistics::new(BufWriter::new(out_file));
 
-        zopfli::compress(&options, &output_type, &data, &mut out_data);
+        zopfli::compress(&options, &output_type, &data, &mut out_file)
+            .unwrap_or_else(|why| panic!("couldn't write to output file {}: {}", out_filename, why));
 
         if options.verbose {
-            let out_size = out_data.len();
+            let out_size = out_file.count;
             println!("Original Size: {}, Compressed: {}, Compression: {}% Removed", filesize, out_size, 100.0 * (filesize - out_size) as f64 / filesize as f64);
         }
+    }
+}
 
-        // Write the `out_data` into the newly created file.
-        out_file.write_all(&out_data)
-            .unwrap_or_else(|why| panic!("couldn't write to output file {}: {}", out_filename, why));
+struct WriteStatistics<W> {
+    inner: W,
+    count: usize,
+}
+
+impl<W> WriteStatistics<W> {
+    fn new(inner: W) -> Self {
+        WriteStatistics {
+            inner: inner,
+            count: 0,
+        }
+    }
+}
+
+impl<W> Write for WriteStatistics<W>
+    where W: Write
+{
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let res = self.inner.write(buf);
+        if let Ok(size) = res {
+            self.count += size;
+        }
+        res
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.inner.flush()
     }
 }

src/zlib.rs

@@ -1,26 +1,24 @@
-use std::io;
-
+use std::io::{self, Write};
 use adler32::adler32;
+use byteorder::{LittleEndian, WriteBytesExt};
 
 use deflate::{deflate, BlockType};
 use Options;
 
-pub fn zlib_compress(options: &Options, in_data: &[u8], out: &mut Vec<u8>) {
-    let checksum = adler32(io::Cursor::new(&in_data)).expect("Error with adler32");
+pub fn zlib_compress<W>(options: &Options, in_data: &[u8], mut out: W) -> io::Result<()>
+    where W: Write
+{
     let cmf = 120;  /* CM 8, CINFO 7. See zlib spec.*/
     let flevel = 3;
     let fdict = 0;
     let mut cmfflg = 256 * cmf + fdict * 32 + flevel * 64;
     let fcheck = 31 - cmfflg % 31;
     cmfflg += fcheck;
 
-    out.push((cmfflg / 256) as u8);
-    out.push((cmfflg % 256) as u8);
+    try!(out.by_ref().write_u16::<LittleEndian>(cmfflg));
 
-    deflate(options, BlockType::Dynamic, in_data, out);
+    try!(deflate(options, BlockType::Dynamic, in_data, out.by_ref()));
 
-    out.push(((checksum >> 24) % 256) as u8);
-    out.push(((checksum >> 16) % 256) as u8);
-    out.push(((checksum >> 8) % 256) as u8);
-    out.push((checksum % 256) as u8);
+    let checksum = adler32(io::Cursor::new(&in_data)).expect("Error with adler32");
+    out.write_u32::<LittleEndian>(checksum)
 }