nescript/src/main.rs

use clap::Parser;
use std::path::{Path, PathBuf};

use nescript::analyzer;
use nescript::assets;
use nescript::codegen::IrCodeGen;
use nescript::errors::render_diagnostics;
use nescript::ir;
use nescript::linker::{render_mlb, render_source_map, Linker, PrgBank};
use nescript::optimizer;
use nescript::parser::ast::BankType;

#[derive(Parser)]
#[command(name = "nescript", about = "NEScript compiler — NES game development")]
enum Cli {
    /// Compile a .ne source file into a .nes ROM
    Build {
        /// Input source file
        input: PathBuf,

        /// Output ROM file (default: input with .nes extension)
        #[arg(short, long)]
        output: Option<PathBuf>,

        /// Enable debug mode (runtime checks, debug.log)
        #[arg(long)]
        debug: bool,

        /// Dump generated 6502 assembly to stdout
        #[arg(long)]
        asm_dump: bool,

        /// Dump the lowered IR program to stdout (after optimization)
        #[arg(long)]
        dump_ir: bool,

        /// Dump a human-readable memory map of variable allocations
        /// to stdout.
        #[arg(long)]
        memory_map: bool,

        /// Dump a call graph showing which functions call which.
        #[arg(long)]
        call_graph: bool,

        /// Skip the IR optimizer pass. Useful for bisecting
        /// optimizer-introduced miscompiles: if a program misbehaves
        /// with the optimizer on but works with `--no-opt`, the bug
        /// lives in `src/optimizer/`.
        #[arg(long)]
        no_opt: bool,

        /// Write a Mesen-compatible symbol file (`.mlb`) next to the
        /// ROM. Contains one `<type>:<address>:<label>` entry per
        /// function, state handler, and user variable. Enables
        /// symbol-level debugging in Mesen / fceux without manual
        /// address lookups.
        #[arg(long, value_name = "PATH")]
        symbols: Option<PathBuf>,

        /// Write a plain-text source map (`.map`) next to the ROM.
        /// Each line has the form `<rom_offset_hex> <file_id>
        /// <line> <col>` and records the position of every IR-level
        /// statement in the assembled fixed bank. Useful for
        /// reverse-mapping a crash address back to the source.
        #[arg(long, value_name = "PATH")]
        source_map: Option<PathBuf>,
    },
    /// Type-check a source file without building
    Check {
        /// Input source file
        input: PathBuf,
    },
}

fn main() {
    let cli = Cli::parse();

    match cli {
        Cli::Build {
            input,
            output,
            debug,
            asm_dump,
            dump_ir,
            memory_map,
            call_graph,
            no_opt,
            symbols,
            source_map,
        } => {
            let output = output.unwrap_or_else(|| input.with_extension("nes"));
            match compile(
                &input,
                &CompileOptions {
                    debug,
                    asm_dump,
                    dump_ir,
                    memory_map,
                    call_graph,
                    no_opt,
                    symbols: symbols.clone(),
                    source_map: source_map.clone(),
                },
            ) {
                Ok(rom) => {
                    std::fs::write(&output, rom).unwrap_or_else(|e| {
                        eprintln!("error: failed to write {}: {e}", output.display());
                        std::process::exit(1);
                    });
                    println!(
                        "compiled {} -> {} ({} bytes)",
                        input.display(),
                        output.display(),
                        std::fs::metadata(&output).map(|m| m.len()).unwrap_or(0)
                    );
                }
                Err(()) => std::process::exit(1),
            }
        }
        Cli::Check { input } => match check(&input) {
            Ok(()) => println!("no errors found in {}", input.display()),
            Err(()) => std::process::exit(1),
        },
    }
}

/// Print a human-readable memory map of variable allocations.
/// Entries are sorted by address and labelled with their scope
/// (zero-page vs RAM).
fn print_memory_map(analysis: &nescript::analyzer::AnalysisResult) {
    let mut allocs: Vec<_> = analysis.var_allocations.iter().collect();
    allocs.sort_by_key(|a| a.address);

    println!("=== NEScript Memory Map ===");
    println!("Zero Page ($00-$FF):");
    println!("  $00-$0F  [SYSTEM]  reserved (frame flag, input, state, params, scratch)");
    for a in allocs.iter().filter(|a| a.address < 0x100) {
        if a.size == 1 {
            println!("  ${:04X}    [USER]    {} (u8)", a.address, a.name);
        } else {
            println!(
                "  ${:04X}-${:04X}  [USER]  {} ({} bytes)",
                a.address,
                a.address + a.size - 1,
                a.name,
                a.size
            );
        }
    }

    let ram_allocs: Vec<_> = allocs.iter().filter(|a| a.address >= 0x100).collect();
    if !ram_allocs.is_empty() {
        println!("\nRAM ($0200-$07FF):");
        println!("  $0200-$02FF  [SYSTEM]  OAM shadow buffer");
        for a in &ram_allocs {
            if a.size == 1 {
                println!("  ${:04X}        [USER]    {} (u8)", a.address, a.name);
            } else {
                println!(
                    "  ${:04X}-${:04X}  [USER]  {} ({} bytes)",
                    a.address,
                    a.address + a.size - 1,
                    a.name,
                    a.size
                );
            }
        }
    }

    // Summary line.
    let zp_used: u16 = allocs
        .iter()
        .filter(|a| a.address < 0x80)
        .map(|a| a.size)
        .sum();
    let ram_used: u16 = allocs
        .iter()
        .filter(|a| a.address >= 0x300)
        .map(|a| a.size)
        .sum();
    println!();
    println!("Zero Page: {zp_used}/128 bytes used");
    println!("Main RAM:  {ram_used}/1280 bytes used");
}

/// Print a human-readable call graph of the analyzed program.
/// Entries show the max call depth reached from each entry point
/// (state handler) and the transitive callees.
fn print_call_graph(analysis: &nescript::analyzer::AnalysisResult) {
    use std::collections::BTreeMap;

    let sorted: BTreeMap<_, _> = analysis
        .call_graph
        .iter()
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();
    let max_depth = analysis.max_depths.values().copied().max().unwrap_or(0);

    println!("=== Call Graph (max depth: {max_depth} / 8) ===");
    if sorted.is_empty() {
        println!("  (no functions or handlers)");
        return;
    }
    for (caller, callees) in &sorted {
        if let Some(depth) = analysis.max_depths.get(caller) {
            println!("{caller} (max depth {depth})");
        } else {
            println!("{caller}");
        }
        if callees.is_empty() {
            println!("  └── (leaf)");
        } else {
            let count = callees.len();
            for (i, callee) in callees.iter().enumerate() {
                let branch = if i + 1 == count {
                    "└──"
                } else {
                    "├──"
                };
                println!("  {branch} {callee}");
            }
        }
    }
}

fn dump_asm(instructions: &[nescript::asm::Instruction]) {
    use nescript::asm::{AddressingMode, Opcode};
    for inst in instructions {
        // A bare `NOP` with a `Label` operand is a label *definition*
        // (the pseudo-instruction the codegen emits when marking a
        // position). Any other opcode with `Label` mode is an actual
        // instruction like `JSR foo` or `JMP bar`, so we show the
        // opcode + target.
        if inst.opcode == Opcode::NOP {
            if let AddressingMode::Label(name) = &inst.mode {
                println!("{name}:");
                continue;
            }
        }
        println!("    {:?} {:?}", inst.opcode, inst.mode);
    }
}

#[allow(clippy::struct_excessive_bools)]
struct CompileOptions {
    debug: bool,
    asm_dump: bool,
    dump_ir: bool,
    memory_map: bool,
    call_graph: bool,
    no_opt: bool,
    symbols: Option<PathBuf>,
    source_map: Option<PathBuf>,
}

fn compile(input: &PathBuf, opts: &CompileOptions) -> Result<Vec<u8>, ()> {
    let debug = opts.debug;
    let asm_dump = opts.asm_dump;
    let dump_ir = opts.dump_ir;
    let memory_map = opts.memory_map;
    let call_graph = opts.call_graph;
    let no_opt = opts.no_opt;
    let symbols_path = opts.symbols.as_ref();
    let source_map_path = opts.source_map.as_ref();
    let raw_source = std::fs::read_to_string(input).map_err(|e| {
        eprintln!("error: failed to read {}: {e}", input.display());
    })?;

    // Preprocess: inline include directives
    let source = nescript::parser::preprocess_source(&raw_source, Some(input)).map_err(|e| {
        eprintln!("error: {e}");
    })?;

    let filename = input.to_string_lossy();

    // Parse
    let (program, parse_diags) = nescript::parser::parse(&source);
    if !parse_diags.is_empty() {
        render_diagnostics(&source, &filename, &parse_diags);
    }
    if parse_diags
        .iter()
        .any(nescript::errors::Diagnostic::is_error)
    {
        return Err(());
    }
    let program = program.ok_or(())?;

    // Analyze
    let analysis = analyzer::analyze(&program);
    if !analysis.diagnostics.is_empty() {
        render_diagnostics(&source, &filename, &analysis.diagnostics);
    }
    if analysis
        .diagnostics
        .iter()
        .any(nescript::errors::Diagnostic::is_error)
    {
        return Err(());
    }

    // IR lowering and (optionally) optimization. `--no-opt` skips
    // the IR optimizer pass entirely so optimizer-introduced
    // miscompiles can be bisected against the unoptimized output.
    let mut ir_program = ir::lower(&program, &analysis);
    if !no_opt {
        optimizer::optimize(&mut ir_program);
    }

    if dump_ir {
        print!("{}", ir_program.pretty());
    }

    if memory_map {
        print_memory_map(&analysis);
    }

    if call_graph {
        print_call_graph(&analysis);
    }

    // Resolve sprite assets (CHR data + tile indices) relative to the
    // source file's directory, so `@binary` / `@chr` paths work naturally.
    let source_dir = input.parent().unwrap_or_else(|| Path::new("."));
    let sprites = assets::resolve_sprites(&program, source_dir).map_err(|e| {
        eprintln!("error: {e}");
    })?;

    // Resolve audio assets: user-declared sfx/music plus any
    // builtins referenced via `play foo` / `start_music bar` for
    // names that aren't in the program's sfx/music declarations.
    let sfx = assets::resolve_sfx(&program).map_err(|e| {
        eprintln!("error: {e}");
    })?;
    let music = assets::resolve_music(&program).map_err(|e| {
        eprintln!("error: {e}");
    })?;

    // Resolve palette and background declarations into fixed-size
    // ROM data blobs. These are purely compile-time — the byte
    // arrays came from the parser and all the analyzer validation
    // has already run.
    let palettes = assets::resolve_palettes(&program);
    let backgrounds = assets::resolve_backgrounds(&program);

    // IR-based code generation. Lower → optimize → emit 6502.
    //
    // We hold on to the codegen after `generate()` because it
    // carries the source-location marker list — one entry per
    // `SourceLoc` IR op — which the CLI reads to emit a source
    // map. Dropping the codegen before then would throw that
    // metadata away. Source-marker emission is opt-in (the label
    // pseudo-ops shift peephole block boundaries, which would
    // flip release-mode ROM bytes if it was always on) — so we
    // only enable it when the user actually asked for a source
    // map on the command line.
    let emit_source_map = source_map_path.is_some();
    let mut codegen = IrCodeGen::new(&analysis.var_allocations, &ir_program)
        .with_sprites(&sprites)
        .with_audio(&sfx, &music)
        .with_debug(debug)
        .with_source_map(emit_source_map);
    let mut instructions = codegen.generate(&ir_program);

    // Peephole pass: cleans up the IR codegen's temp-heavy output —
    // dead stores, redundant loads, short-branch folds, etc.
    nescript::codegen::peephole::optimize(&mut instructions);

    if asm_dump {
        dump_asm(&instructions);
    }

    // Link into ROM with both graphic assets (sprite CHR) and audio
    // assets (sfx envelopes, music note streams) spliced in. We use
    // `Linker::with_mapper` so the iNES header's mapper byte
    // reflects the source's `mapper:` declaration — without this
    // the CLI always shipped mapper 0 (NROM) regardless of whether
    // the program actually needed MMC1/MMC3 bank switching.
    //
    // For banked mappers (MMC1, UxROM, MMC3) we collect the
    // declared `bank X: prg` entries and turn each into an empty
    // 16 KB switchable slot. User code currently still lives in
    // the fixed bank — the declared banks exist so programs that
    // outgrow 16 KB have real ROM space to grow into and so
    // mapper-specific fixtures (vectors, trampolines, bank-select
    // helpers) land in the right place.
    let linker = Linker::with_mapper(program.game.mirroring, program.game.mapper)
        .with_header(program.game.header);
    let switchable_banks: Vec<PrgBank> = program
        .banks
        .iter()
        .filter(|b| b.bank_type == BankType::Prg)
        .map(|b| PrgBank::empty(&b.name))
        .collect();
    let link_result = linker.link_banked_with_ppu_detailed(
        &instructions,
        &sprites,
        &sfx,
        &music,
        &palettes,
        &backgrounds,
        &switchable_banks,
    );

    if let Some(path) = symbols_path {
        let mlb = render_mlb(&link_result, &analysis.var_allocations);
        std::fs::write(path, mlb).map_err(|e| {
            eprintln!("error: failed to write symbol file {}: {e}", path.display());
        })?;
    }
    if let Some(path) = source_map_path {
        let map = render_source_map(&link_result, codegen.source_locs(), &source);
        std::fs::write(path, map).map_err(|e| {
            eprintln!("error: failed to write source map {}: {e}", path.display());
        })?;
    }

    Ok(link_result.rom)
}

fn check(input: &PathBuf) -> Result<(), ()> {
    let raw_source = std::fs::read_to_string(input).map_err(|e| {
        eprintln!("error: failed to read {}: {e}", input.display());
    })?;

    let source = nescript::parser::preprocess_source(&raw_source, Some(input)).map_err(|e| {
        eprintln!("error: {e}");
    })?;

    let filename = input.to_string_lossy();

    let (program, parse_diags) = nescript::parser::parse(&source);
    if !parse_diags.is_empty() {
        render_diagnostics(&source, &filename, &parse_diags);
    }
    if parse_diags
        .iter()
        .any(nescript::errors::Diagnostic::is_error)
    {
        return Err(());
    }
    let program = program.ok_or(())?;

    let analysis = analyzer::analyze(&program);
    if !analysis.diagnostics.is_empty() {
        render_diagnostics(&source, &filename, &analysis.diagnostics);
    }
    if analysis
        .diagnostics
        .iter()
        .any(nescript::errors::Diagnostic::is_error)
    {
        return Err(());
    }

    Ok(())
}