state.rs

    _padding: u8,

    io_registers: [u8; 128],

    ram: BessBuffer,
    vram: BessBuffer,
    mbc_ram: BessBuffer,
    oam: BessBuffer,
    hram: BessBuffer,
    background_palettes: BessBuffer,
    object_palettes: BessBuffer,
}

impl BessCore {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        model: String,
        pc: u16,
        af: u16,
        bc: u16,
        de: u16,
        hl: u16,
        sp: u16,
        ime: bool,
        ie: u8,
        execution_mode: u8,
        io_registers: [u8; 128],
    ) -> Self {
        Self {
            header: BessBlockHeader::new(
                String::from("CORE"),
                ((size_of::<u16>() * 2)
                    + size_of::<u32>()
                    + (size_of::<u16>() * 6)
                    + (size_of::<u8>() * 4)
                    + (size_of::<u8>() * 128)
                    + ((size_of::<u32>() + size_of::<u32>()) * 7)) as u32,
            ),
            major: 1,
            minor: 1,
            model,
            pc,
            af,
            bc,
            de,
            hl,
            sp,
            ime,
            ie,
            execution_mode,
            _padding: 0,
            io_registers,
            ram: BessBuffer::default(),
            vram: BessBuffer::default(),
            mbc_ram: BessBuffer::default(),
            oam: BessBuffer::default(),
            hram: BessBuffer::default(),
            background_palettes: BessBuffer::default(),
            object_palettes: BessBuffer::default(),
        }
    }

    pub fn from_data(data: &mut Cursor<Vec<u8>>) -> Self {
        let mut instance = Self::default();
        instance.read(data);
        instance
    }

    pub fn verify(&self) -> Result<(), Error> {
        if self.header.magic != "CORE" {
            return Err(Error::CustomError(String::from("Invalid magic")));
        }
        if self.oam.size != 0xa0 {
            return Err(Error::CustomError(String::from("Invalid OAM size")));
        }
        if self.hram.size != 0x7f {
            return Err(Error::CustomError(String::from("Invalid HRAM size")));
        }
        if (self.is_cgb() && self.background_palettes.size != 0x40)
            || (self.is_dmg() && self.background_palettes.size != 0x00)
        {
            return Err(Error::CustomError(String::from(
                "Invalid background palettes size",
            )));
        }
        if (self.is_cgb() && self.object_palettes.size != 0x40)
            || (self.is_dmg() && self.object_palettes.size != 0x00)
        {
            return Err(Error::CustomError(String::from(
                "Invalid object palettes size",
            )));
        }
        Ok(())
    }

    /// Obtains the BESS (Game Boy) model string using the
    /// provided `GameBoy` instance.
    fn bess_model(gb: &GameBoy) -> String {
        let mut buffer = [0x00_u8; 4];

        if gb.is_dmg() {
            buffer[0] = b'G';
        } else if gb.is_cgb() {
            buffer[0] = b'C';
        } else if gb.is_sgb() {
            buffer[0] = b'S';
        } else {
            buffer[0] = b' ';
        }

        if gb.is_dmg() {
            buffer[1] = b'D';
        } else if gb.is_cgb() {
            buffer[1] = b'C';
        } else if gb.is_sgb() {
            buffer[1] = b'N';
        } else {
            buffer[1] = b' ';
        }

        if gb.is_dmg() {
            buffer[2] = b'B';
        } else if gb.is_cgb() {
            buffer[2] = b'A';
        } else {
            buffer[2] = b' ';
        }

        buffer[3] = b' ';

        String::from_utf8(Vec::from(buffer)).unwrap()
    }

    fn is_dmg(&self) -> bool {
        if let Some(first_char) = self.model.chars().next() {
            return first_char == 'G';
        }
        false
    }

    fn is_cgb(&self) -> bool {
        if let Some(first_char) = self.model.chars().next() {
            return first_char == 'C';
        }
        false
    }
}

impl Serialize for BessCore {
    fn write(&mut self, buffer: &mut Cursor<Vec<u8>>) {
        self.header.write(buffer);

        buffer.write_all(&self.major.to_le_bytes()).unwrap();
        buffer.write_all(&self.minor.to_le_bytes()).unwrap();

        buffer.write_all(self.model.as_bytes()).unwrap();

        buffer.write_all(&self.pc.to_le_bytes()).unwrap();
        buffer.write_all(&self.af.to_le_bytes()).unwrap();
        buffer.write_all(&self.bc.to_le_bytes()).unwrap();
        buffer.write_all(&self.de.to_le_bytes()).unwrap();
        buffer.write_all(&self.hl.to_le_bytes()).unwrap();
        buffer.write_all(&self.sp.to_le_bytes()).unwrap();

        buffer.write_all(&(self.ime as u8).to_le_bytes()).unwrap();
        buffer.write_all(&self.ie.to_le_bytes()).unwrap();
        buffer
            .write_all(&self.execution_mode.to_le_bytes())
            .unwrap();
        buffer.write_all(&self._padding.to_le_bytes()).unwrap();

        buffer.write_all(&self.io_registers).unwrap();

        self.ram.write(buffer);
        self.vram.write(buffer);
        self.mbc_ram.write(buffer);
        self.oam.write(buffer);
        self.hram.write(buffer);
        self.background_palettes.write(buffer);
        self.object_palettes.write(buffer);
    }

    fn read(&mut self, data: &mut Cursor<Vec<u8>>) {
        self.header.read(data);

        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.major = u16::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.minor = u16::from_le_bytes(buffer);

        let mut buffer = [0x00; 4];
        data.read_exact(&mut buffer).unwrap();
        self.model = String::from_utf8(Vec::from(buffer)).unwrap();

        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.pc = u16::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.af = u16::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.bc = u16::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.de = u16::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.hl = u16::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u16>()];
        data.read_exact(&mut buffer).unwrap();
        self.sp = u16::from_le_bytes(buffer);

        let mut buffer = [0x00; 1];
        data.read_exact(&mut buffer).unwrap();
        self.ime = buffer[0] != 0;
        let mut buffer = [0x00; size_of::<u8>()];
        data.read_exact(&mut buffer).unwrap();
        self.ie = u8::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u8>()];
        data.read_exact(&mut buffer).unwrap();
        self.execution_mode = u8::from_le_bytes(buffer);
        let mut buffer = [0x00; size_of::<u8>()];
        data.read_exact(&mut buffer).unwrap();
        self._padding = u8::from_le_bytes(buffer);

        data.read_exact(&mut self.io_registers).unwrap();

        self.ram.read(data);
        self.vram.read(data);
        self.mbc_ram.read(data);
        self.oam.read(data);
        self.hram.read(data);
        self.background_palettes.read(data);
        self.object_palettes.read(data);
    }
}

impl State for BessCore {
    fn from_gb(gb: &mut GameBoy) -> Result<Self, Error> {
        let mut core = Self::new(
            Self::bess_model(gb),
            gb.cpu_i().pc(),
            gb.cpu_i().af(),
            gb.cpu_i().bc(),
            gb.cpu_i().de(),
            gb.cpu_i().hl(),
            gb.cpu_i().sp(),
            gb.cpu_i().ime(),
            gb.mmu_i().ie,
            u8::from(gb.cpu().halted()),
            // @TODO: these registers cannot be completely retrieved
            // and because of that some audio noise is played
            // The loading of the registers should be done in a much
            // more manual way like SameBoy does here
            // https://github.com/LIJI32/SameBoy/blob/7e6f1f866e89430adaa6be839aecc4a2ccabd69c/Core/save_state.c#L673
            gb.mmu().read_many_unsafe(0xff00, 128).try_into().unwrap(),
        );
        core.ram.fill_buffer(gb.mmu().ram());
        core.vram.fill_buffer(gb.ppu().vram_device());
        core.mbc_ram.fill_buffer(gb.rom_i().ram_data());
        core.oam.fill_buffer(&gb.mmu().read_many(0xfe00, 0x00a0));
        core.hram.fill_buffer(&gb.mmu().read_many(0xff80, 0x007f));
        if gb.is_cgb() {
            core.background_palettes
                .fill_buffer(&gb.ppu_i().palettes_color()[0]);
            core.object_palettes
                .fill_buffer(&gb.ppu_i().palettes_color()[1]);
        }
        Ok(core)
    }

    fn to_gb(&self, gb: &mut GameBoy) -> Result<(), Error> {
        gb.cpu().set_pc(self.pc);
        gb.cpu().set_af(self.af);
        gb.cpu().set_bc(self.bc);
        gb.cpu().set_de(self.de);
        gb.cpu().set_hl(self.hl);
        gb.cpu().set_sp(self.sp);

        gb.cpu().set_ime(self.ime);
        gb.mmu().ie = self.ie;

        match self.execution_mode {
            0 => gb.cpu().set_halted(false),
            1 => gb.cpu().set_halted(true),
            2 => gb.cpu().stop(),
            _ => unimplemented!(),
        }

        // @TODO: we need to be careful about this writing and
        // should make this a bit more robust, to handle this
        // special case/situations
        // The registers should be handled in a more manual manner
        // to avoid unwanted side effects
        // https://github.com/LIJI32/SameBoy/blob/7e6f1f866e89430adaa6be839aecc4a2ccabd69c/Core/save_state.c#L1003
        gb.mmu().write_many_unsafe(0xff00, &self.io_registers);

        gb.mmu().set_ram(self.ram.buffer.to_vec());
        gb.ppu().set_vram(&self.vram.buffer);
        gb.rom().set_ram_data(&self.mbc_ram.buffer);
        gb.mmu().write_many(0xfe00, &self.oam.buffer);
        gb.mmu().write_many(0xff80, &self.hram.buffer);

        if gb.is_cgb() {
            // updates the internal palettes for the CGB with the values
            // stored in the BESS state
            gb.ppu().set_palettes_color([
                self.background_palettes.buffer.to_vec().try_into().unwrap(),
                self.object_palettes.buffer.to_vec().try_into().unwrap(),
            ]);

            // updates the speed of the CGB according to the KEY1 register
            let is_double = self.io_registers[0x4d_usize] & 0x80 == 0x80;
            gb.mmu().set_speed(if is_double {
                GameBoySpeed::Double
            } else {
                GameBoySpeed::Normal
            });

            // need to disable DMA transfer to avoid unwanted
            // DMA transfers when loading the state
            gb.dma().set_active_hdma(false);
        }

        Ok(())
    }
}

impl Default for BessCore {
    fn default() -> Self {
        Self::new(
            String::from("GD  "),
            0x0000_u16,
            0x0000_u16,
            0x0000_u16,
            0x0000_u16,
            0x0000_u16,
            0x0000_u16,
            false,
            0x00,
            0,
            [0x00; 128],
        )
    }
}

pub struct BessMbrRegister {
    address: u16,
    value: u8,
}

impl BessMbrRegister {
    pub fn new(address: u16, value: u8) -> Self {
        Self { address, value }
    }
}

pub struct BessMbc {
    header: BessBlockHeader,
    registers: Vec<BessMbrRegister>,
}

impl BessMbc {
    pub fn new(registers: Vec<BessMbrRegister>) -> Self {
        Self {
            header: BessBlockHeader::new(
                String::from("MBC "),
                ((size_of::<u8>() + size_of::<u16>()) * registers.len()) as u32,
            ),
            registers,
        }
    }

    pub fn from_data(data: &mut Cursor<Vec<u8>>) -> Self {
        let mut instance = Self::default();
        instance.read(data);
        instance
    }
}

impl Serialize for BessMbc {
    fn write(&mut self, buffer: &mut Cursor<Vec<u8>>) {
        self.header.write(buffer);
        for register in self.registers.iter() {
            buffer.write_all(&register.address.to_le_bytes()).unwrap();
            buffer.write_all(&register.value.to_le_bytes()).unwrap();
        }
    }

    fn read(&mut self, data: &mut Cursor<Vec<u8>>) {
        self.header.read(data);
        for _ in 0..(self.header.size / 3) {
            let mut buffer = [0x00; size_of::<u16>()];
            data.read_exact(&mut buffer).unwrap();
            let address = u16::from_le_bytes(buffer);
            let mut buffer = [0x00; size_of::<u8>()];
            data.read_exact(&mut buffer).unwrap();
            let value = u8::from_le_bytes(buffer);
            self.registers.push(BessMbrRegister::new(address, value));
        }
    }
}

impl State for BessMbc {
    fn from_gb(gb: &mut GameBoy) -> Result<Self, Error> {
        let mut registers = vec![];
        match gb.cartridge().rom_type().mbc_type() {
            MbcType::NoMbc => (),
            MbcType::Mbc1 => {
                registers.push(BessMbrRegister::new(
                    0x0000,
                    if gb.rom().ram_enabled() {
                        0x0a_u8
                    } else {
                        0x00_u8
                    },
                ));
                registers.push(BessMbrRegister::new(
                    0x2000,
                    gb.rom().rom_bank() as u8 & 0x1f,
                ));
                registers.push(BessMbrRegister::new(0x4000, gb.rom().ram_bank()));
                registers.push(BessMbrRegister::new(0x6000, 0x00_u8));
            }
            MbcType::Mbc3 => {
                registers.push(BessMbrRegister::new(
                    0x0000,
                    if gb.rom().ram_enabled() {
                        0x0a_u8
                    } else {
                        0x00_u8
                    },
                ));
                registers.push(BessMbrRegister::new(0x2000, gb.rom().rom_bank() as u8));
                registers.push(BessMbrRegister::new(0x4000, gb.rom().ram_bank()));
            }
            MbcType::Mbc5 => {
                registers.push(BessMbrRegister::new(
                    0x0000,
                    if gb.rom().ram_enabled() {
                        0x0a_u8
                    } else {
                        0x00_u8
                    },
                ));
                registers.push(BessMbrRegister::new(0x2000, gb.rom().rom_bank() as u8));
                registers.push(BessMbrRegister::new(
                    0x3000,
                    (gb.rom().rom_bank() >> 8) as u8 & 0x01,
                ));
                registers.push(BessMbrRegister::new(0x4000, gb.rom().ram_bank()));
            }
            _ => unimplemented!(),
        }

        Ok(Self::new(registers))
    }

    fn to_gb(&self, gb: &mut GameBoy) -> Result<(), Error> {
        for register in self.registers.iter() {
            gb.mmu().write(register.address, register.value);
        }
        Ok(())
    }
}

impl Default for BessMbc {
    fn default() -> Self {
        Self::new(vec![])
    }
}

/// Top level manager structure containing the
/// entrypoint static methods for saving and loading
/// [BESS](https://github.com/LIJI32/SameBoy/blob/master/BESS.md) state
/// files and buffers for the Game Boy.
#[cfg_attr(feature = "wasm", wasm_bindgen)]
pub struct StateManager;

impl StateManager {
    pub fn save_file(
        file_path: &str,
        gb: &mut GameBoy,
        format: Option<SaveStateFormat>,
    ) -> Result<(), Error> {
        let mut file = File::create(file_path)
            .map_err(|_| Error::CustomError(format!("Failed to create file: {}", file_path)))?;
        let data = Self::save(gb, format)?;
        file.write_all(&data)
            .map_err(|_| Error::CustomError(format!("Failed to write to file: {}", file_path)))?;
        file.flush()
            .map_err(|_| Error::CustomError(format!("Failed to flush file: {}", file_path)))?;
        Ok(())
    }

    pub fn load_file(
        file_path: &str,
        gb: &mut GameBoy,
        format: Option<SaveStateFormat>,
    ) -> Result<(), Error> {
        let mut file = File::open(file_path)
            .map_err(|_| Error::CustomError(format!("Failed to open file: {}", file_path)))?;
        let mut data = vec![];
        file.read_to_end(&mut data)
            .map_err(|_| Error::CustomError(format!("Failed to read from file: {}", file_path)))?;
        Self::load(&data, gb, format)?;
        Ok(())
    }
}

impl StateManager {
    pub fn save(gb: &mut GameBoy, format: Option<SaveStateFormat>) -> Result<Vec<u8>, Error> {
        let mut data = Cursor::new(vec![]);
        match format {
            Some(SaveStateFormat::Bos) | None => {
                let mut state = BosState::from_gb(gb)?;
                state.write(&mut data);
            }
            Some(SaveStateFormat::Bess) => {
                let mut state = BessState::from_gb(gb)?;
                state.write(&mut data);
            }
        }
        Ok(data.into_inner())
    }

    pub fn load(
        data: &[u8],
        gb: &mut GameBoy,
        format: Option<SaveStateFormat>,
    ) -> Result<(), Error> {
        let data = &mut Cursor::new(data.to_vec());
        let format = match format {
            Some(format) => format,
            None => {
                if BosState::is_bos(data) {
                    SaveStateFormat::Bos
                } else if BessState::is_bess(data) {
                    SaveStateFormat::Bess
                } else {
                    return Err(Error::CustomError(String::from(
                        "Unknown save state file format",
                    )));
                }
            }
        };
        match format {
            SaveStateFormat::Bos => {
                let mut state = BosState::default();
                state.read(data);
                state.to_gb(gb)?;
            }
            SaveStateFormat::Bess => {
                let mut state = BessState::default();
                state.read(data);
                state.to_gb(gb)?;
            }
        }
        Ok(())
    }

    pub fn read_bos(data: &[u8]) -> Result<BosState, Error> {
        let data = &mut Cursor::new(data.to_vec());
        let mut state = BosState::default();
        state.read(data);
        Ok(state)
    }

    pub fn read_bess(data: &[u8]) -> Result<BessState, Error> {
        let data = &mut Cursor::new(data.to_vec());
        let mut state = BessState::default();
        state.read(data);
        Ok(state)
    }

    /// Obtains the thumbnail of the save state file, this thumbnail is
    /// stored in raw RGB format.
    ///
    /// This operation is currently only supported for the BOS format.
    pub fn thumbnail(data: &[u8], format: Option<SaveStateFormat>) -> Result<Vec<u8>, Error> {
        let data = &mut Cursor::new(data.to_vec());
        let format = match format {
            Some(format) => format,
            None => {
                if BosState::is_bos(data) {
                    SaveStateFormat::Bos
                } else if BessState::is_bess(data) {
                    SaveStateFormat::Bess
                } else {
                    return Err(Error::CustomError(String::from(
                        "Unknown save state file format",
                    )));
                }
            }
        };
        match format {
            SaveStateFormat::Bos => {
                let mut state = BosState::default();
                state.read(data);
                Ok(state.image_buffer.unwrap().image.to_vec())
            }
            SaveStateFormat::Bess => Err(Error::CustomError(String::from(
                "Format foes not support thumbnail",
            ))),
        }
    }
}

#[cfg(feature = "wasm")]
#[cfg_attr(feature = "wasm", wasm_bindgen)]
impl StateManager {
    pub fn save_wa(gb: &mut GameBoy, format: Option<SaveStateFormat>) -> Result<Vec<u8>, String> {
        Self::save(gb, format).map_err(|e| e.to_string())
    }

    pub fn load_wa(
        data: &[u8],
        gb: &mut GameBoy,
        format: Option<SaveStateFormat>,
    ) -> Result<(), String> {
        Self::load(data, gb, format).map_err(|e| e.to_string())
    }

    pub fn read_bos_wa(data: &[u8]) -> Result<BosState, String> {
        Self::read_bos(data).map_err(|e| e.to_string())
    }

    pub fn read_bess_wa(data: &[u8]) -> Result<BessState, String> {
        Self::read_bess(data).map_err(|e| e.to_string())
    }

    pub fn thumbnail_wa(data: &[u8], format: Option<SaveStateFormat>) -> Result<Vec<u8>, String> {
        Self::thumbnail(data, format).map_err(|e| e.to_string())
    }
}

#[cfg(test)]
mod tests {
    use boytacean_encoding::zippy::{decode_zippy, encode_zippy};

    use crate::{gb::GameBoy, state::State};

    use super::{BessCore, SaveStateFormat, StateManager};

    #[test]
    fn test_bess_core() {
        let mut gb = GameBoy::default();
        gb.load(true).unwrap();
        gb.load_rom_file("res/roms/test/firstwhite.gb", None)
            .unwrap();
        let bess_core = BessCore::from_gb(&mut gb).unwrap();
        assert_eq!(bess_core.model, "GDB ");
        assert_eq!(bess_core.pc, 0x0000);
        assert_eq!(bess_core.af, 0x0000);
        assert_eq!(bess_core.bc, 0x0000);
        assert_eq!(bess_core.de, 0x0000);
        assert_eq!(bess_core.hl, 0x0000);
        assert_eq!(bess_core.sp, 0x0000);
        assert_eq!(bess_core.ime, false);
        assert_eq!(bess_core.ie, 0x00);
        assert_eq!(bess_core.execution_mode, 0);
        assert_eq!(bess_core.io_registers.len(), 128);
        assert_eq!(
            bess_core.io_registers,
            [
                63, 0, 0, 255, 0, 0, 0, 248, 255, 255, 255, 255, 255, 255, 255, 224, 128, 63, 0,
                255, 191, 255, 63, 0, 255, 191, 127, 255, 159, 255, 191, 255, 255, 0, 0, 191, 0, 0,
                240, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 134, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 126, 255, 254, 0, 255, 255,
                255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
                255, 255, 255, 255, 255, 0, 0, 0, 0, 255, 255, 255, 255, 249, 255, 255, 255, 255,
                255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255
            ]
        );
        assert_eq!(bess_core.ram.size, 0x2000);
        assert_eq!(bess_core.vram.size, 0x2000);
        assert_eq!(bess_core.mbc_ram.size, 0x2000);
        assert_eq!(bess_core.oam.size, 0x00a0);
        assert_eq!(bess_core.hram.size, 0x007f);
        assert_eq!(bess_core.background_palettes.size, 0x0000);
        assert_eq!(bess_core.object_palettes.size, 0x0000);
    }

    #[test]
    fn test_load_bos() {
        let mut gb = GameBoy::default();
        gb.load(true).unwrap();
        gb.load_rom_file("res/roms/test/firstwhite.gb", None)
            .unwrap();
        let data = StateManager::save(&mut gb, Some(SaveStateFormat::Bos)).unwrap();
        StateManager::load(&data, &mut gb, Some(SaveStateFormat::Bos)).unwrap();
        StateManager::load(&data, &mut gb, None).unwrap();
    }

    #[test]
    fn test_load_bess() {
        let mut gb = GameBoy::default();
        gb.load(true).unwrap();
        gb.load_rom_file("res/roms/test/firstwhite.gb", None)
            .unwrap();
        let data = StateManager::save(&mut gb, Some(SaveStateFormat::Bess)).unwrap();
        StateManager::load(&data, &mut gb, Some(SaveStateFormat::Bess)).unwrap();
        StateManager::load(&data, &mut gb, None).unwrap();
    }

    #[test]
    fn test_compression() {
        let mut gb = GameBoy::default();
        gb.load(true).unwrap();
        gb.load_rom_file("res/roms/test/firstwhite.gb", None)
            .unwrap();
        gb.step_to(0x0100);
        let data = StateManager::save(&mut gb, Some(SaveStateFormat::Bess)).unwrap();
        let encoded = encode_zippy(&data).unwrap();
        let decoded = decode_zippy(&encoded).unwrap();
        assert_eq!(data, decoded);
        assert_eq!(encoded.len(), 811);
        assert_eq!(decoded.len(), 25154);
    }
}