assetbundle struct to abstract over unityfs and serializedfile

9 files changed, 95 insertions, 89 deletions

diff --git a/exporter/src/bin/material_stats.rs b/exporter/src/bin/material_stats.rs
index 3eb4dc0..b853442 100644
--- a/exporter/src/bin/material_stats.rs
+++ b/exporter/src/bin/material_stats.rs
@@ -1,48 +1,22 @@
-use anyhow::anyhow;
 use std::{env::args, fs::File, io::BufReader};
-use unity_tools::{classes::material::Material, serialized_file::SerializedFile, unityfs::UnityFS};
+use unity_tools::{assetbundle::AssetBundle, classes::material::Material};
 
 fn main() -> anyhow::Result<()> {
     env_logger::init_from_env("LOG");
     let file = BufReader::new(File::open(args().nth(1).unwrap()).unwrap());
-    let fs = UnityFS::open(file)?;
+    let mut bundle = AssetBundle::open(file)?;
 
     let mode = args().nth(2).unwrap();
 
-    let cabfile = fs
-        .find_main_file()
-        .ok_or(anyhow!("no CAB file found"))?
-        .to_owned();
-
-    let cab = fs.read(&cabfile)?;
-    let mut file = SerializedFile::read(cab)?;
-    let mut sharedassets = file
-        .find_fs_shared_assets(&fs)
-        .map(|n| {
-            let f = fs.read(&n)?;
-            let f = SerializedFile::read(f)?;
-            Ok::<_, anyhow::Error>(f)
-        })
-        .transpose()?;
-
-    for ob in file.objects.clone() {
-        if file.get_object_type_tree(&ob)?.type_string != "Material" {
-            continue;
-        }
-        let mat = file.read_object(ob)?.parse::<Material>()?;
+    for ob in bundle.all_toplevel_of_class("Material").collect::<Vec<_>>() {
+        let mat = ob.load(&mut bundle)?.parse::<Material>()?;
 
         match mode.as_str() {
             "material" => {
                 println!("{}", mat.name)
             }
             "shader" => {
-                println!(
-                    "{}",
-                    mat.shader
-                        .load(&mut file, sharedassets.as_mut())?
-                        .parsed
-                        .name
-                )
+                println!("{}", mat.shader.load(&mut bundle)?.parsed.name)
             }
             x => panic!("unknown mode {x:?}"),
         }
diff --git a/exporter/src/bin/meshes.rs b/exporter/src/bin/meshes.rs
index db8f2f8..4cd86e8 100644
--- a/exporter/src/bin/meshes.rs
+++ b/exporter/src/bin/meshes.rs
@@ -1,53 +1,35 @@
 #![feature(array_chunks)]
-use anyhow::anyhow;
+use glam::Vec3;
 use std::{
     env::args,
     fs::{File, create_dir_all},
     io::{BufReader, BufWriter, Write},
 };
 use unity_tools::{
+    assetbundle::AssetBundle,
     classes::mesh::{Mesh, VertexDataChannel},
-    object::parser::FromValue,
-    serialized_file::SerializedFile,
-    unityfs::UnityFS,
 };
 
 fn main() -> anyhow::Result<()> {
     env_logger::init_from_env("LOG");
     let file = BufReader::new(File::open(args().nth(1).unwrap()).unwrap());
-    let fs = UnityFS::open(file)?;
+    let mut bundle = AssetBundle::open(file)?;
 
     let mut i = 0;
     create_dir_all("/tmp/a").unwrap();
 
-    let cabfile = fs
-        .header
-        .nodes()
-        .iter()
-        .find(|n| !n.name.ends_with(".resource") && !n.name.ends_with(".resS"))
-        .ok_or(anyhow!("no CAB file found"))?
-        .to_owned();
-
-    let mut cab = fs.read(&cabfile)?;
-    let mut file = SerializedFile::read(&mut cab)?;
-    for ob in file.objects.clone() {
-        if file.get_object_type_tree(&ob)?.type_string != "Mesh" {
-            continue;
-        }
-        let value = file.read_object(ob)?;
-        let mesh = Mesh::from_value(value).unwrap();
+    for ob in bundle.all_toplevel_of_class("Mesh").collect::<Vec<_>>() {
+        let mesh = ob.load(&mut bundle)?.parse::<Mesh>()?;
         let mut obj = BufWriter::new(File::create(format!(
             "/tmp/a/{}_{i}.obj",
             mesh.name.replace("/", "-").replace(".", "-")
         ))?);
 
-        let (pos_dims, positions) = mesh
+        let positions = mesh
             .vertex_data
-            .read_channel(VertexDataChannel::Position)
+            .read_channel_vec::<Vec3>(VertexDataChannel::Position)?
             .unwrap();
-        assert_eq!(pos_dims, 3);
-
-        for [x, y, z] in positions.array_chunks() {
+        for Vec3 { x, y, z } in positions {
             writeln!(obj, "v {x} {y} {z}")?;
         }
         for [a, b, c] in mesh.read_indecies() {
diff --git a/exporter/src/bin/textures.rs b/exporter/src/bin/textures.rs
index a3b6cec..56cfc9b 100644
--- a/exporter/src/bin/textures.rs
+++ b/exporter/src/bin/textures.rs
@@ -1,38 +1,26 @@
-use anyhow::anyhow;
 use log::warn;
 use std::{
     env::args,
     fs::{File, create_dir_all},
     io::BufReader,
 };
-use unity_tools::{
-    classes::texture2d::Texture2D, object::parser::FromValue, serialized_file::SerializedFile,
-    unityfs::UnityFS,
-};
+use unity_tools::{assetbundle::AssetBundle, classes::texture2d::Texture2D};
 
 fn main() -> anyhow::Result<()> {
     env_logger::init_from_env("LOG");
     let file = BufReader::new(File::open(args().nth(1).unwrap()).unwrap());
-    let mut fs = UnityFS::open(file)?;
+    let mut bundle = AssetBundle::open(file)?;
 
     let mut i = 0;
     create_dir_all("/tmp/a").unwrap();
 
-    let cabfile = fs
-        .find_main_file()
-        .ok_or(anyhow!("no CAB file found"))?
-        .to_owned();
-
-    let mut cab = fs.read(&cabfile)?;
-    let mut file = SerializedFile::read(&mut cab)?;
-    for ob in file.objects.clone() {
-        if file.get_object_type_tree(&ob)?.type_string != "Texture2D" {
-            continue;
-        }
-        let value = file.read_object(ob)?;
-        let mut texture = Texture2D::from_value(value)?;
+    for ob in bundle
+        .all_toplevel_of_class("Texture2D")
+        .collect::<Vec<_>>()
+    {
+        let mut texture = ob.load(&mut bundle)?.parse::<Texture2D>()?;
         if texture.image_data.is_empty() {
-            texture.image_data = texture.stream_data.read(&mut fs)?;
+            texture.image_data = texture.stream_data.read(&bundle.fs)?;
         }
         let path = format!(
             "/tmp/a/{}_{i}.png",
diff --git a/src/assetbundle.rs b/src/assetbundle.rs
new file mode 100644
index 0000000..e10a8f3
--- /dev/null
+++ b/src/assetbundle.rs
@@ -0,0 +1,56 @@
+use crate::{
+    classes::pptr::PPtr,
+    serialized_file::SerializedFile,
+    unityfs::{NodeReader, UnityFS, block_reader::BlockReader, multi_reader::MultiReader},
+};
+use anyhow::{Context, Result, anyhow};
+use std::{
+    io::{Read, Seek},
+    marker::PhantomData,
+};
+
+/// High-level wrapper around UnityFS, SerializedFile and all the classes.
+pub struct AssetBundle<T> {
+    pub fs: UnityFS<T>,
+    pub(crate) main: SerializedFile<NodeReader<BlockReader<MultiReader<T>>>>,
+    pub(crate) shared_assets: Option<SerializedFile<NodeReader<BlockReader<MultiReader<T>>>>>,
+}
+
+impl<T: Read + Seek> AssetBundle<T> {
+    pub fn open(inner: T) -> Result<Self> {
+        let fs = UnityFS::open(inner).context("opening UnityFS")?;
+        let main_ni = fs
+            .find_main_file()
+            .ok_or(anyhow!("AssetBundle seems to lack main file"))?;
+        let main = SerializedFile::read(fs.read(main_ni)?)?;
+        let shared_assets = if let Some(n) = main.find_fs_shared_assets(&fs) {
+            Some(SerializedFile::read(fs.read(&n)?)?)
+        } else {
+            None
+        };
+        Ok(Self {
+            fs,
+            main,
+            shared_assets,
+        })
+    }
+
+    pub fn all_toplevel_of_class(&self, class_name: &str) -> impl Iterator<Item = PPtr> {
+        self.main
+            .all_objects_of_class(class_name)
+            .map(|o| (0, o))
+            .chain(
+                self.shared_assets
+                    .as_ref()
+                    .map(|e| e.all_objects_of_class(class_name).map(|o| (1, o)))
+                    .into_iter()
+                    .flatten(),
+            )
+            .map(|(fi, o)| PPtr {
+                class: class_name.to_owned(),
+                file_id: fi,
+                path_id: o.path_id,
+                _class: PhantomData,
+            })
+    }
+}
diff --git a/src/classes/pptr.rs b/src/classes/pptr.rs
index 9b54cbb..366e17c 100644
--- a/src/classes/pptr.rs
+++ b/src/classes/pptr.rs
@@ -1,6 +1,6 @@
 use crate::{
+    assetbundle::AssetBundle,
     object::{Value, parser::FromValue},
-    serialized_file::SerializedFile,
 };
 use anyhow::{Result, anyhow, bail};
 use log::debug;
@@ -13,7 +13,7 @@ use std::{
 #[derive(Debug, Serialize)]
 pub struct PPtr<T = Value> {
     #[serde(skip, default)]
-    _class: PhantomData<T>,
+    pub(crate) _class: PhantomData<T>,
     pub class: String,
     pub file_id: i32,
     pub path_id: i64,
@@ -54,27 +54,26 @@ impl<T: FromValue> PPtr<T> {
     pub fn is_null(&self) -> bool {
         self.path_id == 0 && self.file_id == 0
     }
-    pub fn load(
-        &self,
-        file: &mut SerializedFile<impl Read + Seek>,
-        shared_assets: Option<&mut SerializedFile<impl Read + Seek>>,
-    ) -> Result<T> {
+    pub fn load(&self, bundle: &mut AssetBundle<impl Read + Seek>) -> Result<T> {
         debug!(
             "loading PPtr<{}> file_id={} path_id={}",
             self.class, self.file_id, self.path_id
         );
         match self.file_id {
             0 => {
-                let ob = file
+                let ob = bundle
+                    .main
                     .objects
                     .iter()
                     .find(|o| o.path_id == self.path_id)
                     .ok_or(anyhow!("object with path id {} not found", self.path_id))?
                     .to_owned();
-                file.read_object(ob)?.parse()
+                bundle.main.read_object(ob)?.parse()
             }
             1 => {
-                let file = shared_assets.unwrap();
+                let file = bundle.shared_assets.as_mut().ok_or(anyhow!(
+                    "shared assets referenced but not included in bundle"
+                ))?;
                 let ob = file
                     .objects
                     .iter()
diff --git a/src/classes/streaminginfo.rs b/src/classes/streaminginfo.rs
index 21029f4..e308f1c 100644
--- a/src/classes/streaminginfo.rs
+++ b/src/classes/streaminginfo.rs
@@ -26,7 +26,7 @@ impl FromValue for StreamingInfo {
 }
 
 impl StreamingInfo {
-    pub fn read(&self, fs: &mut UnityFS<impl Read + Seek>) -> Result<Vec<u8>> {
+    pub fn read(&self, fs: &UnityFS<impl Read + Seek>) -> Result<Vec<u8>> {
         if !self.path.starts_with("archive:") {
             bail!("StreamingInfo path does not start on archive:")
         }
diff --git a/src/lib.rs b/src/lib.rs
index f01eaae..68e80f4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,3 +5,4 @@ pub mod helper;
 pub mod object;
 pub mod serialized_file;
 pub mod unityfs;
+pub mod assetbundle;
diff --git a/src/scene/mod.rs b/src/scene/mod.rs
deleted file mode 100644
index 8b13789..0000000
--- a/src/scene/mod.rs
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/src/serialized_file.rs b/src/serialized_file.rs
index a6514ee..aba72b9 100644
--- a/src/serialized_file.rs
+++ b/src/serialized_file.rs
@@ -345,6 +345,13 @@ impl<T: Read + Seek> SerializedFile<T> {
             .ok_or(anyhow!("type tree missing"))?;
         Value::read(typetree, self.endianness, &mut self.file)
     }
+
+    pub fn all_objects_of_class(&self, class_name: &str) -> impl Iterator<Item = &ObjectInfo> {
+        self.objects.iter().filter(move |o| {
+            self.get_object_type_tree(&o)
+                .map_or(false, |t| t.type_string == class_name)
+        })
+    }
 }
 
 impl TypeTreeNode {