rlox/interpreter/src/resolver/resolve.rs

use std::collections::HashMap;

use rlox2_frontend::parser::{Expr, Stmt};

use crate::Runtime;
use crate::{LoxError, ResolverError};

/*====================================================================================================================*/

type ResolverScope = HashMap<String, ResolveStatus>;
type ResolverFrame = Vec<ResolverScope>;

/*====================================================================================================================*/

#[derive(Debug, Clone, Copy)]
enum ResolveStatus {
    Declared,
    Defined,
}

/*====================================================================================================================*/

pub fn resolve(statement: &mut Stmt, runtime: &mut Runtime) -> Result<(), LoxError> {
    let global_names = runtime.globals().keys().cloned();

    let mut resolver = Resolver::new(global_names);

    resolver.resolve_stmt(statement)?;

    Ok(())
}

#[derive(Debug)]
struct Resolver {
    global_scope: ResolverScope,
    // local_scopes: Vec<ResolverScope>,
    frames: Vec<ResolverFrame>,

    closure_vars: HashMap<String, usize>,
}

impl Resolver {
    fn new(global_names: impl Iterator<Item = String>) -> Self {
        let mut global_scope = ResolverScope::new();

        for name in global_names {
            global_scope.insert(name, ResolveStatus::Defined);
        }

        Resolver {
            global_scope,
            frames: vec![ResolverFrame::new()],
            closure_vars: HashMap::new(),
        }
    }

    fn local_scopes(&self) -> &ResolverFrame {
        self.frames.last().unwrap()
    }

    fn local_scopes_mut(&mut self) -> &mut ResolverFrame {
        self.frames.last_mut().unwrap()
    }

    fn enter_scope(&mut self) {
        self.local_scopes_mut().push(ResolverScope::new());
    }

    fn exit_scope(&mut self) {
        self.local_scopes_mut()
            .pop()
            .expect("Tried to pop from empty ResolverFrame");
    }

    fn push_frame(&mut self) {
        self.frames.push(vec![ResolverScope::new()]);
    }

    fn pop_frame(&mut self) {
        self.frames.pop().expect("Tried to pop non-existant ResolverFrame");
    }

    fn declare(&mut self, name: &str) {
        let name = name.to_owned();
        if let Some(local_scope) = self.local_scopes_mut().last_mut() {
            local_scope.insert(name, ResolveStatus::Declared);
        } else {
            self.global_scope.insert(name, ResolveStatus::Declared);
        }
    }

    fn define(&mut self, name: &str) {
        let name = name.to_owned();
        if let Some(local_scope) = self.local_scopes_mut().last_mut() {
            local_scope.insert(name, ResolveStatus::Defined);
        } else {
            self.global_scope.insert(name, ResolveStatus::Declared);
        }
    }

    fn check(&self, name: &str) -> Option<ResolveStatus> {
        if let Some(local_scope) = self.local_scopes().last() {
            local_scope.get(name).cloned()
        } else {
            self.global_scope.get(name).cloned()
        }
    }

    fn resolve_var(&mut self, name: &str) -> Result<Expr, ResolverError> {
        let name = name.to_owned();

        let mut level = 0;

        for scope in self.local_scopes().iter().rev() {
            if scope.contains_key(&name) {
                return Ok(Expr::LocalVariable { name, level });
            }
            // DEBUG
            level += 1;
        }

        for frame in self.frames.iter().rev().skip(1) {
            for scope in frame.iter().rev() {
                if scope.contains_key(&name) {
                    if !self.closure_vars.contains_key(&name) {
                        // the level at which the closed-over variable will be collected from
                        // from the perspective of the parameter/closure scope i.e. the outmost of the function
                        self.closure_vars
                            .insert(name.clone(), level - self.local_scopes().len());
                    }
                    return Ok(Expr::LocalVariable {
                        name,
                        // distance from actual variable refernce to parameter/closure scope
                        level: self.local_scopes().len() - 1,
                    });
                }
                level += 1;
            }
        }

        if self.global_scope.contains_key(&name) {
            return Ok(Expr::GlobalVariable { name });
        }

        Err(ResolverError::UnresolvableVariable { name })
    }

    fn resolve_stmt(&mut self, stmt: &mut Stmt) -> Result<(), ResolverError> {
        match stmt {
            Stmt::Print { expr } => self.resolve_expr(expr),
            Stmt::IfStmt {
                condition,
                then_branch,
                else_branch,
            } => {
                self.resolve_expr(condition)?;
                self.resolve_stmt(then_branch)?;
                if let Some(else_branch) = else_branch {
                    self.resolve_stmt(else_branch)?;
                }
                Ok(())
            }
            Stmt::While { condition, body } => {
                self.resolve_expr(condition)?;
                self.resolve_stmt(body)?;
                Ok(())
            }
            Stmt::VarDecl { name, initializer } => {
                self.declare(name);
                self.resolve_expr(initializer)?;
                self.define(name);
                Ok(())
            }
            Stmt::Block { statements } => {
                self.enter_scope();
                for statement in statements.iter_mut() {
                    self.resolve_stmt(statement)?;
                }
                self.exit_scope();
                Ok(())
            }
            Stmt::Class { name, methods: _ } => {
                self.declare(name);
                self.define(name);
                Ok(())
            }
            Stmt::ExprStmt { expr } => self.resolve_expr(expr),
            Stmt::Break => Ok(()),
            Stmt::Return { expr } => self.resolve_expr(expr),
        }
    }

    fn resolve_expr(&mut self, expr: &mut Expr) -> Result<(), ResolverError> {
        match expr {
            Expr::Literal { literal: _ } => Ok(()),
            Expr::Unary { op: _, expr } => self.resolve_expr(expr),
            Expr::Binary { left, op: _, right } => {
                self.resolve_expr(left)?;
                self.resolve_expr(right)?;
                Ok(())
            }
            Expr::Logical { left, op: _, right } => {
                self.resolve_expr(left)?;
                self.resolve_expr(right)?;
                Ok(())
            }
            Expr::Grouping { expr } => self.resolve_expr(expr),
            Expr::Variable { name } => {
                if let Some(ResolveStatus::Declared) = self.check(name) {
                    let name = name.clone();
                    return Err(ResolverError::VarInOwnInitializer { name });
                }

                *expr = self.resolve_var(name)?;

                Ok(())
            }
            Expr::LocalVariable { name, level: _ } => {
                panic!("Tried to resolve variable {name} twice")
            }
            Expr::GlobalVariable { name } => {
                panic!("Tried to resolve variable {name} twice")
            }

            Expr::Assignment { target, value } => {
                self.resolve_expr(value)?;

                let target = target.as_mut();
                if let Expr::Variable { name } = target {
                    *target = self.resolve_var(name)?;
                } else {
                    panic!("Invalid assignment target {target}");
                }

                Ok(())
            }
            Expr::Call { callee, args } => {
                self.resolve_expr(callee)?;
                for arg in args.iter_mut() {
                    self.resolve_expr(arg)?;
                }
                Ok(())
            }
            Expr::Function {
                name,
                param_names,
                closure_vars,
                body,
            } => {
                let old_closure_names = std::mem::take(&mut self.closure_vars);

                self.push_frame();

                self.declare(name);
                self.define(name);

                for param_name in param_names {
                    self.declare(param_name);
                    self.define(param_name);
                }

                self.resolve_stmt(body)?;

                self.pop_frame();

                let closure_names = std::mem::replace(&mut self.closure_vars, old_closure_names);

                for closure_var in closure_names {
                    closure_vars.push(closure_var);
                }

                Ok(())
            }
        }
    }
}