Add multiple jobs support for import task

Cargo.toml

@@ -33,7 +33,7 @@ markup5ever_rcdom = { version = "0.35", optional = true }
 memchr = { version = "2.7", optional = true }
 mozjpeg = { version = "0.10", optional = true }
 msg_tool_macro = { version = "0.2.9" }
-num_cpus = { version = "1.17", optional = true }
+num_cpus = "1.17"
 overf = "0.1"
 parse-size = { version = "1.1", optional = true }
 pelite = { version = "0.10", optional = true }
@@ -67,7 +67,7 @@ artemis-arc = ["artemis", "msg_tool_macro/artemis-arc", "sha1"]
 bgi = ["fancy-regex"]
 bgi-arc = ["bgi", "rand", "utils-bit-stream"]
 bgi-audio = ["bgi"]
-bgi-img = ["bgi", "image", "rand", "utils-threadpool", "utils-bit-stream"]
+bgi-img = ["bgi", "image", "rand", "utils-bit-stream"]
 cat-system = ["fancy-regex", "flate2", "int-enum"]
 cat-system-arc = ["cat-system", "pelite", "utils-blowfish", "utils-crc32"]
 cat-system-img = ["cat-system", "flate2", "image", "mozjpeg", "utils-bit-stream"]
@@ -86,7 +86,7 @@ hexen-haus = ["memchr", "utils-str"]
 hexen-haus-arc = ["hexen-haus"]
 hexen-haus-img = ["hexen-haus", "image"]
 kirikiri = ["emote-psb", "fancy-regex", "flate2", "json", "lz4", "utils-escape"]
-kirikiri-arc = ["kirikiri", "adler", "fastcdc", "flate2", "parse-size", "sha2", "xp3", "zstd", "utils-threadpool"]
+kirikiri-arc = ["kirikiri", "adler", "fastcdc", "flate2", "parse-size", "sha2", "xp3", "zstd"]
 kirikiri-img = ["kirikiri", "image", "libtlg-rs"]
 silky = []
 softpal = ["int-enum"]
@@ -97,9 +97,9 @@ will-plus-img = ["will-plus", "image"]
 yaneurao = []
 yaneurao-itufuru = ["yaneurao"]
 # basic feature
-image = ["png", "utils-threadpool"]
+image = ["png"]
 image-jpg = ["mozjpeg"]
-image-jxl = ["image", "jpegxl-sys", "utils-threadpool"]
+image-jxl = ["image", "jpegxl-sys"]
 image-webp = ["webp"]
 lossless-audio = ["utils-pcm"]
 audio-flac = ["libflac-sys", "utils-pcm"]
@@ -112,7 +112,6 @@ utils-crc32 = []
 utils-escape = ["fancy-regex"]
 utils-pcm = []
 utils-str = []
-utils-threadpool = ["num_cpus"]
 
 [target.'cfg(windows)'.dependencies]
 windows-sys = { version = "0.61", features = ["Win32_Globalization", "Win32_System_Diagnostics_Debug"] }

src/args.rs

@@ -72,7 +72,7 @@ fn parse_jxl_distance(s: &str) -> Result<f32, String> {
 }
 
 /// Tools for export and import scripts
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 #[clap(
     group = ArgGroup::new("encodingg").multiple(false),
     group = ArgGroup::new("output_encodingg").multiple(false),
@@ -530,7 +530,7 @@ pub struct Arg {
     pub command: Command,
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 #[clap(group = ArgGroup::new("patched_encodingg").multiple(false), group = ArgGroup::new("patched_archive_encodingg").multiple(false))]
 pub struct ImportArgs {
     /// Input script file or directory
@@ -593,9 +593,12 @@ pub struct ImportArgs {
     pub replacement_json: Option<String>,
     #[arg(long, action = ArgAction::SetTrue)]
     pub warn_when_output_file_not_found: bool,
+    #[arg(short = 'j', long, default_value_t = 1)]
+    /// Workers count for import scripts in parallel.
+    pub jobs: usize,
 }
 
-#[derive(Subcommand, Debug)]
+#[derive(Subcommand, Debug, Clone)]
 /// Commands
 pub enum Command {
     /// Extract from script

src/main.rs

@@ -2603,6 +2603,7 @@ fn main() {
         std::process::exit(1);
     });
     let arg = args::parse_args();
+    let argn = std::sync::Arc::new(arg.clone());
     if arg.backtrace {
         unsafe { std::env::set_var("RUST_LIB_BACKTRACE", "1") };
     }
@@ -2870,7 +2871,7 @@ fn main() {
                 }
                 None => None,
             };
-            let repl = match &args.replacement_json {
+            let repl = std::sync::Arc::new(match &args.replacement_json {
                 Some(replacement_json) => {
                     let b = utils::files::read_file(replacement_json).unwrap();
                     let s = String::from_utf8(b).unwrap();
@@ -2878,7 +2879,7 @@ fn main() {
                     Some(table)
                 }
                 None => None,
-            };
+            });
             let (scripts, is_dir) =
                 utils::files::collect_files(&args.input, arg.recursive, false).unwrap();
             if is_dir {
@@ -2897,25 +2898,76 @@ fn main() {
             } else {
                 None
             };
+            let workers = if args.jobs > 1 {
+                Some(
+                    utils::threadpool::ThreadPool::<()>::new(
+                        args.jobs,
+                        Some("import-worker-"),
+                        true,
+                    )
+                    .unwrap(),
+                )
+            } else {
+                None
+            };
             for script in scripts.iter() {
-                let re = import_script(
+                if let Some(workers) = workers.as_ref() {
-                    &script,
+                    let arg = argn.clone();
-                    &arg,
+                    let cfg = cfg.clone();
-                    cfg.clone(),
+                    let script = script.clone();
-                    args,
+                    let name_csv = name_csv.as_ref().map(|s| s.clone());
-                    root_dir,
+                    let repl = repl.clone();
-                    name_csv.as_ref(),
+                    let root_dir = root_dir.map(|s| s.to_path_buf());
-                    repl.as_ref(),
+                    let args = args.clone();
-                );
+                    if let Err(e) = workers.execute(
-                match re {
+                        move |_| {
-                    Ok(s) => {
+                            let re = import_script(
-                        COUNTER.inc(s);
+                                &script,
-                    }
+                                &arg,
-                    Err(e) => {
+                                cfg,
+                                &args,
+                                root_dir.as_ref().map(|s| s.as_path()),
+                                name_csv.as_ref(),
+                                (*repl).as_ref(),
+                            );
+                            match re {
+                                Ok(s) => {
+                                    COUNTER.inc(s);
+                                }
+                                Err(e) => {
+                                    COUNTER.inc_error();
+                                    eprintln!("Error exporting {}: {}", script, e);
+                                    if arg.backtrace {
+                                        eprintln!("Backtrace: {}", e.backtrace());
+                                    }
+                                }
+                            }
+                        },
+                        true,
+                    ) {
                         COUNTER.inc_error();
-                        eprintln!("Error exporting {}: {}", script, e);
+                        eprintln!("Error executing import worker: {}", e);
-                        if arg.backtrace {
+                    }
-                            eprintln!("Backtrace: {}", e.backtrace());
+                } else {
+                    let re = import_script(
+                        &script,
+                        &arg,
+                        cfg.clone(),
+                        args,
+                        root_dir,
+                        name_csv.as_ref(),
+                        (*repl).as_ref(),
+                    );
+                    match re {
+                        Ok(s) => {
+                            COUNTER.inc(s);
+                        }
+                        Err(e) => {
+                            COUNTER.inc_error();
+                            eprintln!("Error exporting {}: {}", script, e);
+                            if arg.backtrace {
+                                eprintln!("Backtrace: {}", e.backtrace());
+                            }
                         }
                     }
                 }

src/types.rs

@@ -1063,7 +1063,6 @@ impl AsRef<str> for LosslessAudioFormat {
     }
 }
 
-#[cfg(feature = "utils-threadpool")]
 #[allow(unused)]
 pub(crate) fn get_default_threads() -> usize {
     num_cpus::get().max(2) / 2

src/utils/mod.rs

@@ -28,7 +28,6 @@ pub mod pcm;
 #[cfg(feature = "utils-str")]
 pub mod str;
 pub mod struct_pack;
-#[cfg(feature = "utils-threadpool")]
 pub mod threadpool;
 
 #[cfg(windows)]