openbirch-rs/src/lib/environment.rs

use std::{collections::HashMap, rc::Rc, thread::current};

use crate::node::{
    NodeEnum,
    function::{self, Function, FunctionType, NativeFunctionType},
};

pub type EnvironmentInternalSymbolKey = u16;
// pub type Environment = HashMap<EnvironmentInternalSymbolKey, Rc<NodeEnum>>;

struct Scope {
    pub stack_len: usize,
    pub shadow_len: usize,
}

pub struct Environment {
    map: HashMap<EnvironmentInternalSymbolKey, Rc<NodeEnum>>,
    stack_depth: usize,
    max_stack_depth: usize,

    stack_sizes: Vec<Scope>,
    stack: Vec<EnvironmentInternalSymbolKey>,
    stack_shadows: Vec<(EnvironmentInternalSymbolKey, Rc<NodeEnum>)>,

    hidden: Vec<EnvironmentInternalSymbolKey>,
    enable_closures: bool,

    float_precision: u64,

    // stack: Vec<Vec<EnvironmentInternalSymbolKey>>,
    symbol_to_id: HashMap<String, EnvironmentInternalSymbolKey>,
    id_to_symbol: HashMap<EnvironmentInternalSymbolKey, String>,
    unique_keys: EnvironmentInternalSymbolKey,
}

impl Environment {
    pub fn new() -> Self {
        Self {
            map: HashMap::new(),
            stack_depth: 0,
            max_stack_depth: 5000,

            stack_sizes: vec![],
            stack: vec![],
            stack_shadows: vec![],
            hidden: vec![],
            enable_closures: false,

            float_precision: 64,

            symbol_to_id: HashMap::new(),
            id_to_symbol: HashMap::new(),
            unique_keys: 0,
        }
    }

    pub fn hide(&mut self, key: EnvironmentInternalSymbolKey) {
        match self.hidden.binary_search(&key) {
            Ok(_) => {}
            Err(pos) => self.hidden.insert(pos, key),
        }
    }

    pub fn unhide(&mut self, key: EnvironmentInternalSymbolKey) {
        match self.hidden.binary_search(&key) {
            Ok(pos) => {
                self.hidden.remove(pos);
            }
            Err(_) => {}
        }
    }

    pub fn get_float_precision(&self) -> u64 {
        self.float_precision
    }

    pub fn set_float_precision(&mut self, precision: u64) {
        self.float_precision = precision;
    }

    pub fn disable_closures(&mut self) {
        self.enable_closures = true;
    }

    pub fn enable_closures(&mut self) {
        self.enable_closures = false;
    }

    pub fn closures_enabled(&self) -> bool {
        self.enable_closures
    }

    pub fn push_stack(&mut self) -> Result<(), String> {
        if self.stack_depth == self.max_stack_depth {
            self.unwind_stack();
            return Err("Max Stack depth exceeded".to_owned());
        }

        self.stack_depth += 1;
        self.stack_sizes.push(Scope {
            stack_len: 0,
            shadow_len: 0,
        });

        // debug_assert!(self.stack_depth == self.stack.len() - 1);

        Ok(())
    }

    pub fn pop_stack(&mut self) -> Result<(), String> {
        if self.stack_depth == 0 {
            return Err("Trying to pop empty stack".to_owned());
        }

        self.stack_depth -= 1;
        let scope = self.stack_sizes.pop().unwrap();
        for _ in 0..scope.stack_len {
            self.map.remove(&self.stack.pop().unwrap());
        }
        for _ in 0..scope.shadow_len {
            let (id, value) = self.stack_shadows.pop().unwrap();
            self.map.insert(id, value);
        }

        Ok(())
    }

    fn unwind_stack(&mut self) {
        for scope in &self.stack_sizes {
            for _ in 0..scope.stack_len {
                self.map.remove(&self.stack.pop().unwrap());
            }
        }

        self.stack_sizes = vec![];
        self.stack_shadows = vec![];
        self.stack_depth = 0;
    }

    pub fn define(&mut self, name: impl Into<String>, value: Rc<NodeEnum>) {
        let name = name.into();
        let id = if let Some(value) = self.str_to_id(&name) {
            *value
        } else {
            self.get_new_id()
        };

        self.map.insert(id, value);
        self.insert_id_to_str(id, name);
    }

    pub fn undefine(&mut self, key: EnvironmentInternalSymbolKey) {
        self.map.remove(&key);
    }

    pub fn define_native_function(&mut self, name: &'static str, func: NativeFunctionType) {
        let f = Function::new(FunctionType::Native(name, func));

        self.define(name.to_owned(), f);
    }

    pub fn get(&self, key: &EnvironmentInternalSymbolKey) -> Option<&Rc<NodeEnum>> {
        if self.hidden.binary_search(key).is_ok() {
            // This value is hidden, so we pretend its not defined
            return None;
        }
        self.map.get(key)
    }

    fn get_current_scope(&mut self) -> &[u16] {
        let len = self.stack.len();
        let current_scope_size = self.stack_sizes.last().unwrap().stack_len;
        let start = len - current_scope_size;
        self.stack.get(start..len).unwrap()
    }

    pub fn insert(&mut self, key: EnvironmentInternalSymbolKey, value: Rc<NodeEnum>) {
        if self.stack_depth != 0 {
            let mut shadow = false;
            if self.map.get(&key).is_some() {
                if !self.get_current_scope().contains(&key) {
                    // We need to shadow this variable
                    self.stack_shadows
                        .push((key, self.map.insert(key, value.clone()).unwrap()));
                    shadow = true;
                }
            }
            let scope = unsafe { self.stack_sizes.get_unchecked_mut(self.stack_depth - 1) };
            scope.stack_len += 1;
            self.stack.push(key);
            if shadow {
                scope.shadow_len += 1;
                return;
            }
        }
        self.map.insert(key, value);
    }

    pub fn str_to_id(&self, value: &String) -> Option<&EnvironmentInternalSymbolKey> {
        self.symbol_to_id.get(value)
    }
    pub fn id_to_str(&self, value: &EnvironmentInternalSymbolKey) -> Option<&String> {
        self.id_to_symbol.get(value)
    }
    pub fn insert_str_to_id(&mut self, key: String, value: EnvironmentInternalSymbolKey) {
        self.symbol_to_id.insert(key.clone(), value);
        self.id_to_symbol.insert(value, key);
    }
    pub fn insert_id_to_str(&mut self, key: EnvironmentInternalSymbolKey, value: String) {
        self.id_to_symbol.insert(key, value.clone());
        self.symbol_to_id.insert(value, key);
    }
    pub fn get_new_id(&mut self) -> EnvironmentInternalSymbolKey {
        self.unique_keys += 1;
        self.unique_keys
    }
}