lemur/src/emulator/shrooms_vb_core.rs

use std::{ffi::c_void, ptr, slice};

use anyhow::{anyhow, Result};
use bitflags::bitflags;
use num_derive::{FromPrimitive, ToPrimitive};
use serde::{Deserialize, Serialize};

use super::address_set::AddressSet;

#[repr(C)]
struct VB {
    _data: [u8; 0],
}

#[allow(non_camel_case_types)]
type c_int = i32;
#[allow(non_camel_case_types)]
type c_uint = u32;

#[repr(u32)]
#[derive(FromPrimitive, ToPrimitive)]
enum VBDataType {
    S8 = 0,
    U8 = 1,
    S16 = 2,
    U16 = 3,
    S32 = 4,
    F32 = 5,
}

#[repr(i32)]
#[derive(FromPrimitive, ToPrimitive)]
enum VBOption {
    PseudoHalt = 0,
}

bitflags! {
    #[repr(transparent)]
    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
    pub struct VBKey: u16 {
        const PWR = 0x0001;
        const SGN = 0x0002;
        const A   = 0x0004;
        const B   = 0x0008;
        const RT  = 0x0010;
        const LT  = 0x0020;
        const RU  = 0x0040;
        const RR  = 0x0080;
        const LR  = 0x0100;
        const LL  = 0x0200;
        const LD  = 0x0400;
        const LU  = 0x0800;
        const STA = 0x1000;
        const SEL = 0x2000;
        const RL  = 0x4000;
        const RD  = 0x8000;
    }
}

#[derive(Clone, Copy, Debug)]
pub enum VBRegister {
    Program(u32),
    System(u32),
    PC,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum VBWatchpointType {
    Read,
    Write,
    Access,
}

type OnExecute =
    extern "C" fn(sim: *mut VB, address: u32, code: *const u16, length: c_int) -> c_int;
type OnFrame = extern "C" fn(sim: *mut VB) -> c_int;
type OnRead = extern "C" fn(
    sim: *mut VB,
    address: u32,
    type_: VBDataType,
    value: *mut i32,
    cycles: *mut u32,
) -> c_int;
type OnWrite = extern "C" fn(
    sim: *mut VB,
    address: u32,
    type_: VBDataType,
    value: *mut i32,
    cycles: *mut u32,
    cancel: *mut c_int,
) -> c_int;

#[link(name = "vb")]
extern "C" {
    #[link_name = "vbEmulate"]
    fn vb_emulate(sim: *mut VB, cycles: *mut u32) -> c_int;
    #[link_name = "vbEmulateEx"]
    fn vb_emulate_ex(sims: *mut *mut VB, count: c_uint, cycles: *mut u32) -> c_int;
    #[link_name = "vbGetCartRAM"]
    fn vb_get_cart_ram(sim: *mut VB, size: *mut u32) -> *mut c_void;
    #[link_name = "vbGetCartROM"]
    fn vb_get_cart_rom(sim: *mut VB, size: *mut u32) -> *mut c_void;
    #[link_name = "vbGetPixels"]
    fn vb_get_pixels(
        sim: *mut VB,
        left: *mut c_void,
        left_stride_x: c_int,
        left_stride_y: c_int,
        right: *mut c_void,
        right_stride_x: c_int,
        right_stride_y: c_int,
    );
    #[link_name = "vbGetProgramCounter"]
    fn vb_get_program_counter(sim: *mut VB) -> u32;
    #[link_name = "vbGetProgramRegister"]
    fn vb_get_program_register(sim: *mut VB, index: c_uint) -> i32;
    #[link_name = "vbGetSamples"]
    fn vb_get_samples(
        sim: *mut VB,
        typ_: *mut VBDataType,
        capacity: *mut c_uint,
        position: *mut c_uint,
    ) -> *mut c_void;
    #[link_name = "vbGetSystemRegister"]
    fn vb_get_system_register(sim: *mut VB, index: c_uint) -> i32;
    #[link_name = "vbGetUserData"]
    fn vb_get_user_data(sim: *mut VB) -> *mut c_void;
    #[link_name = "vbInit"]
    fn vb_init(sim: *mut VB) -> *mut VB;
    #[link_name = "vbRead"]
    fn vb_read(sim: *mut VB, address: u32, typ_: VBDataType) -> i32;
    #[link_name = "vbReset"]
    fn vb_reset(sim: *mut VB);
    #[link_name = "vbSetCartRAM"]
    fn vb_set_cart_ram(sim: *mut VB, sram: *mut c_void, size: u32) -> c_int;
    #[link_name = "vbSetCartROM"]
    fn vb_set_cart_rom(sim: *mut VB, rom: *mut c_void, size: u32) -> c_int;
    #[link_name = "vbSetExecuteCallback"]
    fn vb_set_execute_callback(sim: *mut VB, callback: Option<OnExecute>) -> Option<OnExecute>;
    #[link_name = "vbSetFrameCallback"]
    fn vb_set_frame_callback(sim: *mut VB, callback: Option<OnFrame>) -> Option<OnFrame>;
    #[link_name = "vbSetKeys"]
    fn vb_set_keys(sim: *mut VB, keys: u16) -> u16;
    #[link_name = "vbSetOption"]
    fn vb_set_option(sim: *mut VB, key: VBOption, value: c_int);
    #[link_name = "vbSetPeer"]
    fn vb_set_peer(sim: *mut VB, peer: *mut VB);
    #[link_name = "vbSetProgramCounter"]
    fn vb_set_program_counter(sim: *mut VB, value: u32) -> u32;
    #[link_name = "vbSetProgramRegister"]
    fn vb_set_program_register(sim: *mut VB, index: c_uint, value: i32) -> i32;
    #[link_name = "vbSetReadCallback"]
    fn vb_set_read_callback(sim: *mut VB, callback: Option<OnRead>) -> Option<OnRead>;
    #[link_name = "vbSetSamples"]
    fn vb_set_samples(
        sim: *mut VB,
        samples: *mut c_void,
        typ_: VBDataType,
        capacity: c_uint,
    ) -> c_int;
    #[link_name = "vbSetSystemRegister"]
    fn vb_set_system_register(sim: *mut VB, index: c_uint, value: u32) -> u32;
    #[link_name = "vbSetUserData"]
    fn vb_set_user_data(sim: *mut VB, tag: *mut c_void);
    #[link_name = "vbSetWriteCallback"]
    fn vb_set_write_callback(sim: *mut VB, callback: Option<OnWrite>) -> Option<OnWrite>;
    #[link_name = "vbSizeOf"]
    fn vb_size_of() -> usize;
    #[link_name = "vbWrite"]
    fn vb_write(sim: *mut VB, address: u32, _type: VBDataType, value: i32) -> i32;
}

#[no_mangle]
extern "C" fn on_frame(sim: *mut VB) -> c_int {
    // SAFETY: the *mut VB owns its userdata.
    // There is no way for the userdata to be null or otherwise invalid.
    let data: &mut VBState = unsafe { &mut *vb_get_user_data(sim).cast() };
    data.frame_seen = true;
    1
}

#[no_mangle]
extern "C" fn on_execute(sim: *mut VB, address: u32, _code: *const u16, _length: c_int) -> c_int {
    // SAFETY: the *mut VB owns its userdata.
    // There is no way for the userdata to be null or otherwise invalid.
    let data: &mut VBState = unsafe { &mut *vb_get_user_data(sim).cast() };

    let mut stopped = data.stop_reason.is_some();
    if data.step_from.is_some_and(|s| s != address) {
        data.step_from = None;
        data.stop_reason = Some(StopReason::Stepped);
        stopped = true;
    }
    if data.breakpoints.binary_search(&address).is_ok() {
        data.stop_reason = Some(StopReason::Breakpoint);
        stopped = true;
    }

    if stopped {
        1
    } else {
        0
    }
}

#[no_mangle]
extern "C" fn on_read(
    sim: *mut VB,
    address: u32,
    _type: VBDataType,
    _value: *mut i32,
    _cycles: *mut u32,
) -> c_int {
    // SAFETY: the *mut VB owns its userdata.
    // There is no way for the userdata to be null or otherwise invalid.
    let data: &mut VBState = unsafe { &mut *vb_get_user_data(sim).cast() };

    if let Some(start) = data.read_watchpoints.start_of_range_containing(address) {
        let watch = if data.write_watchpoints.contains(address) {
            VBWatchpointType::Access
        } else {
            VBWatchpointType::Read
        };
        data.stop_reason = Some(StopReason::Watchpoint(watch, start));
    }

    // Don't stop here, the debugger expects us to break after the memory access.
    // We'll stop in on_execute instead.
    0
}

#[no_mangle]
extern "C" fn on_write(
    sim: *mut VB,
    address: u32,
    _type: VBDataType,
    _value: *mut i32,
    _cycles: *mut u32,
    _cancel: *mut c_int,
) -> c_int {
    // SAFETY: the *mut VB owns its userdata.
    // There is no way for the userdata to be null or otherwise invalid.
    let data: &mut VBState = unsafe { &mut *vb_get_user_data(sim).cast() };

    if let Some(start) = data.write_watchpoints.start_of_range_containing(address) {
        let watch = if data.read_watchpoints.contains(address) {
            VBWatchpointType::Access
        } else {
            VBWatchpointType::Write
        };
        data.stop_reason = Some(StopReason::Watchpoint(watch, start));
    }

    // Don't stop here, the debugger expects us to break after the memory access.
    // We'll stop in on_execute instead.
    0
}

const AUDIO_CAPACITY_SAMPLES: usize = 834 * 4;
const AUDIO_CAPACITY_FLOATS: usize = AUDIO_CAPACITY_SAMPLES * 2;
pub const EXPECTED_FRAME_SIZE: usize = 834 * 2;

struct VBState {
    frame_seen: bool,
    stop_reason: Option<StopReason>,
    step_from: Option<u32>,
    breakpoints: Vec<u32>,
    read_watchpoints: AddressSet,
    write_watchpoints: AddressSet,
}

impl VBState {
    fn needs_execute_callback(&self) -> bool {
        self.step_from.is_some()
            || !self.breakpoints.is_empty()
            || !self.read_watchpoints.is_empty()
            || !self.write_watchpoints.is_empty()
    }
}

pub enum StopReason {
    Breakpoint,
    Watchpoint(VBWatchpointType, u32),
    Stepped,
}

#[repr(transparent)]
pub struct Sim {
    sim: *mut VB,
}

impl Sim {
    pub fn new() -> Self {
        // init the VB instance itself
        let size = unsafe { vb_size_of() };
        // allocate a vec of u64 so that this memory is 8-byte aligned
        let memory = vec![0u64; size.div_ceil(4)];
        let sim: *mut VB = Box::into_raw(memory.into_boxed_slice()).cast();
        unsafe { vb_init(sim) };
        // pseudohalt is disabled due to breaking red alarm
        unsafe { vb_set_option(sim, VBOption::PseudoHalt, 0) };
        unsafe { vb_set_keys(sim, VBKey::SGN.bits()) };
        unsafe { vb_reset(sim) };

        // set up userdata
        let state = VBState {
            frame_seen: false,
            stop_reason: None,
            step_from: None,
            breakpoints: vec![],
            read_watchpoints: AddressSet::new(),
            write_watchpoints: AddressSet::new(),
        };
        unsafe { vb_set_user_data(sim, Box::into_raw(Box::new(state)).cast()) };
        unsafe { vb_set_frame_callback(sim, Some(on_frame)) };

        // set up audio buffer
        let audio_buffer = vec![0.0f32; AUDIO_CAPACITY_FLOATS];
        let samples: *mut c_void = Box::into_raw(audio_buffer.into_boxed_slice()).cast();
        unsafe { vb_set_samples(sim, samples, VBDataType::F32, AUDIO_CAPACITY_SAMPLES as u32) };

        Sim { sim }
    }

    pub fn reset(&mut self) {
        unsafe { vb_reset(self.sim) };
    }

    pub fn load_cart(&mut self, mut rom: Vec<u8>, mut sram: Vec<u8>) -> Result<()> {
        self.unload_cart();

        rom.shrink_to_fit();
        sram.shrink_to_fit();

        let size = rom.len() as u32;
        let rom = Box::into_raw(rom.into_boxed_slice()).cast();
        let status = unsafe { vb_set_cart_rom(self.sim, rom, size) };
        if status != 0 {
            let _: Vec<u8> =
                unsafe { Vec::from_raw_parts(rom.cast(), size as usize, size as usize) };
            return Err(anyhow!("Invalid ROM size of {} bytes", size));
        }

        let size = sram.len() as u32;
        let sram = Box::into_raw(sram.into_boxed_slice()).cast();
        let status = unsafe { vb_set_cart_ram(self.sim, sram, size) };
        if status != 0 {
            let _: Vec<u8> =
                unsafe { Vec::from_raw_parts(sram.cast(), size as usize, size as usize) };
            return Err(anyhow!("Invalid SRAM size of {} bytes", size));
        }

        Ok(())
    }

    fn unload_cart(&mut self) {
        let mut size = 0;
        let rom = unsafe { vb_get_cart_rom(self.sim, &mut size) };
        unsafe { vb_set_cart_rom(self.sim, ptr::null_mut(), 0) };
        if !rom.is_null() {
            let _: Vec<u8> =
                unsafe { Vec::from_raw_parts(rom.cast(), size as usize, size as usize) };
        }

        let sram = unsafe { vb_get_cart_ram(self.sim, &mut size) };
        unsafe { vb_set_cart_ram(self.sim, ptr::null_mut(), 0) };
        if !sram.is_null() {
            let _: Vec<u8> =
                unsafe { Vec::from_raw_parts(sram.cast(), size as usize, size as usize) };
        }
    }

    pub fn read_sram(&mut self, buffer: &mut [u8]) {
        let mut size = 0;
        let sram = unsafe { vb_get_cart_ram(self.sim, &mut size) };
        if sram.is_null() {
            return;
        }
        let bytes = unsafe { slice::from_raw_parts(sram.cast(), size as usize) };
        buffer.copy_from_slice(bytes);
    }

    pub fn link(&mut self, peer: &mut Sim) {
        unsafe { vb_set_peer(self.sim, peer.sim) };
    }

    pub fn unlink(&mut self) {
        unsafe { vb_set_peer(self.sim, ptr::null_mut()) };
    }

    pub fn emulate(&mut self) {
        let mut cycles = 20_000_000;
        unsafe { vb_emulate(self.sim, &mut cycles) };
    }

    pub fn emulate_many(sims: &mut [Sim]) {
        let mut cycles = 20_000_000;
        let count = sims.len() as c_uint;
        let sims = sims.as_mut_ptr().cast();
        unsafe { vb_emulate_ex(sims, count, &mut cycles) };
    }

    pub fn read_pixels(&mut self, buffers: &mut [u8]) -> bool {
        let data = self.get_state();
        if !data.frame_seen {
            return false;
        }
        data.frame_seen = false;

        // the buffer must be big enough for our data
        assert!(buffers.len() >= 384 * 224 * 2);
        unsafe {
            vb_get_pixels(
                self.sim,
                buffers.as_mut_ptr().cast(),
                2,
                384 * 2,
                buffers.as_mut_ptr().offset(1).cast(),
                2,
                384 * 2,
            );
        };
        true
    }

    pub fn read_samples(&mut self, samples: &mut Vec<f32>, weight: f32) {
        let mut position = 0;
        let ptr =
            unsafe { vb_get_samples(self.sim, ptr::null_mut(), ptr::null_mut(), &mut position) };
        // SAFETY: position is an offset in a buffer of (f32, f32). so, position * 2 is an offset in a buffer of f32.
        let read_samples: &[f32] =
            unsafe { slice::from_raw_parts(ptr.cast(), position as usize * 2) };
        samples.resize(read_samples.len(), 0.0);
        for (index, sample) in read_samples.iter().enumerate() {
            samples[index] += sample * weight;
        }

        unsafe {
            vb_set_samples(
                self.sim,
                ptr,
                VBDataType::F32,
                AUDIO_CAPACITY_SAMPLES as u32,
            )
        };
    }

    pub fn set_keys(&mut self, keys: VBKey) {
        unsafe { vb_set_keys(self.sim, keys.bits()) };
    }

    pub fn read_register(&mut self, register: VBRegister) -> u32 {
        match register {
            VBRegister::Program(index) => unsafe {
                vb_get_program_register(self.sim, index) as u32
            },
            VBRegister::System(index) => unsafe { vb_get_system_register(self.sim, index) as u32 },
            VBRegister::PC => unsafe { vb_get_program_counter(self.sim) },
        }
    }

    pub fn write_register(&mut self, register: VBRegister, value: u32) {
        match register {
            VBRegister::Program(index) => unsafe {
                vb_set_program_register(self.sim, index, value as i32);
            },
            VBRegister::System(index) => unsafe {
                vb_set_system_register(self.sim, index, value);
            },
            VBRegister::PC => unsafe {
                vb_set_program_counter(self.sim, value);
            },
        }
    }

    pub fn read_memory(&mut self, start: u32, length: usize, into: &mut Vec<u8>) {
        let mut address = start;
        for _ in 0..length {
            let byte = unsafe { vb_read(self.sim, address, VBDataType::U8) };
            into.push(byte as u8);
            address = address.wrapping_add(1);
        }
    }

    pub fn write_memory(&mut self, start: u32, buffer: &[u8]) {
        let mut address = start;
        for byte in buffer {
            unsafe { vb_write(self.sim, address, VBDataType::U8, *byte as i32) };
            address = address.wrapping_add(1);
        }
    }

    pub fn add_breakpoint(&mut self, address: u32) {
        let data = self.get_state();
        if let Err(index) = data.breakpoints.binary_search(&address) {
            data.breakpoints.insert(index, address);
        }

        unsafe {
            vb_set_execute_callback(self.sim, Some(on_execute));
        }
    }

    pub fn remove_breakpoint(&mut self, address: u32) {
        let data = self.get_state();
        if let Ok(index) = data.breakpoints.binary_search(&address) {
            data.breakpoints.remove(index);
            if !data.needs_execute_callback() {
                unsafe { vb_set_execute_callback(self.sim, None) };
            }
        }
    }

    pub fn add_watchpoint(&mut self, address: u32, length: usize, watch: VBWatchpointType) {
        match watch {
            VBWatchpointType::Read => self.add_read_watchpoint(address, length),
            VBWatchpointType::Write => self.add_write_watchpoint(address, length),
            VBWatchpointType::Access => {
                self.add_read_watchpoint(address, length);
                self.add_write_watchpoint(address, length);
            }
        }
    }

    fn add_read_watchpoint(&mut self, address: u32, length: usize) {
        let state = self.get_state();
        state.read_watchpoints.add(address, length);
        if !state.read_watchpoints.is_empty() {
            unsafe { vb_set_read_callback(self.sim, Some(on_read)) };
            unsafe { vb_set_execute_callback(self.sim, Some(on_execute)) };
        }
    }

    fn add_write_watchpoint(&mut self, address: u32, length: usize) {
        let state = self.get_state();
        state.write_watchpoints.add(address, length);
        if !state.write_watchpoints.is_empty() {
            unsafe { vb_set_write_callback(self.sim, Some(on_write)) };
            unsafe { vb_set_execute_callback(self.sim, Some(on_execute)) };
        }
    }

    pub fn remove_watchpoint(&mut self, address: u32, length: usize, watch: VBWatchpointType) {
        match watch {
            VBWatchpointType::Read => self.remove_read_watchpoint(address, length),
            VBWatchpointType::Write => self.remove_write_watchpoint(address, length),
            VBWatchpointType::Access => {
                self.remove_read_watchpoint(address, length);
                self.remove_write_watchpoint(address, length);
            }
        }
    }

    fn remove_read_watchpoint(&mut self, address: u32, length: usize) {
        let state = self.get_state();
        state.read_watchpoints.remove(address, length);
        let needs_execute = state.needs_execute_callback();
        if state.read_watchpoints.is_empty() {
            unsafe { vb_set_read_callback(self.sim, None) };
            if !needs_execute {
                unsafe { vb_set_execute_callback(self.sim, None) };
            }
        }
    }

    fn remove_write_watchpoint(&mut self, address: u32, length: usize) {
        let state = self.get_state();
        state.write_watchpoints.remove(address, length);
        let needs_execute = state.needs_execute_callback();
        if state.write_watchpoints.is_empty() {
            unsafe { vb_set_write_callback(self.sim, None) };
            if !needs_execute {
                unsafe { vb_set_execute_callback(self.sim, None) };
            }
        }
    }

    pub fn step(&mut self) {
        let current_pc = unsafe { vb_get_program_counter(self.sim) };
        let data = self.get_state();
        data.step_from = Some(current_pc);
        unsafe {
            vb_set_execute_callback(self.sim, Some(on_execute));
        }
    }

    pub fn clear_debug_state(&mut self) {
        let data = self.get_state();
        data.step_from = None;
        data.breakpoints.clear();
        data.read_watchpoints.clear();
        data.write_watchpoints.clear();
        unsafe { vb_set_read_callback(self.sim, None) };
        unsafe { vb_set_write_callback(self.sim, None) };
        unsafe { vb_set_execute_callback(self.sim, None) };
    }

    pub fn stop_reason(&mut self) -> Option<StopReason> {
        let data = self.get_state();
        let reason = data.stop_reason.take();
        if !data.needs_execute_callback() {
            unsafe { vb_set_execute_callback(self.sim, None) };
        }
        reason
    }

    fn get_state(&mut self) -> &mut VBState {
        // SAFETY: the *mut VB owns its userdata.
        // There is no way for the userdata to be null or otherwise invalid.
        unsafe { &mut *vb_get_user_data(self.sim).cast() }
    }
}

impl Drop for Sim {
    fn drop(&mut self) {
        let ptr =
            unsafe { vb_get_samples(self.sim, ptr::null_mut(), ptr::null_mut(), ptr::null_mut()) };
        // SAFETY: the audio buffer originally came from a Vec<u32>
        let floats: Vec<f32> = unsafe {
            Vec::from_raw_parts(ptr.cast(), AUDIO_CAPACITY_FLOATS, AUDIO_CAPACITY_FLOATS)
        };
        drop(floats);

        // SAFETY: the *mut VB owns its userdata.
        // There is no way for the userdata to be null or otherwise invalid.
        let ptr: *mut VBState = unsafe { vb_get_user_data(self.sim).cast() };
        // SAFETY: we made this pointer ourselves, we can for sure free it
        unsafe { drop(Box::from_raw(ptr)) };

        // If we're linked to another sim, unlink from them.
        unsafe { vb_set_peer(self.sim, ptr::null_mut()) };

        let len = unsafe { vb_size_of() }.div_ceil(4);
        // SAFETY: the sim's memory originally came from a Vec<u64>
        let bytes: Vec<u64> = unsafe { Vec::from_raw_parts(self.sim.cast(), len, len) };
        drop(bytes);
    }
}