nescript/src/rom/mod.rs

#[cfg(test)]
mod tests;

use crate::parser::ast::{Mapper, Mirroring};

/// iNES header magic bytes
const INES_MAGIC: [u8; 4] = [0x4E, 0x45, 0x53, 0x1A]; // "NES\x1A"

/// PRG ROM bank size (16 KB)
const PRG_BANK_SIZE: usize = 16384;

/// CHR ROM bank size (8 KB)
const CHR_BANK_SIZE: usize = 8192;

/// Build a complete iNES ROM file.
///
/// Supports both single-bank (NROM) and multi-bank (MMC1, `UxROM`,
/// MMC3) layouts. When you call [`RomBuilder::set_prg`] the builder
/// treats the bytes as a single 16 KB bank and pads out. When you
/// call [`RomBuilder::set_prg_banks`] the builder writes each bank
/// back-to-back in the order you provided. The iNES header's PRG
/// bank count always reflects the actual number of 16 KB slots.
pub struct RomBuilder {
    /// One Vec per 16 KB PRG bank, in physical order. An empty
    /// outer Vec means no PRG has been set yet; a single inner Vec
    /// means a classic NROM layout.
    prg_banks: Vec<Vec<u8>>,
    chr_data: Vec<u8>,
    mapper: u8,
    mirroring: Mirroring,
}

impl RomBuilder {
    pub fn new(mirroring: Mirroring) -> Self {
        Self {
            prg_banks: Vec::new(),
            chr_data: Vec::new(),
            mapper: 0, // NROM
            mirroring,
        }
    }

    #[allow(dead_code)]
    pub fn set_mapper(&mut self, mapper: u8) {
        self.mapper = mapper;
    }

    /// Set the PRG ROM data as a single bank. Will be padded to fill
    /// 16 KB or 32 KB. Equivalent to calling `set_prg_banks` with a
    /// one- or two-element Vec depending on whether the data crosses
    /// the 16 KB boundary.
    pub fn set_prg(&mut self, data: Vec<u8>) {
        // Preserve the legacy single-bank behaviour: if the data is
        // <= 16 KB we emit exactly one 16 KB bank; if it's larger
        // (but still <= 32 KB) we split into two consecutive banks
        // so the iNES header byte 4 reflects 2, matching the old
        // `set_prg` contract used by all current NROM tests.
        if data.len() <= PRG_BANK_SIZE {
            self.prg_banks = vec![data];
        } else {
            let mut first = data;
            let second = first.split_off(PRG_BANK_SIZE);
            self.prg_banks = vec![first, second];
        }
    }

    /// Set the PRG ROM data as a list of 16 KB banks. Each bank will
    /// be padded with $FF to fill its 16 KB slot. The banks are
    /// written in the order provided — for mapper-specific layouts
    /// the caller (usually the Linker) is responsible for placing
    /// the fixed bank last.
    ///
    /// # Panics
    /// Panics if any single bank exceeds 16 KB, which would indicate
    /// a compiler bug (the allocator is expected to overflow-check
    /// before calling the rom builder).
    pub fn set_prg_banks(&mut self, banks: Vec<Vec<u8>>) {
        for (i, bank) in banks.iter().enumerate() {
            assert!(
                bank.len() <= PRG_BANK_SIZE,
                "PRG bank {i} exceeds 16 KB ({} bytes)",
                bank.len()
            );
        }
        self.prg_banks = banks;
    }

    /// Set the CHR ROM data. Will be padded to fill 8 KB.
    pub fn set_chr(&mut self, data: Vec<u8>) {
        self.chr_data = data;
    }

    /// Build the complete .nes file.
    pub fn build(self) -> Vec<u8> {
        // Determine PRG size. If no banks were set, fall back to a
        // single empty bank so the ROM has a valid 16 KB PRG slot.
        let prg_banks_vec: Vec<Vec<u8>> = if self.prg_banks.is_empty() {
            vec![Vec::new()]
        } else {
            self.prg_banks
        };
        let prg_banks = prg_banks_vec.len();
        let prg_size = prg_banks * PRG_BANK_SIZE;

        // CHR: 1 bank (8 KB), or 0 if using CHR RAM
        let chr_banks = usize::from(!self.chr_data.is_empty());
        let chr_size = chr_banks * CHR_BANK_SIZE;

        let mut rom = Vec::with_capacity(16 + prg_size + chr_size);

        // iNES header (16 bytes)
        rom.extend_from_slice(&INES_MAGIC);
        rom.push(prg_banks as u8); // PRG ROM banks (16 KB units)
        rom.push(chr_banks as u8); // CHR ROM banks (8 KB units)

        // Flags 6: mirroring, mapper low nibble
        let mut flags6 = match self.mirroring {
            Mirroring::Horizontal => 0,
            Mirroring::Vertical => 1,
        };
        flags6 |= (self.mapper & 0x0F) << 4;
        rom.push(flags6);

        // Flags 7: mapper high nibble
        let flags7 = self.mapper & 0xF0;
        rom.push(flags7);

        // Bytes 8-15: padding zeros
        rom.extend_from_slice(&[0u8; 8]);

        // PRG ROM data: each bank padded to 16 KB, concatenated.
        for mut bank in prg_banks_vec {
            bank.resize(PRG_BANK_SIZE, 0xFF);
            rom.extend_from_slice(&bank);
        }

        // CHR ROM data (padded to fill)
        if chr_banks > 0 {
            let mut chr = self.chr_data;
            chr.resize(chr_size, 0x00);
            rom.extend_from_slice(&chr);
        }

        rom
    }
}

/// Validate that a byte slice looks like a valid iNES ROM.
pub fn validate_ines(data: &[u8]) -> Result<RomInfo, &'static str> {
    if data.len() < 16 {
        return Err("file too small for iNES header");
    }
    if data[0..4] != INES_MAGIC {
        return Err("invalid iNES magic bytes");
    }

    let prg_banks = data[4] as usize;
    let chr_banks = data[5] as usize;
    let expected_size = 16 + prg_banks * PRG_BANK_SIZE + chr_banks * CHR_BANK_SIZE;

    if data.len() < expected_size {
        return Err("file too small for declared ROM banks");
    }

    let mirroring = if data[6] & 1 == 0 {
        Mirroring::Horizontal
    } else {
        Mirroring::Vertical
    };

    let mapper = (data[6] >> 4) | (data[7] & 0xF0);

    Ok(RomInfo {
        prg_banks,
        chr_banks,
        mapper,
        mirroring,
    })
}

#[derive(Debug)]
pub struct RomInfo {
    pub prg_banks: usize,
    pub chr_banks: usize,
    pub mapper: u8,
    pub mirroring: Mirroring,
}

/// Map a `Mapper` enum variant to its iNES mapper number.
pub fn mapper_number(mapper: Mapper) -> u8 {
    match mapper {
        Mapper::NROM => 0,
        Mapper::MMC1 => 1,
        Mapper::UxROM => 2,
        Mapper::MMC3 => 4,
    }
}