Add reset methods to decoders as well

While there's no inherent reset method on the underlying streams we can at least reuse the buffer.
Byron · Sep 9, 2015 · a930f18 · a930f18
1 parent 0e66bab
commit a930f18
Show file tree

Hide file tree

Showing 4 changed files with 138 additions and 14 deletions.
diff --git a/src/deflate.rs b/src/deflate.rs
@@ -2,7 +2,6 @@
 
 use std::io::prelude::*;
 use std::io;
-use std::iter::repeat;
 
 use raw;
 
@@ -118,7 +117,7 @@ impl<R: Read> DecoderReader<R> {
     /// Creates a new decoder which will decompress data read from the given
     /// stream.
     pub fn new(r: R) -> DecoderReader<R> {
-        DecoderReader::new_with_buf(r, repeat(0).take(32 * 1024).collect())
+        DecoderReader::new_with_buf(r, vec![0; 32 * 1024])
     }
 
     /// Same as `new`, but the intermediate buffer for data is specified.
@@ -129,6 +128,17 @@ impl<R: Read> DecoderReader<R> {
         DecoderReader { inner: raw::DecoderReader::new(r, true, buf) }
     }
 
+    /// Resets the state of this decoder entirely, swapping out the input
+    /// stream for another.
+    ///
+    /// This will reset the internal state of this decoder and replace the
+    /// input stream with the one provided, returning the previous input
+    /// stream. Future data read from this decoder will be the decompressed
+    /// version of `r`'s data.
+    pub fn reset(&mut self, r: R) -> R {
+        self.inner.reset(r, true)
+    }
+
     /// Consumes this decoder, returning the underlying reader.
     pub fn into_inner(self) -> R {
         self.inner.into_inner()
@@ -153,6 +163,22 @@ impl<W: Write> DecoderWriter<W> {
         }
     }
 
+    /// Resets the state of this decoder entirely, swapping out the output
+    /// stream for another.
+    ///
+    /// This function will finish encoding the current stream into the current
+    /// output stream before swapping out the two output streams. If the stream
+    /// cannot be finished an error is returned.
+    ///
+    /// This will then reset the internal state of this decoder and replace the
+    /// output stream with the one provided, returning the previous output
+    /// stream. Future data written to this decoder will be decompressed into
+    /// the output stream `w`.
+    pub fn reset(&mut self, w: W) -> io::Result<W> {
+        try!(self.inner.finish());
+        Ok(self.inner.reset(w, true))
+    }
+
     /// Consumes this encoder, flushing the output stream.
     ///
     /// This will flush the underlying data stream and then return the contained
@@ -242,4 +268,38 @@ mod tests {
         r.read_to_end(&mut c).unwrap();
         assert!(a == b && b == c);
     }
+
+    #[test]
+    fn reset_decoder() {
+        let v = thread_rng().gen_iter::<u8>().take(1024 * 1024)
+                            .collect::<Vec<_>>();
+        let mut w = EncoderWriter::new(Vec::new(), Default);
+        w.write_all(&v).unwrap();
+        let data = w.finish().unwrap();
+
+        {
+            let (mut a, mut b, mut c) = (Vec::new(), Vec::new(), Vec::new());
+            let mut r = DecoderReader::new(&data[..]);
+            r.read_to_end(&mut a).unwrap();
+            r.reset(&data);
+            r.read_to_end(&mut b).unwrap();
+
+            let mut r = DecoderReader::new(&data[..]);
+            r.read_to_end(&mut c).unwrap();
+            assert!(a == b && b == c && c == v);
+        }
+
+        {
+            let mut w = DecoderWriter::new(Vec::new());
+            w.write_all(&data).unwrap();
+            let a = w.reset(Vec::new()).unwrap();
+            w.write_all(&data).unwrap();
+            let b = w.finish().unwrap();
+
+            let mut w = DecoderWriter::new(Vec::new());
+            w.write_all(&data).unwrap();
+            let c = w.finish().unwrap();
+            assert!(a == b && b == c && c == v);
+        }
+    }
 }
diff --git a/src/gz.rs b/src/gz.rs
@@ -7,7 +7,6 @@ use std::env;
 use std::ffi::CString;
 use std::io::prelude::*;
 use std::io;
-use std::iter::repeat;
 use std::mem;
 
 use Compression;
@@ -127,8 +126,7 @@ impl Builder {
     pub fn read<R: Read>(self, r: R, lvl: Compression) -> EncoderReader<R> {
         let crc = CrcReader::new(r);
         EncoderReader {
-            inner: raw::EncoderReader::new(crc, lvl, true,
-                                           repeat(0).take(32 * 1024).collect()),
+            inner: raw::EncoderReader::new(crc, lvl, true, vec![0; 32 * 1024]),
             header: self.into_header(lvl),
             pos: 0,
             eof: false,
@@ -138,7 +136,7 @@ impl Builder {
     fn into_header(self, lvl: Compression) -> Vec<u8> {
         let Builder { extra, filename, comment, mtime } = self;
         let mut flg = 0;
-        let mut header = repeat(0u8).take(10).collect::<Vec<_>>();
+        let mut header = vec![0u8; 10];
         match extra {
             Some(v) => {
                 flg |= FEXTRA;
@@ -338,7 +336,7 @@ impl<R: Read> DecoderReader<R> {
 
         let extra = if flg & FEXTRA != 0 {
             let xlen = try!(read_le_u16(&mut crc_reader));
-            let mut extra = repeat(0).take(xlen as usize).collect::<Vec<_>>();
+            let mut extra = vec![0; xlen as usize];
             try!(fill(&mut crc_reader, &mut extra));
             Some(extra)
         } else {
@@ -376,7 +374,7 @@ impl<R: Read> DecoderReader<R> {
         }
 
         let flate = raw::DecoderReader::new(crc_reader.into_inner(), true,
-                                            repeat(0).take(32 * 1024).collect());
+                                            vec![0; 32 * 1024]);
         return Ok(DecoderReader {
             inner: CrcReader::new(flate),
             header: Header {

diff --git a/src/raw.rs b/src/raw.rs
@@ -106,6 +106,12 @@ impl<W: Write> DecoderWriter<W> {
         })
     }
 
+    pub fn reset(&mut self, w: W, raw: bool) -> W {
+        self.0.stream = Stream::new_decompress(raw);
+        self.0.buf.truncate(0);
+        mem::replace(&mut self.0.inner, Some(w)).unwrap()
+    }
+
     pub fn finish(&mut self) -> io::Result<()> {
         self.0.finish(&mut |stream, inner| {
             stream.decompress_vec(&[], inner, Flush::Finish)
@@ -215,6 +221,13 @@ impl<R: Read> DecoderReader<R> {
         })
     }
 
+    pub fn reset(&mut self, r: R, raw: bool) -> R {
+        self.0.stream = Stream::new_decompress(raw);
+        self.0.cap = 0;
+        self.0.pos = 0;
+        mem::replace(&mut self.0.inner, r)
+    }
+
     pub fn into_inner(self) -> R { self.0.inner }
 
     pub fn read_raw(&mut self, buf: &mut [u8]) -> io::Result<usize> {

diff --git a/src/zlib.rs b/src/zlib.rs
@@ -2,7 +2,6 @@
 
 use std::io::prelude::*;
 use std::io;
-use std::iter::repeat;
 
 use raw;
 
@@ -87,8 +86,7 @@ impl<R: Read> EncoderReader<R> {
     /// stream and emit the compressed stream.
     pub fn new(r: R, level: ::Compression) -> EncoderReader<R> {
         EncoderReader {
-            inner: raw::EncoderReader::new(r, level, false,
-                                           repeat(0).take(32 * 1024).collect())
+            inner: raw::EncoderReader::new(r, level, false, vec![0; 32 * 1024]),
         }
     }
 
@@ -119,7 +117,7 @@ impl<R: Read> DecoderReader<R> {
     /// Creates a new decoder which will decompress data read from the given
     /// stream.
     pub fn new(r: R) -> DecoderReader<R> {
-        DecoderReader::new_with_buf(r, repeat(0).take(32 * 1024).collect())
+        DecoderReader::new_with_buf(r, vec![0; 32 * 1024])
     }
 
     /// Same as `new`, but the intermediate buffer for data is specified.
@@ -130,6 +128,17 @@ impl<R: Read> DecoderReader<R> {
         DecoderReader { inner: raw::DecoderReader::new(r, false, buf) }
     }
 
+    /// Resets the state of this decoder entirely, swapping out the input
+    /// stream for another.
+    ///
+    /// This will reset the internal state of this decoder and replace the
+    /// input stream with the one provided, returning the previous input
+    /// stream. Future data read from this decoder will be the decompressed
+    /// version of `r`'s data.
+    pub fn reset(&mut self, r: R) -> R {
+        self.inner.reset(r, false)
+    }
+
     /// Consumes this decoder, returning the underlying reader.
     pub fn into_inner(self) -> R {
         self.inner.into_inner()
@@ -145,7 +154,7 @@ impl<R: Read> Read for DecoderReader<R> {
 impl<W: Write> DecoderWriter<W> {
     /// Creates a new decoder which will write uncompressed data to the stream.
     ///
-    /// When this encoder is dropped or unwrapped the final pieces of data will
+    /// When this decoder is dropped or unwrapped the final pieces of data will
     /// be flushed.
     pub fn new(w: W) -> DecoderWriter<W> {
         DecoderWriter {
@@ -154,6 +163,18 @@ impl<W: Write> DecoderWriter<W> {
         }
     }
 
+    /// Resets the state of this decoder entirely, swapping out the output
+    /// stream for another.
+    ///
+    /// This will reset the internal state of this decoder and replace the
+    /// output stream with the one provided, returning the previous output
+    /// stream. Future data written to this decoder will be decompressed into
+    /// the output stream `w`.
+    pub fn reset(&mut self, w: W) -> io::Result<W> {
+        try!(self.inner.finish());
+        Ok(self.inner.reset(w, false))
+    }
+
     /// Consumes this encoder, flushing the output stream.
     ///
     /// This will flush the underlying data stream and then return the contained
@@ -212,6 +233,38 @@ mod tests {
         let w = w.finish().unwrap().finish().unwrap();
         assert!(w == v);
     }
-}
 
+    #[test]
+    fn reset_decoder() {
+        let v = thread_rng().gen_iter::<u8>().take(1024 * 1024)
+                            .collect::<Vec<_>>();
+        let mut w = EncoderWriter::new(Vec::new(), Default);
+        w.write_all(&v).unwrap();
+        let data = w.finish().unwrap();
+
+        {
+            let (mut a, mut b, mut c) = (Vec::new(), Vec::new(), Vec::new());
+            let mut r = DecoderReader::new(&data[..]);
+            r.read_to_end(&mut a).unwrap();
+            r.reset(&data);
+            r.read_to_end(&mut b).unwrap();
 
+            let mut r = DecoderReader::new(&data[..]);
+            r.read_to_end(&mut c).unwrap();
+            assert!(a == b && b == c && c == v);
+        }
+
+        {
+            let mut w = DecoderWriter::new(Vec::new());
+            w.write_all(&data).unwrap();
+            let a = w.reset(Vec::new()).unwrap();
+            w.write_all(&data).unwrap();
+            let b = w.finish().unwrap();
+
+            let mut w = DecoderWriter::new(Vec::new());
+            w.write_all(&data).unwrap();
+            let c = w.finish().unwrap();
+            assert!(a == b && b == c && c == v);
+        }
+    }
+}