horrible pptr external files system

1 files changed, 90 insertions, 42 deletions

diff --git a/src/assetbundle.rs b/src/assetbundle.rs
index 8f6af15..0af9895 100644
--- a/src/assetbundle.rs
+++ b/src/assetbundle.rs
@@ -1,81 +1,129 @@
 use crate::{
     classes::pptr::PPtr,
     serialized_file::SerializedFile,
-    unityfs::{NodeReader, UnityFS, block_reader::BlockReader, multi_reader::MultiReader},
+    unityfs::{
+        NodeReader, UnityFS, block_reader::BlockReader, header::NodeInfo, multi_reader::MultiReader,
+    },
 };
 use anyhow::{Context, Result, anyhow};
 use log::debug;
 use std::{
-    io::{Read, Seek},
+    collections::HashMap,
+    fs::File,
+    io::{BufReader, Read, Seek},
     marker::PhantomData,
+    path::Path,
+    sync::{Arc, Mutex},
 };
 
 /// 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>>>>>,
+    pub ser_files:
+        HashMap<String, Arc<Mutex<SerializedFile<NodeReader<BlockReader<MultiReader<T>>>>>>>,
+    pub main_file: NodeInfo,
+    pub default_resources: SerializedFile<BufReader<File>>,
 }
 
 impl<T: Read + Seek> AssetBundle<T> {
-    pub fn open(inner: T) -> Result<Self> {
+    pub fn open(inner: T, support_dir: impl AsRef<Path>) -> Result<Self> {
         let fs = UnityFS::open(inner).context("opening UnityFS")?;
-        let main_ni = fs
+        let main_file = fs
             .find_main_file()
-            .ok_or(anyhow!("AssetBundle seems to lack main file"))?;
-        debug!("detected {:?} as main file", main_ni.name);
-        let main = SerializedFile::read(fs.read(main_ni)?, 0)?;
-        let shared_assets = if let Some(n) = main.find_fs_shared_assets(&fs) {
-            debug!("detect {:?} as shared assets", n.name);
-            Some(SerializedFile::read(fs.read(&n)?, 1)?)
-        } else {
-            None
-        };
+            .ok_or(anyhow!("AssetBundle seems to lack main file"))?
+            .clone();
+        debug!("opening default resource file");
+        let default_resources = SerializedFile::read(
+            BufReader::new(File::open(
+                support_dir.as_ref().join("unity default resources"),
+            )?),
+            "Library/unity default resources".to_owned(),
+        )?;
+        debug!("detected {:?} as main file", main_file.name);
         Ok(Self {
             fs,
-            main,
-            shared_assets,
+            main_file,
+            default_resources,
+            ser_files: HashMap::new(),
         })
     }
 
-    pub fn all_toplevel(&self) -> impl Iterator<Item = PPtr> {
-        self.main
+    pub fn get_fs_file(
+        &mut self,
+        node: &NodeInfo,
+    ) -> Result<Arc<Mutex<SerializedFile<NodeReader<BlockReader<MultiReader<T>>>>>>> {
+        if !self.ser_files.contains_key(&node.name) {
+            let file =
+                SerializedFile::read(self.fs.read(node)?, format!("archive:/{}", node.name))?;
+            self.ser_files
+                .insert(node.name.clone(), Arc::new(Mutex::new(file)));
+        }
+        Ok(self.ser_files.get(&node.name).unwrap().clone())
+    }
+
+    pub fn all_toplevel(&mut self) -> Vec<PPtr> {
+        let main = self.get_fs_file(&self.main_file.clone()).unwrap();
+        let main = main.lock().unwrap();
+        let shared_assets = main.find_fs_shared_assets(&self.fs);
+        let main_obs = main
             .objects
             .iter()
-            .map(|o| (0, o))
-            .chain(
-                self.shared_assets
-                    .as_ref()
-                    .map(|e| e.objects.iter().map(|o| (1, o)))
-                    .into_iter()
-                    .flatten(),
-            )
-            .map(|(si, o)| PPtr {
+            .map(|o| (main.ecx.clone(), o.to_owned()))
+            .collect::<Vec<_>>();
+        let shared_obs = if let Some(shared_assets) = shared_assets {
+            let shared = self.get_fs_file(&shared_assets).unwrap();
+            let shared = shared.lock().unwrap();
+            shared
+                .objects
+                .iter()
+                .map(|o| (shared.ecx.clone(), o.to_owned()))
+                .collect()
+        } else {
+            Vec::new()
+        };
+
+        main_obs
+            .into_iter()
+            .chain(shared_obs)
+            .map(|(ecx, o)| PPtr {
                 class: "".to_string(),
-                source_file: si,
+                ecx,
                 file_id: 0,
                 path_id: o.path_id,
                 _class: PhantomData,
             })
+            .collect()
     }
 
-    pub fn all_toplevel_of_class(&self, class_name: &str) -> impl Iterator<Item = PPtr> {
-        self.main
+    pub fn all_toplevel_of_class(&mut self, class_name: &str) -> Vec<PPtr> {
+        let main = self.get_fs_file(&self.main_file.clone()).unwrap();
+        let main = main.lock().unwrap();
+        let shared_assets = main.find_fs_shared_assets(&self.fs);
+        let main_obs = 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(|(si, o)| PPtr {
-                class: class_name.to_owned(),
-                source_file: si,
+            .map(|o| (main.ecx.clone(), o.to_owned()))
+            .collect::<Vec<_>>();
+        let shared_obs = if let Some(shared_assets) = shared_assets {
+            let shared = self.get_fs_file(&shared_assets).unwrap();
+            let shared = shared.lock().unwrap();
+            shared
+                .all_objects_of_class(class_name)
+                .map(|o| (shared.ecx.clone(), o.to_owned()))
+                .collect()
+        } else {
+            Vec::new()
+        };
+
+        main_obs
+            .into_iter()
+            .chain(shared_obs)
+            .map(|(ecx, o)| PPtr {
+                class: "".to_string(),
+                ecx,
                 file_id: 0,
                 path_id: o.path_id,
                 _class: PhantomData,
             })
+            .collect()
     }
 }