allow preloading files

1 files changed, 23 insertions, 14 deletions

diff --git a/src/assetbundle.rs b/src/assetbundle.rs
index 0af9895..f87d7b8 100644
--- a/src/assetbundle.rs
+++ b/src/assetbundle.rs
@@ -1,16 +1,14 @@
 use crate::{
     classes::pptr::PPtr,
     serialized_file::SerializedFile,
-    unityfs::{
-        NodeReader, UnityFS, block_reader::BlockReader, header::NodeInfo, multi_reader::MultiReader,
-    },
+    unityfs::{UnityFS, header::NodeInfo},
 };
 use anyhow::{Context, Result, anyhow};
 use log::debug;
 use std::{
     collections::HashMap,
     fs::File,
-    io::{BufReader, Read, Seek},
+    io::{BufReader, Cursor, Read, Seek},
     marker::PhantomData,
     path::Path,
     sync::{Arc, Mutex},
@@ -19,13 +17,16 @@ use std::{
 /// High-level wrapper around UnityFS, SerializedFile and all the classes.
 pub struct AssetBundle<T> {
     pub fs: UnityFS<T>,
-    pub ser_files:
-        HashMap<String, Arc<Mutex<SerializedFile<NodeReader<BlockReader<MultiReader<T>>>>>>>,
+    pub ser_files: HashMap<String, Arc<Mutex<SerializedFile<Box<dyn ReadSeek>>>>>,
     pub main_file: NodeInfo,
     pub default_resources: SerializedFile<BufReader<File>>,
+    pub full_decomp: bool,
 }
 
-impl<T: Read + Seek> AssetBundle<T> {
+pub trait ReadSeek: Read + Seek + 'static {}
+impl<T: Read + Seek + 'static> ReadSeek for T {}
+
+impl<T: Read + Seek + 'static> AssetBundle<T> {
     pub fn open(inner: T, support_dir: impl AsRef<Path>) -> Result<Self> {
         let fs = UnityFS::open(inner).context("opening UnityFS")?;
         let main_file = fs
@@ -44,21 +45,29 @@ impl<T: Read + Seek> AssetBundle<T> {
             fs,
             main_file,
             default_resources,
+            full_decomp: false,
             ser_files: HashMap::new(),
         })
     }
 
     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))?;
+        nodeinfo: &NodeInfo,
+    ) -> Result<Arc<Mutex<SerializedFile<Box<dyn ReadSeek>>>>> {
+        if !self.ser_files.contains_key(&nodeinfo.name) {
+            let mut node = self.fs.read(nodeinfo)?;
+            let node_box = if self.full_decomp {
+                let mut buf = Vec::new();
+                node.read_to_end(&mut buf)?;
+                Box::new(Cursor::new(buf)) as Box<dyn ReadSeek>
+            } else {
+                Box::new(node) as Box<dyn ReadSeek>
+            };
+            let file = SerializedFile::read(node_box, format!("archive:/{}", nodeinfo.name))?;
             self.ser_files
-                .insert(node.name.clone(), Arc::new(Mutex::new(file)));
+                .insert(nodeinfo.name.clone(), Arc::new(Mutex::new(file)));
         }
-        Ok(self.ser_files.get(&node.name).unwrap().clone())
+        Ok(self.ser_files.get(&nodeinfo.name).unwrap().clone())
     }
 
     pub fn all_toplevel(&mut self) -> Vec<PPtr> {