From fede62e46b1aa8ecd49375bb4cddd60b8dc1f1dc Mon Sep 17 00:00:00 2001
From: CheerfulPianissimo <sooraj2642@gmail.com>
Date: Sun, 24 Mar 2024 18:19:32 +0530
Subject: [PATCH] feat(clipboard): reduce buffer allocations for encoding image

---
 wayshot/src/wayshot.rs | 77 +++++++++++++++++++++++-------------------
 1 file changed, 43 insertions(+), 34 deletions(-)

diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs
index 04b20649..c6f446b1 100644
--- a/wayshot/src/wayshot.rs
+++ b/wayshot/src/wayshot.rs
@@ -127,45 +127,54 @@ fn main() -> Result<()> {
         }
     };
 
-    let mut buffer = Cursor::new(Vec::new());
-    image_buffer.write_to(&mut buffer, requested_encoding)?;
-
     if let Some(file) = file {
-        let mut file = File::create(file)?;
-        file.write_all(buffer.get_ref())?;
-    } else {
-        if stdout_print {
-            let stdout = stdout();
-            let mut writer = BufWriter::new(stdout.lock());
-            writer.write_all(buffer.get_ref())?;
+        image_buffer.save(file)?;
+    } else if stdout_print {
+        let mut buffer = Cursor::new(Vec::new());
+        image_buffer.write_to(&mut buffer, requested_encoding)?;
+        let stdout = stdout();
+        let mut writer = BufWriter::new(stdout.lock());
+        writer.write_all(buffer.get_ref())?;
+        if cli.clipboard {
+            clipboard_daemonize(buffer)?;
         }
     }
 
-    if cli.clipboard {
-        let mut opts = Options::new();
-        match unsafe { fork() } {
-            // Having the image persistently available on the clipboard requires a wayshot process to be alive.
-            // Fork the process with a child detached from the main process and have the parent exit
-            Ok(ForkResult::Parent { .. }) => {
-                return Ok(());
-            }
-            Ok(ForkResult::Child) => {
-                opts.foreground(true); // Offer the image till something else is available on the clipboard
-                opts.copy(
-                    Source::Bytes(buffer.into_inner().into()),
-                    MimeType::Autodetect,
-                )?;
-            }
-            Err(e) => {
-                tracing::warn!("Fork failed with error: {e}, couldn't offer image on the clipboard persistently.
-                 Use a clipboard manager to record screenshot.");
-                opts.copy(
-                    Source::Bytes(buffer.into_inner().into()),
-                    MimeType::Autodetect,
-                )?;
-            }
-        }
+    if !stdout_print && cli.clipboard {
+        let mut buffer = Cursor::new(Vec::new());
+        image_buffer.write_to(&mut buffer, requested_encoding)?;
+        clipboard_daemonize(buffer)?;
     }
 
     Ok(())
 }
+
+/// Daemonize and copy the given buffer containing the encoded image to the clipboard
+fn clipboard_daemonize(buffer: Cursor<Vec<u8>>) -> Result<()> {
+    let mut opts = Options::new();
+    match unsafe { fork() } {
+        // Having the image persistently available on the clipboard requires a wayshot process to be alive.
+        // Fork the process with a child detached from the main process and have the parent exit
+        Ok(ForkResult::Parent { .. }) => {
+            return Ok(());
+        }
+        Ok(ForkResult::Child) => {
+            opts.foreground(true); // Offer the image till something else is available on the clipboard
+            opts.copy(
+                Source::Bytes(buffer.into_inner().into()),
+                MimeType::Autodetect,
+            )?;
+        }
+        Err(e) => {
+            tracing::warn!(
+                "Fork failed with error: {e}, couldn't offer image on the clipboard persistently.
+                 Use a clipboard manager to record screenshot."
+            );
+            opts.copy(
+                Source::Bytes(buffer.into_inner().into()),
+                MimeType::Autodetect,
+            )?;
+        }
+    }
+    Ok(())
+}