1 files changed, 83 insertions, 40 deletions

diff --git a/shared/src/respack.rs b/shared/src/respack.rs
index 8c0b9d8..dee0cd0 100644
--- a/shared/src/respack.rs
+++ b/shared/src/respack.rs
@@ -18,6 +18,7 @@ use crate::{packets::Resource, resources::RespackEntry, store::ResourceStore};
 use anyhow::{Result, bail};
 use log::{info, warn};
 use std::{
+    collections::HashMap,
     io::{Read, Seek, SeekFrom, Write},
     marker::PhantomData,
 };
@@ -49,59 +50,101 @@ pub fn save_respack(
     Ok(())
 }
 
+/// Copies the entire respack to the store. This is usually a dumb idea because the pack supportes on-demand reading.
 pub fn load_respack(
-    mut input: impl Read + Seek,
+    input: impl Read + Seek,
     store: &ResourceStore,
 ) -> Result<Option<Resource<RespackEntry>>> {
-    let mut magic = [0u8; MAGIC.len()];
-    input.read_exact(&mut magic)?;
-    if magic != *MAGIC {
-        bail!("wrong magic bytes");
-    }
-    let mut entry = [0u8; 32];
-    input.read_exact(&mut entry)?;
-    let entry = if entry != [0u8; 32] {
-        Some(Resource(entry, PhantomData))
-    } else {
-        None
-    };
-
-    let mut count = [0u8; size_of::<u64>()];
-    input.read_exact(&mut count)?;
-    let count = u64::from_be_bytes(count);
-
+    let mut pack = RespackReader::open(input)?;
     let mut load_queue = Vec::new();
-    let mut found_entry = false;
-    for _ in 0..count {
-        let mut res = [0u8; 32];
-        let mut off = [0u8; size_of::<u64>()];
-        let mut size = [0u8; size_of::<u64>()];
-        input.read_exact(&mut res)?;
-        input.read_exact(&mut off)?;
-        input.read_exact(&mut size)?;
-
-        found_entry |= Some(Resource(res, PhantomData)) == entry;
-        if store.get_raw_size(Resource(res, PhantomData))?.is_none() {
-            load_queue.push((res, u64::from_be_bytes(off), u64::from_be_bytes(size)))
+    for res in pack.index.keys() {
+        if store.get_raw_size(*res)?.is_none() {
+            load_queue.push(*res);
         }
     }
-    if !found_entry && entry.is_some() {
-        warn!("respack does not contain its entry resource")
-    }
     info!(
-        "loading {} of {count} resources from pack",
+        "loading {} of {} resources from pack",
         load_queue.len(),
+        pack.index.len()
     );
-
-    for (res, off, size) in load_queue {
-        input.seek(SeekFrom::Start(off))?;
+    for res in load_queue {
         let mut buf = Vec::new();
-        input.by_ref().take(size).read_to_end(&mut buf)?;
+        pack.read(res)?.unwrap().read_to_end(&mut buf)?;
         let key = store.set_raw(&buf)?;
-        if key.0 != res {
+        if key != res {
             warn!("respack containes mislabeled resources")
         }
     }
 
-    Ok(entry)
+    Ok(pack.entry)
+}
+
+pub struct RespackReader<T> {
+    entry: Option<Resource<RespackEntry>>,
+    index: HashMap<Resource, (u64, u64)>,
+    inner: T,
+}
+impl<T: Read + Seek> RespackReader<T> {
+    pub fn open(mut inner: T) -> Result<Self> {
+        let mut magic = [0u8; MAGIC.len()];
+        inner.read_exact(&mut magic)?;
+        if magic != *MAGIC {
+            bail!("wrong magic bytes");
+        }
+        let mut entry = [0u8; 32];
+        inner.read_exact(&mut entry)?;
+        let entry = if entry != [0u8; 32] {
+            Some(Resource(entry, PhantomData))
+        } else {
+            None
+        };
+
+        let mut count = [0u8; size_of::<u64>()];
+        inner.read_exact(&mut count)?;
+        let count = u64::from_be_bytes(count);
+
+        let mut index = HashMap::new();
+        let mut found_entry = false;
+        for _ in 0..count {
+            let mut res = [0u8; 32];
+            let mut off = [0u8; size_of::<u64>()];
+            let mut size = [0u8; size_of::<u64>()];
+            inner.read_exact(&mut res)?;
+            inner.read_exact(&mut off)?;
+            inner.read_exact(&mut size)?;
+
+            found_entry |= Some(Resource(res, PhantomData)) == entry;
+            index.insert(
+                Resource(res, PhantomData),
+                (u64::from_be_bytes(off), u64::from_be_bytes(size)),
+            );
+        }
+        if !found_entry && entry.is_some() {
+            warn!("respack does not contain its entry resource")
+        }
+        info!("opened respack with {} resources", index.len());
+        Ok(Self {
+            entry,
+            index,
+            inner,
+        })
+    }
+    pub fn entry(&self) -> Option<Resource<RespackEntry>> {
+        self.entry.clone()
+    }
+    pub fn get_size(&self, key: Resource) -> Option<usize> {
+        self.index.get(&key).map(|(_off, size)| *size as usize)
+    }
+    pub fn iter(&self, mut cb: impl FnMut(Resource, usize)) {
+        for (r, (_, s)) in &self.index {
+            cb(*r, *s as usize)
+        }
+    }
+    pub fn read(&mut self, key: Resource) -> Result<Option<impl Read>> {
+        let Some((off, size)) = self.index.get(&key).copied() else {
+            return Ok(None);
+        };
+        self.inner.seek(SeekFrom::Start(off))?;
+        Ok(Some(self.inner.by_ref().take(size)))
+    }
 }