From dbbdf92bdd9a375b8f66496711982d6b95623e40 Mon Sep 17 00:00:00 2001
From: Luke Street <luke.street@encounterpc.com>
Date: Thu, 9 Feb 2023 16:38:51 -0500
Subject: [PATCH] Support LZSS compressed files

---
 src/cmd/pak.rs   | 46 +++++++++++++++++++++++++++++-----------------
 src/util/lzss.rs | 38 ++++++++++++++++++++++++++++++++++++++
 src/util/mod.rs  |  1 +
 3 files changed, 68 insertions(+), 17 deletions(-)
 create mode 100644 src/util/lzss.rs

diff --git a/src/cmd/pak.rs b/src/cmd/pak.rs
index e532832..54abb14 100644
--- a/src/cmd/pak.rs
+++ b/src/cmd/pak.rs
@@ -1,18 +1,18 @@
 use std::{
+    borrow::Cow,
     fmt::{Debug, Display, Formatter, Write},
     fs,
-    fs::File,
     io::{Cursor, Read, Seek, SeekFrom},
     path::PathBuf,
 };
 
 use anyhow::{anyhow, bail, ensure, Result};
 use argh::FromArgs;
-use binrw::{BinRead, BinReaderExt, BinResult, Endian, NullString};
-use binrw_derive::{binrw, BinWrite};
+use binrw::{BinReaderExt, BinResult, Endian};
+use binrw_derive::binrw;
 use uuid::Uuid;
 
-use crate::util::file::map_file;
+use crate::util::{file::map_file, lzss::decompress};
 
 #[derive(FromArgs, PartialEq, Debug)]
 /// process PAK files
@@ -180,8 +180,8 @@ pub struct AssetDirectoryEntry {
     pub unk1: u32,
     pub unk2: u32,
     pub offset: u64,
+    pub decompressed_size: u64,
     pub size: u64,
-    pub unk3: u64,
 }
 
 #[binrw]
@@ -306,18 +306,30 @@ fn extract(args: ExtractArgs) -> Result<()> {
                 .unwrap_or_else(|| format!("{}", asset_entry.asset_id));
             let filename = format!("{}.{}", name, asset_entry.asset_type);
             let path = args.output.join(filename);
-            // log::info!("Extracting {}", path.display());
-            let mut data = &data
-                [asset_entry.offset as usize..(asset_entry.offset + asset_entry.size) as usize];
-            if asset_entry.asset_type == *b"FMV0" {
-                let (hdr, fmv_data, _) = FormDescriptor::slice(data, Endian::Little)?;
-                // log::info!("{:?}", hdr);
-                data = fmv_data;
-            } else if asset_entry.asset_type == *b"ROOM" {
-                let mut file = File::create("ROOM-format.txt")?;
-                dump_rfrm(&mut file, data, 0)?;
-            }
-            fs::write(path, data)?;
+            log::info!("Extracting {} ({:?})", path.display(), asset_entry);
+            let data: Cow<[u8]> = if asset_entry.size != asset_entry.decompressed_size {
+                let compression_bytes =
+                    &data[asset_entry.offset as usize..asset_entry.offset as usize + 4];
+                let compression_type = u32::from_le_bytes(compression_bytes.try_into().unwrap());
+                log::info!("Decompressing {}", compression_type);
+                let mut out = vec![0u8; asset_entry.decompressed_size as usize];
+                let data = &data[asset_entry.offset as usize + 4
+                    ..(asset_entry.offset + asset_entry.size) as usize];
+                match compression_type {
+                    1 => decompress::<1>(data, &mut out),
+                    2 => decompress::<2>(data, &mut out),
+                    3 => decompress::<3>(data, &mut out),
+                    _ => bail!("Unsupported compression mode {}", compression_type),
+                }
+                Cow::Owned(out)
+            } else {
+                Cow::Borrowed(
+                    &data[asset_entry.offset as usize
+                        ..(asset_entry.offset + asset_entry.size) as usize],
+                )
+            };
+            // TODO strip RFRM header?
+            fs::write(path, &data)?;
         }
     } else {
         bail!("Failed to locate ADIR chunk");
diff --git a/src/util/lzss.rs b/src/util/lzss.rs
new file mode 100644
index 0000000..621f0f5
--- /dev/null
+++ b/src/util/lzss.rs
@@ -0,0 +1,38 @@
+use anyhow::Result;
+
+/// https://wiki.axiodl.com/w/LZSS_Compression
+pub fn decompress<const M: u8>(mut input: &[u8], output: &mut [u8]) {
+    let group_len = 2usize.pow(M as u32 - 1);
+    let mut out_cur = 0usize;
+
+    let mut header_byte = 0u8;
+    let mut group = 0u8;
+    while !input.is_empty() {
+        if group == 0 {
+            header_byte = input[0];
+            input = &input[1..];
+            group = 8;
+        }
+
+        if header_byte & 0x80 == 0 {
+            output[out_cur..group_len + out_cur].copy_from_slice(&input[..group_len]);
+            input = &input[group_len..];
+            out_cur += group_len;
+        } else {
+            let count = (input[0] as usize >> 4) + (4 - M as usize);
+            let length = (((input[0] as usize & 0xF) << 0x8) | input[1] as usize) << (M - 1);
+            input = &input[2..];
+
+            let seek = out_cur - length;
+            for n in 0..count * group_len {
+                output[out_cur + n] = output[seek + n];
+            }
+            out_cur += count * group_len;
+        }
+
+        header_byte <<= 1;
+        group -= 1;
+    }
+
+    debug_assert_eq!(out_cur, output.len());
+}
diff --git a/src/util/mod.rs b/src/util/mod.rs
index 2e172cd..fc06744 100644
--- a/src/util/mod.rs
+++ b/src/util/mod.rs
@@ -1 +1,2 @@
 pub mod file;
+pub mod lzss;