pack-bundle command + write callback

Author: yungcomputerchair

src/bin/ffbuildtool.rs

@@ -27,6 +27,8 @@ enum Commands {
     ReadBundle(ReadBundleArgs),
     #[cfg(feature = "lzma")]
     ExtractBundle(ExtractBundleArgs),
+    #[cfg(feature = "lzma")]
+    PackBundle(PackBundleArgs),
 }
 
 #[derive(Args, Debug)]
@@ -101,8 +103,8 @@ struct ValidateBuildArgs {
 #[derive(Args, Debug)]
 struct ReadBundleArgs {
     /// Path to the compressed asset bundle
-    #[clap(short = 'b', long)]
-    bundle_path: String,
+    #[clap(short = 'i', long)]
+    input_bundle: String,
 
     /// Whether to calculate the hashes of each file in the bundle
     #[clap(short = 'c', long, action)]
@@ -113,14 +115,30 @@ struct ReadBundleArgs {
 #[derive(Args, Debug)]
 struct ExtractBundleArgs {
     /// Path to the compressed asset bundle
-    #[clap(short = 'b', long)]
-    bundle_path: String,
+    #[clap(short = 'i', long)]
+    input_bundle: String,
 
     /// Path to the output directory. If not specified, will be extracted to a directory named after the bundle.
     #[clap(short = 'o', long)]
     output_dir: Option<String>,
 }
 
+#[cfg(feature = "lzma")]
+#[derive(Args, Debug)]
+struct PackBundleArgs {
+    /// Path to the input directory
+    #[clap(short = 'i', long)]
+    input_dir: String,
+
+    /// Path to the output bundle
+    #[clap(short = 'o', long)]
+    output_bundle: String,
+
+    /// Compression level to use
+    #[clap(short = 'l', long, default_value = "6")]
+    compression_level: u32,
+}
+
 #[derive(Debug, PartialEq, Eq)]
 enum ItemState {
     Downloading,
@@ -236,6 +254,8 @@ async fn main() -> Result<(), String> {
         Commands::ReadBundle(args) => read_bundle(args).await,
         #[cfg(feature = "lzma")]
         Commands::ExtractBundle(args) => extract_bundle(args).await,
+        #[cfg(feature = "lzma")]
+        Commands::PackBundle(args) => pack_bundle(args).await,
     }
 }
 
@@ -359,8 +379,10 @@ async fn validate_build(args: ValidateBuildArgs) -> Result<(), String> {
 async fn read_bundle(args: ReadBundleArgs) -> Result<(), String> {
     use std::time::Instant;
 
+    use ffbuildtool::bundle::AssetBundle;
+
     let start = Instant::now();
-    let (header, mut bundle) = ffbuildtool::bundle::AssetBundle::from_file(&args.bundle_path)?;
+    let (header, mut bundle) = AssetBundle::from_file(&args.input_bundle)?;
     println!("Bundle read in {}ms", start.elapsed().as_millis());
 
     if args.calculate_hashes {
@@ -380,26 +402,71 @@ async fn read_bundle(args: ReadBundleArgs) -> Result<(), String> {
 async fn extract_bundle(args: ExtractBundleArgs) -> Result<(), String> {
     use std::{path::PathBuf, time::Instant};
 
+    use ffbuildtool::{bundle::AssetBundle, util};
+
     let start = Instant::now();
-    let (_, bundle) = ffbuildtool::bundle::AssetBundle::from_file(&args.bundle_path)?;
+    let (header, bundle) = AssetBundle::from_file(&args.input_bundle)?;
     println!("Bundle read in {}ms", start.elapsed().as_millis());
+    println!(
+        "------------------------\n{}\n------------------------\n{}",
+        header, bundle
+    );
 
     let output_dir = args.output_dir.unwrap_or({
-        let bundle_name = ffbuildtool::util::get_file_name_without_parent(&args.bundle_path);
-        let bundle_name_url_encoded = ffbuildtool::util::url_encode(bundle_name);
-        let bundle_path = PathBuf::from(&args.bundle_path);
+        let bundle_name = util::get_file_name_without_parent(&args.input_bundle);
+        let bundle_name_url_encoded = util::url_encode(bundle_name);
+        let bundle_path = PathBuf::from(&args.input_bundle);
         bundle_path
             .parent()
             .unwrap_or(&bundle_path)
             .join(bundle_name_url_encoded)
             .to_string_lossy()
             .to_string()
     });
-    println!("Extracting bundle {} to {}", args.bundle_path, output_dir);
+    println!("Extracting bundle {} to {}", args.input_bundle, output_dir);
 
     let start = Instant::now();
     bundle.extract_files(&output_dir)?;
     println!("Bundle extracted in {}ms", start.elapsed().as_millis());
 
     Ok(())
 }
+
+#[cfg(feature = "lzma")]
+async fn pack_bundle(args: PackBundleArgs) -> Result<(), String> {
+    use std::{sync::LazyLock, time::Instant};
+
+    use ffbuildtool::bundle::AssetBundle;
+
+    fn cb(level_idx: usize, file: usize, total_files: usize, current_file_name: String) {
+        static PBS: OnceLock<Mutex<HashMap<usize, ProgressBar>>> = OnceLock::new();
+        static PB_TEMPLATE: LazyLock<ProgressStyle> = LazyLock::new(|| {
+            ProgressStyle::default_bar()
+                .template("{bar:40} {pos} / {len} {msg}")
+                .unwrap()
+        });
+
+        let mut pbs = PBS
+            .get_or_init(|| Mutex::new(HashMap::new()))
+            .lock()
+            .unwrap();
+        let pb = pbs.entry(level_idx).or_insert_with(|| {
+            let pb = ProgressBar::new(total_files as u64);
+            pb.set_style(PB_TEMPLATE.clone());
+            pb
+        });
+
+        pb.set_position(file as u64);
+        pb.set_message(current_file_name);
+    }
+
+    let start = Instant::now();
+    let bundle = AssetBundle::from_directory(&args.input_dir)?;
+    println!("Files read in {}ms", start.elapsed().as_millis());
+
+    let start = Instant::now();
+    bundle.to_file(&args.output_bundle, args.compression_level, Some(cb))?;
+    println!("Bundle created in {}ms", start.elapsed().as_millis());
+
+    Ok(())
+}

src/bundle.rs

@@ -11,6 +11,9 @@ use lzma::{LzmaReader, LzmaWriter};
 
 use crate::{util, Error, FileInfo};
 
+// level index, file index, total files, file name
+pub type CompressionCallback = fn(usize, usize, usize, String);
+
 fn read_u32<T: Read>(reader: &mut T) -> Result<u32, Error> {
     let mut buf = [0; 4];
     reader.read_exact(&mut buf)?;
@@ -92,7 +95,12 @@ impl std::fmt::Display for AssetBundleHeader {
                 i, level.compressed_end, level.uncompressed_end
             )?;
         }
-        write!(f, "Bundle size: {}", self.bundle_size)
+        write!(
+            f,
+            "Bundle size: {} ({} bytes)",
+            util::bytes_to_human_readable(self.bundle_size),
+            self.bundle_size
+        )
     }
 }
 impl AssetBundleHeader {
@@ -292,15 +300,15 @@ impl std::fmt::Display for LevelFile {
                 f,
                 "{} - {} ({} bytes) - {}",
                 self.name,
-                util::bytes_to_human_readable(self.data.len() as u64),
+                util::bytes_to_human_readable(self.data.len() as u32),
                 self.data.len(),
                 hash
             ),
             None => write!(
                 f,
                 "{} - {} ({} bytes)",
                 self.name,
-                util::bytes_to_human_readable(self.data.len() as u64),
+                util::bytes_to_human_readable(self.data.len() as u32),
                 self.data.len()
             ),
         }
@@ -336,15 +344,29 @@ impl Level {
         Ok(Self { files })
     }
 
-    fn write<W: Write>(&self, writer: &mut W, compression: u32) -> Result<usize, Error> {
+    fn write<W: Write>(
+        &self,
+        writer: &mut W,
+        compression: u32,
+        level_idx: usize,
+        callback: Option<CompressionCallback>,
+    ) -> Result<usize, Error> {
         let mut writer = Counter::new(LzmaWriter::new_compressor(writer, compression)?);
         let header = self.gen_header();
         header.write(&mut writer)?;
 
+        let num_files = header.files.len();
         for (idx, file) in header.files.iter().enumerate() {
             let padding_size = file.offset as usize - writer.writer_bytes();
             writer.write_all(&vec![0; padding_size])?;
             writer.write_all(&self.files[idx].data)?;
+            if let Some(callback) = callback {
+                callback(level_idx, idx, num_files, file.name.clone());
+            }
+        }
+
+        if let Some(callback) = callback {
+            callback(level_idx, num_files, num_files, "Done".to_string());
         }
 
         // pad to 4 bytes
@@ -442,14 +464,20 @@ impl AssetBundle {
         Ok((header, Self { levels }))
     }
 
-    fn write<W: Write>(&self, writer: &mut W, compression: u32) -> Result<(), Error> {
+    fn write<W: Write>(
+        &self,
+        writer: &mut W,
+        compression: u32,
+        callback: Option<CompressionCallback>,
+    ) -> Result<(), Error> {
         let mut buf = Vec::new();
         let mut buf_writer = Counter::new(&mut buf);
         let mut uncompressed_bytes_written = 0;
 
         let mut level_ends = Vec::new();
-        for level in &self.levels {
-            uncompressed_bytes_written += level.write(&mut buf_writer, compression)?;
+        for (idx, level) in self.levels.iter().enumerate() {
+            uncompressed_bytes_written +=
+                level.write(&mut buf_writer, compression, idx, callback)?;
             let uncompressed_end = uncompressed_bytes_written as u32;
             let compressed_end = buf_writer.writer_bytes() as u32;
             level_ends.push(LevelEnds {
@@ -492,11 +520,57 @@ impl AssetBundle {
             .map_err(|e| format!("Couldn't read bundle: {}", e))
     }
 
-    pub fn to_file(&self, path: &str, compression_level: u32) -> Result<(), String> {
+    pub fn from_directory(path: &str) -> Result<Self, String> {
+        // each subdirectory with the name `levelX` contains the files for that level.
+        // they must be in order-- starting from level0-- for their files to be included.
+        // all loose files get put at the end of level0.
+        let root_path = PathBuf::from(path);
+        if !root_path.is_dir() {
+            return Err(format!("Invalid root directory: {}", path));
+        }
+
+        let mut levels = Vec::new();
+        for i in 0.. {
+            let level_dir = root_path.join(format!("level{}", i));
+            if !level_dir.as_path().is_dir() {
+                break;
+            }
+
+            let Ok(files) = Self::get_level_files_from_dir(&level_dir) else {
+                return Err(format!(
+                    "Couldn't read files in dir: {}",
+                    level_dir.display()
+                ));
+            };
+
+            levels.push(Level { files });
+        }
+
+        let Ok(loose_files) = Self::get_level_files_from_dir(&root_path) else {
+            return Err(format!(
+                "Couldn't read files in dir: {}",
+                root_path.display()
+            ));
+        };
+        if levels.is_empty() {
+            levels.push(Level { files: loose_files });
+        } else {
+            levels[0].files.extend(loose_files);
+        }
+
+        Ok(Self { levels })
+    }
+
+    pub fn to_file(
+        &self,
+        path: &str,
+        compression_level: u32,
+        callback: Option<CompressionCallback>,
+    ) -> Result<(), String> {
         let file =
             File::create(path).map_err(|e| format!("Couldn't create file {}: {}", path, e))?;
         let mut writer = BufWriter::new(file);
-        self.write(&mut writer, compression_level)
+        self.write(&mut writer, compression_level, callback)
             .map_err(|e| format!("Couldn't write bundle: {}", e))?;
         writer
             .flush()
@@ -553,4 +627,11 @@ impl AssetBundle {
 
         Ok(result)
     }
+
+    pub fn get_num_files(&self, level: usize) -> Result<usize, Error> {
+        if level >= self.levels.len() {
+            return Err(format!("Level {} does not exist", level).into());
+        }
+        Ok(self.levels[level].files.len())
+    }
 }

src/util.rs

@@ -266,7 +266,7 @@ pub fn file_path_to_uri(file_path: &str) -> String {
     format!("file:///{}", path)
 }
 
-pub fn bytes_to_human_readable(bytes: u64) -> String {
+pub fn bytes_to_human_readable(bytes: u32) -> String {
     let units = ["B", "KB", "MB", "GB"];
     let mut bytes = bytes as f64;
     let mut unit = 0;