From 4141f361e363cf227ce582a5308f9f70035da2b8 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Fri, 4 Jul 2025 13:55:56 +0800 Subject: [PATCH] Add create support for DSC file --- Cargo.toml | 2 +- src/ext/vec.rs | 18 +++ src/scripts/bgi/archive/dsc.rs | 268 ++++++++++++++++++++++++++++++++- src/utils/bit_stream.rs | 5 +- 4 files changed, 289 insertions(+), 4 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b3bbe40..880050f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,7 +29,7 @@ utf16string = "0.2" [features] default = ["bgi", "bgi-arc", "bgi-img", "cat-system", "cat-system-arc", "cat-system-img", "circus", "escude", "escude-arc", "kirikiri", "kirikiri-img", "yaneurao", "yaneurao-itufuru"] bgi = [] -bgi-arc = ["bgi", "utils-bit-stream"] +bgi-arc = ["bgi", "rand", "utils-bit-stream"] bgi-img = ["bgi", "image", "utils-bit-stream"] cat-system = [] cat-system-arc = ["cat-system", "blowfish", "utils-crc32"] diff --git a/src/ext/vec.rs b/src/ext/vec.rs index 33f20d9..1b1c8af 100644 --- a/src/ext/vec.rs +++ b/src/ext/vec.rs @@ -24,3 +24,21 @@ impl VecExt for Vec { } } } + +pub trait SliceExt { + fn rfind(&self, pattern: &[T]) -> Option; +} + +impl SliceExt for [T] { + fn rfind(&self, pattern: &[T]) -> Option { + if pattern.is_empty() || self.len() < pattern.len() { + return None; + } + for i in (0..=self.len() - pattern.len()).rev() { + if &self[i..i + pattern.len()] == pattern { + return Some(i); + } + } + None + } +} diff --git a/src/scripts/bgi/archive/dsc.rs b/src/scripts/bgi/archive/dsc.rs index f9fbaf4..dab1f96 100644 --- a/src/scripts/bgi/archive/dsc.rs +++ b/src/scripts/bgi/archive/dsc.rs @@ -4,7 +4,9 @@ use crate::scripts::base::*; use crate::types::*; use crate::utils::bit_stream::*; use anyhow::Result; -use std::io::Write; +use rand::Rng; +use std::collections::BinaryHeap; +use std::io::{Seek, Write}; #[derive(Debug)] struct HuffmanCode { @@ -208,6 +210,240 @@ impl<'a> DscDecoder<'a> { } } +#[derive(Debug, Clone, Copy)] +enum LzssOp { + Literal(u8), + Match { len: u16, offset: u16 }, +} + +#[derive(Debug)] +struct FreqNode { + freq: u32, + symbol: Option, + left: Option>, + right: Option>, +} +impl PartialEq for FreqNode { + fn eq(&self, other: &Self) -> bool { + self.freq == other.freq + } +} +impl Eq for FreqNode {} +impl PartialOrd for FreqNode { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} +impl Ord for FreqNode { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + other.freq.cmp(&self.freq) + } +} + +fn calculate_huffman_depths(freqs: &[u32]) -> Vec { + let mut heap = BinaryHeap::new(); + for (symbol, &freq) in freqs.iter().enumerate() { + if freq > 0 { + heap.push(FreqNode { + freq, + symbol: Some(symbol as u16), + left: None, + right: None, + }); + } + } + + if heap.len() <= 1 { + let mut depths = vec![0; 512]; + if let Some(node) = heap.pop() { + depths[node.symbol.unwrap() as usize] = 1; + } + return depths; + } + + while heap.len() > 1 { + let node1 = heap.pop().unwrap(); + let node2 = heap.pop().unwrap(); + let new_node = FreqNode { + freq: node1.freq + node2.freq, + symbol: None, + left: Some(Box::new(node1)), + right: Some(Box::new(node2)), + }; + heap.push(new_node); + } + + let mut depths = vec![0; 512]; + if let Some(root) = heap.pop() { + fn traverse(node: &FreqNode, depth: u8, depths: &mut [u8]) { + if let Some(symbol) = node.symbol { + if depth == 0 { + depths[symbol as usize] = 1; + } else { + depths[symbol as usize] = depth; + } + } else { + if let Some(ref left) = node.left { + traverse(left, depth + 1, depths); + } + if let Some(ref right) = node.right { + traverse(right, depth + 1, depths); + } + } + } + traverse(&root, 0, &mut depths); + } + depths +} + +fn generate_canonical_codes(depths: &[u8]) -> Vec> { + let mut codes_with_depths = vec![]; + for (symbol, &depth) in depths.iter().enumerate() { + if depth > 0 { + codes_with_depths.push((symbol as u16, depth)); + } + } + codes_with_depths.sort_by(|a, b| { + let depth_cmp = a.1.cmp(&b.1); + if depth_cmp == std::cmp::Ordering::Equal { + a.0.cmp(&b.0) + } else { + depth_cmp + } + }); + + let mut huffman_codes = vec![None; 512]; + let mut current_code = 0u16; + let mut last_depth = 0u8; + + for &(symbol, depth) in &codes_with_depths { + if last_depth != 0 { + current_code <<= depth - last_depth; + } + huffman_codes[symbol as usize] = Some((current_code, depth)); + current_code += 1; + last_depth = depth; + } + + huffman_codes +} + +pub struct DscEncoder<'a, T: Write + Seek> { + stream: MsbBitWriter<'a, T>, + magic: u32, + key: u32, + dec_count: u32, +} + +impl<'a, T: Write + Seek> DscEncoder<'a, T> { + pub fn new(writer: &'a mut T) -> Self { + let stream = MsbBitWriter::new(writer); + DscEncoder { + stream, + magic: 0x5344 << 16, // "DS" + key: rand::rng().random(), + dec_count: 0, + } + } + + pub fn pack(mut self, data: &[u8]) -> Result<()> { + // LZSS compression + let mut ops = vec![]; + let mut pos = 0; + while pos < data.len() { + let mut best_len = 0; + let mut best_offset = 0; + let max_len = (data.len() - pos).min(257); + + if max_len >= 2 { + let search_start = if pos > 4097 { pos - 4097 } else { 0 }; + let lookbehind = &data[search_start..pos]; + for len in (2..=max_len).rev() { + if let Some(found_idx) = lookbehind.rfind(&data[pos..pos + len]) { + let offset = lookbehind.len() - found_idx; + if offset >= 2 { + best_len = len; + best_offset = offset; + break; + } + } + } + } + + if best_len >= 2 { + ops.push(LzssOp::Match { + len: best_len as u16, + offset: best_offset as u16, + }); + pos += best_len; + } else { + ops.push(LzssOp::Literal(data[pos])); + pos += 1; + } + } + + let symbols: Vec = ops + .iter() + .map(|op| match op { + LzssOp::Literal(byte) => *byte as u16, + LzssOp::Match { len, .. } => 256 + (len - 2), + }) + .collect(); + self.dec_count = symbols.len() as u32; + + let mut freqs = vec![0u32; 512]; + for &s in &symbols { + freqs[s as usize] += 1; + } + + let depths = calculate_huffman_depths(&freqs); + let huffman_codes = generate_canonical_codes(&depths); + + self.stream.writer.write_all(b"DSC FORMAT 1.00\0")?; + self.stream.writer.seek(std::io::SeekFrom::Start(0x10))?; + self.stream.writer.write_u32(self.key)?; + self.stream.writer.write_u32(data.len() as u32)?; + self.stream.writer.write_u32(self.dec_count)?; + self.stream.writer.seek(std::io::SeekFrom::Start(0x20))?; + + for depth in depths.iter() { + let key = self.update_key(); + self.stream.writer.write_u8(depth.overflowing_add(key).0)?; + } + + for op in &ops { + match op { + LzssOp::Literal(byte) => { + let symbol = *byte as u16; + let (code, len) = huffman_codes[symbol as usize].unwrap(); + self.stream.put_bits(code as u32, len)?; + } + LzssOp::Match { len, offset } => { + let symbol = 256 + (len - 2); + let (code, huff_len) = huffman_codes[symbol as usize].unwrap(); + self.stream.put_bits(code as u32, huff_len)?; + self.stream.put_bits((*offset - 2) as u32, 12)?; + } + } + } + self.stream.flush()?; + Ok(()) + } + + fn update_key(&mut self) -> u8 { + let v0 = 20021 * (self.key & 0xffff); + let mut v1 = self.magic | (self.key >> 16); + v1 = v1 + .overflowing_mul(20021) + .0 + .overflowing_add(self.key.overflowing_mul(346).0) + .0; + v1 = (v1 + (v0 >> 16)) & 0xffff; + self.key = (v1 << 16) + (v0 & 0xffff) + 1; + v1 as u8 + } +} + #[derive(Debug)] pub struct DscBuilder {} @@ -251,6 +487,23 @@ impl ScriptBuilder for DscBuilder { } None } + + fn can_create_file(&self) -> bool { + true + } + + fn create_file<'a>( + &'a self, + filename: &'a str, + mut writer: Box, + _encoding: Encoding, + _file_encoding: Encoding, + ) -> Result<()> { + let encoder = DscEncoder::new(&mut writer); + let data = crate::utils::files::read_file(filename)?; + encoder.pack(&data)?; + Ok(()) + } } #[derive(Debug)] @@ -291,4 +544,17 @@ impl Script for Dsc { f.write_all(&self.data)?; Ok(()) } + + fn custom_import<'a>( + &'a self, + custom_filename: &'a str, + mut file: Box, + _encoding: Encoding, + _output_encoding: Encoding, + ) -> Result<()> { + let encoder = DscEncoder::new(&mut file); + let data = crate::utils::files::read_file(custom_filename)?; + encoder.pack(&data)?; + Ok(()) + } } diff --git a/src/utils/bit_stream.rs b/src/utils/bit_stream.rs index 98c360c..8dba11f 100644 --- a/src/utils/bit_stream.rs +++ b/src/utils/bit_stream.rs @@ -42,7 +42,7 @@ impl MsbBitStream { } pub struct MsbBitWriter<'a, T: Write> { - writer: &'a mut T, + pub writer: &'a mut T, buffer: u32, buffer_size: u32, } @@ -58,7 +58,8 @@ impl<'a, T: Write> MsbBitWriter<'a, T> { pub fn flush(&mut self) -> Result<()> { if self.buffer_size > 0 { - self.writer.write_u8((self.buffer & 0xFF) as u8)?; + self.writer + .write_u8(((self.buffer << (8 - self.buffer_size)) & 0xFF) as u8)?; self.buffer = 0; self.buffer_size = 0; }