From a930f183166bd77a3ee0574f6cd512fb7bbeef73 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Wed, 9 Sep 2015 16:20:11 -0700 Subject: [PATCH] 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. --- src/deflate.rs | 64 +++++++++++++++++++++++++++++++++++++++++++++++-- src/gz.rs | 10 ++++---- src/raw.rs | 13 ++++++++++ src/zlib.rs | 65 +++++++++++++++++++++++++++++++++++++++++++++----- 4 files changed, 138 insertions(+), 14 deletions(-) diff --git a/src/deflate.rs b/src/deflate.rs index 504807a38..7d22bf88e 100644 --- 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 DecoderReader { /// Creates a new decoder which will decompress data read from the given /// stream. pub fn new(r: R) -> DecoderReader { - 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 DecoderReader { 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 DecoderWriter { } } + /// 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 { + 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::().take(1024 * 1024) + .collect::>(); + 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 index 8d5a5fc2a..9dcd6033c 100644 --- 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(self, r: R, lvl: Compression) -> EncoderReader { 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 { let Builder { extra, filename, comment, mtime } = self; let mut flg = 0; - let mut header = repeat(0u8).take(10).collect::>(); + let mut header = vec![0u8; 10]; match extra { Some(v) => { flg |= FEXTRA; @@ -338,7 +336,7 @@ impl DecoderReader { 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::>(); + let mut extra = vec![0; xlen as usize]; try!(fill(&mut crc_reader, &mut extra)); Some(extra) } else { @@ -376,7 +374,7 @@ impl DecoderReader { } 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 index 0118960f4..6d78f64b3 100644 --- a/src/raw.rs +++ b/src/raw.rs @@ -106,6 +106,12 @@ impl DecoderWriter { }) } + 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 DecoderReader { }) } + 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 { diff --git a/src/zlib.rs b/src/zlib.rs index b61fff5a6..ebdc7aaf4 100644 --- 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 EncoderReader { /// stream and emit the compressed stream. pub fn new(r: R, level: ::Compression) -> EncoderReader { 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 DecoderReader { /// Creates a new decoder which will decompress data read from the given /// stream. pub fn new(r: R) -> DecoderReader { - 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 DecoderReader { 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 Read for DecoderReader { impl DecoderWriter { /// 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 { DecoderWriter { @@ -154,6 +163,18 @@ impl DecoderWriter { } } + /// 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 { + 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::().take(1024 * 1024) + .collect::>(); + 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); + } + } +}