From 6c30607e782c835150d76283a8192cf1787bae7a Mon Sep 17 00:00:00 2001
From: ef3d0c3e <ef3d0c3e@pundalik.org>
Date: Tue, 20 Aug 2024 09:52:46 +0200
Subject: [PATCH] Better image sizes

---
 src/png_data/header.rs |  1 -
 src/png_data/main.rs   | 18 +++++++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/png_data/header.rs b/src/png_data/header.rs
index b487ad6..9b9121f 100644
--- a/src/png_data/header.rs
+++ b/src/png_data/header.rs
@@ -1,6 +1,5 @@
 use crc::Crc;
 
-
 pub trait Encode {
 	/// Encode the data into a vector
 	fn encode(&self, vec: &mut Vec<u8>);
diff --git a/src/png_data/main.rs b/src/png_data/main.rs
index 5496f5a..0c63131 100644
--- a/src/png_data/main.rs
+++ b/src/png_data/main.rs
@@ -98,6 +98,22 @@ fn best_layout(size: u64, bits_per_pixel: u8) -> (u32, u32) {
 	(width as u32, sz.div_ceil(width) as u32)
 }
 
+/// Gets the minimum image buffer size in bytes
+fn minimum_size(color: ColorType, depth: BitDepth, width: u32, height: u32) -> usize
+{
+	let samples = width as usize * color.samples();
+	(match depth {
+		BitDepth::Sixteen => samples * 2,
+		BitDepth::Eight => samples,
+		subbyte => {
+			let samples_per_byte = 8 / subbyte as usize;
+			let whole = samples / samples_per_byte;
+			let fract = usize::from(samples % samples_per_byte > 0);
+			whole + fract
+		}
+	}) * height as usize
+}
+
 fn encode(input: String, output: String, layout: String, matches: Matches) -> Result<(), String> {
 	let layout = str_to_layout(layout.as_str())?;
 	let comment = matches.opt_str("c");
@@ -139,7 +155,7 @@ fn encode(input: String, output: String, layout: String, matches: Matches) -> Re
 		.map_err(|err| format!("Failed to write png header: {err}"))?;
 
 	// Image byte length
-	let byte_len = ((width as usize) * (height as usize) * (bits_per_pixel as usize)).div_ceil(8);
+	let byte_len = minimum_size(layout.0, layout.1, width, height);
 	data.reserve(byte_len);
 
 	data.extend_from_slice(input_data.as_slice());