From 0f88f11b7ec92a7195c5e06140d2e61c2aebd989 Mon Sep 17 00:00:00 2001
From: Samrat Man Singh <samratmansingh@gmail.com>
Date: Fri, 15 Dec 2023 08:18:29 +0545
Subject: [PATCH] Pre-allocate buffer in `remove_extra_data`

---
 src/darwin.rs      |  1 +
 src/image_utils.rs | 18 ++++++++++++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/darwin.rs b/src/darwin.rs
index 4fdfc4d..be16623 100644
--- a/src/darwin.rs
+++ b/src/darwin.rs
@@ -20,6 +20,7 @@ fn capture(display_info: &DisplayInfo, cg_rect: CGRect) -> Result<RgbaImage> {
     let height = cg_image.height();
     let clean_buf = remove_extra_data(
         width,
+        height,
         cg_image.bytes_per_row(),
         Vec::from(cg_image.data().bytes()),
     );
diff --git a/src/image_utils.rs b/src/image_utils.rs
index 43d277e..d71a7f8 100644
--- a/src/image_utils.rs
+++ b/src/image_utils.rs
@@ -24,10 +24,19 @@ pub fn bgra_to_rgba_image(width: u32, height: u32, buf: Vec<u8>) -> Result<RgbaI
 /// https://github.com/nashaofu/screenshots-rs/issues/29
 /// https://github.com/nashaofu/screenshots-rs/issues/38
 #[cfg(any(target_os = "macos", test))]
-pub fn remove_extra_data(width: usize, bytes_per_row: usize, buf: Vec<u8>) -> Vec<u8> {
-    buf.chunks_exact(bytes_per_row)
-        .flat_map(|row| row.split_at(width * 4).0.to_owned())
-        .collect()
+pub fn remove_extra_data(
+    width: usize,
+    height: usize,
+    bytes_per_row: usize,
+    buf: Vec<u8>,
+) -> Vec<u8> {
+    let extra_bytes_per_row = bytes_per_row - width * 4;
+    let mut result = Vec::with_capacity(buf.len() - extra_bytes_per_row * height);
+    for row in buf.chunks_exact(bytes_per_row) {
+        result.extend_from_slice(&row[..width * 4]);
+    }
+
+    result
 }
 
 #[cfg(any(target_os = "linux", test))]
@@ -61,6 +70,7 @@ mod tests {
     #[test]
     fn extra_data() {
         let clean = remove_extra_data(
+            2,
             2,
             9,
             vec![