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

use crate::name_analysis::symbol::class_member_symbol::ClassMemberSymbol;
use crate::name_analysis::symbol::function_symbol::FunctionSymbol;
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::use_symbol::{ConcreteUseSymbol, StarUseSymbol};
use crate::name_analysis::symbol::variable_symbol::VariableSymbol;
use crate::name_analysis::symbol::*;
use crate::name_analysis::symbol_table::SymbolInsertError::SymbolAlreadyDefined;
use crate::name_analysis::symbol_table::SymbolLookupError::NoDefinition;
use scope::Scope;
use std::fmt::Display;
use std::ops::Deref;

mod scope;

#[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 {
        Self {
            scopes: vec![Scope::new(None, 0, String::from("GlobalScope"))],
            current_scope_id: 0,
        }
    }

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

    pub fn push_scope(&mut self, debug_name: &str) {
        let id = self.scopes.len();
        self.scopes
            .push(Scope::new(Some(self.current_scope_id), 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;
        }
    }

    fn current_scope(&self) -> &Scope {
        self.scopes.last().unwrap()
    }

    fn current_scope_mut(&mut self) -> &mut Scope {
        self.scopes.last_mut().unwrap()
    }

    fn find_current_scope_concrete_use_symbol(
        &self,
        declared_name: &str,
    ) -> Option<&ConcreteUseSymbol> {
        self.current_scope()
            .concrete_use_symbols()
            .get(declared_name)
    }

    fn find_current_scope_star_use_symbol(&self, base_fqn: &str) -> Option<&StarUseSymbol> {
        self.current_scope().star_use_symbols().get(base_fqn)
    }

    fn find_current_scope_type_symbol(&self, declared_name: &str) -> Option<&TypeSymbol> {
        self.current_scope().type_symbols().get(declared_name)
    }

    fn find_current_scope_module_symbol(&self, declared_name: &str) -> Option<&ModuleSymbol> {
        self.current_scope().module_symbols().get(declared_name)
    }

    fn find_current_scope_parameter_symbol(&self, declared_name: &str) -> Option<&ParameterSymbol> {
        self.current_scope().parameter_symbols().get(declared_name)
    }

    fn find_current_scope_variable_symbol(&self, declared_name: &str) -> Option<&VariableSymbol> {
        self.current_scope().variable_symbols().get(declared_name)
    }

    fn find_current_scope_class_member_symbol(
        &self,
        declared_name: &str,
    ) -> Option<&ClassMemberSymbol> {
        self.current_scope()
            .class_member_symbols()
            .get(declared_name)
    }

    fn find_current_scope_variable_or_parameter_symbol(
        &self,
        declared_name: &str,
    ) -> Option<Symbol> {
        self.find_current_scope_variable_symbol(declared_name)
            .map(|variable_symbol| Symbol::Variable(variable_symbol.clone()))
            .or_else(|| {
                self.find_current_scope_parameter_symbol(declared_name)
                    .map(|parameter_symbol| Symbol::Parameter(parameter_symbol.clone()))
            })
    }

    fn find_current_scope_usable_symbol(&self, declared_name: &str) -> Option<Symbol> {
        self.find_current_scope_concrete_use_symbol(declared_name)
            .map(|concrete_use_symbol| Symbol::ConcreteUse(concrete_use_symbol.clone()))
            .or_else(|| {
                self.find_current_scope_type_symbol(declared_name)
                    .map(|type_symbol| Symbol::Type(type_symbol.clone()))
            })
            .or_else(|| {
                self.find_current_scope_module_symbol(declared_name)
                    .map(|module_symbol| Symbol::Module(module_symbol.clone()))
            })
    }

    pub fn insert_concrete_use_symbol(
        &mut self,
        concrete_use_symbol: ConcreteUseSymbol,
    ) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_usable_symbol(concrete_use_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            self.current_scope_mut().concrete_use_symbols_mut().insert(
                concrete_use_symbol.declared_name().to_string(),
                concrete_use_symbol,
            );
            Ok(())
        }
    }

    pub fn insert_star_use_symbol(
        &mut self,
        star_use_symbol: StarUseSymbol,
    ) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_star_use_symbol(star_use_symbol.base_fqn())
        {
            Err(SymbolAlreadyDefined(Symbol::StarUse(
                defined_symbol.clone(),
            )))
        } else {
            self.current_scope_mut()
                .star_use_symbols_mut()
                .insert(star_use_symbol.base_fqn().to_string(), star_use_symbol);
            Ok(())
        }
    }

    pub fn insert_module_symbol(
        &mut self,
        module_symbol: ModuleSymbol,
    ) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_usable_symbol(module_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            self.current_scope_mut()
                .module_symbols_mut()
                .insert(module_symbol.declared_name().to_string(), module_symbol);
            Ok(())
        }
    }

    pub fn insert_type_symbol(&mut self, type_symbol: TypeSymbol) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_usable_symbol(type_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            self.current_scope_mut()
                .type_symbols_mut()
                .insert(type_symbol.declared_name().to_string(), type_symbol);
            Ok(())
        }
    }

    pub fn insert_function_symbol(
        &mut self,
        function_symbol: FunctionSymbol,
    ) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_usable_symbol(function_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            self.current_scope_mut()
                .function_symbols_mut()
                .insert(function_symbol.declared_name().to_string(), function_symbol);
            Ok(())
        }
    }

    pub fn insert_parameter_symbol(
        &mut self,
        parameter_symbol: ParameterSymbol,
    ) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_parameter_symbol(parameter_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(Symbol::Parameter(
                defined_symbol.clone(),
            )))
        } else {
            self.current_scope_mut().parameter_symbols_mut().insert(
                parameter_symbol.declared_name().to_string(),
                parameter_symbol,
            );
            Ok(())
        }
    }

    pub fn insert_variable_symbol(
        &mut self,
        variable_symbol: VariableSymbol,
    ) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_variable_or_parameter_symbol(variable_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(defined_symbol))
        } else {
            self.current_scope_mut()
                .variable_symbols_mut()
                .insert(variable_symbol.declared_name().to_string(), variable_symbol);
            Ok(())
        }
    }

    pub fn insert_class_member_symbol(
        &mut self,
        class_member_symbol: ClassMemberSymbol,
    ) -> Result<(), SymbolInsertError> {
        if let Some(defined_symbol) =
            self.find_current_scope_class_member_symbol(class_member_symbol.declared_name())
        {
            Err(SymbolAlreadyDefined(Symbol::ClassMember(
                defined_symbol.clone(),
            )))
        } else {
            self.current_scope_mut().class_member_symbols_mut().insert(
                class_member_symbol.declared_name().to_string(),
                class_member_symbol,
            );
            Ok(())
        }
    }

    fn lookup_type_in_scope_by_declared_name<'a>(
        scope: &'a Scope,
        declared_name: &str,
    ) -> Option<&'a TypeSymbol> {
        scope.type_symbols().get(declared_name)
    }

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

    fn lookup_type_in_scope_by_fqn<'a>(scope: &'a Scope, fqn: &str) -> Option<&'a TypeSymbol> {
        for type_symbol in scope.type_symbols().values() {
            match type_symbol {
                TypeSymbol::Concrete(concrete_type_symbol) => {
                    if concrete_type_symbol.fqn() == fqn {
                        return Some(type_symbol);
                    }
                }
                TypeSymbol::Generic(_) => {}
            }
        }
        None
    }

    pub fn lookup_type_by_fqn(
        &self,
        fqn: &str,
        scope_id: usize,
    ) -> Result<&TypeSymbol, SymbolLookupError> {
        let mut scope_opt = Some(&self.scopes[scope_id]);
        while let Some(scope) = scope_opt {
            if let Some(type_symbol) = Self::lookup_type_in_scope_by_fqn(scope, fqn) {
                return Ok(type_symbol);
            }
            scope_opt = if let Some(parent_id) = scope.parent() {
                Some(&self.scopes[parent_id])
            } else {
                None
            }
        }
        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 scope in &self.scopes {
            writeln!(f, "{}", scope)?;
        }
        Ok(())
    }
}