From 7be43815d17a8adefd503ae052bf5df21ca3fd34 Mon Sep 17 00:00:00 2001
From: Moritz Borcherding <moritz.borcherding@web.de>
Date: Mon, 23 Dec 2024 13:53:34 +0100
Subject: [PATCH] fix take_drain to do what it says on the tin

---
 src/bin/zstd.rs                  | 9 ++++-----
 src/encoding/frame_compressor.rs | 4 ++--
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/bin/zstd.rs b/src/bin/zstd.rs
index 471389d..bdd80bd 100644
--- a/src/bin/zstd.rs
+++ b/src/bin/zstd.rs
@@ -159,8 +159,7 @@ fn main() {
 
     if flags.is_empty() {
         let mut encoder = FrameCompressor::new(CompressionLevel::Fastest);
-        let mut output = Vec::new();
-        encoder.set_drain(&mut output);
+        encoder.set_drain(Vec::new());
 
         for path in file_paths {
             let start_instant = Instant::now();
@@ -173,10 +172,8 @@ fn main() {
                 last_percent: 0,
             };
             encoder.set_source(file);
-
-            encoder.drain_mut().unwrap().clear();
             encoder.compress();
-            let output = encoder.drain_mut().unwrap();
+            let mut output: Vec<_> = encoder.take_drain().unwrap();
             println!(
                 "Compressed {path:} from {} to {} ({}%) took {}ms",
                 input_len,
@@ -188,6 +185,8 @@ fn main() {
                 },
                 start_instant.elapsed().as_millis()
             );
+            output.clear();
+            encoder.set_drain(output);
         }
     } else {
         decompress(&flags, &file_paths);
diff --git a/src/encoding/frame_compressor.rs b/src/encoding/frame_compressor.rs
index ac13559..29cdfad 100644
--- a/src/encoding/frame_compressor.rs
+++ b/src/encoding/frame_compressor.rs
@@ -211,8 +211,8 @@ impl<R: Read, W: Write, M: Matcher> FrameCompressor<R, W, M> {
     }
 
     /// Retrieve the drain
-    pub fn take_drain(&mut self) -> Option<&mut W> {
-        self.compressed_data.as_mut()
+    pub fn take_drain(&mut self) -> Option<W> {
+        self.compressed_data.take()
     }
 
     /// Before calling [FrameCompressor::compress] you can replace the matcher