deimos-lang/src/name_analysis/symbol_table.rs

use std::cell::RefCell;
use crate::name_analysis::symbol::{
    FunctionSymbol, ModuleSymbol, ParameterSymbol, Symbol, SymbolInner, TypeSymbol,
    UseStatementSymbol, VariableSymbol,
};
use crate::name_analysis::symbol_table::SymbolInsertError::SymbolAlreadyDefined;
use crate::name_analysis::symbol_table::SymbolLookupError::NoDefinition;
use std::collections::HashMap;
use std::fmt::Display;
use std::rc::Rc;
/* Scope */

#[derive(Debug)]
struct Scope {
    parent: Option<usize>,
    use_statement_symbols: HashMap<String, Rc<RefCell<UseStatementSymbol>>>,
    module_symbols: HashMap<String, Rc<ModuleSymbol>>,
    type_symbols: HashMap<String, Rc<TypeSymbol>>,
    function_symbols: HashMap<String, Rc<FunctionSymbol>>,
    parameter_symbols: HashMap<String, Rc<ParameterSymbol>>,
    variable_symbols: HashMap<String, Rc<VariableSymbol>>,
    debug_name: String,
}

impl Scope {
    pub fn new(parent: Option<usize>, debug_name: String) -> Scope {
        Scope {
            parent,
            use_statement_symbols: HashMap::new(),
            module_symbols: HashMap::new(),
            type_symbols: HashMap::new(),
            function_symbols: HashMap::new(),
            parameter_symbols: HashMap::new(),
            variable_symbols: HashMap::new(),
            debug_name,
        }
    }

    fn get_any_symbol(&self, name: &str) -> Option<Symbol> {
        self.variable_symbols.get(name)
            .map(|s| Symbol::Variable(s.clone()))
            .or_else(|| self.parameter_symbols.get(name).map(|s| Symbol::Parameter(s.clone())))
            .or_else(|| self.function_symbols.get(name).map(|s| Symbol::Function(s.clone())))
            .or_else(|| self.type_symbols.get(name).map(|ts| Symbol::Type(ts.clone())))
            .or_else(|| self.module_symbols.get(name).map(|ms| Symbol::Module(ms.clone())))
            .or_else(|| self.use_statement_symbols.get(name).map(|us| Symbol::UseStatement(us.clone())))
    }

    fn get_module_symbol_by_declared_name(&self, name: &str) -> Option<Rc<ModuleSymbol>> {
        for module_symbol in self.module_symbols.values() {
            if module_symbol.declared_name() == name {
                return Some(module_symbol.clone());
            }
        }
        None
    }

    fn get_usable_symbol_by_fqn(&self, fqn: &str) -> Option<Symbol> {
        for function_symbol in self.function_symbols.values() {
            if function_symbol.fqn() == fqn {
                return Some(Symbol::Function(function_symbol.clone()));
            }
        }
        for type_symbol in self.type_symbols.values() {
            if type_symbol.fqn() == fqn {
                return Some(Symbol::Type(type_symbol.clone()));
            }
        }
        None
    }

    fn get_usable_symbol_by_declared_name(&self, declared_name: &str) -> Option<Symbol> {
        for function_symbol in self.function_symbols.values() {
            if function_symbol.declared_name() == declared_name {
                return Some(Symbol::Function(function_symbol.clone()));
            }
        }
        for type_symbol in self.type_symbols.values() {
            if type_symbol.declared_name() == declared_name {
                return Some(Symbol::Type(type_symbol.clone()));
            }
        }
        None
    }

    fn get_value_symbol_by_declared_name(&self, declared_name: &str) -> Option<Symbol> {
        for variable_symbol in self.variable_symbols.values() {
            if variable_symbol.declared_name() == declared_name {
                return Some(Symbol::Variable(variable_symbol.clone()));
            }
        }
        for parameter_symbol in self.parameter_symbols.values() {
            if parameter_symbol.declared_name() == declared_name {
                return Some(Symbol::Parameter(parameter_symbol.clone()));
            }
        }
        None
    }
}

/* Symbol table */

#[derive(Debug)]
pub enum SymbolInsertError {
    SymbolAlreadyDefined(Symbol),
}

#[derive(Debug)]
pub enum SymbolLookupError {
    NoDefinition,
}

#[derive(Debug)]
pub struct SymbolTable {
    scopes: Vec<Scope>,
    current_scope_id: usize,
}

/// Contains a vec of scopes, like a flattened tree
impl SymbolTable {
    pub fn new() -> Self {
        let mut t = SymbolTable {
            scopes: vec![Scope::new(None, String::from("GlobalScope"))],
            current_scope_id: 0,
        };
        t
    }

    pub fn current_scope_id(&self) -> usize {
        self.current_scope_id
    }

    pub fn scopes(&self) -> &Vec<Scope> {
        &self.scopes
    }

    pub fn push_scope(&mut self, debug_name: &str) {
        let id = self.scopes.len();
        self.scopes.push(Scope::new(
            Some(self.current_scope_id),
            debug_name.to_string(),
        ));
        self.current_scope_id = id;
    }

    pub fn pop_scope(&mut self) {
        if let Some(parent_id) = self.scopes[self.current_scope_id].parent {
            self.current_scope_id = parent_id;
        }
    }

    pub fn insert_use_statement_symbol(
        &mut self,
        use_statement_symbol: UseStatementSymbol,
    ) -> Result<(), SymbolInsertError> {
        let current_scope = self.scopes.get_mut(self.current_scope_id).unwrap();
        if let Some(defined_symbol) =
            current_scope.get_usable_symbol_by_declared_name(use_statement_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            current_scope.use_statement_symbols.insert(
                use_statement_symbol.declared_name().to_string(),
                Rc::new(RefCell::new(use_statement_symbol)),
            );
            Ok(())
        }
    }

    pub fn insert_module_symbol(
        &mut self,
        module_symbol: ModuleSymbol,
    ) -> Result<(), SymbolInsertError> {
        let current_scope = self.scopes.get_mut(self.current_scope_id).unwrap();
        if let Some(defined_symbol) =
            current_scope.get_module_symbol_by_declared_name(module_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(Symbol::Module(defined_symbol.clone())))
        } else {
            current_scope.module_symbols.insert(
                module_symbol.declared_name().to_string(),
                Rc::new(module_symbol),
            );
            Ok(())
        }
    }

    pub fn insert_type_symbol(&mut self, type_symbol: TypeSymbol) -> Result<(), SymbolInsertError> {
        let current_scope = self.scopes.get_mut(self.current_scope_id).unwrap();
        if let Some(defined_symbol) =
            current_scope.get_usable_symbol_by_declared_name(type_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            current_scope.type_symbols.insert(
                type_symbol.declared_name().to_string(),
                Rc::new(type_symbol),
            );
            Ok(())
        }
    }

    pub fn insert_function_symbol(
        &mut self,
        function_symbol: FunctionSymbol,
    ) -> Result<(), SymbolInsertError> {
        let current_scope = self.scopes.get_mut(self.current_scope_id).unwrap();
        if let Some(defined_symbol) =
            current_scope.get_usable_symbol_by_declared_name(function_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            current_scope.function_symbols.insert(
                function_symbol.declared_name().to_string(),
                Rc::new(function_symbol),
            );
            Ok(())
        }
    }

    pub fn insert_parameter_symbol(
        &mut self,
        parameter_symbol: ParameterSymbol,
    ) -> Result<(), SymbolInsertError> {
        let current_scope = self.scopes.get_mut(self.current_scope_id).unwrap();
        if let Some(defined_symbol) =
            current_scope.get_value_symbol_by_declared_name(parameter_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            current_scope.parameter_symbols.insert(
                parameter_symbol.declared_name().to_string(),
                Rc::new(parameter_symbol),
            );
            Ok(())
        }
    }

    pub fn insert_variable_symbol(
        &mut self,
        variable_symbol: VariableSymbol,
    ) -> Result<(), SymbolInsertError> {
        let current_scope = self.scopes.get_mut(self.current_scope_id).unwrap();
        if let Some(defined_symbol) =
            current_scope.get_value_symbol_by_declared_name(variable_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            current_scope.variable_symbols.insert(
                variable_symbol.declared_name().to_string(),
                Rc::new(variable_symbol),
            );
            Ok(())
        }
    }

    pub fn lookup(&self, name: &str, scope_id: usize) -> Result<Symbol, SymbolLookupError> {
        let mut scope_opt = Some(&self.scopes[scope_id]);
        while let Some(scope) = scope_opt {
            if let Some(symbol) = scope.get_any_symbol(name) {
                return Ok(symbol);
            }
            scope_opt = if let Some(parent_id) = scope.parent {
                Some(&self.scopes[parent_id])
            } else {
                None
            };
        }
        Err(NoDefinition)
    }

    pub fn lookup_usable_by_fqn(
        &self,
        fully_qualified_name: &str,
        scope_id: usize,
    ) -> Result<Symbol, SymbolLookupError> {
        for scope in &self.scopes {
            if let Some(symbol) = scope.get_usable_symbol_by_fqn(fully_qualified_name) {
                return Ok(symbol);
            }
        }
        Err(NoDefinition)
    }
}

impl Display for SymbolTable {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        writeln!(f, "SymbolTable(current_scope = {})", self.current_scope_id)?;
        for (i, scope) in self.scopes.iter().enumerate() {
            writeln!(f, "Scope {} {}", i, scope.debug_name)?;
            for (name, symbol) in &scope.use_statement_symbols {
                writeln!(f, "  {}({})", name, symbol.borrow())?;
            }
            for (name, symbol) in &scope.module_symbols {
                writeln!(f, "  {}({})", name, symbol)?;
            }
            for (name, symbol) in &scope.type_symbols {
                writeln!(f, "  {}({})", name, symbol)?;
            }
            for (name, symbol) in &scope.function_symbols {
                writeln!(f, "  {}({})", name, symbol)?;
            }
            for (name, symbol) in &scope.variable_symbols {
                writeln!(f, "  {}({})", name, symbol)?;
            }
        }
        Ok(())
    }
}