deimos-lang/src/name_analysis/symbol_table/mod.rs

use crate::name_analysis::symbol::function_symbol::FunctionSymbol;
use crate::name_analysis::symbol::generic_type_symbol::GenericTypeSymbol;
use crate::name_analysis::symbol::module_symbol::ModuleSymbol;
use crate::name_analysis::symbol::parameter_symbol::ParameterSymbol;
use crate::name_analysis::symbol::type_symbol::TypeSymbol;
use crate::name_analysis::symbol::usable_symbol::UsableSymbol;
use crate::name_analysis::symbol::use_symbol::{ConcreteUseSymbol, StarUseSymbol};
use crate::name_analysis::symbol::variable_symbol::VariableSymbol;
use crate::name_analysis::symbol::{ExpressibleSymbol, LVSymbol, Symbol};
use crate::name_analysis::symbol_table::fqn_context::FqnContext;
use crate::name_analysis::symbol_table::symbol_tree::SymbolTree;
use crate::name_analysis::symbol_table::SymbolInsertError::SymbolAlreadyDefined;
use crate::name_analysis::util::join_fqn_parts;
use scope::Scope;
use std::cell::RefCell;
use std::fmt::Display;
use std::ops::Deref;
use std::rc::Rc;

pub(self) mod fqn_context;
mod scope;
pub mod symbol_tree;

#[derive(Debug)]
pub enum SymbolInsertError {
    SymbolAlreadyDefined(Rc<RefCell<dyn Symbol>>),
}

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

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

/// Contains a vec of scopes, like a flattened tree
impl SymbolTable {
    pub fn new() -> Self {
        Self {
            scopes: vec![Scope::new(None, 0, String::from("GlobalScope"))],
            current_scope_id: 0,
            symbol_tree: Box::new(SymbolTree::new("<DEFAULT>")),
            fqn_context: Box::new(FqnContext::new()),
        }
    }

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

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

    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 set_current_scope(&mut self, id: usize) {
        self.current_scope_id = id;
    }

    pub fn set_current_fqn(&mut self, names: &[&str]) {
        self.fqn_context = Box::new(FqnContext::new());
        for name in names {
            self.fqn_context.push(*name);
        }
    }

    pub fn push_fqn_part(&mut self, part: &str) {
        self.fqn_context.push(part);
    }

    pub fn pop_fqn_part(&mut self) {
        self.fqn_context.pop();
    }

    pub fn current_fqn(&self) -> Vec<&str> {
        self.fqn_context.current_fqn()
    }

    pub fn current_fqn_owned(&self) -> Vec<Rc<str>> {
        self.fqn_context.current_fqn_owned()
    }

    pub fn resolve_fqn(&self, part: Rc<str>) -> Vec<Rc<str>> {
        let mut parts = self.current_fqn_owned();
        parts.push(part);
        parts
    }

    pub fn find_usable_symbol(&self, fqn_parts: &[Rc<str>]) -> Option<UsableSymbol> {
        self.symbol_tree
            .find_interface(fqn_parts)
            .map(|interface_symbol| UsableSymbol::Interface(interface_symbol))
            .or_else(|| {
                self.symbol_tree
                    .find_class(fqn_parts)
                    .map(|class_symbol| UsableSymbol::Class(class_symbol))
            })
            .or_else(|| {
                self.symbol_tree
                    .find_function(fqn_parts)
                    .map(|function_symbol| UsableSymbol::Function(function_symbol))
            })
    }

    pub fn register_module(&mut self, fqn_parts: &[&str]) {
        self.symbol_tree.register_module_by_fqn_parts(fqn_parts);
    }

    pub fn register_function_symbol(
        &mut self,
        function_symbol: FunctionSymbol,
    ) -> Rc<RefCell<FunctionSymbol>> {
        let as_rc = Rc::new(RefCell::new(function_symbol));
        self.symbol_tree.register_function(as_rc.clone());
        as_rc
    }

    pub fn resolve_global_imports(&mut self) {
        let global_scope = &self.scopes[0];
        for concrete_use_symbol in global_scope.concrete_use_symbols() {
            match self
                .symbol_tree
                .find_usable_symbol(concrete_use_symbol.borrow().fqn_parts())
            {
                Ok(usable_symbol) => {
                    concrete_use_symbol.borrow_mut().set_resolved_symbol(usable_symbol);
                }
                Err(symbol_lookup_error) => {
                    // panic because this is definitely a compiler/configuration error
                    panic!(
                        "{:?}: Cannot resolve global import {}",
                        symbol_lookup_error,
                        join_fqn_parts(concrete_use_symbol.borrow().fqn_parts())
                    );
                }
            }
        }
        
        for star_use_symbol in global_scope.star_use_symbols() {
            let mut star_use_symbol_ref_mut = star_use_symbol.borrow_mut();
            match self.symbol_tree.find_all_by_base_fqn(star_use_symbol_ref_mut.fqn_parts()) {
                Ok(usable_symbols) => {
                    star_use_symbol_ref_mut.set_resolved_symbols(usable_symbols);
                }
                Err(symbol_lookup_error) => {
                    // again, panic because this is a compiler/configuration error
                    panic!(
                        "{:?}: Cannot resolve global star import {}::*",
                        symbol_lookup_error,
                        join_fqn_parts(star_use_symbol.borrow().fqn_parts())
                    );
                }
            }
        }
    }

    pub fn find_usable_symbols_by_base_fqn(
        &self,
        fqn_parts: &[Rc<str>],
    ) -> Result<Vec<UsableSymbol>, SymbolLookupError> {
        self.symbol_tree.find_all_by_base_fqn(fqn_parts)
    }

    fn current_scope(&self) -> &Scope {
        &self.scopes[self.current_scope_id]
    }

    fn current_scope_mut(&mut self) -> &mut Scope {
        &mut self.scopes[self.current_scope_id]
    }

    pub fn insert_concrete_use_symbol(
        &mut self,
        concrete_use_symbol: ConcreteUseSymbol,
    ) -> Result<Rc<RefCell<ConcreteUseSymbol>>, SymbolInsertError> {
        if let Some(defined_symbol) = self
            .current_scope()
            .find_module_level_symbol(concrete_use_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol.to_symbol()))
        } else {
            let inserted = self
                .current_scope_mut()
                .insert_concrete_use_symbol(concrete_use_symbol);
            Ok(inserted)
        }
    }

    pub fn insert_star_use_symbol(
        &mut self,
        star_use_symbol: StarUseSymbol,
    ) -> Result<Rc<RefCell<StarUseSymbol>>, SymbolInsertError> {
        let inserted = self
            .current_scope_mut()
            .insert_star_use_symbol(star_use_symbol);
        Ok(inserted)
    }

    pub fn insert_module_symbol(
        &mut self,
        module_symbol: ModuleSymbol,
    ) -> Result<Rc<RefCell<ModuleSymbol>>, SymbolInsertError> {
        if let Some(defined_symbol) = self
            .current_scope()
            .find_module_level_symbol(module_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol.to_symbol()))
        } else {
            let inserted = self.current_scope_mut().insert_module_symbol(module_symbol);
            self.symbol_tree.register_module(inserted.clone());
            Ok(inserted)
        }
    }

    pub fn insert_function_symbol(
        &mut self,
        function_symbol: FunctionSymbol,
    ) -> Result<Rc<RefCell<FunctionSymbol>>, SymbolInsertError> {
        if let Some(defined_symbol) = self
            .current_scope()
            .find_module_level_symbol(function_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol.to_symbol()))
        } else {
            let inserted = self
                .current_scope_mut()
                .insert_function_symbol(function_symbol);
            self.symbol_tree.register_function(inserted.clone());
            Ok(inserted)
        }
    }

    pub fn insert_generic_type_symbol(
        &mut self,
        generic_type_symbol: GenericTypeSymbol,
    ) -> Result<Rc<RefCell<GenericTypeSymbol>>, SymbolInsertError> {
        if let Some(defined_symbol) = self
            .current_scope()
            .find_type_symbol(generic_type_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol.to_symbol()))
        } else {
            let inserted = self
                .current_scope_mut()
                .insert_generic_type_symbol(generic_type_symbol);
            Ok(inserted)
        }
    }

    pub fn insert_parameter_symbol(
        &mut self,
        parameter_symbol: ParameterSymbol,
    ) -> Result<Rc<RefCell<ParameterSymbol>>, SymbolInsertError> {
        if let Some(defined_symbol) = self
            .current_scope()
            .find_lv_symbol(parameter_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol.to_symbol()))
        } else {
            let inserted = self
                .current_scope_mut()
                .insert_parameter_symbol(parameter_symbol);
            Ok(inserted)
        }
    }

    pub fn insert_variable_symbol(
        &mut self,
        variable_symbol: VariableSymbol,
    ) -> Result<Rc<RefCell<VariableSymbol>>, SymbolInsertError> {
        if let Some(defined_symbol) = self
            .current_scope()
            .find_lv_symbol(variable_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol.to_symbol()))
        } else {
            let inserted = self
                .current_scope_mut()
                .insert_variable_symbol(variable_symbol);
            Ok(inserted)
        }
    }

    pub fn lookup_type(&self, declared_name: &str) -> Result<TypeSymbol, SymbolLookupError> {
        let mut current_scope: Option<&Scope> = Some(self.current_scope());
        while let Some(scope) = current_scope.take() {
            if let Some(type_symbol) = scope.find_type_symbol(declared_name) {
                return Ok(type_symbol);
            } else {
                current_scope = scope
                    .parent()
                    .and_then(|parent_id| self.scopes.get(parent_id));
            }
        }
        Err(SymbolLookupError::NoDefinition)
    }

    pub fn lookup_type_by_fqn(&self, fqn_parts: &[&str]) -> Result<TypeSymbol, SymbolLookupError> {
        todo!()
    }

    pub fn lookup_function_symbol(
        &self,
        declared_name: &str,
    ) -> Result<Rc<RefCell<FunctionSymbol>>, SymbolLookupError> {
        let mut current_scope: Option<&Scope> = Some(self.current_scope());
        while let Some(scope) = current_scope.take() {
            if let Some(function_symbol) = scope.find_function_symbol(declared_name) {
                return Ok(function_symbol.clone());
            } else {
                current_scope = scope
                    .parent()
                    .and_then(|parent_id| self.scopes.get(parent_id));
            }
        }
        Err(SymbolLookupError::NoDefinition)
    }

    pub fn lookup_lv_symbol(&self, declared_name: &str) -> Result<LVSymbol, SymbolLookupError> {
        let mut current_scope: Option<&Scope> = Some(self.current_scope());
        while let Some(scope) = current_scope.take() {
            if let Some(lv_symbol) = scope.find_lv_symbol(declared_name) {
                return Ok(lv_symbol);
            } else {
                current_scope = scope
                    .parent()
                    .and_then(|parent_id| self.scopes.get(parent_id));
            }
        }
        Err(SymbolLookupError::NoDefinition)
    }

    pub fn lookup_expressible_symbol(
        &self,
        declared_name: &str,
    ) -> Result<ExpressibleSymbol, SymbolLookupError> {
        let mut current_scope: Option<&Scope> = Some(self.current_scope());
        while let Some(scope) = current_scope.take() {
            if let Some(expressible_symbol) = scope.find_expressible_symbol(declared_name) {
                return Ok(expressible_symbol);
            } else {
                current_scope = scope
                    .parent()
                    .and_then(|parent_id| self.scopes.get(parent_id));
            }
        }
        Err(SymbolLookupError::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 scope in &self.scopes {
            writeln!(f, "{}", scope)?;
        }
        writeln!(f, "---SymbolTable: SymbolTree---")?;
        writeln!(f, "{}", self.symbol_tree)?;
        Ok(())
    }
}