From 03e955d500506ff018fd3671c2f0fb502b688b19 Mon Sep 17 00:00:00 2001
From: asolana <110843012+ksolana@users.noreply.github.com>
Date: Thu, 9 Jan 2025 11:40:25 -0800
Subject: [PATCH] Make the buffer size configurable (#307)

* Add new_with_capacity to allow users to create Writers with custom capacity.

* Add test for buffer_with_capacity
---
 src/stream/zio/writer.rs | 48 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/stream/zio/writer.rs b/src/stream/zio/writer.rs
index 8484089e..16011c8a 100644
--- a/src/stream/zio/writer.rs
+++ b/src/stream/zio/writer.rs
@@ -43,13 +43,20 @@ where
     W: Write,
     D: Operation,
 {
-    /// Creates a new `Writer`.
+    /// Creates a new `Writer` with a fixed buffer capacity of 32KB
     ///
     /// 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
+        new_with_capacity(W, D, 32 * 1024)
+    }
+
+    /// Creates a new `Writer` with user defined capacity.
+    ///
+    /// All output from the given operation will be forwarded to `writer`.
+    pub fn new_with_capacity(writer: W, operation: D, capacity: usize) -> Self {
         Self::with_output_buffer(
-            Vec::with_capacity(32 * 1024),
+            Vec::with_capacity(capacity),
             writer,
             operation,
         )
@@ -314,6 +321,25 @@ mod tests {
         assert_eq!(&decoded, input);
     }
 
+    #[test]
+    fn test_compress_with_capacity() {
+        use crate::stream::raw::Encoder;
+
+        let input = b"AbcdefghAbcdefgh.";
+
+        // Test writer
+        let mut output = Vec::new();
+        {
+            let mut writer =
+                Writer::new_with_capacity(&mut output, Encoder::new(1).unwrap(), 64);
+            assert_eq!(writer.buffer().capacity() == 64);
+            writer.write_all(input).unwrap();
+            writer.finish().unwrap();
+        }
+        let decoded = crate::decode_all(&output[..]).unwrap();
+        assert_eq!(&decoded, input);
+    }
+
     #[test]
     fn test_decompress() {
         use crate::stream::raw::Decoder;
@@ -331,4 +357,22 @@ mod tests {
         // println!("Output: {:?}", output);
         assert_eq!(&output, input);
     }
+
+    #[test]
+    fn test_decompress_with_capacity() {
+        use crate::stream::raw::Decoder;
+
+        let input = b"AbcdefghAbcdefgh.";
+        let compressed = crate::encode_all(&input[..], 1).unwrap();
+
+        // Test writer
+        let mut output = Vec::new();
+        {
+            let mut writer = Writer::new(&mut output, Decoder::new().unwrap(), 64);
+            assert_eq!(writer.buffer().capacity() == 64);
+            writer.write_all(&compressed).unwrap();
+            writer.finish().unwrap();
+        }
+        assert_eq!(&output, input);
+    }
 }