Add zio::Writer::with_output_buffer

src/stream/write/mod.rs

@@ -187,23 +187,29 @@ impl<W: Write> Encoder<'static, W> {
         dictionary: &[u8],
     ) -> io::Result<Self> {
         let encoder = raw::Encoder::with_dictionary(level, dictionary)?;
-        let writer = zio::Writer::new(writer, encoder);
-        Ok(Encoder { writer })
+        Ok(Self::with_encoder(writer, encoder))
     }
 }
 
 impl<'a, W: Write> Encoder<'a, W> {
+    /// Creates a new encoder from a prepared zio writer.
+    pub fn with_writer(writer: zio::Writer<W, raw::Encoder<'a>>) -> Self {
+        Self { writer }
+    }
+
+    /// Creates a new encoder from the given `Write` and raw encoder.
+    pub fn with_encoder(writer: W, encoder: raw::Encoder<'a>) -> Self {
+        let writer = zio::Writer::new(writer, encoder);
+        Self::with_writer(writer)
+    }
+
     /// Creates an encoder that uses the provided context to compress a stream.
     pub fn with_context(
         writer: W,
         context: &'a mut zstd_safe::CCtx<'static>,
     ) -> Self {
-        Self {
-            writer: zio::Writer::new(
-                writer,
-                raw::Encoder::with_context(context),
-            ),
-        }
+        let encoder = raw::Encoder::with_context(context);
+        Self::with_encoder(writer, encoder)
     }
 
     /// Creates a new encoder, using an existing prepared `EncoderDictionary`.
@@ -218,8 +224,7 @@ impl<'a, W: Write> Encoder<'a, W> {
         'b: 'a,
     {
         let encoder = raw::Encoder::with_prepared_dictionary(dictionary)?;
-        let writer = zio::Writer::new(writer, encoder);
-        Ok(Encoder { writer })
+        Ok(Self::with_encoder(writer, encoder))
     }
 
     /// Creates a new encoder, using a ref prefix
@@ -232,8 +237,7 @@ impl<'a, W: Write> Encoder<'a, W> {
         'b: 'a,
     {
         let encoder = raw::Encoder::with_ref_prefix(level, ref_prefix)?;
-        let writer = zio::Writer::new(writer, encoder);
-        Ok(Encoder { writer })
+        Ok(Self::with_encoder(writer, encoder))
     }
 
     /// Returns a wrapper around `self` that will finish the stream on drop.
@@ -340,12 +344,32 @@ impl<W: Write> Decoder<'static, W> {
     /// but requires the dictionary to be present during decompression.)
     pub fn with_dictionary(writer: W, dictionary: &[u8]) -> io::Result<Self> {
         let decoder = raw::Decoder::with_dictionary(dictionary)?;
-        let writer = zio::Writer::new(writer, decoder);
-        Ok(Decoder { writer })
+        Ok(Self::with_decoder(writer, decoder))
     }
 }
 
 impl<'a, W: Write> Decoder<'a, W> {
+    /// Creates a new decoder around the given prepared zio writer.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// fn wrap<W: std::io::Write>(writer: W) -> zstd::stream::write::Decoder<'static, W> {
+    ///   let decoder = zstd::stream::raw::Decoder::new().unwrap();
+    ///   let writer = zstd::stream::zio::Writer::new(writer, decoder);
+    ///   zstd::stream::write::Decoder::with_writer(writer)
+    /// }
+    /// ```
+    pub fn with_writer(writer: zio::Writer<W, raw::Decoder<'a>>) -> Self {
+        Decoder { writer }
+    }
+
+    /// Creates a new decoder around the given `Write` and raw decoder.
+    pub fn with_decoder(writer: W, decoder: raw::Decoder<'a>) -> Self {
+        let writer = zio::Writer::new(writer, decoder);
+        Decoder { writer }
+    }
+
     /// Creates a new decoder, using an existing prepared `DecoderDictionary`.
     ///
     /// (Provides better compression ratio for small files,
@@ -358,8 +382,7 @@ impl<'a, W: Write> Decoder<'a, W> {
         'b: 'a,
     {
         let decoder = raw::Decoder::with_prepared_dictionary(dictionary)?;
-        let writer = zio::Writer::new(writer, decoder);
-        Ok(Decoder { writer })
+        Ok(Self::with_decoder(writer, decoder))
     }
 
     /// Acquires a reference to the underlying writer.

src/stream/zio/writer.rs

@@ -11,16 +11,30 @@ use crate::stream::raw::{InBuffer, Operation, OutBuffer};
 /// It can be used with either compression or decompression, and forwards the
 /// output to a wrapped `Write`.
 pub struct Writer<W, D> {
-    writer: W,
+    /// Either an encoder or a decoder.
     operation: D,
 
+    /// Where we send the output of the operation.
+    writer: W,
+
+    /// Offset into the buffer
+    ///
+    /// Only things after this matter. Things before have already been sent to the writer.
     offset: usize,
+
+    /// Output buffer
+    ///
+    /// Where the operation writes, before it gets flushed to the writer
     buffer: Vec<u8>,
 
     // When `true`, indicates that nothing should be added to the buffer.
     // All that's left if to empty the buffer.
     finished: bool,
 
+    /// When `true`, the operation just finished a frame.
+    ///
+    /// Only happens when decompressing.
+    /// The context needs to be re-initialized to process the next frame.
     finished_frame: bool,
 }
 
@@ -33,13 +47,31 @@ where
     ///
     /// All output from the given operation will be forwarded to `writer`.
     pub fn new(writer: W, operation: D) -> Self {
+        // 32KB buffer? That's what flate2 uses
+        Self::with_output_buffer(
+            Vec::with_capacity(32 * 1024),
+            writer,
+            operation,
+        )
+    }
+
+    /// Creates a new `Writer` using the given output buffer.
+    ///
+    /// The output buffer _must_ have pre-allocated capacity (its capacity will not be changed after).
+    ///
+    /// Usually you would use `Vec::with_capacity(desired_buffer_size)`.
+    pub fn with_output_buffer(
+        output_buffer: Vec<u8>,
+        writer: W,
+        operation: D,
+    ) -> Self {
         Writer {
             writer,
             operation,
 
             offset: 0,
             // 32KB buffer? That's what flate2 uses
-            buffer: Vec::with_capacity(32 * 1024),
+            buffer: output_buffer,
 
             finished: false,
             finished_frame: false,