6use bincode::{Decode, Encode};
7use jgenesis_common::num::{GetBit, U16Ext, U24Ext};
8use std::{array, mem};

10#[rustfmt::skip]
11const EVOLUTION_PROBABILITY: &[u8; 53] = &[
12    90,37,17, 8, 3, 1,90,63,44,32,23,17,12, 9, 7, 5, 4, 3, 2,
13    90,72,58,46,38,31,25,21,17,14,11, 9, 8, 7, 5, 4, 4, 3, 2,
14    2 ,88,77,67,59,52,46,41,37,86,79,71,65,60,55
15];
16
17#[rustfmt::skip]
18const EVOLUTION_NEXT_LPS: &[u8; 53] = &[
19    1 , 6, 8,10,12,15, 7,19,21,22,23,25,26,28,29,31,32,34,35,
20    20,39,40,42,44,45,46,25,26,26,27,28,29,30,31,33,33,34,35,
21    36,39,47,48,49,50,51,44,45,47,47,48,49,50,51
22];
23
24#[rustfmt::skip]
25const EVOLUTION_NEXT_MPS: &[u8; 53] = &[
26    1 , 2, 3, 4, 5, 5, 7, 8, 9,10,11,12,13,14,15,16,17,18, 5,
27    20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,
28    5 ,40,41,42,43,44,45,46,24,48,49,50,51,52,43
29];

32const MODE_2_CONTEXT_TABLE: &[u8; 15] = &[0, 0, 0, 15, 17, 19, 21, 23, 25, 25, 25, 25, 25, 27, 29];

34#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Encode, Decode)]
35enum DecompressionMode {
36    // Miscellaneous data
37    #[default]
38    Zero,
39    // 2bpp graphical data
40    One,
41    // 4bpp graphical data
42    Two,
43}
44
45impl DecompressionMode {
46    fn bpp(self) -> u32 {
47        match self {
48            Self::Zero => 1,
49            Self::One => 2,
50            Self::Two => 3,
51        }
52    }
53}
54
55impl DecompressionMode {
56    fn from_byte(byte: u8) -> Self {
57        match byte {
58            0x00 => Self::Zero,
59            0x01 => Self::One,
60            0x02 => Self::Two,
61            _ => {
62                log::warn!("Unexpected SPC7110 decompression mode, defaulting to 0: {byte:02X}");
63                Self::Zero
64            }
65        }
66    }
67}
68
69#[derive(Debug, Clone, Default, Encode, Decode)]
70struct DecompressionState {
71    initialized: bool,
72    mode: DecompressionMode,
73    source: u32,
74    out: u32,
75    decoded: u8,
76    in_count: u8,
77    buffer_index: u8,
78    a: u8,
79    b: u8,
80    c: u8,
81    context: u8,
82    top: u8,
83    input: u16,
84    plane1: u8,
85    plane_buffer: [u8; 16],
86    pixel_order: [u8; 16],
87    real_order: [u8; 16],
88    context_index: [u8; 32],
89    context_invert: [u8; 32],
90}
91
92impl DecompressionState {
93    fn input_msb(&self) -> u8 {
94        (self.input >> 8) as u8
95    }
96}
97
98#[derive(Debug, Clone, Default, Encode, Decode)]
99pub struct Spc7110Decompressor {
100    pub rom_directory_base: u32,
101    pub rom_directory_index: u8,
102    pub target_offset: u16,
103    pub length_counter: u16,
104    pub skip_enabled: bool,
105    state: DecompressionState,
106}
107
108impl Spc7110Decompressor {
109    pub fn new() -> Self {
110        Self::default()
111    }
112
113    pub fn write_rom_directory_base_low(&mut self, value: u8) {
114        self.rom_directory_base.set_low_byte(value);
115    }
116
117    pub fn write_rom_directory_base_mid(&mut self, value: u8) {
118        self.rom_directory_base.set_mid_byte(value);
119    }
120
121    pub fn write_rom_directory_base_high(&mut self, value: u8) {
122        self.rom_directory_base.set_high_byte(value);
123    }
124
125    pub fn write_target_offset_low(&mut self, value: u8) {
126        self.target_offset.set_lsb(value);
127    }
128
129    pub fn write_target_offset_high(&mut self, value: u8, data_rom: &[u8]) {
130        self.target_offset.set_msb(value);
131
132        // Writing offset MSB initializes decompressor
133        self.initialize(data_rom);
134    }
135
136    pub fn write_length_counter_low(&mut self, value: u8) {
137        self.length_counter.set_lsb(value);
138    }
139
140    pub fn write_length_counter_high(&mut self, value: u8) {
141        self.length_counter.set_msb(value);
142    }
143
144    pub fn write_mode(&mut self, value: u8) {
145        // Supposedly, $02 causes the decompressor to skip <offset> rows of pixels after initialization,
146        // while $00 causes it to ignore the target offset
147        self.skip_enabled = value == 0x02;
148    }
149
150    pub fn read_mode(&self) -> u8 {
151        if self.skip_enabled { 0x02 } else { 0x00 }
152    }
153
154    pub fn read_status(&self) -> u8 {
155        u8::from(self.state.initialized) << 7
156    }
157
158    pub fn next_byte(&mut self, data_rom: &[u8]) -> u8 {
159        if !self.state.initialized {
160            return 0;
161        }
162
163        self.length_counter = self.length_counter.wrapping_sub(1);
164
165        match self.state.mode {
166            DecompressionMode::Zero => self.next_byte_mode_0(data_rom),
167            DecompressionMode::One => self.next_byte_mode_1(data_rom),
168            DecompressionMode::Two => self.next_byte_mode_2(data_rom),
169        }
170    }
171
172    fn next_byte_mode_0(&mut self, data_rom: &[u8]) -> u8 {
173        // Decompress 8 bits and output them
174        self.state.decoded = 0;
175        for ctx_offset in [0, 1, 3, 7] {
176            self.state.context = ctx_offset + self.state.decoded;
177            self.decompress_bit(data_rom);
178        }
179
180        self.state.out = (self.state.out << 4)
181            ^ (((self.state.out >> 12) ^ u32::from(self.state.decoded)) & 0xF);
182
183        self.state.decoded = 0;
184        for ctx_offset in [0, 1, 3, 7] {
185            self.state.context = 15 + ctx_offset + self.state.decoded;
186            self.decompress_bit(data_rom);
187        }
188
189        self.state.out = (self.state.out << 4)
190            ^ (((self.state.out >> 12) ^ u32::from(self.state.decoded)) & 0xF);
191
192        self.state.out as u8
193    }
194
195    fn next_byte_mode_1(&mut self, data_rom: &[u8]) -> u8 {
196        let byte = if !self.state.buffer_index.bit(0) {
197            // Decompress the next 16 bits
198            for _ in 0..8 {
199                self.state.a = ((self.state.out >> 2) & 0x03) as u8;
200                self.state.b = ((self.state.out >> 14) & 0x03) as u8;
201
202                self.state.decoded = 0;
203                self.state.context = get_context(self.state.a, self.state.b, self.state.c);
204                self.decompress_bit(data_rom);
205
206                self.state.context = 2 * self.state.context + 5 + self.state.decoded;
207                self.decompress_bit(data_rom);
208
209                self.adjust_pixel_order(2);
210            }
211
212            // Deinterleave into 2bpp bitplanes and return bitplane 0
213            let (plane1, plane0) = deinterleave_bits(self.state.out);
214            self.state.plane1 = plane1 as u8;
215
216            plane0 as u8
217        } else {
218            // Return bitplane 1 from the last call
219            self.state.plane1
220        };
221
222        self.state.buffer_index = self.state.buffer_index.wrapping_add(1);
223
224        byte
225    }
226
227    fn next_byte_mode_2(&mut self, data_rom: &[u8]) -> u8 {
228        let byte = if self.state.buffer_index & 0x11 == 0 {
229            // Decompress the next 32 bits
230            for _ in 0..8 {
231                self.state.a = (self.state.out & 0xF) as u8;
232                self.state.b = ((self.state.out >> 28) & 0xF) as u8;
233
234                self.state.decoded = 0;
235                self.state.context = 0;
236                self.decompress_bit(data_rom);
237
238                self.state.context = self.state.decoded + 1;
239                self.decompress_bit(data_rom);
240
241                self.state.context = if self.state.context == 2 {
242                    self.state.decoded + 11
243                } else {
244                    get_context(self.state.a, self.state.b, self.state.c)
245                        + 3
246                        + 5 * self.state.decoded
247                };
248                self.decompress_bit(data_rom);
249
250                self.state.context =
251                    MODE_2_CONTEXT_TABLE[self.state.context as usize] + (self.state.decoded & 0x01);
252                self.decompress_bit(data_rom);
253
254                self.adjust_pixel_order(4);
255            }
256
257            // Deinterleave into 4bpp bitplanes and return bitplane 0
258            // This is designed for SNES tile data, so $00-$0F alternate between planes 0 and 1
259            // and $10-$1F alternate between planes 2 and 3
260            let (even_bits, odd_bits) = deinterleave_bits(self.state.out);
261            let (plane2, plane0) = deinterleave_bits(odd_bits.into());
262            let (plane3, plane1) = deinterleave_bits(even_bits.into());
263            self.state.plane1 = plane1 as u8;
264            self.state.plane_buffer[(self.state.buffer_index & 0xF) as usize] = plane2 as u8;
265            self.state.plane_buffer[((self.state.buffer_index + 1) & 0xF) as usize] = plane3 as u8;
266
267            plane0 as u8
268        } else if self.state.buffer_index & 0x10 == 0 {
269            // Return bitplane 1 from the last call
270            self.state.plane1
271        } else {
272            // Return bitplane 2 or 3 from the call 16-17 bytes ago
273            self.state.plane_buffer[(self.state.buffer_index & 0xF) as usize]
274        };
275
276        self.state.buffer_index = self.state.buffer_index.wrapping_add(1);
277
278        byte
279    }
280
281    fn initialize(&mut self, data_rom: &[u8]) {
282        self.state.initialized = true;
283
284        // Directory entries are 4 bytes
285        // Byte 0 contains the decompression mode (0/1/2)
286        // Bytes 1-3 contain the data ROM address, in big endian (unlike everything else in this chip)
287        let directory_addr = self.rom_directory_base + 4 * u32::from(self.rom_directory_index);
288        self.state.mode = DecompressionMode::from_byte(rom_get(data_rom, directory_addr));
289        self.state.source = u32::from_be_bytes([
290            0,
291            rom_get(data_rom, directory_addr + 1),
292            rom_get(data_rom, directory_addr + 2),
293            rom_get(data_rom, directory_addr + 3),
294        ]);
295
296        self.state.buffer_index = 0;
297        self.state.out = 0;
298        self.state.top = 255;
299        self.state.c = 0;
300
301        let input_msb = rom_get(data_rom, self.state.source);
302        self.state.input = u16::from_be_bytes([input_msb, 0]);
303        self.state.source += 1;
304        self.state.in_count = 0;
305
306        self.state.pixel_order = array::from_fn(|i| i as u8);
307        self.state.context_index.fill(0);
308        self.state.context_invert.fill(0);
309
310        if self.skip_enabled {
311            // Skip the next N rows of pixels, where N = target offset
312            // Not sure this is right, but not multiplying by bpp causes graphical glitches in
313            // Super Power League 4
314            let skip_bytes = self.state.mode.bpp() * u32::from(self.target_offset);
315            for _ in 0..skip_bytes {
316                self.next_byte(data_rom);
317            }
318
319            self.target_offset = 0;
320        }
321    }
322
323    fn decompress_bit(&mut self, data_rom: &[u8]) {
324        let context = self.state.context as usize;
325
326        self.state.decoded = (self.state.decoded << 1) | self.state.context_invert[context];
327
328        let evolution = self.state.context_index[context] as usize;
329        self.state.top -= EVOLUTION_PROBABILITY[evolution];
330
331        if self.state.input_msb() > self.state.top {
332            // Output LPS, and possibly swap LPS and MPS for this context
333            let input_msb = self.state.input_msb() - 1 - self.state.top;
334            self.state.input = (self.state.input & 0x00FF) | (u16::from(input_msb) << 8);
335
336            self.state.top = EVOLUTION_PROBABILITY[evolution] - 1;
337            if self.state.top > 79 {
338                self.state.context_invert[context] ^= 1;
339            }
340
341            self.state.decoded ^= 1;
342
343            self.state.context_index[context] = EVOLUTION_NEXT_LPS[evolution];
344        } else {
345            // Output MPS
346            if self.state.top <= 126 {
347                self.state.context_index[context] = EVOLUTION_NEXT_MPS[evolution];
348            }
349        }
350
351        while self.state.top <= 126 {
352            if self.state.in_count == 0 {
353                let input_lsb = rom_get(data_rom, self.state.source);
354                self.state.input = (self.state.input & 0xFF00) | u16::from(input_lsb);
355                self.state.source += 1;
356                self.state.in_count = 8;
357            }
358
359            self.state.top = (self.state.top << 1) | 1;
360            self.state.input <<= 1;
361            self.state.in_count -= 1;
362        }
363    }
364
365    fn adjust_pixel_order(&mut self, bpp: u8) {
366        let mut x = self.state.a;
367        for m in 0.. {
368            mem::swap(&mut x, &mut self.state.pixel_order[m]);
369            if x == self.state.a {
370                break;
371            }
372        }
373
374        for m in 0..1 << bpp {
375            self.state.real_order[m] = self.state.pixel_order[m];
376        }
377
378        x = self.state.c;
379        for m in 0.. {
380            mem::swap(&mut x, &mut self.state.real_order[m]);
381            if x == self.state.c {
382                break;
383            }
384        }
385
386        x = self.state.b;
387        for m in 0.. {
388            mem::swap(&mut x, &mut self.state.real_order[m]);
389            if x == self.state.b {
390                break;
391            }
392        }
393
394        x = self.state.a;
395        for m in 0.. {
396            mem::swap(&mut x, &mut self.state.real_order[m]);
397            if x == self.state.a {
398                break;
399            }
400        }
401
402        self.state.out =
403            (self.state.out << bpp) + u32::from(self.state.real_order[self.state.decoded as usize]);
404
405        self.state.c = self.state.b;
406    }
407}
408
409fn rom_get(data_rom: &[u8], address: u32) -> u8 {
410    data_rom.get(address as usize).copied().unwrap_or(0)
411}
412
413fn get_context(a: u8, b: u8, c: u8) -> u8 {
414    if a == b && b == c {
415        0
416    } else if a == b {
417        1
418    } else if b == c {
419        2
420    } else if a == c {
421        3
422    } else {
423        4
424    }
425}

430fn deinterleave_bits(n: u32) -> (u16, u16) {
431    let mut n: u64 = n.into();
432
433    n = (n & 0x0000_0000_5555_5555) | ((n << 31) & 0x5555_5555_0000_0000);
434    n = (n | (n >> 1)) & 0x3333_3333_3333_3333;
435    n = (n | (n >> 2)) & 0x0F0F_0F0F_0F0F_0F0F;
436    n = (n | (n >> 4)) & 0x00FF_00FF_00FF_00FF;
437    n = (n | (n >> 8)) & 0x0000_FFFF_0000_FFFF;
438
439    (n as u16, (n >> 32) as u16)
440}

442#[cfg(test)]
443mod tests {
444    use super::*;
445
446    #[test]
447    fn deinterleave() {
448        let (even, odd) = deinterleave_bits(0xFFFFFFFF);
449        assert_eq!(even, 0xFFFF);
450        assert_eq!(odd, 0xFFFF);
451
452        let (even, odd) = deinterleave_bits(0x5555AAAA);
453        assert_eq!(even, 0xFF00);
454        assert_eq!(odd, 0x00FF);
455
456        let (even, odd) = deinterleave_bits(0x12345678);
457        assert_eq!(even, 0x46EC);
458        assert_eq!(odd, 0x1416);
459    }
460
461    const DECOMPRESSED: &[u8] =
462        "Test123.ABCDABCDAAAAAAAAaaaabbbbccccdddd7654321076543210.Test123".as_bytes();
463
464    const MODE_0_COMPRESSED: &[u8; 45] = &[
465        0x68, 0x91, 0x36, 0x15, 0xF8, 0xBF, 0x42, 0x35, 0x2F, 0x67, 0x3D, 0xB7, 0xAA, 0x05, 0xB4,
466        0xF7, 0x70, 0x7A, 0x26, 0x20, 0xEA, 0x58, 0x2C, 0x09, 0x61, 0x00, 0xC5, 0x00, 0x8C, 0x6F,
467        0xFF, 0xD1, 0x42, 0x9D, 0xEE, 0x7F, 0x72, 0x87, 0xDF, 0xD6, 0x5F, 0x92, 0x65, 0x00, 0x00,
468    ];
469
470    const MODE_1_COMPRESSED: &[u8; 47] = &[
471        0x4B, 0xF6, 0x80, 0x1E, 0x3A, 0x4C, 0x42, 0x6C, 0xDA, 0x16, 0x0F, 0xC6, 0x44, 0xED, 0x64,
472        0x10, 0x77, 0xAF, 0x50, 0x00, 0x05, 0xC0, 0x01, 0x27, 0x22, 0xB0, 0x83, 0x51, 0x05, 0x32,
473        0x4A, 0x1E, 0x74, 0x93, 0x08, 0x76, 0x07, 0xE5, 0x32, 0x12, 0xB4, 0x99, 0x9E, 0x55, 0xA3,
474        0xF8, 0x00,
475    ];
476
477    const MODE_2_COMPRESSED: &[u8; 52] = &[
478        0x13, 0xB3, 0x27, 0xA6, 0xF4, 0x5C, 0xD8, 0xED, 0x6C, 0x6D, 0xF8, 0x76, 0x80, 0xA7, 0x87,
479        0x20, 0x39, 0x4B, 0x37, 0x1A, 0xCC, 0x3F, 0xE4, 0x3D, 0xBE, 0x65, 0x2D, 0x89, 0x7E, 0x0B,
480        0x0A, 0xD3, 0x46, 0xD5, 0x0C, 0x1F, 0xD3, 0x81, 0xF3, 0xAD, 0xDD, 0xE8, 0x5C, 0xC0, 0xBD,
481        0x62, 0xAA, 0xCB, 0xF8, 0xB5, 0x38, 0x00,
482    ];
483
484    fn perform_decompression(rom: &[u8]) -> Vec<u8> {
485        let mut decompressor = Spc7110Decompressor::new();
486        decompressor.rom_directory_base = 0;
487        decompressor.rom_directory_index = 0x50 / 4;
488        decompressor.length_counter = 64;
489        decompressor.target_offset = 0;
490
491        let mut decompressed = Vec::with_capacity(64);
492        decompressor.initialize(rom);
493        for _ in 0..64 {
494            decompressed.push(decompressor.next_byte(rom));
495        }
496
497        assert_eq!(decompressor.length_counter, 0);
498
499        decompressed
500    }
501
502    #[test]
503    fn mode_0() {
504        let mut rom = vec![0_u8; 0x100];
505        rom[..MODE_0_COMPRESSED.len()].copy_from_slice(MODE_0_COMPRESSED);
506        rom[0x50..0x54].copy_from_slice(&[0x00, 0x00, 0x00, 0x00]);
507
508        let decompressed = perform_decompression(&rom);
509
510        assert_eq!(decompressed.as_slice(), DECOMPRESSED);
511    }
512
513    #[test]
514    fn mode_1() {
515        let mut rom = vec![0_u8; 0x100];
516        rom[..MODE_1_COMPRESSED.len()].copy_from_slice(MODE_1_COMPRESSED);
517        rom[0x50..0x54].copy_from_slice(&[0x01, 0x00, 0x00, 0x00]);
518
519        let decompressed = perform_decompression(&rom);
520
521        assert_eq!(decompressed.as_slice(), DECOMPRESSED);
522    }
523
524    #[test]
525    fn mode_2() {
526        let mut rom = vec![0_u8; 0x100];
527        rom[..MODE_2_COMPRESSED.len()].copy_from_slice(MODE_2_COMPRESSED);
528        rom[0x50..0x54].copy_from_slice(&[0x02, 0x00, 0x00, 0x00]);
529
530        let decompressed = perform_decompression(&rom);
531
532        assert_eq!(decompressed.as_slice(), DECOMPRESSED);
533    }
534}