lemur/src/emulator.rs

use std::{
    collections::HashMap,
    fmt::Display,
    fs::{self, File},
    io::{Read, Seek, SeekFrom, Write},
    path::{Path, PathBuf},
    sync::{
        atomic::{AtomicBool, Ordering},
        mpsc::{self, RecvError, TryRecvError},
        Arc, Weak,
    },
};

use anyhow::Result;
use atomic::Atomic;
use bytemuck::NoUninit;
use egui_toast::{Toast, ToastKind, ToastOptions};
use tracing::{error, warn};

use crate::{
    audio::Audio,
    graphics::TextureSink,
    memory::{MemoryRange, MemoryRegion},
};
use shrooms_vb_core::{Sim, StopReason, EXPECTED_FRAME_SIZE};
pub use shrooms_vb_core::{VBKey, VBRegister, VBWatchpointType};

mod address_set;
mod shrooms_vb_core;

#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub enum SimId {
    Player1,
    Player2,
}
impl SimId {
    pub const fn values() -> [Self; 2] {
        [Self::Player1, Self::Player2]
    }
    pub const fn to_index(self) -> usize {
        match self {
            Self::Player1 => 0,
            Self::Player2 => 1,
        }
    }
}
impl Display for SimId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Self::Player1 => "Player 1",
            Self::Player2 => "Player 2",
        })
    }
}

struct Cart {
    rom_path: PathBuf,
    rom: Vec<u8>,
    sram_file: File,
    sram: Vec<u8>,
}

impl Cart {
    fn load(rom_path: &Path, sim_id: SimId) -> Result<Self> {
        let rom = fs::read(rom_path)?;

        let mut sram_file = File::options()
            .read(true)
            .write(true)
            .create(true)
            .truncate(false)
            .open(sram_path(rom_path, sim_id))?;
        sram_file.set_len(8 * 1024)?;

        let mut sram = vec![];
        sram_file.read_to_end(&mut sram)?;
        Ok(Cart {
            rom_path: rom_path.to_path_buf(),
            rom,
            sram_file,
            sram,
        })
    }
}

fn sram_path(rom_path: &Path, sim_id: SimId) -> PathBuf {
    match sim_id {
        SimId::Player1 => rom_path.with_extension("p1.sram"),
        SimId::Player2 => rom_path.with_extension("p2.sram"),
    }
}

pub struct EmulatorBuilder {
    rom: Option<PathBuf>,
    commands: mpsc::Receiver<EmulatorCommand>,
    sim_state: Arc<[Atomic<SimState>; 2]>,
    state: Arc<Atomic<EmulatorState>>,
    audio_on: Arc<[AtomicBool; 2]>,
    linked: Arc<AtomicBool>,
    start_paused: bool,
}

impl EmulatorBuilder {
    pub fn new() -> (Self, EmulatorClient) {
        let (queue, commands) = mpsc::channel();
        let builder = Self {
            rom: None,
            commands,
            sim_state: Arc::new([
                Atomic::new(SimState::Uninitialized),
                Atomic::new(SimState::Uninitialized),
            ]),
            state: Arc::new(Atomic::new(EmulatorState::Paused)),
            audio_on: Arc::new([AtomicBool::new(true), AtomicBool::new(true)]),
            linked: Arc::new(AtomicBool::new(false)),
            start_paused: false,
        };
        let client = EmulatorClient {
            queue,
            sim_state: builder.sim_state.clone(),
            state: builder.state.clone(),
            audio_on: builder.audio_on.clone(),
            linked: builder.linked.clone(),
        };
        (builder, client)
    }

    pub fn with_rom(self, path: &Path) -> Self {
        Self {
            rom: Some(path.into()),
            ..self
        }
    }

    pub fn start_paused(self, paused: bool) -> Self {
        Self {
            start_paused: paused,
            ..self
        }
    }

    pub fn build(self) -> Result<Emulator> {
        let mut emulator = Emulator::new(
            self.commands,
            self.sim_state,
            self.state,
            self.audio_on,
            self.linked,
        )?;
        if let Some(path) = self.rom {
            emulator.load_cart(SimId::Player1, &path)?;
        }
        if self.start_paused {
            emulator.pause_sims()?;
        }
        Ok(emulator)
    }
}

pub struct Emulator {
    sims: Vec<Sim>,
    carts: [Option<Cart>; 2],
    audio: Audio,
    commands: mpsc::Receiver<EmulatorCommand>,
    sim_state: Arc<[Atomic<SimState>; 2]>,
    state: Arc<Atomic<EmulatorState>>,
    audio_on: Arc<[AtomicBool; 2]>,
    linked: Arc<AtomicBool>,
    renderers: HashMap<SimId, TextureSink>,
    messages: HashMap<SimId, mpsc::Sender<Toast>>,
    debuggers: HashMap<SimId, DebugInfo>,
    watched_regions: HashMap<MemoryRange, Weak<MemoryRegion>>,
    eye_contents: Vec<u8>,
    audio_samples: Vec<f32>,
    buffer: Vec<u8>,
}

impl Emulator {
    fn new(
        commands: mpsc::Receiver<EmulatorCommand>,
        sim_state: Arc<[Atomic<SimState>; 2]>,
        state: Arc<Atomic<EmulatorState>>,
        audio_on: Arc<[AtomicBool; 2]>,
        linked: Arc<AtomicBool>,
    ) -> Result<Self> {
        Ok(Self {
            sims: vec![],
            carts: [None, None],
            audio: Audio::init()?,
            commands,
            sim_state,
            state,
            audio_on,
            linked,
            renderers: HashMap::new(),
            messages: HashMap::new(),
            debuggers: HashMap::new(),
            watched_regions: HashMap::new(),
            eye_contents: vec![0u8; 384 * 224 * 2],
            audio_samples: Vec::with_capacity(EXPECTED_FRAME_SIZE),
            buffer: vec![],
        })
    }

    pub fn load_cart(&mut self, sim_id: SimId, path: &Path) -> Result<()> {
        let cart = Cart::load(path, sim_id)?;
        self.reset_sim(sim_id, Some(cart))?;
        Ok(())
    }

    pub fn start_second_sim(&mut self, rom: Option<PathBuf>) -> Result<()> {
        let rom_path = if let Some(path) = rom {
            Some(path)
        } else {
            self.carts[0].as_ref().map(|c| c.rom_path.clone())
        };
        let cart = match rom_path {
            Some(rom_path) => Some(Cart::load(&rom_path, SimId::Player2)?),
            None => None,
        };
        self.reset_sim(SimId::Player2, cart)?;
        self.link_sims();
        Ok(())
    }

    fn reset_sim(&mut self, sim_id: SimId, new_cart: Option<Cart>) -> Result<()> {
        self.save_sram(sim_id)?;

        let index = sim_id.to_index();
        while self.sims.len() <= index {
            self.sims.push(Sim::new());
            self.sim_state[index].store(SimState::NoGame, Ordering::Release);
        }
        let sim = &mut self.sims[index];
        sim.reset();
        if let Some(cart) = new_cart {
            sim.load_cart(cart.rom.clone(), cart.sram.clone())?;
            self.carts[index] = Some(cart);
            self.sim_state[index].store(SimState::Ready, Ordering::Release);
        }
        if self.sim_state[index].load(Ordering::Acquire) == SimState::Ready {
            self.resume_sims();
        }
        Ok(())
    }

    fn link_sims(&mut self) {
        let (first, second) = self.sims.split_at_mut(1);
        let Some(first) = first.first_mut() else {
            return;
        };
        let Some(second) = second.first_mut() else {
            return;
        };
        first.link(second);
        self.linked.store(true, Ordering::Release);
    }

    fn unlink_sims(&mut self) {
        let Some(first) = self.sims.first_mut() else {
            return;
        };
        first.unlink();
        self.linked.store(false, Ordering::Release);
    }

    fn pause_sims(&mut self) -> Result<()> {
        if self
            .state
            .compare_exchange(
                EmulatorState::Running,
                EmulatorState::Paused,
                Ordering::AcqRel,
                Ordering::Relaxed,
            )
            .is_ok()
        {
            for sim_id in SimId::values() {
                self.save_sram(sim_id)?;
            }
        }
        Ok(())
    }

    fn resume_sims(&mut self) {
        let _ = self.state.compare_exchange(
            EmulatorState::Paused,
            EmulatorState::Running,
            Ordering::AcqRel,
            Ordering::Relaxed,
        );
    }

    fn save_sram(&mut self, sim_id: SimId) -> Result<()> {
        let sim = self.sims.get_mut(sim_id.to_index());
        let cart = self.carts[sim_id.to_index()].as_mut();
        if let (Some(sim), Some(cart)) = (sim, cart) {
            sim.read_sram(&mut cart.sram);
            cart.sram_file.seek(SeekFrom::Start(0))?;
            cart.sram_file.write_all(&cart.sram)?;
        }
        Ok(())
    }

    pub fn stop_second_sim(&mut self) -> Result<()> {
        self.save_sram(SimId::Player2)?;
        self.renderers.remove(&SimId::Player2);
        self.sims.truncate(1);
        self.sim_state[SimId::Player2.to_index()].store(SimState::Uninitialized, Ordering::Release);
        self.stop_debugging(SimId::Player2);
        self.linked.store(false, Ordering::Release);
        Ok(())
    }

    fn start_debugging(&mut self, sim_id: SimId, sender: DebugSender) {
        if self.sim_state[sim_id.to_index()].load(Ordering::Acquire) != SimState::Ready {
            // Can't debug unless a game is connected
            return;
        }
        let debug = DebugInfo {
            sender,
            stop_reason: Some(DebugStopReason::Paused),
        };
        self.debuggers.insert(sim_id, debug);
        self.state
            .store(EmulatorState::Debugging, Ordering::Release);
    }

    fn stop_debugging(&mut self, sim_id: SimId) {
        if let Some(sim) = self.sims.get_mut(sim_id.to_index()) {
            sim.clear_debug_state();
        }
        self.debuggers.remove(&sim_id);
        if self.debuggers.is_empty() {
            let _ = self.state.compare_exchange(
                EmulatorState::Debugging,
                EmulatorState::Running,
                Ordering::AcqRel,
                Ordering::Relaxed,
            );
        }
    }

    fn debug_interrupt(&mut self, sim_id: SimId) {
        self.debug_stop(sim_id, DebugStopReason::Paused);
    }

    fn debug_stop(&mut self, sim_id: SimId, reason: DebugStopReason) {
        let Some(debugger) = self.debuggers.get_mut(&sim_id) else {
            self.stop_debugging(sim_id);
            return;
        };
        if debugger.stop_reason != Some(reason) {
            debugger.stop_reason = Some(reason);
            if debugger.sender.send(DebugEvent::Stopped(reason)).is_err() {
                self.stop_debugging(sim_id);
            }
        }
    }

    fn debug_continue(&mut self, sim_id: SimId) -> bool {
        let Some(debugger) = self.debuggers.get_mut(&sim_id) else {
            self.stop_debugging(sim_id);
            return false;
        };
        debugger.stop_reason = None;
        true
    }
    fn debug_step(&mut self, sim_id: SimId) {
        if self.debug_continue(sim_id) {
            let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                return;
            };
            sim.step();
        }
    }

    fn watch_memory(&mut self, range: MemoryRange, region: Weak<MemoryRegion>) {
        self.watched_regions.insert(range, region);
    }

    pub fn run(&mut self) {
        loop {
            let idle = self.tick();
            if idle {
                // The game is paused, and we have output all the video/audio we have.
                // Block the thread until a new command comes in.
                match self.commands.recv() {
                    Ok(command) => self.handle_command(command),
                    Err(RecvError) => {
                        return;
                    }
                }
            }
            loop {
                match self.commands.try_recv() {
                    Ok(command) => self.handle_command(command),
                    Err(TryRecvError::Empty) => {
                        break;
                    }
                    Err(TryRecvError::Disconnected) => {
                        return;
                    }
                }
            }
            self.watched_regions.retain(|range, region| {
                let Some(region) = region.upgrade() else {
                    return false;
                };
                let Some(sim) = self.sims.get_mut(range.sim.to_index()) else {
                    return false;
                };
                self.buffer.clear();
                sim.read_memory(range.start, range.length, &mut self.buffer);
                region.update(&self.buffer);
                true
            });
        }
    }

    // returns true if the emulator is "idle" (i.e. this didn't output anything)
    pub fn tick(&mut self) -> bool {
        let p1_state = self.sim_state[SimId::Player1.to_index()].load(Ordering::Acquire);
        let p2_state = self.sim_state[SimId::Player2.to_index()].load(Ordering::Acquire);
        let state = self.state.load(Ordering::Acquire);

        // Emulation
        // Don't emulate if the state is "paused", or if any sim is paused in the debugger
        let running = match state {
            EmulatorState::Paused => false,
            EmulatorState::Running => true,
            EmulatorState::Debugging => self.debuggers.values().all(|d| d.stop_reason.is_none()),
        };
        let p1_running = running && p1_state == SimState::Ready;
        let p2_running = running && p2_state == SimState::Ready;
        let mut idle = !p1_running && !p2_running;
        if p1_running && p2_running {
            Sim::emulate_many(&mut self.sims);
        } else if p1_running {
            self.sims[SimId::Player1.to_index()].emulate();
        } else if p2_running {
            self.sims[SimId::Player2.to_index()].emulate();
        }

        // Debug state
        if state == EmulatorState::Debugging {
            for sim_id in SimId::values() {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    continue;
                };
                if let Some(reason) = sim.stop_reason() {
                    let stop_reason = match reason {
                        StopReason::Stepped => DebugStopReason::Trace,
                        StopReason::Watchpoint(watch, address) => {
                            DebugStopReason::Watchpoint(watch, address)
                        }
                        StopReason::Breakpoint => DebugStopReason::Breakpoint,
                    };
                    self.debug_stop(sim_id, stop_reason);
                }
            }
        }

        // Video
        for sim_id in SimId::values() {
            let Some(renderer) = self.renderers.get_mut(&sim_id) else {
                continue;
            };
            let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                continue;
            };
            if sim.read_pixels(&mut self.eye_contents) {
                idle = false;
                if renderer.queue_render(&self.eye_contents).is_err() {
                    self.renderers.remove(&sim_id);
                }
            }
        }

        // Audio
        // Audio playback speed is how we keep the emulator running in real time.
        // Even if we're muted, call `read_samples` to know how many frames of silence to play.
        let p1_audio =
            p1_running && self.audio_on[SimId::Player1.to_index()].load(Ordering::Acquire);
        let p2_audio =
            p2_running && self.audio_on[SimId::Player2.to_index()].load(Ordering::Acquire);
        let (p1_weight, p2_weight) = match (p1_audio, p2_audio) {
            (true, true) => (0.5, 0.5),
            (true, false) => (1.0, 0.0),
            (false, true) => (0.0, 1.0),
            (false, false) => (0.0, 0.0),
        };
        if let Some(sim) = self.sims.get_mut(SimId::Player1.to_index()) {
            sim.read_samples(&mut self.audio_samples, p1_weight);
        }
        if let Some(sim) = self.sims.get_mut(SimId::Player2.to_index()) {
            sim.read_samples(&mut self.audio_samples, p2_weight);
        }
        if !self.audio_samples.is_empty() {
            idle = false;
        }
        self.audio.update(&self.audio_samples);
        self.audio_samples.clear();
        idle
    }

    fn handle_command(&mut self, command: EmulatorCommand) {
        match command {
            EmulatorCommand::ConnectToSim(sim_id, renderer, messages) => {
                self.renderers.insert(sim_id, renderer);
                self.messages.insert(sim_id, messages);
            }
            EmulatorCommand::LoadGame(sim_id, path) => {
                if let Err(error) = self.load_cart(sim_id, &path) {
                    self.report_error(sim_id, format!("Error loading rom: {error}"));
                }
            }
            EmulatorCommand::StartSecondSim(path) => {
                if let Err(error) = self.start_second_sim(path) {
                    self.report_error(
                        SimId::Player2,
                        format!("Error starting second sim: {error}"),
                    );
                }
            }
            EmulatorCommand::StopSecondSim => {
                if let Err(error) = self.stop_second_sim() {
                    self.report_error(
                        SimId::Player2,
                        format!("Error stopping second sim: {error}"),
                    );
                }
            }
            EmulatorCommand::Pause => {
                if let Err(error) = self.pause_sims() {
                    self.report_error(SimId::Player1, format!("Error pausing: {error}"));
                }
            }
            EmulatorCommand::Resume => {
                self.resume_sims();
            }
            EmulatorCommand::StartDebugging(sim_id, debugger) => {
                self.start_debugging(sim_id, debugger);
            }
            EmulatorCommand::StopDebugging(sim_id) => {
                self.stop_debugging(sim_id);
            }
            EmulatorCommand::DebugInterrupt(sim_id) => {
                self.debug_interrupt(sim_id);
            }
            EmulatorCommand::DebugContinue(sim_id) => {
                self.debug_continue(sim_id);
            }
            EmulatorCommand::DebugStep(sim_id) => {
                self.debug_step(sim_id);
            }
            EmulatorCommand::ReadRegister(sim_id, register, done) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                let value = sim.read_register(register);
                let _ = done.send(value);
            }
            EmulatorCommand::WriteRegister(sim_id, register, value) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                sim.write_register(register, value);
            }
            EmulatorCommand::ReadMemory(sim_id, start, length, mut buffer, done) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                sim.read_memory(start, length, &mut buffer);
                let _ = done.send(buffer);
            }
            EmulatorCommand::WriteMemory(sim_id, start, buffer, done) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                sim.write_memory(start, &buffer);
                let _ = done.send(buffer);
            }
            EmulatorCommand::WatchMemory(range, region) => {
                self.watch_memory(range, region);
            }
            EmulatorCommand::AddBreakpoint(sim_id, address) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                sim.add_breakpoint(address);
            }
            EmulatorCommand::RemoveBreakpoint(sim_id, address) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                sim.remove_breakpoint(address);
            }
            EmulatorCommand::AddWatchpoint(sim_id, address, length, watch) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                sim.add_watchpoint(address, length, watch);
            }
            EmulatorCommand::RemoveWatchpoint(sim_id, address, length, watch) => {
                let Some(sim) = self.sims.get_mut(sim_id.to_index()) else {
                    return;
                };
                sim.remove_watchpoint(address, length, watch);
            }
            EmulatorCommand::SetAudioEnabled(p1, p2) => {
                self.audio_on[SimId::Player1.to_index()].store(p1, Ordering::Release);
                self.audio_on[SimId::Player2.to_index()].store(p2, Ordering::Release);
            }
            EmulatorCommand::Link => {
                self.link_sims();
            }
            EmulatorCommand::Unlink => {
                self.unlink_sims();
            }
            EmulatorCommand::Reset(sim_id) => {
                if let Err(error) = self.reset_sim(sim_id, None) {
                    self.report_error(sim_id, format!("Error resetting sim: {error}"));
                }
            }
            EmulatorCommand::SetKeys(sim_id, keys) => {
                if let Some(sim) = self.sims.get_mut(sim_id.to_index()) {
                    sim.set_keys(keys);
                }
            }
            EmulatorCommand::Exit(done) => {
                for sim_id in SimId::values() {
                    if let Err(error) = self.save_sram(sim_id) {
                        self.report_error(sim_id, format!("Error saving sram on exit: {error}"));
                    }
                }
                let _ = done.send(());
            }
        }
    }

    fn report_error(&self, sim_id: SimId, message: String) {
        let messages = self
            .messages
            .get(&sim_id)
            .or_else(|| self.messages.get(&SimId::Player1));
        if let Some(msg) = messages {
            let toast = Toast::new()
                .kind(ToastKind::Error)
                .options(ToastOptions::default().duration_in_seconds(5.0))
                .text(&message);
            if msg.send(toast).is_ok() {
                return;
            }
        }
        error!("{}", message);
    }
}

#[derive(Debug)]
pub enum EmulatorCommand {
    ConnectToSim(SimId, TextureSink, mpsc::Sender<Toast>),
    LoadGame(SimId, PathBuf),
    StartSecondSim(Option<PathBuf>),
    StopSecondSim,
    Pause,
    Resume,
    StartDebugging(SimId, DebugSender),
    StopDebugging(SimId),
    DebugInterrupt(SimId),
    DebugContinue(SimId),
    DebugStep(SimId),
    ReadRegister(SimId, VBRegister, oneshot::Sender<u32>),
    WriteRegister(SimId, VBRegister, u32),
    ReadMemory(SimId, u32, usize, Vec<u8>, oneshot::Sender<Vec<u8>>),
    WriteMemory(SimId, u32, Vec<u8>, oneshot::Sender<Vec<u8>>),
    WatchMemory(MemoryRange, Weak<MemoryRegion>),
    AddBreakpoint(SimId, u32),
    RemoveBreakpoint(SimId, u32),
    AddWatchpoint(SimId, u32, usize, VBWatchpointType),
    RemoveWatchpoint(SimId, u32, usize, VBWatchpointType),
    SetAudioEnabled(bool, bool),
    Link,
    Unlink,
    Reset(SimId),
    SetKeys(SimId, VBKey),
    Exit(oneshot::Sender<()>),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, NoUninit)]
#[repr(usize)]
pub enum SimState {
    Uninitialized,
    NoGame,
    Ready,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, NoUninit)]
#[repr(usize)]
pub enum EmulatorState {
    Paused,
    Running,
    Debugging,
}

type DebugSender = tokio::sync::mpsc::UnboundedSender<DebugEvent>;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DebugStopReason {
    // We are stepping
    Trace,
    // We hit a breakpoint
    Breakpoint,
    // We hit a watchpoint
    Watchpoint(VBWatchpointType, u32),
    // The debugger told us to pause
    Paused,
}

struct DebugInfo {
    sender: DebugSender,
    stop_reason: Option<DebugStopReason>,
}

pub enum DebugEvent {
    Stopped(DebugStopReason),
}

#[derive(Clone)]
pub struct EmulatorClient {
    queue: mpsc::Sender<EmulatorCommand>,
    sim_state: Arc<[Atomic<SimState>; 2]>,
    state: Arc<Atomic<EmulatorState>>,
    audio_on: Arc<[AtomicBool; 2]>,
    linked: Arc<AtomicBool>,
}

impl EmulatorClient {
    pub fn sim_state(&self, sim_id: SimId) -> SimState {
        self.sim_state[sim_id.to_index()].load(Ordering::Acquire)
    }
    pub fn emulator_state(&self) -> EmulatorState {
        self.state.load(Ordering::Acquire)
    }
    pub fn is_audio_enabled(&self, sim_id: SimId) -> bool {
        self.audio_on[sim_id.to_index()].load(Ordering::Acquire)
    }
    pub fn are_sims_linked(&self) -> bool {
        self.linked.load(Ordering::Acquire)
    }
    pub fn send_command(&self, command: EmulatorCommand) -> bool {
        match self.queue.send(command) {
            Ok(()) => true,
            Err(err) => {
                warn!(
                    "could not send command {:?} as emulator is shut down",
                    err.0
                );
                false
            }
        }
    }
}