From 889642fd75120bd5d14528af86ca9792a04756a2 Mon Sep 17 00:00:00 2001
From: LegionMammal978 <mattlloydhouse@gmail.com>
Date: Sat, 3 Aug 2024 17:15:30 -0400
Subject: [PATCH] Add option for Zopfli iteration count

---
 benches/zopfli.rs            |  4 ++--
 src/cli.rs                   |  9 +++++++--
 src/deflate/mod.rs           |  4 ++--
 src/deflate/zopfli_oxipng.rs |  6 +++---
 src/main.rs                  | 10 ++++++----
 tests/flags.rs               |  4 ++--
 6 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/benches/zopfli.rs b/benches/zopfli.rs
index 912782fe..8c5b1818 100644
--- a/benches/zopfli.rs
+++ b/benches/zopfli.rs
@@ -3,13 +3,13 @@
 extern crate oxipng;
 extern crate test;
 
-use std::{num::NonZeroU8, path::PathBuf};
+use std::{num::NonZeroU64, path::PathBuf};
 
 use oxipng::{internal_tests::*, *};
 use test::Bencher;
 
 // SAFETY: trivially safe. Stopgap solution until const unwrap is stabilized.
-const DEFAULT_ZOPFLI_ITERATIONS: NonZeroU8 = unsafe { NonZeroU8::new_unchecked(15) };
+const DEFAULT_ZOPFLI_ITERATIONS: NonZeroU64 = unsafe { NonZeroU64::new_unchecked(15) };
 
 #[bench]
 fn zopfli_16_bits_strategy_0(b: &mut Bencher) {
diff --git a/src/cli.rs b/src/cli.rs
index 3c06c766..fe4e37c2 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -332,10 +332,15 @@ be processed successfully.")
                 .help("Use the much slower but stronger Zopfli compressor")
                 .long_help("\
 Use the much slower but stronger Zopfli compressor for main compression trials. \
-Recommended use is with '-o max' and '--fast'.")
+Recommended use is with '-o max' and '--fast'. The default number of iterations is 15. \
+Using fewer iterations may speed up compression for large files.")
                 .short('Z')
                 .long("zopfli")
-                .action(ArgAction::SetTrue),
+                .num_args(0..=1)
+                .require_equals(true)
+                .value_name("iterations")
+                .default_missing_value("15")
+                .value_parser(1..),
         )
         .arg(
             Arg::new("timeout")
diff --git a/src/deflate/mod.rs b/src/deflate/mod.rs
index 18f4812d..74ca09b4 100644
--- a/src/deflate/mod.rs
+++ b/src/deflate/mod.rs
@@ -1,6 +1,6 @@
 mod deflater;
 #[cfg(feature = "zopfli")]
-use std::num::NonZeroU8;
+use std::num::NonZeroU64;
 use std::{fmt, fmt::Display};
 
 pub use deflater::{crc32, deflate, inflate};
@@ -25,7 +25,7 @@ pub enum Deflaters {
         /// The number of compression iterations to do. 15 iterations are fine
         /// for small files, but bigger files will need to be compressed with
         /// less iterations, or else they will be too slow.
-        iterations: NonZeroU8,
+        iterations: NonZeroU64,
     },
 }
 
diff --git a/src/deflate/zopfli_oxipng.rs b/src/deflate/zopfli_oxipng.rs
index cb6e6eb4..9a776429 100644
--- a/src/deflate/zopfli_oxipng.rs
+++ b/src/deflate/zopfli_oxipng.rs
@@ -1,11 +1,11 @@
-use std::num::NonZeroU8;
+use std::num::NonZeroU64;
 
 use crate::{PngError, PngResult};
 
-pub fn deflate(data: &[u8], iterations: NonZeroU8) -> PngResult<Vec<u8>> {
+pub fn deflate(data: &[u8], iterations: NonZeroU64) -> PngResult<Vec<u8>> {
     let mut output = Vec::with_capacity(data.len());
     let options = zopfli::Options {
-        iteration_count: iterations.into(),
+        iteration_count: iterations,
         ..Default::default()
     };
     match zopfli::compress(options, zopfli::Format::Zlib, data, &mut output) {
diff --git a/src/main.rs b/src/main.rs
index 33b793a7..cd8417c6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -17,7 +17,7 @@
 mod rayon;
 
 #[cfg(feature = "zopfli")]
-use std::num::NonZeroU8;
+use std::num::NonZeroU64;
 use std::{ffi::OsString, fs::DirBuilder, io::Write, path::PathBuf, process::exit, time::Duration};
 
 use clap::ArgMatches;
@@ -324,10 +324,12 @@ fn parse_opts_into_struct(
         opts.strip = StripChunks::Safe;
     }
 
-    if matches.get_flag("zopfli") {
+    if let Some(&iterations) = matches.get_one::<i64>("zopfli") {
         #[cfg(feature = "zopfli")]
-        if let Some(iterations) = NonZeroU8::new(15) {
-            opts.deflate = Deflaters::Zopfli { iterations };
+        {
+            opts.deflate = Deflaters::Zopfli {
+                iterations: NonZeroU64::new(iterations as u64).unwrap(),
+            };
         }
     } else if let (Deflaters::Libdeflater { compression }, Some(x)) =
         (&mut opts.deflate, matches.get_one::<i64>("compression"))
diff --git a/tests/flags.rs b/tests/flags.rs
index 887da87c..a4e8ed9a 100644
--- a/tests/flags.rs
+++ b/tests/flags.rs
@@ -1,7 +1,7 @@
 #[cfg(feature = "filetime")]
 use std::cell::RefCell;
 #[cfg(feature = "zopfli")]
-use std::num::NonZeroU8;
+use std::num::NonZeroU64;
 #[cfg(feature = "filetime")]
 use std::ops::Deref;
 use std::{
@@ -681,7 +681,7 @@ fn zopfli_mode() {
     let input = PathBuf::from("tests/files/zopfli_mode.png");
     let (output, mut opts) = get_opts(&input);
     opts.deflate = Deflaters::Zopfli {
-        iterations: NonZeroU8::new(15).unwrap(),
+        iterations: NonZeroU64::new(15).unwrap(),
     };
 
     test_it_converts(