1use crate::ext::atomic::*;
3use crate::ext::io::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::bit_stream::*;
7use crate::utils::img::*;
8use crate::utils::struct_pack::*;
9use crate::utils::threadpool::*;
10use anyhow::Result;
11use msg_tool_macro::*;
12use std::io::{Read, Seek, Write};
13use std::sync::atomic::AtomicBool;
14use std::sync::{Arc, Mutex};
15
16#[derive(Debug)]
17pub struct BgiCBGBuilder {}
19
20impl BgiCBGBuilder {
21 pub const fn new() -> Self {
23 BgiCBGBuilder {}
24 }
25}
26
27impl ScriptBuilder for BgiCBGBuilder {
28 fn default_encoding(&self) -> Encoding {
29 Encoding::Cp932
30 }
31
32 fn build_script(
33 &self,
34 data: Vec<u8>,
35 _filename: &str,
36 _encoding: Encoding,
37 _archive_encoding: Encoding,
38 config: &ExtraConfig,
39 _archive: Option<&Box<dyn Script>>,
40 ) -> Result<Box<dyn Script>> {
41 Ok(Box::new(BgiCBG::new(data, config)?))
42 }
43
44 fn extensions(&self) -> &'static [&'static str] {
45 &[]
46 }
47
48 fn script_type(&self) -> &'static ScriptType {
49 &ScriptType::BGICbg
50 }
51
52 fn is_image(&self) -> bool {
53 true
54 }
55
56 fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
57 if buf_len >= 0x10 && buf.starts_with(b"CompressedBG___") {
58 return Some(255);
59 }
60 None
61 }
62
63 fn can_create_image_file(&self) -> bool {
64 true
65 }
66
67 fn create_image_file<'a>(
68 &'a self,
69 data: ImageData,
70 _filename: &str,
71 mut writer: Box<dyn WriteSeek + 'a>,
72 _options: &ExtraConfig,
73 ) -> Result<()> {
74 let encoder = CbgEncoder::new(data)?;
75 let data = encoder.encode()?;
76 writer.write_all(&data)?;
77 Ok(())
78 }
79}
80
81#[derive(Debug, StructPack, StructUnpack)]
82struct BgiCBGHeader {
83 width: u16,
84 height: u16,
85 bpp: u32,
86 _unk: u64,
87 intermediate_length: u32,
88 key: u32,
89 enc_length: u32,
90 check_sum: u8,
91 check_xor: u8,
92 version: u16,
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96enum CbgColorType {
97 Bgra32,
98 Bgr24,
99 Grayscale,
100 Bgr565,
101}
102
103fn convert_bgr565_to_bgr24(input: Vec<u8>, width: u16, height: u16) -> ImageData {
104 let pixel_count = width as usize * height as usize;
105 let mut output = Vec::with_capacity(pixel_count * 3);
106
107 for chunk in input.chunks_exact(2) {
108 let pixel = u16::from_le_bytes([chunk[0], chunk[1]]);
109
110 let blue_5bit = (pixel & 0x1) as u8;
111 let green_6bit = ((pixel >> 5) & 0x3) as u8;
112 let red_5bit = ((pixel >> 11) & 0x1) as u8;
113
114 let blue = ((blue_5bit as u16 * 255) / 31) as u8;
115 let green = ((green_6bit as u16 * 255) / 63) as u8;
116 let red = ((red_5bit as u16 * 255) / 31) as u8;
117
118 output.push(blue);
119 output.push(green);
120 output.push(red);
121 }
122
123 ImageData {
124 width: width as u32,
125 height: height as u32,
126 color_type: ImageColorType::Bgr,
127 depth: 8,
128 data: output,
129 }
130}
131
132#[derive(Debug)]
133pub struct BgiCBG {
135 header: BgiCBGHeader,
136 data: MemReader,
137 color_type: CbgColorType,
138 decode_workers: usize,
139}
140
141impl BgiCBG {
142 pub fn new(data: Vec<u8>, config: &ExtraConfig) -> Result<Self> {
147 let mut reader = MemReader::new(data);
148 let mut magic = [0u8; 16];
149 reader.read_exact(&mut magic)?;
150 if !magic.starts_with(b"CompressedBG___") {
151 return Err(anyhow::anyhow!("Invalid magic: {:?}", magic));
152 }
153 let header = BgiCBGHeader::unpack(&mut reader, false, Encoding::Cp932, &None)?;
154 if header.version > 2 {
155 return Err(anyhow::anyhow!("Unsupported version: {}", header.version));
156 }
157 let color_type = match header.bpp {
158 32 => CbgColorType::Bgra32,
159 24 => CbgColorType::Bgr24,
160 8 => CbgColorType::Grayscale,
161 16 => {
162 if header.version == 2 {
163 return Err(anyhow::anyhow!("Unsupported BPP 16 in version 2"));
164 }
165 CbgColorType::Bgr565
166 }
167 _ => return Err(anyhow::anyhow!("Unsupported BPP: {}", header.bpp)),
168 };
169 Ok(BgiCBG {
170 header,
171 data: reader,
172 color_type,
173 decode_workers: config.bgi_img_workers.max(1),
174 })
175 }
176}
177
178impl Script for BgiCBG {
179 fn default_output_script_type(&self) -> OutputScriptType {
180 OutputScriptType::Json
181 }
182
183 fn default_format_type(&self) -> FormatOptions {
184 FormatOptions::None
185 }
186
187 fn is_image(&self) -> bool {
188 true
189 }
190
191 fn export_image(&self) -> Result<ImageData> {
192 let decoder = CbgDecoder::new(
193 self.data.to_ref(),
194 &self.header,
195 self.color_type,
196 self.decode_workers,
197 )?;
198 Ok(decoder.unpack()?)
199 }
200
201 fn import_image<'a>(
202 &'a self,
203 data: ImageData,
204 _filename: &str,
205 mut file: Box<dyn WriteSeek + 'a>,
206 ) -> Result<()> {
207 let encoder = CbgEncoder::new(data)?;
208 let encoded_data = encoder.encode()?;
209 file.write_all(&encoded_data)?;
210 Ok(())
211 }
212}
213
214struct CbgDecoder<'a> {
215 stream: MsbBitStream<MemReaderRef<'a>>,
216 info: &'a BgiCBGHeader,
217 color_type: CbgColorType,
218 key: u32,
219 magic: u32,
220 pixel_size: u8,
221 stride: usize,
222 workers: usize,
223}
224
225impl<'a> CbgDecoder<'a> {
226 fn new(
227 reader: MemReaderRef<'a>,
228 info: &'a BgiCBGHeader,
229 color_type: CbgColorType,
230 workers: usize,
231 ) -> Result<Self> {
232 let magic = 0;
233 let key = info.key;
234 let stream = MsbBitStream::new(reader);
235 let pixel_size = info.bpp as u8 / 8;
236 let stride = info.width as usize * (info.bpp as usize / 8);
237 Ok(CbgDecoder {
238 stream,
239 info,
240 key,
241 magic,
242 color_type,
243 pixel_size,
244 stride,
245 workers,
246 })
247 }
248
249 fn unpack(mut self) -> Result<ImageData> {
250 self.stream.m_input.pos = 0x30;
251 if self.info.version < 2 {
252 return self.unpack_v1();
253 } else if self.info.version == 2 {
254 if self.info.enc_length < 0x80 {
255 return Err(anyhow::anyhow!(
256 "Invalid encoded length: {}",
257 self.info.enc_length
258 ));
259 }
260 return self.unpack_v2();
261 }
262 Err(anyhow::anyhow!("Unknown version: {}", self.info.version))
263 }
264
265 fn unpack_v1(&mut self) -> Result<ImageData> {
266 let leaf_nodes_weight = {
267 let stream = MemReader::new(self.read_encoded()?);
268 let mut stream_ref = stream.to_ref();
269 Self::read_weight_table(&mut stream_ref, 0x100)?
270 };
271 let tree = HuffmanTree::new(&leaf_nodes_weight, false);
272 let mut packed = Vec::with_capacity(self.info.intermediate_length as usize);
273 packed.resize(self.info.intermediate_length as usize, 0);
274 self.huffman_decompress(&tree, &mut packed)?;
275 let buf_size = self.stride * self.info.height as usize;
276 let mut output = Vec::with_capacity(buf_size);
277 output.resize(buf_size, 0);
278 Self::unpack_zeros(&packed, &mut output);
279 self.reverse_average_sampling(&mut output);
280 let color_type = match self.color_type {
281 CbgColorType::Bgra32 => ImageColorType::Bgra,
282 CbgColorType::Bgr24 => ImageColorType::Bgr,
283 CbgColorType::Grayscale => ImageColorType::Grayscale,
284 CbgColorType::Bgr565 => {
285 return Ok(convert_bgr565_to_bgr24(
286 output,
287 self.info.width,
288 self.info.height,
289 ));
290 }
291 };
292 Ok(ImageData {
293 width: self.info.width as u32,
294 height: self.info.height as u32,
295 color_type,
296 depth: 8,
297 data: output,
298 })
299 }
300
301 fn unpack_v2(&mut self) -> Result<ImageData> {
302 let dct_data = self.read_encoded()?;
303 let mut dct = [[0.0f32; 64]; 2];
304 for i in 0..0x80 {
305 dct[i >> 6][i & 0x3f] = dct_data[i] as f32 * DCT_TABLE[i & 0x3f];
306 }
307
308 let base_offset = self.stream.m_input.pos;
309 let tree1 = HuffmanTree::new(
310 &Self::read_weight_table(&mut self.stream.m_input, 0x10)?,
311 true,
312 );
313 let tree2 = HuffmanTree::new(
314 &Self::read_weight_table(&mut self.stream.m_input, 0xB0)?,
315 true,
316 );
317
318 let width = ((self.info.width as i32 + 7) & -8) as i32;
319 let height = ((self.info.height as i32 + 7) & -8) as i32;
320 let y_blocks = height / 8;
321
322 let mut offsets = Vec::with_capacity((y_blocks + 1) as usize);
323 let input_base =
324 (self.stream.m_input.pos + ((y_blocks + 1) as usize * 4) - base_offset) as i32;
325
326 for _ in 0..=y_blocks {
327 let offset = self.stream.m_input.read_i32()?;
328 offsets.push(offset - input_base);
329 }
330
331 let input = self.stream.m_input.data[self.stream.m_input.pos..].to_vec();
332 let pad_skip = ((width >> 3) + 7) >> 3;
333
334 let output_size = (width * height * 4) as usize;
335 let output = vec![0u8; output_size];
336 let output_mutex = Mutex::new(output);
337
338 let decoder = Arc::new(ParallelCbgDecoder {
339 input,
340 output: output_mutex,
341 bpp: self.info.bpp as i32,
342 width,
343 height,
344 tree1,
345 tree2,
346 dct,
347 has_alpha: AtomicBool::new(false),
348 });
349
350 let thread_pool = ThreadPool::new(self.workers, Some("cbg-decoder-worker-"), false)?;
351 let mut dst = 0i32;
352
353 for i in 0..y_blocks {
354 let block_offset = offsets[i as usize] + pad_skip;
355 let next_offset = if i + 1 == y_blocks {
356 decoder.input.len() as i32
357 } else {
358 offsets[(i + 1) as usize]
359 };
360 let closure_dst = dst;
361 let decoder_ref = Arc::clone(&decoder);
362
363 thread_pool.execute(
364 move |_| {
365 decoder_ref.unpack_block(block_offset, next_offset - block_offset, closure_dst)
366 },
367 true,
368 )?;
369 dst += width * 32;
370 }
371
372 if self.info.bpp == 32 {
373 let decoder_ref = Arc::clone(&decoder);
374 thread_pool.execute(
375 move |_| decoder_ref.unpack_alpha(offsets[y_blocks as usize]),
376 true,
377 )?;
378 }
379
380 let tasks = thread_pool.into_results();
381
382 for task in tasks {
383 task?;
384 }
385
386 let has_alpha = decoder.has_alpha.qload();
387 let mut output = decoder
388 .output
389 .lock()
390 .map_err(|e| anyhow::anyhow!("Failed to lock output: {}", e))?
391 .clone();
392
393 if !has_alpha {
394 let mut src_idx = 0;
395 let mut dst_idx = 0;
396 for _ in 0..self.info.height {
397 for _ in 0..self.info.width {
398 output[dst_idx] = output[src_idx];
399 output[dst_idx + 1] = output[src_idx + 1];
400 output[dst_idx + 2] = output[src_idx + 2];
401 src_idx += 4;
402 dst_idx += 3;
403 }
404 }
405 output.truncate(dst_idx);
406 }
407
408 let color_type = if has_alpha {
409 ImageColorType::Bgra
410 } else {
411 ImageColorType::Bgr
412 };
413
414 let img = ImageData {
415 width: decoder.width as u32,
416 height: decoder.height as u32,
417 color_type,
418 depth: 8,
419 data: output,
420 };
421
422 if decoder.width != self.info.width as i32 || decoder.height != self.info.height as i32 {
423 return Ok(draw_on_canvas(
424 img,
425 self.info.width as u32,
426 self.info.height as u32,
427 0,
428 0,
429 )?);
430 }
431
432 Ok(img)
433 }
434
435 fn read_encoded(&mut self) -> Result<Vec<u8>> {
436 let mut output = Vec::with_capacity(self.info.enc_length as usize);
437 output.resize(self.info.enc_length as usize, 0);
438 self.stream.m_input.read_exact(&mut output)?;
439 let mut sum = 0u8;
440 let mut xor = 0u8;
441 for i in 0..output.len() {
442 output[i] = output[i].wrapping_sub(self.update_key());
443 sum = sum.wrapping_add(output[i]);
444 xor ^= output[i];
445 }
446 if sum != self.info.check_sum || xor != self.info.check_xor {
447 return Err(anyhow::anyhow!(
448 "Checksum mismatch: sum={}, xor={}",
449 sum,
450 xor
451 ));
452 }
453 Ok(output)
454 }
455
456 fn read_int(input: &mut MemReaderRef<'_>) -> Result<i32> {
457 let mut v = 0;
458 let mut code_length = 0;
459 loop {
460 let code = input.read_i8()?;
461 if code_length >= 32 {
462 return Err(anyhow::anyhow!(
463 "Failed to raed int: code={}, code_length={}",
464 code,
465 code_length
466 ));
467 }
468 v |= ((code & 0x7f) as i32) << code_length;
469 code_length += 7;
470 if code & -128 == 0 {
471 break;
472 }
473 }
474 Ok(v)
475 }
476
477 fn read_weight_table(input: &mut MemReaderRef<'_>, length: usize) -> Result<Vec<u32>> {
478 let mut weights = Vec::with_capacity(length);
479 for _ in 0..length {
480 let weight = Self::read_int(input)? as u32;
481 weights.push(weight);
482 }
483 Ok(weights)
484 }
485
486 fn huffman_decompress(&mut self, tree: &HuffmanTree, output: &mut [u8]) -> Result<()> {
487 for dst in 0..output.len() {
488 output[dst] = tree.decode_token(&mut self.stream)? as u8;
489 }
490 Ok(())
491 }
492
493 fn unpack_zeros(input: &[u8], output: &mut [u8]) {
494 let mut dst = 0;
495 let mut dec_zero = 0;
496 let mut src = 0;
497 while dst < output.len() {
498 let mut code_length = 0;
499 let mut count = 0;
500 let mut code;
501 loop {
502 if src >= input.len() {
503 return;
504 }
505 code = input[src];
506 src += 1;
507 count |= ((code & 0x7f) as usize) << code_length;
508 code_length += 7;
509 if code & 0x80 == 0 {
510 break;
511 }
512 }
513 if dst + count > output.len() {
514 break;
515 }
516 if dec_zero == 0 {
517 if src + count > input.len() {
518 break;
519 }
520 output[dst..dst + count].copy_from_slice(&input[src..src + count]);
521 src += count;
522 } else {
523 for i in 0..count {
524 output[dst + i] = 0;
525 }
526 }
527 dec_zero ^= 1;
528 dst += count;
529 }
530 }
531
532 fn reverse_average_sampling(&self, output: &mut [u8]) {
533 for y in 0..self.info.height {
534 let line = y as usize * self.stride;
535 for x in 0..self.info.width {
536 let pixel = line + x as usize * self.pixel_size as usize;
537 for p in 0..self.pixel_size {
538 let mut avg = 0u32;
539 if x > 0 {
540 avg = avg.wrapping_add(
541 output[pixel + p as usize - self.pixel_size as usize] as u32,
542 );
543 }
544 if y > 0 {
545 avg = avg.wrapping_add(output[pixel + p as usize - self.stride] as u32);
546 }
547 if x > 0 && y > 0 {
548 avg /= 2;
549 }
550 if avg != 0 {
551 output[pixel + p as usize] =
552 output[pixel + p as usize].wrapping_add(avg as u8);
553 }
554 }
555 }
556 }
557 }
558
559 fn update_key(&mut self) -> u8 {
560 let v0 = 20021 * (self.key & 0xffff);
561 let mut v1 = self.magic | (self.key >> 16);
562 v1 = v1
563 .overflowing_mul(20021)
564 .0
565 .overflowing_add(self.key.overflowing_mul(346).0)
566 .0;
567 v1 = (v1 + (v0 >> 16)) & 0xffff;
568 self.key = (v1 << 16) + (v0 & 0xffff) + 1;
569 v1 as u8
570 }
571}
572
573#[derive(Debug)]
574struct HuffmanNode {
575 valid: bool,
576 is_parent: bool,
577 weight: u32,
578 left_index: usize,
579 right_index: usize,
580}
581
582#[derive(Debug)]
583struct HuffmanTree {
584 nodes: Vec<HuffmanNode>,
585}
586
587impl HuffmanTree {
588 fn new(weights: &[u32], v2: bool) -> Self {
589 let mut nodes = Vec::with_capacity(weights.len() * 2);
590 let mut root_node_weight = 0u32;
591 for weight in weights {
592 let node = HuffmanNode {
593 valid: *weight != 0,
594 is_parent: false,
595 weight: *weight,
596 left_index: 0,
597 right_index: 0,
598 };
599 nodes.push(node);
600 root_node_weight = root_node_weight.wrapping_add(*weight);
601 }
602 let mut child_node_index = [0usize; 2];
603 loop {
604 let mut weight = 0u32;
605 for i in 0usize..2usize {
606 let mut min_weight = u32::MAX;
607 child_node_index[i] = usize::MAX;
608 let mut n = 0;
609 if v2 {
610 while n < nodes.len() {
611 if nodes[n].valid {
612 min_weight = nodes[n].weight;
613 child_node_index[i] = n;
614 n += 1;
615 break;
616 }
617 n += 1;
618 }
619 n = n.max(i + 1);
620 }
621 while n < nodes.len() {
622 if nodes[n].valid && nodes[n].weight < min_weight {
623 min_weight = nodes[n].weight;
624 child_node_index[i] = n;
625 }
626 n += 1;
627 }
628 if child_node_index[i] == usize::MAX {
629 continue;
630 }
631 nodes[child_node_index[i]].valid = false;
632 weight = weight.wrapping_add(nodes[child_node_index[i]].weight);
633 }
634 let parent_node = HuffmanNode {
635 valid: true,
636 is_parent: true,
637 left_index: child_node_index[0],
638 right_index: child_node_index[1],
639 weight,
640 };
641 nodes.push(parent_node);
642 if weight >= root_node_weight {
643 break;
644 }
645 }
646 Self { nodes }
647 }
648
649 fn decode_token(&self, stream: &mut MsbBitStream<MemReaderRef<'_>>) -> Result<usize> {
650 let mut node_index = self.nodes.len() - 1;
651 loop {
652 let bit = stream.get_next_bit()?;
653 if !bit {
654 node_index = self.nodes[node_index].left_index;
655 } else {
656 node_index = self.nodes[node_index].right_index;
657 }
658 if !self.nodes[node_index].is_parent {
659 return Ok(node_index);
660 }
661 }
662 }
663
664 fn encode_token(&self, stream: &mut MsbBitWriter<impl Write>, token: usize) -> Result<()> {
665 let mut path = Vec::new();
666 if !self.find_path(self.nodes.len() - 1, token, &mut path) {
667 return Err(anyhow::anyhow!("Token not found in Huffman tree"));
668 }
669 for &bit in path.iter().rev() {
670 stream.put_bit(bit)?;
671 }
672 Ok(())
673 }
674
675 fn find_path(&self, node_index: usize, token: usize, path: &mut Vec<bool>) -> bool {
676 if node_index == usize::MAX {
677 return false;
678 }
679 let node = &self.nodes[node_index];
680 if !node.is_parent {
681 return node_index == token;
682 }
683
684 if self.find_path(node.left_index, token, path) {
685 path.push(false);
686 return true;
687 }
688 if self.find_path(node.right_index, token, path) {
689 path.push(true);
690 return true;
691 }
692 false
693 }
694}
695
696const DCT_TABLE: [f32; 64] = [
697 1.00000000, 1.38703990, 1.30656302, 1.17587554, 1.00000000, 0.78569496, 0.54119611, 0.27589938,
698 1.38703990, 1.92387950, 1.81225491, 1.63098633, 1.38703990, 1.08979023, 0.75066054, 0.38268343,
699 1.30656302, 1.81225491, 1.70710683, 1.53635550, 1.30656302, 1.02655995, 0.70710677, 0.36047992,
700 1.17587554, 1.63098633, 1.53635550, 1.38268340, 1.17587554, 0.92387950, 0.63637930, 0.32442334,
701 1.00000000, 1.38703990, 1.30656302, 1.17587554, 1.00000000, 0.78569496, 0.54119611, 0.27589938,
702 0.78569496, 1.08979023, 1.02655995, 0.92387950, 0.78569496, 0.61731654, 0.42521504, 0.21677275,
703 0.54119611, 0.75066054, 0.70710677, 0.63637930, 0.54119611, 0.42521504, 0.29289323, 0.14931567,
704 0.27589938, 0.38268343, 0.36047992, 0.32442334, 0.27589938, 0.21677275, 0.14931567, 0.07612047,
705];
706
707const BLOCK_FILL_ORDER: [u8; 64] = [
708 0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27, 20,
709 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51, 58, 59,
710 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
711];
712
713struct ParallelCbgDecoder {
714 input: Vec<u8>,
715 output: Mutex<Vec<u8>>,
716 bpp: i32,
717 width: i32,
718 height: i32,
719 tree1: HuffmanTree,
720 tree2: HuffmanTree,
721 dct: [[f32; 64]; 2],
722 has_alpha: AtomicBool,
723}
724
725impl ParallelCbgDecoder {
726 fn unpack_block(&self, offset: i32, length: i32, dst: i32) -> Result<()> {
727 let input = MemReaderRef::new(&self.input[offset as usize..(offset + length) as usize]);
728 let mut reader = MsbBitStream::new(input);
729
730 let block_size = CbgDecoder::read_int(&mut reader.m_input)?;
731 if block_size == -1 {
732 return Ok(());
733 }
734
735 let mut color_data = vec![0i16; block_size as usize];
736 let mut acc = 0i32;
737 let mut i = 0i32;
738
739 while i < block_size && reader.m_input.pos < reader.m_input.data.len() {
740 let count = self.tree1.decode_token(&mut reader)?;
741 if count != 0 {
742 let mut v = reader.get_bits(count as u32)? as i32;
743 if (v >> (count - 1)) == 0 {
744 v = (-1 << count | v) + 1;
745 }
746 acc += v;
747 }
748 color_data[i as usize] = acc as i16;
749 i += 64;
750 }
751
752 if (reader.m_cached_bits & 7) != 0 {
753 reader.get_bits(reader.m_cached_bits & 7)?;
754 }
755
756 i = 0;
757 while i < block_size && reader.m_input.pos < reader.m_input.data.len() {
758 let mut index = 1usize;
759 while index < 64 && reader.m_input.pos < reader.m_input.data.len() {
760 let code = self.tree2.decode_token(&mut reader)?;
761 if code == 0 {
762 break;
763 }
764 if code == 0xf {
765 index += 0x10;
766 continue;
767 }
768 index += code & 0xf;
769 if index >= BLOCK_FILL_ORDER.len() {
770 break;
771 }
772 let bits = code >> 4;
773 let mut v = reader.get_bits(bits as u32)? as i32;
774 if bits != 0 && (v >> (bits - 1)) == 0 {
775 v = (-1 << bits | v) + 1;
776 }
777 color_data[i as usize + BLOCK_FILL_ORDER[index] as usize] = v as i16;
778 index += 1;
779 }
780 i += 64;
781 }
782
783 if self.bpp == 8 {
784 self.decode_grayscale(&color_data, dst)?;
785 } else {
786 self.decode_rgb(&color_data, dst)?;
787 }
788
789 Ok(())
790 }
791
792 fn decode_rgb(&self, data: &[i16], dst: i32) -> Result<()> {
793 let block_count = self.width / 8;
794 let mut dst = dst as usize;
795
796 for i in 0..block_count {
797 let mut src = (i * 64) as usize;
798 let mut ycbcr_block = [[0i16; 3]; 64];
799
800 for channel in 0..3 {
801 self.decode_dct(channel, data, src, &mut ycbcr_block)?;
802 src += (self.width * 8) as usize;
803 }
804
805 let mut output = self
806 .output
807 .lock()
808 .map_err(|e| anyhow::anyhow!("Failed to lock output: {}", e))?;
809
810 for j in 0..64 {
811 let cy = ycbcr_block[j][0] as f32;
812 let cb = ycbcr_block[j][1] as f32;
813 let cr = ycbcr_block[j][2] as f32;
814
815 let r = cy + 1.402f32 * cr - 178.956f32;
817 let g = cy - 0.34414f32 * cb - 0.71414f32 * cr + 135.95984f32;
818 let b = cy + 1.772f32 * cb - 226.316f32;
819
820 let y = j >> 3;
821 let x = j & 7;
822 let p = (y * self.width as usize + x) * 4 + dst;
823
824 output[p] = Self::float_to_byte(b);
825 output[p + 1] = Self::float_to_byte(g);
826 output[p + 2] = Self::float_to_byte(r);
827 }
828 dst += 32;
829 }
830 Ok(())
831 }
832
833 fn decode_grayscale(&self, data: &[i16], dst: i32) -> Result<()> {
834 let mut src = 0;
835 let block_count = self.width / 8;
836 let mut dst = dst as usize;
837
838 for _ in 0..block_count {
839 let mut ycbcr_block = [[0i16; 3]; 64];
840 self.decode_dct(0, data, src, &mut ycbcr_block)?;
841 src += 64;
842
843 let mut output = self
844 .output
845 .lock()
846 .map_err(|e| anyhow::anyhow!("Failed to lock output: {}", e))?;
847
848 for j in 0..64 {
849 let y = j >> 3;
850 let x = j & 7;
851 let p = (y * self.width as usize + x) * 4 + dst;
852 let value = ycbcr_block[j][0] as u8;
853
854 output[p] = value;
855 output[p + 1] = value;
856 output[p + 2] = value;
857 }
858 dst += 32;
859 }
860 Ok(())
861 }
862
863 fn unpack_alpha(&self, offset: i32) -> Result<()> {
864 let mut input = MemReaderRef::new(&self.input[offset as usize..]);
865
866 if input.read_i32()? != 1 {
867 return Ok(());
868 }
869
870 let mut dst = 3;
871 let mut ctl = 1i32 << 1;
872
873 let mut output = self
874 .output
875 .lock()
876 .map_err(|e| anyhow::anyhow!("Failed to lock output: {}", e))?;
877
878 while dst < output.len() {
879 ctl >>= 1;
880 if ctl == 1 {
881 ctl = (input.read_u8()? as i32) | 0x100;
882 }
883
884 if (ctl & 1) != 0 {
885 let v = input.read_u16()? as i32;
886 let mut x = v & 0x3f;
887 if x > 0x1f {
888 x |= -0x40;
889 }
890 let mut y = (v >> 6) & 7;
891 if y != 0 {
892 y |= -8;
893 }
894 let count = ((v >> 9) & 0x7f) + 3;
895
896 let src = dst as isize + (x as isize + y as isize * self.width as isize) * 4;
897 if src < 0 || src >= dst as isize {
898 return Ok(());
899 }
900
901 let mut src = src as usize;
902 for _ in 0..count {
903 output[dst] = output[src];
904 src += 4;
905 dst += 4;
906 }
907 } else {
908 output[dst] = input.read_u8()?;
909 dst += 4;
910 }
911 }
912
913 self.has_alpha.qsave(true);
914 Ok(())
915 }
916
917 fn decode_dct(
918 &self,
919 channel: usize,
920 data: &[i16],
921 src: usize,
922 output: &mut [[i16; 3]; 64],
923 ) -> Result<()> {
924 let d = if channel > 0 { 1 } else { 0 };
925 let mut tmp = [[0f32; 8]; 8];
926
927 for i in 0..8 {
928 if data[src + 8 + i] == 0
930 && data[src + 16 + i] == 0
931 && data[src + 24 + i] == 0
932 && data[src + 32 + i] == 0
933 && data[src + 40 + i] == 0
934 && data[src + 48 + i] == 0
935 && data[src + 56 + i] == 0
936 {
937 let t = data[src + i] as f32 * self.dct[d][i];
938 for row in 0..8 {
939 tmp[row][i] = t;
940 }
941 continue;
942 }
943
944 let v1 = data[src + i] as f32 * self.dct[d][i];
945 let v2 = data[src + 8 + i] as f32 * self.dct[d][8 + i];
946 let v3 = data[src + 16 + i] as f32 * self.dct[d][16 + i];
947 let v4 = data[src + 24 + i] as f32 * self.dct[d][24 + i];
948 let v5 = data[src + 32 + i] as f32 * self.dct[d][32 + i];
949 let v6 = data[src + 40 + i] as f32 * self.dct[d][40 + i];
950 let v7 = data[src + 48 + i] as f32 * self.dct[d][48 + i];
951 let v8 = data[src + 56 + i] as f32 * self.dct[d][56 + i];
952
953 let v10 = v1 + v5;
954 let v11 = v1 - v5;
955 let v12 = v3 + v7;
956 let v13 = (v3 - v7) * 1.414213562f32 - v12;
957 let v1 = v10 + v12;
958 let v7 = v10 - v12;
959 let v3 = v11 + v13;
960 let v5 = v11 - v13;
961 let v14 = v2 + v8;
962 let v15 = v2 - v8;
963 let v16 = v6 + v4;
964 let v17 = v6 - v4;
965 let v8 = v14 + v16;
966 let v11 = (v14 - v16) * 1.414213562f32;
967 let v9 = (v17 + v15) * 1.847759065f32;
968 let v10 = 1.082392200f32 * v15 - v9;
969 let v13 = -2.613125930f32 * v17 + v9;
970 let v6 = v13 - v8;
971 let v4 = v11 - v6;
972 let v2 = v10 + v4;
973
974 tmp[0][i] = v1 + v8;
975 tmp[1][i] = v3 + v6;
976 tmp[2][i] = v5 + v4;
977 tmp[3][i] = v7 - v2;
978 tmp[4][i] = v7 + v2;
979 tmp[5][i] = v5 - v4;
980 tmp[6][i] = v3 - v6;
981 tmp[7][i] = v1 - v8;
982 }
983
984 let mut dst = 0;
985 for i in 0..8 {
986 let v10 = tmp[i][0] + tmp[i][4];
987 let v11 = tmp[i][0] - tmp[i][4];
988 let v12 = tmp[i][2] + tmp[i][6];
989 let v13 = tmp[i][2] - tmp[i][6];
990 let v14 = tmp[i][1] + tmp[i][7];
991 let v15 = tmp[i][1] - tmp[i][7];
992 let v16 = tmp[i][5] + tmp[i][3];
993 let v17 = tmp[i][5] - tmp[i][3];
994
995 let v13 = 1.414213562f32 * v13 - v12;
996 let v1 = v10 + v12;
997 let v7 = v10 - v12;
998 let v3 = v11 + v13;
999 let v5 = v11 - v13;
1000 let v8 = v14 + v16;
1001 let v11 = (v14 - v16) * 1.414213562f32;
1002 let v9 = (v17 + v15) * 1.847759065f32;
1003 let v10 = v9 - v15 * 1.082392200f32;
1004 let v13 = v9 - v17 * 2.613125930f32;
1005 let v6 = v13 - v8;
1006 let v4 = v11 - v6;
1007 let v2 = v10 - v4;
1008
1009 output[dst][channel] = Self::float_to_short(v1 + v8);
1010 output[dst + 1][channel] = Self::float_to_short(v3 + v6);
1011 output[dst + 2][channel] = Self::float_to_short(v5 + v4);
1012 output[dst + 3][channel] = Self::float_to_short(v7 + v2);
1013 output[dst + 4][channel] = Self::float_to_short(v7 - v2);
1014 output[dst + 5][channel] = Self::float_to_short(v5 - v4);
1015 output[dst + 6][channel] = Self::float_to_short(v3 - v6);
1016 output[dst + 7][channel] = Self::float_to_short(v1 - v8);
1017 dst += 8;
1018 }
1019
1020 Ok(())
1021 }
1022
1023 fn float_to_short(f: f32) -> i16 {
1024 let a = 0x80 + ((f as i32) >> 3);
1025 if a <= 0 {
1026 0
1027 } else if a <= 0xff {
1028 a as i16
1029 } else if a < 0x180 {
1030 0xff
1031 } else {
1032 0
1033 }
1034 }
1035
1036 fn float_to_byte(f: f32) -> u8 {
1037 if f >= 255.0 {
1038 0xff
1039 } else if f <= 0.0 {
1040 0
1041 } else {
1042 f as u8
1043 }
1044 }
1045}
1046
1047struct CbgEncoder {
1048 header: BgiCBGHeader,
1049 stream: MemWriter,
1050 img: ImageData,
1051 key: u32,
1052 magic: u32,
1053}
1054
1055impl CbgEncoder {
1056 pub fn new(mut img: ImageData) -> Result<Self> {
1057 if img.depth != 8 {
1058 return Err(anyhow::anyhow!("Unsupported image depth: {}", img.depth));
1059 }
1060 let bpp = match img.color_type {
1061 ImageColorType::Bgr => 24,
1062 ImageColorType::Bgra => 32,
1063 ImageColorType::Grayscale => 8,
1064 ImageColorType::Rgb => {
1065 convert_rgb_to_bgr(&mut img)?;
1066 24
1067 }
1068 ImageColorType::Rgba => {
1069 convert_rgba_to_bgra(&mut img)?;
1070 32
1071 }
1072 };
1073 let key = rand::random();
1074 let header = BgiCBGHeader {
1075 width: img.width as u16,
1076 height: img.height as u16,
1077 bpp,
1078 _unk: 0,
1079 intermediate_length: 0,
1080 key,
1081 enc_length: 0,
1082 check_sum: 0,
1083 check_xor: 0,
1084 version: 1,
1085 };
1086
1087 Ok(CbgEncoder {
1088 header,
1089 stream: MemWriter::new(),
1090 img,
1091 key,
1092 magic: 0,
1093 })
1094 }
1095
1096 pub fn encode(mut self) -> Result<Vec<u8>> {
1097 self.stream.write_all(b"CompressedBG___\0")?;
1098 let header_pos = self.stream.pos;
1099 self.stream.seek(std::io::SeekFrom::Current(0x20))?;
1100
1101 let pixel_size = (self.header.bpp / 8) as usize;
1102 let stride = self.header.width as usize * pixel_size;
1103 let mut sampled_data = self.img.data.clone();
1104 self.average_sampling(&mut sampled_data, stride, pixel_size);
1105
1106 let packed_data = Self::pack_zeros(&sampled_data);
1107 self.header.intermediate_length = packed_data.len() as u32;
1108
1109 let mut frequencies = vec![0u32; 256];
1110 for &byte in &packed_data {
1111 frequencies[byte as usize] += 1;
1112 }
1113 if frequencies.iter().all(|&f| f == 0) {
1114 frequencies[0] = 1;
1115 }
1116
1117 let tree = HuffmanTree::new(&frequencies, false);
1118
1119 let mut weight_writer = MemWriter::new();
1120 for &weight in &frequencies {
1121 Self::write_int(&mut weight_writer, weight as i32)?;
1122 }
1123 let weight_data = weight_writer.into_inner();
1124 self.write_encoded(&weight_data)?;
1125
1126 let mut bit_writer = MsbBitWriter::new(&mut self.stream);
1127 for &byte in &packed_data {
1128 tree.encode_token(&mut bit_writer, byte as usize)?;
1129 }
1130 bit_writer.flush()?;
1131
1132 let final_pos = self.stream.pos;
1133 self.stream.pos = header_pos;
1134 self.header
1135 .pack(&mut self.stream, false, Encoding::Cp932, &None)?;
1136 self.stream.pos = final_pos;
1137
1138 Ok(self.stream.into_inner())
1139 }
1140
1141 fn average_sampling(&self, data: &mut [u8], stride: usize, pixel_size: usize) {
1142 for y in (0..self.header.height as usize).rev() {
1143 let line = y * stride;
1144 for x in (0..self.header.width as usize).rev() {
1145 let pixel = line + x * pixel_size;
1146 for p in 0..pixel_size {
1147 let mut avg = 0u32;
1148 let mut count = 0;
1149 if x > 0 {
1150 avg = avg.wrapping_add(data[pixel + p - pixel_size] as u32);
1151 count += 1;
1152 }
1153 if y > 0 {
1154 avg = avg.wrapping_add(data[pixel + p - stride] as u32);
1155 count += 1;
1156 }
1157 if count > 0 {
1158 avg /= count;
1159 }
1160 if avg != 0 {
1161 data[pixel + p] = data[pixel + p].wrapping_sub(avg as u8);
1162 }
1163 }
1164 }
1165 }
1166 }
1167
1168 fn pack_zeros(input: &[u8]) -> Vec<u8> {
1169 let mut output = Vec::new();
1170 let mut i = 0;
1171 let mut is_zero_run = false;
1172
1173 while i < input.len() {
1174 let mut count = 0;
1175 if is_zero_run {
1176 while i + count < input.len() && input[i + count] == 0 {
1177 count += 1;
1178 }
1179 } else {
1180 while i + count < input.len() && input[i + count] != 0 {
1181 count += 1;
1182 }
1183 }
1184
1185 let mut count_buf = Vec::new();
1186 let mut n = count;
1187 loop {
1188 let mut byte = (n & 0x7f) as u8;
1189 n >>= 7;
1190 if n > 0 {
1191 byte |= 0x80;
1192 }
1193 count_buf.push(byte);
1194 if n == 0 {
1195 break;
1196 }
1197 }
1198 output.extend_from_slice(&count_buf);
1199
1200 if !is_zero_run {
1201 output.extend_from_slice(&input[i..i + count]);
1202 }
1203 i += count;
1204 is_zero_run = !is_zero_run;
1205 }
1206 output
1207 }
1208
1209 fn write_int<W: Write>(writer: &mut W, mut value: i32) -> Result<()> {
1210 loop {
1211 let mut b = (value as u8) & 0x7f;
1212 value >>= 7;
1213 if value != 0 {
1214 b |= 0x80;
1215 }
1216 writer.write_u8(b)?;
1217 if value == 0 {
1218 break;
1219 }
1220 }
1221 Ok(())
1222 }
1223
1224 fn write_encoded(&mut self, data: &[u8]) -> Result<()> {
1225 self.header.enc_length = data.len() as u32;
1226 let mut sum = 0u8;
1227 let mut xor = 0u8;
1228 let mut encoded_data = Vec::with_capacity(data.len());
1229 for &byte in data {
1230 let encrypted_byte = byte.wrapping_add(self.update_key());
1231 sum = sum.wrapping_add(byte);
1232 xor ^= byte;
1233 encoded_data.push(encrypted_byte);
1234 }
1235 self.header.check_sum = sum;
1236 self.header.check_xor = xor;
1237 self.stream.write_all(&encoded_data)?;
1238 Ok(())
1239 }
1240
1241 fn update_key(&mut self) -> u8 {
1242 let v0 = 20021 * (self.key & 0xffff);
1243 let mut v1 = self.magic | (self.key >> 16);
1244 v1 = v1
1245 .overflowing_mul(20021)
1246 .0
1247 .overflowing_add(self.key.overflowing_mul(346).0)
1248 .0;
1249 v1 = (v1 + (v0 >> 16)) & 0xffff;
1250 self.key = (v1 << 16) + (v0 & 0xffff) + 1;
1251 v1 as u8
1252 }
1253}