endorphin-aster-vm/crates/compiler/src/analysis.rs

use std::collections::BTreeSet;

use endor_core::CodeLayout;
use frontend::ast::{AstRoot, Decl, DeclKind, Expr, ExprKind, Function, Spanned, StmtKind};
use frontend::{Diagnostic, DiagnosticKind, Span};

use crate::{CompileOptions, diagnostic};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScopeKind {
    Module,
    Expression,
    Interactive,
    Function,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NameBinding {
    Local(usize),
    Global(usize),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NameUse {
    pub name: String,
    pub span: Span,
    pub binding: NameBinding,
}

#[derive(Debug, Clone)]
pub struct Scope {
    pub kind: ScopeKind,
    pub name: String,
    pub span: Span,
    pub layout: CodeLayout,
    pub references: Vec<NameUse>,
}

#[derive(Debug, Clone)]
pub struct SymbolTable {
    pub root: Scope,
    /// Function and class-method scopes in declaration order. Nested declarations
    /// are rejected until source closure semantics are implemented.
    pub functions: Vec<Scope>,
}

fn error(span: Span, message: impl Into<String>) -> Diagnostic {
    diagnostic(DiagnosticKind::Semantic, span, message)
}

fn declaration_name(declaration: &Decl) -> &frontend::ast::Name {
    match &declaration.kind {
        DeclKind::Function(function) => &function.name,
        DeclKind::Class(class) => &class.name,
    }
}

/// Resolve a closed source unit plus explicitly declared host globals.
/// Function locals are collected before reading expressions; a later let
/// shadows the global for the entire function, but cannot be read early.
pub fn analyze(root: &AstRoot, options: &CompileOptions) -> Result<SymbolTable, Diagnostic> {
    if let AstRoot::Interactive(input) = root
        && input.statements.len() != 1
    {
        return Err(error(
            input.span,
            "interactive input must contain exactly one statement",
        ));
    }
    let mut globals: BTreeSet<String> = options.globals.iter().cloned().collect();
    if globals.contains("") {
        return Err(error(root.span(), "host global names must not be empty"));
    }
    let declarations: Vec<&Decl> = match root {
        AstRoot::Module(program) => program.declarations.iter().collect(),
        AstRoot::Interactive(input) => input
            .statements
            .iter()
            .filter_map(|statement| match &statement.kind {
                StmtKind::Declaration(declaration) => Some(declaration),
                _ => None,
            })
            .collect(),
        AstRoot::Expression(_) => Vec::new(),
    };
    let functions: Vec<&Function> = declarations
        .iter()
        .flat_map(|declaration| match &declaration.kind {
            DeclKind::Function(function) => vec![function],
            DeclKind::Class(class) => class.methods.iter().collect(),
        })
        .collect();
    let (kind, name) = match root {
        AstRoot::Module(_) => (ScopeKind::Module, "<module>"),
        AstRoot::Expression(_) => (ScopeKind::Expression, "<expression>"),
        AstRoot::Interactive(_) => (ScopeKind::Interactive, "<interactive>"),
    };
    let mut root_scope = Builder::new(kind, name, root.span());
    let mut declared = BTreeSet::new();
    for declaration in &declarations {
        let name = declaration_name(declaration);
        declare(&mut declared, &name.text, name.span)?;
        globals.insert(name.text.clone());
        root_scope.intern(&name.text);
    }
    for declaration in &declarations {
        if let DeclKind::Class(class) = &declaration.kind {
            let base_name = class
                .base
                .as_ref()
                .map_or("Object", |base| base.text.as_str());
            if !globals.contains(base_name) {
                return Err(error(
                    class.span,
                    format!("class base is not supplied: {base_name}"),
                ));
            }
            root_scope.intern(base_name);
            let mut members = BTreeSet::new();
            for field in &class.fields {
                declare(&mut members, &field.name.text, field.span)?;
            }
            for method in &class.methods {
                declare(&mut members, &method.name.text, method.span)?;
                if method
                    .parameters
                    .first()
                    .is_none_or(|parameter| parameter.name.text != "self")
                {
                    return Err(error(
                        method.span,
                        "class method requires an explicit first self parameter",
                    ));
                }
            }
        }
    }
    // Analyze all bodies after registering all functions, enabling forward
    // references and recursion without instantiating function objects here.
    let scopes = functions
        .iter()
        .map(|function| function_scope(function, &globals))
        .collect::<Result<Vec<_>, _>>()?;
    match root {
        AstRoot::Module(_) => {}
        AstRoot::Expression(expression) => root_scope.expression(expression, &globals)?,
        AstRoot::Interactive(input) => {
            for statement in &input.statements {
                match &statement.kind {
                    StmtKind::Binding(binding) => {
                        declare(&mut declared, &binding.name.text, binding.name.span)?;
                        root_scope.expression(&binding.value, &globals)?;
                        globals.insert(binding.name.text.clone());
                        root_scope.intern(&binding.name.text);
                    }
                    StmtKind::Assignment(assignment) => {
                        root_scope.expression(&assignment.value, &globals)?;
                        globals.insert(assignment.name.text.clone());
                        root_scope.intern(&assignment.name.text);
                    }
                    StmtKind::MemberAssignment(assignment) => {
                        root_scope.expression(&assignment.object, &globals)?;
                        root_scope.expression(&assignment.value, &globals)?;
                        root_scope.intern(&assignment.name.text);
                    }
                    StmtKind::Declaration(_) => {}
                    StmtKind::Expression(expression) => {
                        root_scope.expression(expression, &globals)?
                    }
                    StmtKind::Return(_) => {
                        return Err(error(statement.span, "return outside a function"));
                    }
                }
            }
        }
    }
    Ok(SymbolTable {
        root: root_scope.scope,
        functions: scopes,
    })
}

fn declare(names: &mut BTreeSet<String>, name: &str, span: Span) -> Result<(), Diagnostic> {
    if name.is_empty() || !names.insert(name.to_owned()) {
        return Err(error(
            span,
            format!("duplicate or empty declaration: {name}"),
        ));
    }
    Ok(())
}

fn function_scope(function: &Function, globals: &BTreeSet<String>) -> Result<Scope, Diagnostic> {
    let mut builder = Builder::new(ScopeKind::Function, &function.name.text, function.span);
    let mut locals = BTreeSet::new();
    for parameter in &function.parameters {
        declare(&mut locals, &parameter.name.text, parameter.name.span)?;
        builder
            .scope
            .layout
            .parameters
            .push(parameter.name.text.clone());
        builder
            .scope
            .layout
            .locals
            .push(parameter.name.text.clone());
        builder.initialized.insert(parameter.name.text.clone());
    }
    for statement in &function.body.statements {
        match &statement.kind {
            StmtKind::Binding(binding) => {
                declare(&mut locals, &binding.name.text, binding.name.span)?;
                builder.scope.layout.locals.push(binding.name.text.clone());
            }
            StmtKind::Assignment(assignment) => {
                if locals.insert(assignment.name.text.clone()) {
                    builder
                        .scope
                        .layout
                        .locals
                        .push(assignment.name.text.clone());
                }
            }
            StmtKind::Declaration(_) => {
                return Err(error(
                    statement.span,
                    "nested declarations are not compiled yet",
                ));
            }
            _ => {}
        }
    }
    let mut returned = false;
    for statement in &function.body.statements {
        if returned {
            return Err(error(
                statement.span,
                "statement after return is unreachable",
            ));
        }
        match &statement.kind {
            StmtKind::Binding(binding) => {
                builder.expression(&binding.value, globals)?;
                builder.initialized.insert(binding.name.text.clone());
            }
            StmtKind::Assignment(assignment) => {
                builder.expression(&assignment.value, globals)?;
                builder.initialized.insert(assignment.name.text.clone());
            }
            StmtKind::Return(statement) => {
                if let Some(value) = &statement.value {
                    builder.expression(value, globals)?;
                }
                returned = true;
            }
            StmtKind::MemberAssignment(assignment) => {
                builder.expression(&assignment.object, globals)?;
                builder.expression(&assignment.value, globals)?;
                builder.intern(&assignment.name.text);
            }
            StmtKind::Expression(expression) => builder.expression(expression, globals)?,
            StmtKind::Declaration(_) => unreachable!("nested declaration rejected above"),
        }
    }
    Ok(builder.scope)
}

struct Builder {
    scope: Scope,
    initialized: BTreeSet<String>,
}

impl Builder {
    fn new(kind: ScopeKind, name: &str, span: Span) -> Self {
        Self {
            scope: Scope {
                kind,
                name: name.into(),
                span,
                layout: CodeLayout::default(),
                references: Vec::new(),
            },
            initialized: BTreeSet::new(),
        }
    }

    fn intern(&mut self, name: &str) -> usize {
        if let Some(index) = self
            .scope
            .layout
            .names
            .iter()
            .position(|candidate| candidate == name)
        {
            index
        } else {
            self.scope.layout.names.push(name.to_owned());
            self.scope.layout.names.len() - 1
        }
    }

    fn expression(
        &mut self,
        expression: &Expr,
        globals: &BTreeSet<String>,
    ) -> Result<(), Diagnostic> {
        match &expression.kind {
            ExprKind::Name(name) => {
                let binding = if let Some(index) = self
                    .scope
                    .layout
                    .locals
                    .iter()
                    .position(|local| local == name)
                {
                    if !self.initialized.contains(name) {
                        return Err(error(
                            expression.span,
                            format!("local variable read before initialization: {name}"),
                        ));
                    }
                    NameBinding::Local(index)
                } else if globals.contains(name) {
                    NameBinding::Global(self.intern(name))
                } else {
                    return Err(error(expression.span, format!("undefined name: {name}")));
                };
                self.scope.references.push(NameUse {
                    name: name.clone(),
                    span: expression.span,
                    binding,
                });
            }
            ExprKind::Unary { operand, .. } => self.expression(operand, globals)?,
            ExprKind::Binary { left, right, .. } => {
                self.expression(left, globals)?;
                self.expression(right, globals)?;
            }
            ExprKind::Member { object, name } => {
                self.expression(object, globals)?;
                self.intern(&name.text);
            }
            ExprKind::Call { callee, arguments } => {
                self.expression(callee, globals)?;
                for argument in arguments {
                    self.expression(argument, globals)?;
                }
            }
            ExprKind::Number(_) | ExprKind::String(_) | ExprKind::Symbol(_) => {}
        }
        Ok(())
    }
}