rlox/src/backend/environment.rs

use crate::{
    common::lox_result::{runtime_error, LoxResult},
    frontend::source_registry::SourceSlice,
};
use std::collections::HashMap;
use std::fmt::Debug;

pub type Environment<T> = HashMap<String, T>;

#[derive(Debug, Clone, PartialEq, Default)]
pub struct EnvironmentStack<T: Clone + Debug + PartialEq> {
    stack: Vec<HashMap<String, T>>,
}

impl<T: Clone + Debug + PartialEq> EnvironmentStack<T> {
    pub fn new() -> Self {
        EnvironmentStack {
            stack: vec![HashMap::new()],
        }
    }

    pub fn is_empty(&self) -> bool {
        self.stack.is_empty()
    }

    pub fn push_new_scope(&mut self) {
        self.stack.push(HashMap::new());
    }

    pub fn push_existing_scope(&mut self, scope: HashMap<String, T>) {
        self.stack.push(scope);
    }

    pub fn pop_scope(&mut self) {
        self.stack.pop();
    }

    pub fn clone_last_scope(&mut self) -> HashMap<String, T> {
        self.stack.last().unwrap().clone()
    }

    pub fn get(&self, name: &str) -> LoxResult<T> {
        let size = self.stack.len();

        for i in (0..size).rev() {
            if let Some(value) = self.stack[i].get(name) {
                return Ok(value.clone());
            }
        }
        runtime_error(
            SourceSlice::synthetic(), // todo change this to the actual source slice
            format!("Undefined variable '{}'", name),
        )
    }

    pub fn declare(&mut self, name: String, value: T) -> LoxResult<T> {
        self.stack.last_mut().unwrap().insert(name, value.clone());
        Ok(value)
    }

    pub fn set(&mut self, name: String, value: T) -> LoxResult<T> {
        let size = self.stack.len();

        for i in (0..size).rev() {
            if let Some(_) = self.stack[i].get(&name) {
                let mut founded = false;
                self.stack[i].entry(name.clone()).and_modify(|v| {
                    *v = value.clone();
                    founded = true;
                });

                if founded {
                    return Ok(value);
                }
            }
        }
        runtime_error(
            SourceSlice::synthetic(), // todo change this to the actual source slice
            format!("Undefined variable '{}'", name),
        )
    }

    pub fn depth(&self) -> usize {
        self.stack.len()
    }
    pub fn scope_contains(&self, index: usize, name: &str) -> bool {
        self.stack[index].contains_key(name)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn declare_and_get() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.declare("a".to_string(), 1).unwrap();
        assert_eq!(env.get("a").unwrap(), 1);
    }

    #[test]
    fn get_undefined_is_error() {
        let env: EnvironmentStack<i32> = EnvironmentStack::new();
        assert!(env.get("nope").is_err());
    }

    #[test]
    fn declare_overwrites_in_same_scope() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.declare("a".to_string(), 1).unwrap();
        env.declare("a".to_string(), 2).unwrap();
        assert_eq!(env.get("a").unwrap(), 2);
    }

    #[test]
    fn inner_scope_shadows_outer_and_unshadows_on_pop() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.declare("a".to_string(), 1).unwrap();
        env.push_new_scope();
        env.declare("a".to_string(), 2).unwrap();
        assert_eq!(env.get("a").unwrap(), 2);
        env.pop_scope();
        assert_eq!(env.get("a").unwrap(), 1);
    }

    #[test]
    fn get_falls_through_to_outer_scope() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.declare("a".to_string(), 1).unwrap();
        env.push_new_scope();
        assert_eq!(env.get("a").unwrap(), 1);
    }

    #[test]
    fn set_updates_existing_value_in_outer_scope() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.declare("a".to_string(), 1).unwrap();
        env.push_new_scope();
        env.set("a".to_string(), 9).unwrap();
        env.pop_scope();
        assert_eq!(env.get("a").unwrap(), 9);
    }

    #[test]
    fn set_undefined_is_error() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        assert!(env.set("a".to_string(), 1).is_err());
    }

    #[test]
    fn pop_scope_discards_inner_declarations() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.push_new_scope();
        env.declare("temp".to_string(), 5).unwrap();
        env.pop_scope();
        assert!(env.get("temp").is_err());
    }

    #[test]
    fn push_existing_scope_makes_values_visible() {
        let mut scope = HashMap::new();
        scope.insert("k".to_string(), 7);
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.push_existing_scope(scope);
        assert_eq!(env.get("k").unwrap(), 7);
    }

    #[test]
    fn clone_last_scope_returns_top_scope() {
        let mut env: EnvironmentStack<i32> = EnvironmentStack::new();
        env.declare("a".to_string(), 1).unwrap();
        let scope = env.clone_last_scope();
        assert_eq!(scope.get("a"), Some(&1));
    }
}