Implement material modification in model unpacking

Cargo.lock

@@ -8,6 +8,15 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
+[[package]]
+name = "aho-corasick"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2969dcb958b36655471fc61f7e416fa76033bdd4bfed0678d8fee1e2d07a1f0"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "anstream"
 version = "0.6.5"
@@ -341,10 +350,18 @@ dependencies = [
  "clap",
  "flate2",
  "gltf",
+ "pathdiff",
  "rayon",
+ "regex",
  "tar",
 ]
 
+[[package]]
+name = "memchr"
+version = "2.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167"
+
 [[package]]
 name = "memoffset"
 version = "0.9.0"
@@ -394,6 +411,12 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "pathdiff"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8835116a5c179084a830efb3adc117ab007512b535bc1a21c991d3b32a6b44dd"
+
 [[package]]
 name = "png"
 version = "0.17.10"
@@ -454,6 +477,35 @@ dependencies = [
  "bitflags 1.3.2",
 ]
 
+[[package]]
+name = "regex"
+version = "1.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "380b951a9c5e80ddfd6136919eef32310721aa4aacd4889a8d39124b026ab343"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5f804c7828047e88b2d32e2d7fe5a105da8ee3264f01902f796c8e067dc2483f"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f"
+
 [[package]]
 name = "rustix"
 version = "0.38.28"

src/unpacker.rs

@@ -1,6 +1,9 @@
 use crate::asset::{Asset, AssetType};
 use flate2::read::GzDecoder;
+use gltf::{json, Document};
 use rayon::prelude::*;
+use std::borrow::Cow;
+
 use std::fs::File;
 use std::path::{Path, PathBuf};
 use std::process::Command;
@@ -15,6 +18,10 @@ pub struct Unpacker {
     pub assets: Vec<Asset>,
 }
 
+fn get_relative_path(path1: &Path, path2: &Path) -> Option<PathBuf> {
+    pathdiff::diff_paths(path1, path2)
+}
+
 impl Unpacker {
     pub fn prepare_environment(&self) {
         let archive_path = Path::new(&self.args.input);
@@ -88,11 +95,10 @@ impl Unpacker {
             prefabs.len(),
             materials.len()
         );
-        let mut counter = 0;
-        let mut result_models : Vec<Asset> = vec![];
-        for prefab in prefabs.iter() {
+
+        prefabs.par_iter().for_each(|prefab|{
             let path = Path::new(&prefab.path);
-            let prefab_content = fs::read_to_string(&path).unwrap();
+            let prefab_content = fs::read_to_string(path).unwrap();
             let matching_materials: Vec<Asset> = materials
                 .clone()
                 .into_iter()
@@ -104,27 +110,69 @@ impl Unpacker {
                 .filter(|a| prefab_content.contains(&a.guid))
                 .collect();
             if matching_materials.len() != 1 || 1 != matching_models.len() {
-                continue;
+                return;
             }
             let material = matching_materials.first().unwrap();
-            let model = matching_models.first().unwrap();
-            if result_models.iter().any(|a| model.guid.eq(&a.guid)) {
-                continue;
-            }
+            let model: &Asset = matching_models.first().unwrap();
             let texture_guid: Option<String> = material.try_get_mat_texture_guid();
-            let texture_asset: &Asset;
-            match &texture_guid {
-                Some(guid) => {
-                    texture_asset = self.assets.iter().find(|a| guid.eq(&a.guid)).unwrap()
-                }
-                None => continue,
-            }
-            // here we should read gltf file and replace material texture with Uri based on texture_asset
             
-            result_models.push(model.clone());
-            counter += 1;
+            let texture_asset: &Asset = match &texture_guid {
+                Some(guid) => {
+                    self.assets.iter().find(|a| guid.eq(&a.guid)).unwrap()
                 }
-        println!("Updated {} models", counter);
+                None => return,
+            };
+            // here we should read gltf file and replace material texture with Uri based on texture_asset
+            let model_path = Path::new(&model.path).with_extension("glb");
+            Self::modify_material(&model_path, Path::new(&texture_asset.path));
+        });
+    }
+
+    fn align_to_multiple_of_four(n: &mut usize) {
+        *n = (*n + 3) & !3;
+    }
+
+    fn modify_material(gltf_path: &Path, texture_asset: &Path) {
+        let file = fs::File::open(gltf_path).unwrap();
+        let reader = io::BufReader::new(file);
+        let mut gltf = gltf::Gltf::from_reader(reader).unwrap();
+        let mut json = gltf.document.into_json();
+        if let Some(rel_path) = get_relative_path(texture_asset, gltf_path) {
+            for image in json.images.iter_mut() {
+                let result = rel_path.file_name().unwrap().to_str().unwrap().to_string();
+                let required_file = gltf_path.with_file_name(&result);
+                if !required_file.exists() {
+                    fs::copy(texture_asset, gltf_path.with_file_name(&result)).unwrap();
+                }
+                println!(
+                    "Image{:?}: {:?} to be replaced with: {}",
+                    image.name, image.uri, &result
+                );
+                image.uri = Some(result);
+            }
+        }
+
+        gltf.document = Document::from_json(json.clone()).unwrap();
+        // Save the modified glTF
+        let json_string = json::serialize::to_string(&json).expect("Serialization error");
+        let mut json_offset = json_string.len();
+        Self::align_to_multiple_of_four(&mut json_offset);
+        let blob = gltf.blob.clone().unwrap_or_default();
+        let buffer_length = blob.len();
+        let glb = gltf::binary::Glb {
+            header: gltf::binary::Header {
+                magic: *b"glTF",
+                version: 2,
+                // N.B., the size of binary glTF file is limited to range of `u32`.
+                length: (json_offset + buffer_length)
+                    .try_into()
+                    .expect("file size exceeds binary glTF limit"),
+            },
+            bin: Some(Cow::Owned(gltf.blob.unwrap_or_default())),
+            json: Cow::Owned(json_string.into_bytes()),
+        };
+        let writer = std::fs::File::create(gltf_path).expect("I/O error");
+        glb.to_writer(writer).expect("glTF binary output error");
     }
 
     pub fn process_data(&self) {
@@ -156,11 +204,11 @@ impl Unpacker {
             if self.args.fbx_to_gltf.is_some() && &asset.asset_type == &AssetType::FbxModel {
                 process_fbx_file(
                     &source_asset,
-                    &path,
+                    path,
                     &self.args.fbx_to_gltf.clone().unwrap(),
                 );
             } else {
-                process_non_fbx_file(&source_asset, &path);
+                process_non_fbx_file(&source_asset, path);
             }
         });