rlox/src/middleend/variable_resolution.rs

use std::collections::HashMap;

use crate::{
    backend::environment::EnvironmentStack,
    common::{
        ast::{AstNode, Expr, NodeId, Stmt},
        lox_result::{LoxError, LoxResult},
    },
    frontend::source_registry::SourceSlice,
    middleend::visit_ast::{walk_expr, walk_stmt, Visitor},
};

struct Resolver {
    scopes: EnvironmentStack<bool>,
    locals: HashMap<NodeId, usize>,
}

impl Resolver {
    pub fn new() -> Self {
        Resolver {
            scopes: EnvironmentStack::new(),
            locals: HashMap::new(),
        }
    }

    fn declare(&mut self, name: &String) {
        if self.scopes.is_empty() {
            return;
        }
        let scope = self.scopes.peek();
        let _ = self.scopes.set(name.clone(), false);
    }

    fn define(&mut self, name: &String) {
        if self.scopes.is_empty() {
            return;
        }
        let _ = self.scopes.set(name.clone(), true);
    }

    fn resolve_local(&mut self, id: NodeId, name: &String) {
        let depth = self.scopes.depth();
        for i in (0..depth).rev() {
            if self.scopes.scope_contains(i, name) {
                // Distance = how many scopes up from the innermost the binding lives.
                self.locals.insert(id, depth - 1 - i);
                return;
            }
        }
        // Not found in any tracked scope: assume global, record nothing.
    }
}

impl Visitor for Resolver {
    fn visit_stmt(&mut self, stmt: &AstNode<Stmt>) -> LoxResult<()> {
        match &stmt.node {
            Stmt::VarDeclaration { name, .. } => {
                self.declare(name);
                // `walk_stmt` resolves the initializer (if present).
                walk_stmt(self, stmt)?;
                self.define(name);
                Ok(())
            }
            // NOTE: assignment is now an `Expr::Assign`, not a statement.
            // Add an `Expr::Assign` arm to `visit_expr` to resolve assignments.
            Stmt::Block { .. } => {
                self.scopes.push_new_scope();
                walk_stmt(self, stmt)?;
                self.scopes.pop_scope();
                Ok(())
            }
            // Expression, Print, Return, If, While, For: default traversal.
            _ => walk_stmt(self, stmt),
        }
    }

    fn visit_expr(&mut self, expr: &AstNode<Expr>) -> LoxResult<()> {
        match &expr.node {
            Expr::Identifier { name, .. } => {
                if !self.scopes.is_empty() && self.scopes.get(name).is_ok() {
                    return Err(LoxError::ParseError {
                        source_slice: SourceSlice::synthetic(),
                        message: "Cant read local varialbe in it own initializer".to_string(),
                    });
                }
                walk_expr(self, expr)
            }
            Expr::Assign { name, .. } => {
                self.resolve_local(expr.id, name);
                walk_expr(self, expr)
            }
            _ => walk_expr(self, expr),
        }
    }
}