3use bincode::{Decode, Encode};
4use jgenesis_common::num::GetBit;
5use jgenesis_proc_macros::{EnumAll, EnumDisplay};
6use std::{array, cmp};

8#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Encode, Decode, EnumDisplay, EnumAll)]
9#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10#[cfg_attr(feature = "clap", derive(jgenesis_proc_macros::CustomValueEnum))]
11pub enum Sn76489Version {
12    #[default]
13    MasterSystem2,
14    Standard, // Sega-customized version used in SMS, Game Gear, Genesis
15    Discrete, // Stock TI version used in SG-1000
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)]
19enum WaveOutput {
20    Negative = 0,
21    Positive = 1,
22}
23
24impl WaveOutput {
25    fn invert(self) -> Self {
26        match self {
27            Self::Negative => Self::Positive,
28            Self::Positive => Self::Negative,
29        }
30    }
31}
32
33impl From<WaveOutput> for f64 {
34    fn from(value: WaveOutput) -> Self {
35        match value {
36            WaveOutput::Negative => 0.0,
37            WaveOutput::Positive => 1.0,
38        }
39    }
40}
41
42#[derive(Debug, Clone, Copy, Encode, Decode)]
43struct SquareWaveGenerator {
44    counter: u16,
45    current_output: WaveOutput,
46    tone: u16,
47    attenuation: u8,
48}

52const ATTENUATION_TO_VOLUME: [f64; 16] = [
53    1.0,
54    0.7943282347242815,
55    0.6309573444801932,
56    0.5011872336272722,
57    0.3981071705534972,
58    0.3162277660168379,
59    0.25118864315095796,
60    0.19952623149688792,
61    0.15848931924611132,
62    0.1258925411794167,
63    0.09999999999999998,
64    0.07943282347242814,
65    0.06309573444801932,
66    0.05011872336272722,
67    0.03981071705534972,
68    0.0,
69];

72const SMS2_ATTENUATION_TO_VOLUME: [f64; 16] = [
73    0.55,
74    0.55,
75    0.55,
76    0.5011872336272722,
77    0.3981071705534972,
78    0.3162277660168379,
79    0.25118864315095796,
80    0.19952623149688792,
81    0.15848931924611132,
82    0.1258925411794167,
83    0.09999999999999998,
84    0.07943282347242814,
85    0.06309573444801932,
86    0.05011872336272722,
87    0.03981071705534972,
88    0.0,
89];

91impl SquareWaveGenerator {
92    fn new() -> Self {
93        Self { counter: 1, current_output: WaveOutput::Negative, tone: 0, attenuation: 0x0F }
94    }
95
96    fn update_tone_low_bits(&mut self, data: u8) {
97        self.tone = (self.tone & 0xFFF0) | u16::from(data & 0x0F);
98    }
99
100    fn update_tone_high_bits(&mut self, data: u8) {
101        self.tone = (self.tone & 0x000F) | (u16::from(data & 0x3F) << 4);
102    }
103
104    fn clock(&mut self) {
105        self.counter -= 1;
106        if self.counter == 0 {
107            // Tone of 0 functions as 1, where the output inverts every PSG clock
108            self.counter = cmp::max(self.tone, 1);
109            self.current_output = self.current_output.invert();
110        }
111    }
112
113    fn sample(self, volume_table: &[f64; 16]) -> f64 {
114        f64::from(self.current_output) * volume_table[self.attenuation as usize]
115    }
116}
117
118#[derive(Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)]
119pub enum NoiseMode {
120    Periodic,
121    White,
122}
123
124#[derive(Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)]
125pub enum NoiseReload {
126    Value(u16),
127    Tone2,
128}
129
130impl NoiseReload {
131    fn from_noise_register(value: u8) -> Self {
132        match value & 0x03 {
133            0x00 => Self::Value(0x10),
134            0x01 => Self::Value(0x20),
135            0x02 => Self::Value(0x40),
136            0x03 => Self::Tone2,
137            _ => unreachable!("value & 0x03 is always <= 0x03"),
138        }
139    }
140
141    fn value(self, tone2: u16) -> u16 {
142        match self {
143            Self::Value(value) => value,
144            Self::Tone2 => tone2,
145        }
146    }
147}
148
149#[derive(Debug, Clone, Encode, Decode)]
150struct NoiseGenerator {
151    counter: u16,
152    current_counter_output: WaveOutput,
153    counter_reload: NoiseReload,
154    lfsr: u16,
155    current_lfsr_output: WaveOutput,
156    noise_type: NoiseMode,
157    attenuation: u8,
158}
159
160const INITIAL_LFSR: u16 = 0x8000;
161
162impl NoiseGenerator {
163    fn new() -> Self {
164        Self {
165            counter: 0,
166            current_counter_output: WaveOutput::Negative,
167            counter_reload: NoiseReload::from_noise_register(0x00),
168            lfsr: INITIAL_LFSR,
169            current_lfsr_output: WaveOutput::Negative,
170            noise_type: NoiseMode::Periodic,
171            attenuation: 0x0F,
172        }
173    }
174
175    fn shift_lfsr(&mut self, version: Sn76489Version) {
176        self.current_lfsr_output =
177            if self.lfsr.bit(0) { WaveOutput::Positive } else { WaveOutput::Negative };
178
179        let input_bit = match self.noise_type {
180            NoiseMode::Periodic => self.lfsr.bit(0),
181            NoiseMode::White => match version {
182                Sn76489Version::Standard | Sn76489Version::MasterSystem2 => {
183                    self.lfsr.bit(0) ^ self.lfsr.bit(3)
184                }
185                Sn76489Version::Discrete => self.lfsr.bit(0) ^ self.lfsr.bit(2) ^ self.lfsr.bit(15),
186            },
187        };
188
189        self.lfsr = (self.lfsr >> 1) | (u16::from(input_bit) << 15);
190    }
191
192    fn write_data(&mut self, data: u8) {
193        self.counter_reload = NoiseReload::from_noise_register(data);
194        self.noise_type = if data.bit(2) { NoiseMode::White } else { NoiseMode::Periodic };
195
196        self.lfsr = INITIAL_LFSR;
197    }
198
199    fn clock(&mut self, tone2: u16, version: Sn76489Version) {
200        self.counter = self.counter.saturating_sub(1);
201        if self.counter != 0 {
202            return;
203        }
204
205        self.counter = self.counter_reload.value(tone2);
206        self.current_counter_output = self.current_counter_output.invert();
207        if self.current_counter_output == WaveOutput::Positive {
208            self.shift_lfsr(version);
209        }
210    }
211
212    fn sample(&self, volume_table: &[f64; 16]) -> f64 {
213        f64::from(self.current_lfsr_output) * volume_table[self.attenuation as usize]
214    }
215}
216
217#[derive(Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)]
218enum Register {
219    Tone0,
220    Tone1,
221    Tone2,
222    Noise,
223    Volume0,
224    Volume1,
225    Volume2,
226    Volume3,
227}
228
229impl Register {
230    fn from_latch_byte(value: u8) -> Self {
231        match value & 0x70 {
232            0x00 => Self::Tone0,
233            0x10 => Self::Volume0,
234            0x20 => Self::Tone1,
235            0x30 => Self::Volume1,
236            0x40 => Self::Tone2,
237            0x50 => Self::Volume2,
238            0x60 => Self::Noise,
239            0x70 => Self::Volume3,
240            _ => unreachable!("value & 0x70 is always one of the above values"),
241        }
242    }
243}
244
245#[derive(Debug, Clone, Copy, PartialEq, Eq)]
246pub enum Sn76489TickEffect {
247    None,
248    Clocked,
249}
250
251#[derive(Debug, Clone, Copy, Encode, Decode)]
252struct StereoControl {
253    square_0_l: bool,
254    square_1_l: bool,
255    square_2_l: bool,
256    noise_l: bool,
257    square_0_r: bool,
258    square_1_r: bool,
259    square_2_r: bool,
260    noise_r: bool,
261}
262
263impl StereoControl {
264    fn write(&mut self, value: u8) {
265        self.square_0_r = value.bit(0);
266        self.square_1_r = value.bit(1);
267        self.square_2_r = value.bit(2);
268        self.noise_r = value.bit(3);
269        self.square_0_l = value.bit(4);
270        self.square_1_l = value.bit(5);
271        self.square_2_l = value.bit(6);
272        self.noise_l = value.bit(7);
273    }
274}
275
276impl Default for StereoControl {
277    fn default() -> Self {
278        Self {
279            square_0_l: true,
280            square_1_l: true,
281            square_2_l: true,
282            noise_l: true,
283            square_0_r: true,
284            square_1_r: true,
285            square_2_r: true,
286            noise_r: true,
287        }
288    }
289}
290
291#[derive(Debug, Clone, Encode, Decode)]
292pub struct Sn76489 {
293    version: Sn76489Version,
294    square_wave_channels: [SquareWaveGenerator; 3],
295    noise_channel: NoiseGenerator,
296    latched_register: Register,
297    stereo_control: StereoControl,
298    divider: u8,
299}
300
301const SN76489_DIVIDER: u8 = 16;
302
303impl Sn76489 {
304    #[must_use]
305    pub fn new(version: Sn76489Version) -> Self {
306        Self {
307            version,
308            square_wave_channels: array::from_fn(|_| SquareWaveGenerator::new()),
309            noise_channel: NoiseGenerator::new(),
310            latched_register: Register::Tone0,
311            stereo_control: StereoControl::default(),
312            divider: SN76489_DIVIDER,
313        }
314    }
315
316    fn write_register_low_bits(&mut self, data: u8) {
317        match self.latched_register {
318            Register::Tone0 => {
319                self.square_wave_channels[0].update_tone_low_bits(data);
320            }
321            Register::Tone1 => {
322                self.square_wave_channels[1].update_tone_low_bits(data);
323            }
324            Register::Tone2 => {
325                self.square_wave_channels[2].update_tone_low_bits(data);
326            }
327            Register::Noise => {
328                self.noise_channel.write_data(data);
329            }
330            Register::Volume0 => {
331                self.square_wave_channels[0].attenuation = data & 0x0F;
332            }
333            Register::Volume1 => {
334                self.square_wave_channels[1].attenuation = data & 0x0F;
335            }
336            Register::Volume2 => {
337                self.square_wave_channels[2].attenuation = data & 0x0F;
338            }
339            Register::Volume3 => {
340                self.noise_channel.attenuation = data & 0x0F;
341            }
342        }
343    }
344
345    fn write_register_high_bits(&mut self, data: u8) {
346        match self.latched_register {
347            Register::Tone0 => {
348                self.square_wave_channels[0].update_tone_high_bits(data);
349            }
350            Register::Tone1 => {
351                self.square_wave_channels[1].update_tone_high_bits(data);
352            }
353            Register::Tone2 => {
354                self.square_wave_channels[2].update_tone_high_bits(data);
355            }
356            _ => {
357                self.write_register_low_bits(data);
358            }
359        }
360    }
361
362    pub fn write(&mut self, value: u8) {
363        if value.bit(7) {
364            // LATCH/DATA byte
365            self.latched_register = Register::from_latch_byte(value);
366            self.write_register_low_bits(value);
367        } else {
368            // DATA byte
369            self.write_register_high_bits(value);
370        }
371    }
372
373    pub fn write_stereo_control(&mut self, value: u8) {
374        self.stereo_control.write(value);
375    }
376
377    #[inline]
378    pub fn tick(&mut self) -> Sn76489TickEffect {
379        self.divider -= 1;
380        if self.divider == 0 {
381            self.divider = SN76489_DIVIDER;
382
383            for channel in &mut self.square_wave_channels {
384                channel.clock();
385            }
386            self.noise_channel.clock(self.square_wave_channels[2].tone, self.version);
387
388            Sn76489TickEffect::Clocked
389        } else {
390            Sn76489TickEffect::None
391        }
392    }
393
394    #[must_use]
395    pub fn version(&self) -> Sn76489Version {
396        self.version
397    }
398
399    pub fn set_version(&mut self, version: Sn76489Version) {
400        self.version = version;
401    }
402
403    #[must_use]
404    pub fn sample(&self) -> (f64, f64) {
405        // TODO rewrite to use integer arithmetic as much as possible
406        let volume_table = match self.version {
407            Sn76489Version::MasterSystem2 => &SMS2_ATTENUATION_TO_VOLUME,
408            Sn76489Version::Standard | Sn76489Version::Discrete => &ATTENUATION_TO_VOLUME,
409        };
410
411        let square_samples = self.square_wave_channels.map(|channel| channel.sample(volume_table));
412        let noise_sample = self.noise_channel.sample(volume_table);
413
414        let sample_l = (f64::from(self.stereo_control.square_0_l) * square_samples[0]
415            + f64::from(self.stereo_control.square_1_l) * square_samples[1]
416            + f64::from(self.stereo_control.square_2_l) * square_samples[2]
417            + f64::from(self.stereo_control.noise_l) * noise_sample)
418            / 4.0;
419        let sample_r = (f64::from(self.stereo_control.square_0_r) * square_samples[0]
420            + f64::from(self.stereo_control.square_1_r) * square_samples[1]
421            + f64::from(self.stereo_control.square_2_r) * square_samples[2]
422            + f64::from(self.stereo_control.noise_r) * noise_sample)
423            / 4.0;
424
425        (sample_l, sample_r)
426    }
427
428    #[must_use]
429    pub fn tone_frequencies(&self) -> [u16; 3] {
430        self.square_wave_channels.map(|channel| channel.tone)
431    }
432
433    #[must_use]
434    pub fn tone_attenuations(&self) -> [u8; 3] {
435        self.square_wave_channels.map(|channel| channel.attenuation)
436    }
437
438    #[must_use]
439    pub fn noise_mode(&self) -> NoiseMode {
440        self.noise_channel.noise_type
441    }
442
443    #[must_use]
444    pub fn noise_reload(&self) -> NoiseReload {
445        self.noise_channel.counter_reload
446    }
447
448    #[must_use]
449    pub fn noise_attenuation(&self) -> u8 {
450        self.noise_channel.attenuation
451    }
452}