more abstraction around unityfs to read multiple files from a single reader

1 files changed, 99 insertions, 0 deletions

diff --git a/src/unityfs/block_reader.rs b/src/unityfs/block_reader.rs
new file mode 100644
index 0000000..8aa18c3
--- /dev/null
+++ b/src/unityfs/block_reader.rs
@@ -0,0 +1,99 @@
+use super::BlockInfo;
+use log::{debug, trace};
+use std::{
+    io::{ErrorKind, Read, Seek, SeekFrom},
+    sync::Arc,
+};
+
+pub struct BlockReader<T> {
+    blocks: Arc<Vec<BlockInfo>>,
+    inner: T,
+    inner_seek_offset: u64,
+    nblock_index: usize,
+    cblock_data: Vec<u8>,
+    cblock_off: usize,
+}
+
+impl<T: Read> BlockReader<T> {
+    pub fn new(blocks: Arc<Vec<BlockInfo>>, inner: T, inner_seek_offset: u64) -> Self {
+        Self {
+            blocks,
+            inner,
+            inner_seek_offset,
+            nblock_index: 0,
+            cblock_data: Vec::new(),
+            cblock_off: 0,
+        }
+    }
+    pub fn load_next_block(&mut self) -> std::io::Result<()> {
+        trace!("loading block {}", self.nblock_index);
+        let block = &self.blocks[self.nblock_index];
+        let mut comp_buf = vec![0; block.comp_size as usize];
+        self.inner.read_exact(&mut comp_buf)?;
+        let decomp_buf = block
+            .comp_scheme
+            .decompress(comp_buf, block.decomp_size as usize)
+            .map_err(|e| {
+                std::io::Error::new(
+                    ErrorKind::InvalidData,
+                    format!("decompression failure: {e}"),
+                )
+            })?;
+        self.nblock_index += 1;
+        self.cblock_data = decomp_buf;
+        self.cblock_off = 0;
+        Ok(())
+    }
+}
+
+impl<T: Read> Read for BlockReader<T> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        if self.cblock_off >= self.cblock_data.len() {
+            self.load_next_block()?;
+        }
+        let size = (self.cblock_data.len() - self.cblock_off).min(buf.len());
+        buf[..size].copy_from_slice(&self.cblock_data[self.cblock_off..self.cblock_off + size]);
+        self.cblock_off += size;
+        Ok(size)
+    }
+}
+impl<T: Seek + Read> Seek for BlockReader<T> {
+    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+        let SeekFrom::Start(pos) = pos else {
+            unimplemented!()
+        };
+        debug!("seek decomp to {pos}");
+        let mut comp_off = self.inner_seek_offset;
+        let mut decomp_off = 0;
+        let mut target_block = None;
+        for (i, b) in self.blocks.iter().enumerate() {
+            if pos <= decomp_off + b.decomp_size as u64 {
+                target_block = Some(i);
+                break;
+            }
+            decomp_off += b.decomp_size as u64;
+            comp_off += b.comp_size as u64;
+        }
+
+        let Some(i) = target_block else {
+            return Err(std::io::Error::new(
+                ErrorKind::UnexpectedEof,
+                "seek out of bounds",
+            ));
+        };
+
+        let block_off = pos - decomp_off;
+        debug!("target is block={i} offset={block_off}");
+        if self.nblock_index == i + 1 {
+            debug!("intra-block seek")
+        } else {
+            debug!("seek comp to {comp_off}");
+            self.inner.seek(SeekFrom::Start(comp_off))?;
+            self.nblock_index = i;
+            self.load_next_block()?;
+        }
+        self.cblock_off = block_off as usize;
+
+        Ok(pos)
+    }
+}