From 4e50a1c649f68041bda58477ffd8603688f4bd00 Mon Sep 17 00:00:00 2001
From: lifegpc <root@lifegpc.com>
Date: Tue, 29 Jul 2025 22:13:42 +0800
Subject: [PATCH] Add option to set png compression level

---
 src/args.rs      |  3 +++
 src/main.rs      | 14 +++++++++++---
 src/types.rs     | 31 +++++++++++++++++++++++++++++++
 src/utils/img.rs |  8 +++++++-
 4 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/src/args.rs b/src/args.rs
index b6c5b8f..7d60099 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -203,6 +203,9 @@ pub struct Arg {
     #[arg(short = 'z', long, global = true, value_name = "LEVEL", value_parser = parse_compression_level)]
     /// Zlib compression level. Default is 6. 0 means no compression, 9 means best compression.
     pub zlib_compression_level: Option<u32>,
+    #[cfg(feature = "image")]
+    #[arg(short = 'g', long, global = true, value_enum, default_value_t = PngCompressionLevel::Fast)]
+    pub png_compression_level: PngCompressionLevel,
     #[command(subcommand)]
     /// Command
     pub command: Command,
diff --git a/src/main.rs b/src/main.rs
index 2dcb7d2..8d066ca 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -417,6 +417,7 @@ pub fn export_script(
                                 img_data.data,
                                 out_type,
                                 &out_path.to_string_lossy(),
+                                config,
                             )?;
                             COUNTER.inc(types::ScriptResult::Ok);
                         }
@@ -449,7 +450,12 @@ pub fn export_script(
                             continue;
                         }
                     }
-                    match utils::img::encode_img(img_data, out_type, &out_path.to_string_lossy()) {
+                    match utils::img::encode_img(
+                        img_data,
+                        out_type,
+                        &out_path.to_string_lossy(),
+                        config,
+                    ) {
                         Ok(_) => {}
                         Err(e) => {
                             eprintln!("Error encoding image: {}", e);
@@ -689,7 +695,7 @@ pub fn export_script(
                         continue;
                     }
                 }
-                utils::img::encode_img(img_data.data, out_type, &f)?;
+                utils::img::encode_img(img_data.data, out_type, &f, config)?;
                 COUNTER.inc(types::ScriptResult::Ok);
             }
             return Ok(types::ScriptResult::Ok);
@@ -720,7 +726,7 @@ pub fn export_script(
                 }
             }
         };
-        utils::img::encode_img(img_data, out_type, &f)?;
+        utils::img::encode_img(img_data, out_type, &f, config)?;
         return Ok(types::ScriptResult::Ok);
     }
     let mut of = match &arg.output_type {
@@ -1451,6 +1457,8 @@ fn main() {
         cat_system_cstl_lang: arg.cat_system_cstl_lang.clone(),
         #[cfg(feature = "flate2")]
         zlib_compression_level: arg.zlib_compression_level,
+        #[cfg(feature = "image")]
+        png_compression_level: arg.png_compression_level,
     };
     match &arg.command {
         args::Command::Export { input, output } => {
diff --git a/src/types.rs b/src/types.rs
index 23f47e9..3ae8db9 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -237,6 +237,8 @@ pub struct ExtraConfig {
     pub cat_system_cstl_lang: Option<String>,
     #[cfg(feature = "flate2")]
     pub zlib_compression_level: Option<u32>,
+    #[cfg(feature = "image")]
+    pub png_compression_level: PngCompressionLevel,
 }
 
 #[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)]
@@ -484,3 +486,32 @@ impl BomType {
         }
     }
 }
+
+#[cfg(feature = "image")]
+#[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)]
+pub enum PngCompressionLevel {
+    #[value(alias = "d")]
+    /// Default level
+    Default,
+    #[value(alias = "f")]
+    /// Fast minimal compression
+    Fast,
+    #[value(alias = "b")]
+    /// Higher compression level
+    ///
+    /// Best in this context isn't actually the highest possible level
+    /// the encoder can do, but is meant to emulate the `Best` setting in the `Flate2`
+    /// library.
+    Best,
+}
+
+#[cfg(feature = "image")]
+impl PngCompressionLevel {
+    pub fn to_compression(&self) -> png::Compression {
+        match self {
+            PngCompressionLevel::Default => png::Compression::Default,
+            PngCompressionLevel::Fast => png::Compression::Fast,
+            PngCompressionLevel::Best => png::Compression::Best,
+        }
+    }
+}
diff --git a/src/utils/img.rs b/src/utils/img.rs
index cd139b9..dfb5953 100644
--- a/src/utils/img.rs
+++ b/src/utils/img.rs
@@ -109,7 +109,12 @@ pub fn convert_rgba_to_bgra(data: &mut ImageData) -> Result<()> {
     Ok(())
 }
 
-pub fn encode_img(mut data: ImageData, typ: ImageOutputType, filename: &str) -> Result<()> {
+pub fn encode_img(
+    mut data: ImageData,
+    typ: ImageOutputType,
+    filename: &str,
+    config: &ExtraConfig,
+) -> Result<()> {
     match typ {
         ImageOutputType::Png => {
             let mut file = crate::utils::files::write_file(filename)?;
@@ -137,6 +142,7 @@ pub fn encode_img(mut data: ImageData, typ: ImageOutputType, filename: &str) ->
             let mut encoder = png::Encoder::new(&mut file, data.width, data.height);
             encoder.set_color(color_type);
             encoder.set_depth(bit_depth);
+            encoder.set_compression(config.png_compression_level.to_compression());
             let mut writer = encoder.write_header()?;
             writer.write_image_data(&data.data)?;
             writer.finish()?;