rlox/src/middleend/variable_resolution.rs

use std::collections::HashMap;

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

struct Resolver {
    scopes: EnvironmentStack<bool>,
    locals: HashMap<SourceSlice, 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 _ = 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, name: &String) {
        let depth = self.scopes.depth();
        for i in (0..depth).rev() {
            if self.scopes.scope_contains(i, name) {
                self.locals.insert(, i);
            }
        }
    }
}

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(())
            }
            Stmt::VarAssigment { name, .. } => {
                match self.scopes.get(name) {
                    Ok(true) => (),
                    Ok(false) => {
                        return Err(LoxError::RuntimeError {
                            source_slice: SourceSlice::default(),
                            message: "Cant read local variable in it own lintilizer".to_string(),
                        })
                    }
                    Err(err) => return Err(err),
                }
                // `walk_stmt` resolves the assigned value.
                walk_stmt(self, stmt)
            }
            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::default(),
                        message: "Cant read local varialbe in it own initializer".to_string(),
                    });
                }
                Ok(())
            }
            _ => walk_expr(self, expr),
        }
    }
}