Better Unity material reading and parsing functionality

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## [0.4.1]
+
+### Fixed
+
+- Improved material parsing, big thanks to [bevity](https://github.com/gamedolphin/bevity).
+
 ## [0.4.0]
 
 ### Added

Cargo.lock

@@ -65,6 +65,12 @@ dependencies = [
  "windows-sys",
 ]
 
+[[package]]
+name = "anyhow"
+version = "1.0.78"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca87830a3e3fb156dc96cfbd31cb620265dd053be734723f22b760d6cc3c3051"
+
 [[package]]
 name = "autocfg"
 version = "1.1.0"
@@ -205,6 +211,12 @@ version = "1.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07"
 
+[[package]]
+name = "equivalent"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
+
 [[package]]
 name = "errno"
 version = "0.3.8"
@@ -285,6 +297,12 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "hashbrown"
+version = "0.14.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "290f1a1d9242c78d09ce40a5e87e7554ee637af1351968159f4952f028f75604"
+
 [[package]]
 name = "heck"
 version = "0.4.1"
@@ -306,6 +324,16 @@ dependencies = [
  "png",
 ]
 
+[[package]]
+name = "indexmap"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d530e1a18b1cb4c484e6e34556a0d948706958449fca0cab753d649f2bce3d1f"
+dependencies = [
+ "equivalent",
+ "hashbrown",
+]
+
 [[package]]
 name = "inflections"
 version = "1.1.1"
@@ -346,11 +374,14 @@ checksum = "c4cd1a83af159aa67994778be9070f0ae1bd732942279cabb14f86f986a21456"
 name = "lwa_unity_unpack"
 version = "0.4.0"
 dependencies = [
+ "anyhow",
  "clap",
  "flate2",
  "gltf",
  "rayon",
  "regex",
+ "serde",
+ "serde_yaml",
  "tar",
 ]
 
@@ -539,6 +570,19 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "serde_yaml"
+version = "0.9.29"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a15e0ef66bf939a7c890a0bf6d5a733c70202225f9888a89ed5c62298b019129"
+dependencies = [
+ "indexmap",
+ "itoa",
+ "ryu",
+ "serde",
+ "unsafe-libyaml",
+]
+
 [[package]]
 name = "simd-adler32"
 version = "0.3.7"
@@ -579,6 +623,12 @@ version = "1.0.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
 
+[[package]]
+name = "unsafe-libyaml"
+version = "0.2.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ab4c90930b95a82d00dc9e9ac071b4991924390d46cbd0dfe566148667605e4b"
+
 [[package]]
 name = "urlencoding"
 version = "2.1.3"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lwa_unity_unpack"
-version = "0.4.0"
+version = "0.4.1"
 edition = "2021"
 repository = "https://github.com/Leinnan/lwa_unity_unpack"
 homepage = "https://github.com/Leinnan/lwa_unity_unpack"
@@ -26,3 +26,6 @@ gltf = "1.4.0"
 rayon = "1.8.0"
 regex = "1.10.2"
 tar = "0.4"
+serde = { version = "1.0", features = ["derive"] }
+serde_yaml = "0.9"
+anyhow = "1.0.78"

src/asset.rs

@@ -1,4 +1,4 @@
-use regex::Regex;
+use crate::primitives::materials::read_single_material;
 use std::ffi::OsStr;
 use std::fs;
 use std::fs::DirEntry;
@@ -31,15 +31,15 @@ impl Asset {
             AssetType::Material => {}
             _ => return None,
         }
-        let file = File::open(&self.path).unwrap();
-        let buf_reader = BufReader::new(file);
-        let search = buf_reader.lines().find(|s| {
-            let ss = s.as_ref().unwrap();
-            ss.contains("m_Texture") && ss.contains("guid: ")
-        });
-        if let Some(line) = search {
-            let line = line.unwrap_or_default();
-            return extract_guid(&line);
+        let content = fs::read_to_string(&self.path).unwrap();
+        let material = read_single_material(&content);
+        if let Ok(mat) = material {
+            return mat
+                .properties
+                .tex_envs
+                .iter()
+                .find_map(|tex| tex.get("_MainTex"))
+                .and_then(|t| t.texture.guid.clone());
         }
         None
     }
@@ -123,9 +123,3 @@ impl Asset {
         }
     }
 }
-
-fn extract_guid(text: &str) -> Option<String> {
-    let re = Regex::new(r"guid: (?P<guid>[A-Za-z0-9]{32})").unwrap();
-    re.captures(text)
-        .and_then(|cap| cap.name("guid").map(|guid| guid.as_str().to_string()))
-}

src/main.rs

@@ -1,7 +1,8 @@
 mod args;
 pub mod asset;
+pub mod primitives;
 mod unpacker;
-
+mod yaml_helpers;
 use clap::Parser;
 
 fn main() {

src/primitives/materials.rs

@@ -0,0 +1,101 @@
+use crate::primitives::reference::FileReference;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+// use bevy::prelude::*;
+use crate::yaml_helpers::parse_unity_yaml;
+use anyhow::{bail, Context, Result};
+
+#[derive(Serialize, Deserialize, Default, Debug)]
+pub struct UnityMaterial {
+    #[serde(alias = "m_Name")]
+    pub name: String,
+
+    #[serde(alias = "m_Shader")]
+    pub shader: FileReference,
+
+    #[serde(alias = "m_SavedProperties")]
+    pub properties: SavedProperties,
+
+    #[serde(default, alias = "stringTagMap")]
+    pub string_tags: HashMap<String, String>,
+}
+
+#[derive(Serialize, Deserialize, Default, Debug)]
+pub struct SavedProperties {
+    #[serde(alias = "serializedVersion")]
+    pub serialized_version: u64,
+
+    #[serde(alias = "m_TexEnvs")]
+    pub tex_envs: Vec<HashMap<String, TextureInfo>>,
+
+    #[serde(alias = "m_Floats")]
+    pub floats: Vec<HashMap<String, f32>>,
+
+    #[serde(alias = "m_Colors")]
+    pub colors: Vec<HashMap<String, UnityColor>>,
+}
+
+#[derive(Serialize, Deserialize, Default, Debug)]
+pub struct TextureInfo {
+    #[serde(alias = "m_Texture")]
+    pub texture: FileReference,
+    #[serde(alias = "m_Scale")]
+    pub scale: UnityVector2,
+    #[serde(alias = "m_Offset")]
+    pub offset: UnityVector2,
+}
+
+#[derive(Serialize, Deserialize, Debug, Default, Copy, Clone)]
+pub struct UnityVector2 {
+    pub x: f32,
+    pub y: f32,
+}
+#[derive(Serialize, Deserialize, Debug, Default, Copy, Clone)]
+pub struct UnityColor {
+    pub r: f32,
+    pub g: f32,
+    pub b: f32,
+    pub a: f32,
+}
+
+// impl From<UnityColor> for Color {
+//     fn from(value: UnityColor) -> Self {
+//         Color::Rgba {
+//             red: value.r,
+//             green: value.g,
+//             blue: value.b,
+//             alpha: value.a,
+//         }
+//     }
+// }
+//
+// impl From<&UnityColor> for Color {
+//     fn from(value: &UnityColor) -> Self {
+//         Color::Rgba {
+//             red: value.r,
+//             green: value.g,
+//             blue: value.b,
+//             alpha: value.a,
+//         }
+//     }
+// }
+
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(tag = "object_type")]
+enum MaterialContainer {
+    Material(UnityMaterial),
+    #[serde(other)]
+    DontCare,
+}
+
+pub fn read_single_material(contents: &str) -> Result<UnityMaterial> {
+    let map = parse_unity_yaml(contents)?;
+
+    let (_, output) = map.into_iter().next().context("0 items in material file")?;
+
+    let MaterialContainer::Material(mat) = output else {
+        bail!("invalid material file");
+    };
+
+    Ok(mat)
+}

src/primitives/mod.rs

@@ -0,0 +1,2 @@
+pub mod materials;
+pub mod reference;

src/primitives/reference.rs

@@ -0,0 +1,53 @@
+use serde::{
+    de::{self, Visitor},
+    Deserialize, Deserializer, Serialize,
+};
+
+#[derive(Serialize, Deserialize, Debug, Default, Clone)]
+pub struct FileReference {
+    #[serde(alias = "fileID")]
+    pub file_id: i64,
+    #[serde(default, deserialize_with = "deserialize_option_string_or_float")]
+    pub guid: Option<String>,
+}
+
+fn deserialize_option_string_or_float<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    struct StringOrFloat;
+
+    impl<'de> Visitor<'de> for StringOrFloat {
+        type Value = Option<String>;
+
+        fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
+            formatter.write_str("a string or a float")
+        }
+
+        fn visit_str<E: de::Error>(self, value: &str) -> Result<Self::Value, E> {
+            Ok(Some(value.to_owned()))
+        }
+
+        fn visit_string<E: de::Error>(self, value: String) -> Result<Self::Value, E> {
+            Ok(Some(value))
+        }
+
+        fn visit_f64<E: de::Error>(self, _: f64) -> Result<Self::Value, E> {
+            Ok(None)
+        }
+
+        fn visit_f32<E: de::Error>(self, _: f32) -> Result<Self::Value, E> {
+            Ok(None)
+        }
+
+        fn visit_none<E: de::Error>(self) -> Result<Self::Value, E> {
+            Ok(None)
+        }
+
+        fn visit_unit<E: de::Error>(self) -> Result<Self::Value, E> {
+            Ok(None)
+        }
+    }
+
+    deserializer.deserialize_any(StringOrFloat)
+}

src/yaml_helpers.rs

@@ -0,0 +1,41 @@
+use anyhow::Result;
+use serde::de::DeserializeOwned;
+use std::collections::HashMap;
+
+pub fn parse_unity_yaml<T: DeserializeOwned>(file: &str) -> Result<HashMap<i64, T>> {
+    let file = cleanup_unity_yaml(file)?;
+    let parse: HashMap<i64, T> = serde_yaml::from_str(&file)?;
+
+    Ok(parse)
+}
+
+fn cleanup_unity_yaml(yaml: &str) -> Result<String> {
+    let lines: Vec<String> = yaml
+        .lines()
+        .filter_map(|line| {
+            if line.starts_with("%YAML") || line.starts_with("%TAG") {
+                // unity specific headers. SKIP!
+                None
+            } else if line.starts_with("--- !u!") {
+                // unity object id declared on this line
+                // --- !u!104 &2 => 104 is object type and 2 is object id
+                let mut splits = line.split_whitespace();
+                let object_id: i64 = splits
+                    .find(|&part| part.starts_with('&'))
+                    .and_then(|num| num[1..].parse().ok())?;
+
+                Some(format!("{}:", object_id))
+            } else if line.starts_with(' ') {
+                Some(line.to_string())
+            } else {
+                Some(format!("  object_type: {}", line.replace(':', "")))
+            }
+        })
+        .collect();
+
+    let mut lines = lines.join("\n");
+
+    lines.push('\n'); // insert new line at the end
+
+    Ok(lines)
+}