rlox/src/backend/variable_resolution.rs

use crate::{
    backend::environment::EnvironmentStack,
    common::{
        ast::{AstNode, AstNodeKind, Expr, Stmt},
        lox_result::{runtime_error, LoxError, LoxResult},
    },
    frontend::source_registry::SourceSlice,
};

use std::fmt::{Debug, Display};
struct Resolver {
    scopes: EnvironmentStack<bool>,
}

impl Resolver {
    pub fn new() -> Self {
        Resolver {
            scopes: EnvironmentStack::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);
    }
}

pub trait StaticAnalyzer<T> {
    fn resolve(&mut self, node: &T) -> LoxResult<bool>;
}

impl<R: AstNodeKind + Clone + Debug + Display> StaticAnalyzer<AstNode<R>> for Resolver
where
    Resolver: StaticAnalyzer<R>,
{
    fn resolve(&mut self, node: &AstNode<R>) -> LoxResult<bool> {
        self.resolve(&node.node)
    }
}

impl StaticAnalyzer<Stmt> for Resolver {
    fn resolve(&mut self, node: &Stmt) -> LoxResult<bool> {
        match node {
            Stmt::Expression { expression, .. } => {
                // Visit the expression
                self.resolve(expression.as_ref())
            }
            Stmt::Print { expression, .. } => {
                // Visit the expression
                self.resolve(expression.as_ref())
            }
            Stmt::VarDeclaration {
                initializer, name, ..
            } => {
                self.declare(name);
                // Visit the initializer if present
                if let Some(init) = initializer {
                    self.resolve(init.as_ref())?;
                }
                self.define(name);
                Ok(true)
            }
            Stmt::VarAssigment { value, name, .. } => {
                match self.scopes.get(name) {
                    Ok(true) => (),
                    Ok(false) => {
                        return Err(LoxError::RuntimeError {
                            source_slice: SourceSlice::default(),
                            message: "Cant read loac variable in it own lintilizer".to_string(),
                        })
                    }
                    Err(err) => return Err(err),
                }
                self.resolve(value.as_ref())
            }
            Stmt::Return { expression, .. } => {
                // Visit the return expression
                self.resolve(expression.as_ref())
            }
            Stmt::Block { statements, .. } => {
                self.scopes.push_new_scope();
                // Visit all statements in the block
                for stmt in statements.iter() {
                    self.resolve(stmt)?;
                }
                self.scopes.pop_scope();
                Ok(true)
            }
            Stmt::If {
                condition,
                then_branch,
                elif_branch,
                else_branch,
                ..
            } => {
                // Visit the condition
                self.resolve(condition.as_ref())?;

                // Visit the then branch
                self.resolve(then_branch.as_ref())?;

                // Visit all elif branches
                for (elif_condition, elif_stmt) in elif_branch.iter() {
                    self.resolve(elif_condition.as_ref())?;
                    self.resolve(elif_stmt.as_ref())?;
                }

                // Visit the else branch if present
                if let Some(else_stmt) = else_branch {
                    self.resolve(else_stmt.as_ref())?;
                }

                Ok(true)
            }
            Stmt::While {
                condition, body, ..
            } => {
                // Visit the condition
                self.resolve(condition.as_ref())?;

                // Visit the body
                self.resolve(body.as_ref())?;
                Ok(true)
            }
            Stmt::For {
                variable,
                condition,
                increment,
                body,
                ..
            } => {
                // Visit the variable initialization
                self.resolve(variable.as_ref())?;

                // Visit the condition
                self.resolve(condition.as_ref())?;

                // Visit the increment
                self.resolve(increment.as_ref())?;

                // Visit the body
                self.resolve(body.as_ref())?;
                Ok(true)
            }
        }
    }
}

impl StaticAnalyzer<Expr> for Resolver {
    fn resolve(&mut self, node: &Expr) -> LoxResult<bool> {
        match node {
            Expr::Literal { .. } => {
                // Leaf node - no children to visit
                Ok(true)
            }
            Expr::Binary { left, right, .. } => {
                // Visit left operand
                self.resolve(left.as_ref())?;

                // Visit right operand
                self.resolve(right.as_ref())
            }
            Expr::Unary { operand, .. } => {
                // Visit the operand
                self.resolve(operand.as_ref())
            }
            Expr::Grouping { expression } => {
                // Visit the grouped expression
                self.resolve(expression.as_ref())
            }
            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(true)
            }
            Expr::Call {
                callee, arguments, ..
            } => {
                // Visit the callee
                self.resolve(callee.as_ref())?;

                // Visit all arguments
                for arg in arguments.iter() {
                    self.resolve(arg)?;
                }

                Ok(true)
            }
        }
    }
}