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

use crate::codegen::ast_util;
use crate::codegen::syscall_map;
use anyhow::{anyhow, Result};
use oxc_ast::ast::*;
use oxc_ast_visit::{walk, Visit};
use oxc_syntax::scope::ScopeFlags;
use oxc_span::GetSpan;

/// AST Visitor that ensures the source code follows the Prometeu subset of JS/TS.
/// 
/// Since the Prometeu Virtual Machine is highly optimized and focused on game logic,
/// it does not support the full ECMA-262 specification (e.g., no `async`, `class`, 
/// `try/catch`, or `generators`).
pub struct Validator {
    /// List of validation errors found during the pass.
    errors: Vec<String>,
    /// Set of function names defined in the project, used to distinguish 
    /// local calls from invalid references.
    local_functions: std::collections::HashSet<String>,
}

impl Validator {
    /// Performs a full validation pass on a program.
    /// 
    /// Returns `Ok(())` if the program is valid, or a combined error message 
    /// listing all violations.
    pub fn validate(program: &Program) -> Result<()> {
        let mut validator = Self { 
            errors: Vec::new(),
            local_functions: std::collections::HashSet::new(),
        };
        
        // 1. Discovery Pass: Collect all function names and imports recursively
        validator.discover_functions(program);
        
        // 2. Traversal Pass: Check every node for compatibility
        validator.visit_program(program);

        if validator.errors.is_empty() {
            Ok(())
        } else {
            Err(anyhow!("Validation errors:\n{}", validator.errors.join("\n")))
        }
    }

    /// Recursively discovers all function declarations in the program.
    fn discover_functions(&mut self, program: &Program) {
        struct FunctionDiscoverer<'a> {
            functions: &'a mut std::collections::HashSet<String>,
        }
        impl<'a, 'b> Visit<'b> for FunctionDiscoverer<'a> {
            fn visit_function(&mut self, f: &Function<'b>, _flags: ScopeFlags) {
                if let Some(ident) = &f.id {
                    self.functions.insert(ident.name.to_string());
                }
                walk::walk_function(self, f, _flags);
            }
            fn visit_import_declaration(&mut self, decl: &ImportDeclaration<'b>) {
                if let Some(specifiers) = &decl.specifiers {
                    for specifier in specifiers {
                        match specifier {
                            ImportDeclarationSpecifier::ImportSpecifier(s) => {
                                self.functions.insert(s.local.name.to_string());
                            }
                            ImportDeclarationSpecifier::ImportDefaultSpecifier(s) => {
                                self.functions.insert(s.local.name.to_string());
                            }
                            ImportDeclarationSpecifier::ImportNamespaceSpecifier(s) => {
                                self.functions.insert(s.local.name.to_string());
                            }
                        }
                    }
                }
            }
        }
        let mut discoverer = FunctionDiscoverer { functions: &mut self.local_functions };
        discoverer.visit_program(program);
    }
}

impl<'a> Visit<'a> for Validator {
    /// Validates that only supported expressions are used.
    fn visit_expression(&mut self, expr: &Expression<'a>) {
        match expr {
            Expression::NumericLiteral(_) |
            Expression::BooleanLiteral(_) |
            Expression::StringLiteral(_) |
            Expression::NullLiteral(_) |
            Expression::Identifier(_) |
            Expression::AssignmentExpression(_) |
            Expression::BinaryExpression(_) |
            Expression::LogicalExpression(_) |
            Expression::UnaryExpression(_) |
            Expression::CallExpression(_) |
            Expression::StaticMemberExpression(_) => {
                // Basic JS logic is supported.
                walk::walk_expression(self, expr);
            }
            _ => {
                self.errors.push(format!("Unsupported expression type at {:?}. Note: Closures, arrow functions, and object literals are not yet supported.", expr.span()));
            }
        }
    }

    fn visit_call_expression(&mut self, expr: &CallExpression<'a>) {
        if let Ok(name) = ast_util::get_callee_name(&expr.callee) {
            if syscall_map::map_syscall(&name).is_none() && !self.local_functions.contains(&name) {
                self.errors.push(format!("Unsupported function call: {} at {:?}", name, expr.span));
            }
        } else {
            self.errors.push(format!("Unsupported callee expression at {:?}", expr.callee.span()));
        }
        walk::walk_call_expression(self, expr);
    }

    fn visit_unary_expression(&mut self, expr: &UnaryExpression<'a>) {
        match expr.operator {
            UnaryOperator::UnaryNegation |
            UnaryOperator::UnaryPlus |
            UnaryOperator::LogicalNot => {
                walk::walk_unary_expression(self, expr);
            }
            _ => {
                self.errors.push(format!("Unsupported unary operator {:?} at {:?}", expr.operator, expr.span));
            }
        }
    }

    fn visit_binary_expression(&mut self, expr: &BinaryExpression<'a>) {
        match expr.operator {
            BinaryOperator::Addition |
            BinaryOperator::Subtraction |
            BinaryOperator::Multiplication |
            BinaryOperator::Division |
            BinaryOperator::Equality |
            BinaryOperator::Inequality |
            BinaryOperator::LessThan |
            BinaryOperator::GreaterThan |
            BinaryOperator::LessEqualThan |
            BinaryOperator::GreaterEqualThan => {
                walk::walk_binary_expression(self, expr);
            }
            _ => {
                self.errors.push(format!("Unsupported binary operator {:?} at {:?}", expr.operator, expr.span));
            }
        }
    }

    fn visit_logical_expression(&mut self, expr: &LogicalExpression<'a>) {
        match expr.operator {
            LogicalOperator::And |
            LogicalOperator::Or => {
                walk::walk_logical_expression(self, expr);
            }
            _ => {
                self.errors.push(format!("Unsupported logical operator {:?} at {:?}", expr.operator, expr.span));
            }
        }
    }

    /// Validates that only supported statements are used.
    fn visit_statement(&mut self, stmt: &Statement<'a>) {
        match stmt {
            Statement::VariableDeclaration(_) |
            Statement::ExpressionStatement(_) |
            Statement::IfStatement(_) |
            Statement::BlockStatement(_) |
            Statement::ExportNamedDeclaration(_) |
            Statement::ImportDeclaration(_) |
            Statement::FunctionDeclaration(_) |
            Statement::ReturnStatement(_) => {
                // These are the only statements the PVM handles currently.
                walk::walk_statement(self, stmt);
            }
            _ => {
                self.errors.push(format!("Unsupported statement type at {:?}. Note: Prometeu does not support while/for loops or classes yet.", stmt.span()));
            }
        }
    }
}