✨ Parallel asset extraction (#46)

* ✨ Parallel asset extraction

* 💸 Add sponsorship information

.github/FUNDING.yaml

@@ -0,0 +1,4 @@
+github: nicks96432
+patreon: nicks96432
+custom:
+  - "https://www.paypal.me/nicks96432"

cli/Cargo.toml

@@ -33,7 +33,7 @@ path = "../manifest"
 
 [features]
 default = ["download", "extract", "manifest"]
-download = ["mltd-asset-download"]
-debug = ["download", "mltd-asset-download/debug"]
+debug = ["mltd-asset-download/debug", "mltd-asset-extract/debug"]
+download = ["dep:mltd-asset-download"]
 extract = ["dep:mltd-asset-extract"]
 manifest = ["dep:mltd-asset-manifest"]

download/src/lib.rs

@@ -11,7 +11,10 @@ use rayon::prelude::{ParallelBridge, ParallelIterator};
 use rayon::ThreadPoolBuilder;
 use ureq::AgentBuilder;
 
-use std::fs::{create_dir_all, File};
+#[cfg(not(feature = "debug"))]
+use std::fs::File;
+
+use std::fs::create_dir_all;
 use std::io::copy;
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicU64, Ordering};
@@ -123,7 +126,7 @@ pub fn download_assets(args: &DownloaderArgs) -> Result<(), DownloadError> {
         let asset_path = output_path.join(filename);
 
         #[cfg(feature = "debug")]
-        let asset_file = Vec::new();
+        let mut asset_file = Vec::new();
 
         #[cfg(not(feature = "debug"))]
         let mut asset_file = match File::create(&asset_path) {

extract/Cargo.toml

@@ -14,6 +14,7 @@ acb = { path = "../acb" }
 # For reading numbers from binary data
 byteorder = "1.5.0"
 
+indicatif = { workspace = true }
 log = { workspace = true }
 num-derive = "0.4.1"
 num-traits = "0.2.17"
@@ -22,6 +23,8 @@ num_cpus = { workspace = true }
 # For unpacking unity assets
 rabex = "0.0.3"
 
+rayon = { workspace = true }
+
 # For decoding Texture2D
 texture2ddecoder = "0.0.5"
 
@@ -48,3 +51,6 @@ env_logger = { workspace = true }
 [dev-dependencies.mltd-utils]
 features = ["log", "rand"]
 path = "../utils"
+
+[features]
+debug = []

extract/src/class/text_asset.rs

@@ -78,5 +78,10 @@ where
         return Ok(());
     }
 
-    ffmpeg(&wav.data, num_cpus::get() / args.parallel, &args.audio_args, output_path)
+    ffmpeg(
+        &wav.data,
+        num_cpus::get() / args.parallel as usize,
+        &args.audio_args,
+        output_path,
+    )
 }

extract/src/class/texture_2d.rs

@@ -804,7 +804,7 @@ where
         let mut buf = Vec::new();
         img.write_with_encoder(PngEncoder::new(Cursor::new(&mut buf)))?;
 
-        ffmpeg(&buf, num_cpus::get() / args.parallel, &args.image_args, output_path)?;
+        ffmpeg(&buf, num_cpus::get() / args.parallel as usize, &args.image_args, output_path)?;
     }
     Ok(())
 }

extract/src/lib.rs

@@ -4,19 +4,26 @@ mod utils;
 mod version;
 
 use std::error::Error;
-use std::fs::{create_dir_all, read_dir, File};
+use std::fs::{read_dir, File};
 use std::io::{Cursor, Seek, SeekFrom};
 use std::path::{Path, PathBuf};
 use std::process::exit;
+use std::sync::atomic::{AtomicU64, Ordering};
 
-use clap::Args;
+#[cfg(not(feature = "debug"))]
+use std::fs::{create_dir_all, write};
+
+use clap::{value_parser, Args};
 use environment::Environment;
+use indicatif::{ProgressBar, ProgressStyle};
 use rabex::config::ExtractionConfig;
 use rabex::files::{BundleFile, SerializedFile};
 use rabex::objects::map;
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use rayon::ThreadPoolBuilder;
 
 use crate::class::asset_bundle::construct_asset_bundle;
-use crate::class::text_asset::extract_acb;
+use crate::class::text_asset::{construct_text_asset, extract_acb};
 use crate::class::texture_2d::extract_texture_2d;
 use crate::environment::{check_file_type, FileType};
 
@@ -54,12 +61,14 @@ pub struct ExtractorArgs {
 
     /// The number of threads to use
     #[arg(short = 'P', long, value_name = "CPUS", display_order = 2)]
-    #[arg(default_value_t = num_cpus::get())]
-    parallel: usize,
+    #[arg(value_parser = value_parser!(u32).range(1..=(num_cpus::get() as i64)))]
+    #[arg(default_value_t = num_cpus::get() as u32)]
+    parallel: u32,
     // TODO: Add option to extract only specific files
 }
 
 pub fn extract_media(args: &ExtractorArgs) -> Result<(), Box<dyn Error>> {
+    #[cfg(not(feature = "debug"))]
     create_dir_all(&args.output)?;
 
     let input_realpath = args.input.canonicalize()?;
@@ -73,9 +82,46 @@ pub fn extract_media(args: &ExtractorArgs) -> Result<(), Box<dyn Error>> {
         exit(1);
     }
 
-    for entry in read_dir(&args.input)? {
-        extract_file(&entry?.path(), args)?;
+    log::debug!("setting progress bar");
+
+    let template = "{msg:60} {eta:4} [{wide_bar:.cyan/blue}] {percent:3}%";
+    let progress_bar_style = match ProgressStyle::with_template(template) {
+        Ok(style) => style,
+        Err(_) => {
+            log::debug!("invalid progress bar template, using default style");
+
+            ProgressStyle::default_bar()
+        }
     }
+    .progress_chars("##-");
+
+    let entries: Vec<_> = read_dir(&args.input)?.collect();
+    let progress_bar = ProgressBar::new(entries.len() as u64).with_style(progress_bar_style);
+
+    let finished_count = AtomicU64::new(0);
+
+    log::debug!("setting thread pool");
+
+    let thread_pool_builder = ThreadPoolBuilder::new().num_threads(args.parallel as usize);
+    thread_pool_builder.build_global()?;
+
+    entries.par_iter().for_each(|entry| {
+        let entry = match entry {
+            Ok(e) => e,
+            Err(e) => {
+                progress_bar.suspend(|| log::warn!("failed to read directory entry: {}", e));
+                return;
+            }
+        };
+
+        if let Err(e) = extract_file(&entry.path(), args) {
+            log::warn!("failed to extract file: {}", e);
+        };
+
+        let cur_finished_count = finished_count.fetch_add(1, Ordering::AcqRel);
+        progress_bar.inc(1);
+        progress_bar.set_message(format!("{}/{}", cur_finished_count, entries.len()));
+    });
 
     Ok(())
 }
@@ -151,6 +197,7 @@ fn extract_object(
 
     let output_dir = args.output.join(asset_path);
 
+    #[cfg(not(feature = "debug"))]
     create_dir_all(&output_dir)?;
 
     for object_info in serialized_file.m_Objects.iter() {
@@ -166,7 +213,19 @@ fn extract_object(
         let data = env.get_object(object_info.m_PathID).unwrap();
 
         match object_info.m_ClassID {
-            map::TextAsset => extract_acb(data, &output_dir, args, serialized_file)?,
+            map::TextAsset => {
+                let text_asset = construct_text_asset(data, serialized_file)?;
+                match text_asset.m_Name.contains("acb") {
+                    true => extract_acb(data, &output_dir, args, serialized_file)?,
+                    false => {
+                        #[cfg(not(feature = "debug"))]
+                        write(
+                            output_dir.join(text_asset.m_Name).with_extension(&args.audio_ext),
+                            text_asset.m_Script.as_bytes(),
+                        )?;
+                    }
+                }
+            }
             map::Texture2D => extract_texture_2d(data, &output_dir, args, serialized_file, env)?,
             map::AssetBundle => {
                 // this class contains some information about the bundle

extract/src/utils/mod.rs

@@ -10,6 +10,7 @@ use std::process::Stdio;
 pub use self::puzzle::*;
 pub use self::read_ext::*;
 
+#[cfg(not(feature = "debug"))]
 pub fn ffmpeg<P>(
     input: &[u8],
     threads: usize,
@@ -21,7 +22,7 @@ where
 {
     #[rustfmt::skip]
     let ffmpeg_args = [
-        "-hide_banner",
+        "-loglevel quiet",
         "-threads", &threads.to_string(),
         "-i", "-",
         "-y",
@@ -45,3 +46,41 @@ where
 
     Ok(())
 }
+
+#[cfg(feature = "debug")]
+pub fn ffmpeg<P>(
+    input: &[u8],
+    threads: usize,
+    args: &str,
+    _output_path: P,
+) -> Result<(), Box<dyn Error>>
+where
+    P: AsRef<Path>,
+{
+    #[rustfmt::skip]
+    let ffmpeg_args = [
+        "-loglevel quiet",
+        "-threads", &threads.to_string(),
+        "-i", "-",
+        "-y",
+        "-f", "data",
+        "-",
+    ];
+
+    let ffmpeg_args = ffmpeg_args.into_iter().chain(args.split_ascii_whitespace());
+
+    let mut ffmpeg = Command::new("ffmpeg")
+        .args(ffmpeg_args)
+        .stdin(Stdio::piped())
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()?;
+
+    let mut stdin = ffmpeg.stdin.take().unwrap();
+    stdin.write_all(input)?;
+    drop(stdin);
+
+    ffmpeg.wait()?;
+
+    Ok(())
+}