assetbundle struct to abstract over unityfs and serializedfile

1 files changed, 56 insertions, 0 deletions

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,
+            })
+    }
+}