rlox/src/middleend/variable_resolution.rs

//! Static variable resolution (Crafting Interpreters, chapter 11).
//!
//! Walks the AST with the shared [`Visitor`] traversal, tracking lexical scopes
//! in a [`ScopeStack`], and records for every variable *reference* how many
//! scopes up its declaration lives (`locals: NodeId -> distance`). It also
//! reports the resolution errors from chapter 11:
//!
//! * reading a local variable in its own initializer (`var a = a;`),
//! * declaring two variables with the same name in one local scope,
//! * `return` outside of any function.
//!
//! Diagnostics accumulate in an [`ErrorSink`] so a single pass surfaces them
//! all. The produced `locals` map is keyed by [`NodeId`] (stable identity),
//! ready for the interpreter to consume via distance-based lookup.

use std::collections::HashMap;

use crate::{
    common::{
        ast::{AstNode, Expr, NodeId, Stmt},
        base_value::{BaseValue, LoxFunction},
        lox_result::{ErrorSink, LoxError, LoxResult},
    },
    middleend::{
        scope_stack::ScopeStack,
        visit_ast::{walk_expr, walk_function, walk_stmt, Visitor},
    },
};

/// Tracks whether resolution is currently inside a function body, so a
/// top-level `return` can be reported.
#[derive(Clone, Copy, PartialEq)]
enum FunctionType {
    None,
    Function,
}

pub struct Resolver {
    scopes: ScopeStack<bool>,
    locals: HashMap<NodeId, usize>,
    errors: ErrorSink,
    current_function: FunctionType,
}

impl Resolver {
    pub fn new() -> Self {
        Resolver {
            scopes: ScopeStack::new(),
            locals: HashMap::new(),
            errors: ErrorSink::new(),
            current_function: FunctionType::None,
        }
    }

    /// Resolve a whole program, returning the per-reference scope distances or
    /// the accumulated resolution errors.
    pub fn resolve_program(
        statements: &[AstNode<Stmt>],
    ) -> Result<HashMap<NodeId, usize>, Vec<LoxError>> {
        let mut resolver = Resolver::new();
        for statement in statements {
            // Visiting only fails for fatal/internal errors, which this pass
            // never produces; user-facing diagnostics go to `errors`.
            let _ = resolver.visit_stmt(statement);
        }
        if resolver.errors.has_errors() {
            Err(resolver.errors.into_errors())
        } else {
            Ok(resolver.locals)
        }
    }

    fn declare(&mut self, name: &str, slice: &crate::frontend::source_registry::SourceSlice) {
        if self.scopes.is_empty() {
            return; // global scope is untracked
        }
        if self.scopes.declared_in_current(name) {
            self.errors.report(LoxError::ParseError {
                source_slice: slice.clone(),
                message: format!("Already a variable named '{}' in this scope.", name),
            });
        }
        self.scopes.declare(name.to_string(), false);
    }

    fn define(&mut self, name: &str) {
        if self.scopes.is_empty() {
            return;
        }
        self.scopes.set_local(name, true);
    }

    fn resolve_local(&mut self, id: NodeId, name: &str) {
        if let Some(distance) = self.scopes.resolve(name) {
            self.locals.insert(id, distance);
        }
        // Not found locally: assume global, record nothing.
    }
}

impl Default for Resolver {
    fn default() -> Self {
        Self::new()
    }
}

impl Visitor for Resolver {
    fn visit_stmt(&mut self, stmt: &AstNode<Stmt>) -> LoxResult<()> {
        match &stmt.node {
            Stmt::VarDeclaration {
                name, initializer, ..
            } => {
                // A function declaration is a var bound to a function literal.
                // Define its name *before* resolving the body so it can recurse;
                // a plain variable is defined *after* its initializer so that
                // `var a = a;` is caught.
                let is_function = matches!(
                    initializer.as_deref().map(|node| &node.node),
                    Some(Expr::Literal {
                        value: BaseValue::Function(_)
                    })
                );
                self.declare(name, &stmt.source_slice);
                if is_function {
                    self.define(name);
                    walk_stmt(self, stmt)?;
                } else {
                    walk_stmt(self, stmt)?; // resolves the initializer, if any
                    self.define(name);
                }
                Ok(())
            }
            Stmt::Block { .. } => {
                self.scopes.begin_scope();
                walk_stmt(self, stmt)?;
                self.scopes.end_scope();
                Ok(())
            }
            Stmt::Return { .. } => {
                if self.current_function == FunctionType::None {
                    self.errors.report(LoxError::ParseError {
                        source_slice: stmt.source_slice.clone(),
                        message: "Can't return from top-level code.".to_string(),
                    });
                }
                walk_stmt(self, stmt) // resolve the returned expression
            }
            // Expression, Print, 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.get_local(name) == Some(&false) {
                    self.errors.report(LoxError::ParseError {
                        source_slice: expr.source_slice.clone(),
                        message: "Can't read local variable in its own initializer.".to_string(),
                    });
                }
                self.resolve_local(expr.id, name);
                Ok(())
            }
            Expr::Assign { name, .. } => {
                walk_expr(self, expr)?; // resolve the assigned value first
                self.resolve_local(expr.id, name);
                Ok(())
            }
            _ => walk_expr(self, expr),
        }
    }

    fn visit_function(&mut self, function: &LoxFunction) -> LoxResult<()> {
        let enclosing = std::mem::replace(&mut self.current_function, FunctionType::Function);
        self.scopes.begin_scope();
        // Parameters carry no source slice of their own; use the body's.
        let param_slice = function.body.source_slice.clone();
        for (param, _ty) in &function.parameters {
            self.declare(param, &param_slice);
            self.define(param);
        }
        walk_function(self, function)?; // resolves guard + body
        self.scopes.end_scope();
        self.current_function = enclosing;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::frontend::lexer::Lexer;
    use crate::frontend::parser::Parser;

    fn parse(src: &str) -> Vec<AstNode<Stmt>> {
        let tokens = Lexer::new(src.to_string(), 0)
            .scans_tokens()
            .expect("source should lex");
        Parser::new(tokens).parse().expect("source should parse")
    }

    fn resolve(src: &str) -> Result<HashMap<NodeId, usize>, Vec<LoxError>> {
        Resolver::resolve_program(&parse(src))
    }

    #[test]
    fn global_variables_are_not_resolved() {
        // Top-level (global) scope is untracked, so nothing is recorded.
        let locals = resolve("x := 1; print x;").expect("should resolve");
        assert!(locals.is_empty());
    }

    #[test]
    fn local_read_resolves_to_distance_zero() {
        let locals = resolve("do x := 1; print x; end").expect("should resolve");
        assert_eq!(locals.len(), 1);
        assert_eq!(*locals.values().next().unwrap(), 0);
    }

    #[test]
    fn nested_scope_resolves_to_outer_distance() {
        let src = "do x := 1; do print x; end end";
        let locals = resolve(src).expect("should resolve");
        assert_eq!(locals.len(), 1);
        // `x` is read one scope above where it is read from.
        assert_eq!(*locals.values().next().unwrap(), 1);
    }

    #[test]
    fn reading_a_variable_in_its_own_initializer_is_an_error() {
        let errors = resolve("do x := x; end").expect_err("should fail");
        assert!(errors
            .iter()
            .any(|e| e.get_message().contains("its own initializer")));
    }

    #[test]
    fn duplicate_declaration_in_same_scope_is_an_error() {
        let errors = resolve("do x := 1; x := 2; end").expect_err("should fail");
        assert!(errors
            .iter()
            .any(|e| e.get_message().contains("Already a variable")));
    }

    #[test]
    fn return_at_top_level_is_an_error() {
        let errors = resolve("return 1;").expect_err("should fail");
        assert!(errors.iter().any(|e| e.get_message().contains("top-level")));
    }

    #[test]
    fn return_inside_a_function_is_allowed() {
        let locals = resolve("f :: fn (n) do return n; end").expect("should resolve");
        // `n` resolves from the body block up to the parameter scope.
        assert_eq!(locals.len(), 1);
        assert_eq!(*locals.values().next().unwrap(), 1);
    }

    #[test]
    fn assignment_target_is_resolved() {
        let locals = resolve("do x := 1; x = 2; end").expect("should resolve");
        // Both the assignment target and... only the target is a reference here.
        assert!(locals.values().all(|&d| d == 0));
        assert!(!locals.is_empty());
    }
}