prometeu-runtime/crates/prometeu-compiler/src/compiler.rs

use std::path::{Path, PathBuf};
use anyhow::{Result, Context, anyhow};
use oxc_allocator::Allocator;
use oxc_parser::Parser;
use oxc_span::SourceType;
use oxc_ast::ast::*;
use crate::codegen::Codegen;
use crate::codegen::validator::Validator;
use std::fs;
use std::collections::{HashMap, VecDeque};
use prometeu_bytecode::disasm::disasm;
use serde::Serialize;

#[derive(Serialize)]
pub struct Symbol {
    pub pc: u32,
    pub file: String,
    pub line: usize,
    pub col: usize,
}

pub struct CompilationUnit {
    pub rom: Vec<u8>,
    pub symbols: Vec<Symbol>,
}

impl CompilationUnit {
    pub fn export(&self, out: &Path, emit_disasm: bool, emit_symbols: bool) -> Result<()> {
        fs::write(out, &self.rom).with_context(|| format!("Failed to write PBC to {:?}", out))?;

        if emit_symbols {
            let symbols_path = out.with_file_name("symbols.json");
            let symbols_json = serde_json::to_string_pretty(&self.symbols)?;
            fs::write(&symbols_path, symbols_json)?;
        }

        if emit_disasm {
            let disasm_path = out.with_extension("disasm.txt");
            
            // Try to parse as PBC, if fails use raw
            let rom_to_disasm = if let Ok(pbc) = prometeu_bytecode::pbc::parse_pbc(&self.rom) {
                pbc.rom
            } else {
                self.rom.clone()
            };

            let instructions = disasm(&rom_to_disasm).map_err(|e| anyhow::anyhow!("Disassembly failed: {}", e))?;
            
            let mut disasm_text = String::new();
            for instr in instructions {
                let symbol = self.symbols.iter().find(|s| s.pc == instr.pc);
                let comment = if let Some(s) = symbol {
                    format!(" ; {}:{}", s.file, s.line)
                } else {
                    "".to_string()
                };

                let operands_str = instr.operands.iter()
                    .map(|o| format!("{:?}", o))
                    .collect::<Vec<_>>()
                    .join(" ");
                
                disasm_text.push_str(&format!("{:08X}  {:?} {}{}\n", instr.pc, instr.opcode, operands_str, comment));
            }
            fs::write(disasm_path, disasm_text)?;
        }

        Ok(())
    }
}

fn resolve_import(base_path: &Path, import_str: &str) -> Result<PathBuf> {
    let mut path = base_path.parent().unwrap().join(import_str);
    if !path.exists() {
        if path.with_extension("ts").exists() {
            path.set_extension("ts");
        } else if path.with_extension("js").exists() {
            path.set_extension("js");
        }
    }
    if !path.exists() {
        return Err(anyhow!("Cannot resolve import '{}' from {:?}", import_str, base_path));
    }
    Ok(path.canonicalize()?)
}

pub fn compile(entry: &Path) -> Result<CompilationUnit> {
    let allocator = Allocator::default();
    let mut modules = HashMap::new();
    let mut queue = VecDeque::new();
    
    let entry_abs = entry.canonicalize()
        .with_context(|| format!("Failed to canonicalize entry path: {:?}", entry))?;
    queue.push_back(entry_abs.clone());

    while let Some(path) = queue.pop_front() {
        let path_str = path.to_string_lossy().to_string();
        if modules.contains_key(&path_str) {
            continue;
        }

        let source_text = fs::read_to_string(&path)
            .with_context(|| format!("Failed to read file: {:?}", path))?;
        
        // Allocate source_text in the allocator to keep it alive
        let source_text_ptr: &str = allocator.alloc_str(&source_text);
        
        let source_type = SourceType::from_path(&path).unwrap_or_default();
        let parser_ret = Parser::new(&allocator, source_text_ptr, source_type).parse();

        if !parser_ret.errors.is_empty() {
            for error in parser_ret.errors {
                eprintln!("{:?}", error);
            }
            return Err(anyhow!("Failed to parse module: {:?}", path));
        }

        // Validate individual module
        Validator::validate(&parser_ret.program)?;

        // Find imports to add to queue
        for item in &parser_ret.program.body {
            if let Statement::ImportDeclaration(decl) = item {
                let import_path = decl.source.value.as_str();
                let resolved = resolve_import(&path, import_path)?;
                queue.push_back(resolved);
            }
        }

        modules.insert(path_str, (source_text_ptr, parser_ret.program));
    }

    let entry_str = entry_abs.to_string_lossy().to_string();
    let mut program_list = Vec::new();
    
    // Add entry program first
    let entry_data = modules.get(&entry_str).ok_or_else(|| anyhow!("Entry module not found after loading"))?;
    program_list.push((entry_str.clone(), entry_data.0.to_string(), &entry_data.1));
    
    // Add all other programs
    for (path, (source, program)) in &modules {
        if path != &entry_str {
            program_list.push((path.clone(), source.to_string(), program));
        }
    }

    let mut codegen = Codegen::new(entry_str.clone(), entry_data.0.to_string());
    let rom = codegen.compile_programs(program_list)?;

    Ok(CompilationUnit {
        rom,
        symbols: codegen.symbols,
    })
}