deimos-lang/src/name_analysis/symbol_table.rs

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::cell::RefCell;
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<RefCell<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_type_symbol_by_declared_name(&self, declared_name: &str) -> Option<Rc<TypeSymbol>> {
        if let Some(type_symbol) = self.type_symbols.get(declared_name) {
            Some(type_symbol.clone())
        } else {
            for use_statement_symbol in self.use_statement_symbols.values() {
                let borrowed = use_statement_symbol.borrow();
                if borrowed.declared_name() == declared_name {
                    if let Some(referenced_symbol) = borrowed.referenced_symbol() {
                        match *referenced_symbol {
                            Symbol::Type(type_symbol) => return Some(type_symbol.clone()),
                            _ => continue,
                        }
                    }
                }
            }
            None
        }
    }

    fn get_type_symbol_by_fqn(&self, fqn: &str) -> Option<Rc<TypeSymbol>> {
        self.type_symbols.values().find(|s| s.fqn() == fqn).cloned()
    }

    fn get_usable_symbol_by_fqn(&self, fqn: &str) -> Option<Symbol> {
        for function_symbol in self.function_symbols.values() {
            if function_symbol.borrow().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.borrow().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
    }

    fn get_expressible_by_declared_name(&self, declared_name: &str) -> Option<Symbol> {
        self.variable_symbols
            .get(declared_name)
            .map(|s| Symbol::Variable(s.clone()))
            .or_else(|| {
                self.parameter_symbols
                    .get(declared_name)
                    .map(|p| Symbol::Parameter(p.clone()))
            })
            .or_else(|| {
                self.function_symbols
                    .get(declared_name)
                    .map(|f| Symbol::Function(f.clone()))
            })
            .or_else(|| {
                self.type_symbols
                    .get(declared_name)
                    .map(|t| Symbol::Type(t.clone()))
            })
            .or_else(|| {
                self.use_statement_symbols
                    .get(declared_name)
                    .map(|us| Symbol::UseStatement(us.clone()))
            })
    }

    fn get_expressible_by_fqn(&self, fqn: &str) -> Option<Symbol> {
        self.function_symbols
            .values()
            .find(|fs| fs.borrow().fqn() == fqn)
            .map(|f| Symbol::Function(f.clone()))
            .or_else(|| {
                self.type_symbols
                    .values()
                    .find(|ts| ts.fqn() == fqn)
                    .map(|ts| Symbol::Type(ts.clone()))
            })
    }
}

/* 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<Rc<RefCell<UseStatementSymbol>>, 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 {
            let declared_name = use_statement_symbol.declared_name().to_string();
            let to_insert = Rc::new(RefCell::new(use_statement_symbol));
            let to_return = to_insert.clone();
            current_scope
                .use_statement_symbols
                .insert(declared_name, to_insert);
            Ok(to_return)
        }
    }

    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<Rc<RefCell<FunctionSymbol>>, 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 {
            let declared_name = function_symbol.declared_name().to_string();
            let to_insert = Rc::new(RefCell::new(function_symbol));
            let to_return = to_insert.clone();
            current_scope
                .function_symbols
                .insert(declared_name, to_insert);
            Ok(to_return)
        }
    }

    pub fn insert_parameter_symbol(
        &mut self,
        parameter_symbol: ParameterSymbol,
    ) -> Result<Rc<ParameterSymbol>, 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 {
            let to_insert = Rc::new(parameter_symbol);
            let to_return = to_insert.clone();
            current_scope
                .parameter_symbols
                .insert(to_insert.declared_name().to_string(), to_insert);
            Ok(to_return)
        }
    }

    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(())
        }
    }

    #[deprecated(note = "Use more specific lookup methods.")]
    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_type_by_declared_name(
        &self,
        declared_name: &str,
        scope_id: usize,
    ) -> Result<Rc<TypeSymbol>, SymbolLookupError> {
        let mut scope_opt = Some(&self.scopes[scope_id]);
        while let Some(scope) = scope_opt {
            if let Some(symbol) = scope.get_type_symbol_by_declared_name(declared_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_type_by_fqn(
        &self,
        fqn: &str,
        scope_id: usize,
    ) -> Result<Rc<TypeSymbol>, SymbolLookupError> {
        let mut scope_opt = Some(&self.scopes[scope_id]);
        while let Some(scope) = scope_opt {
            if let Some(symbol) = scope.get_type_symbol_by_fqn(fqn) {
                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)
    }

    pub fn lookup_expressible_by_declared_name(
        &self,
        declared_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_expressible_by_declared_name(declared_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_expressible_by_fqn(
        &self,
        fqn: &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_expressible_by_fqn(fqn) {
                return Ok(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 (i, scope) in self.scopes.iter().enumerate() {
            writeln!(f, "----Scope {} {}----", i, scope.debug_name)?;
            for symbol in scope.use_statement_symbols.values() {
                writeln!(f, "{:#?}", symbol.borrow())?;
            }
            for symbol in scope.module_symbols.values() {
                writeln!(f, "{:#?}", symbol)?;
            }
            for symbol in scope.type_symbols.values() {
                writeln!(f, "{:#?}", symbol)?;
            }
            for symbol in scope.function_symbols.values() {
                writeln!(f, "{:#?}", symbol.borrow())?;
            }
            for symbol in scope.parameter_symbols.values() {
                writeln!(f, "{:#?}", symbol)?;
            }
            for symbol in scope.variable_symbols.values() {
                writeln!(f, "{:#?}", symbol)?;
            }
        }
        Ok(())
    }
}