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

use crate::name_analysis::symbol::class_symbol::ClassSymbol;
use crate::name_analysis::symbol::function_symbol::FunctionSymbol;
use crate::name_analysis::symbol::interface_symbol::InterfaceSymbol;
use crate::name_analysis::symbol::module_symbol::ModuleSymbol;
use crate::name_analysis::symbol::usable_symbol::UsableSymbol;
use crate::name_analysis::symbol_table::SymbolLookupError;
use crate::util::indent_writer::IndentWriter;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::rc::Rc;

#[derive(Debug)]
pub struct SymbolTree {
    debug_name: String,
    children: Box<HashMap<Rc<str>, SymbolTree>>,
    classes: Box<HashMap<Rc<str>, Rc<RefCell<ClassSymbol>>>>,
    interfaces: Box<HashMap<Rc<str>, Rc<RefCell<InterfaceSymbol>>>>,
    functions: Box<HashMap<Rc<str>, Rc<RefCell<FunctionSymbol>>>>,
}

impl SymbolTree {
    pub fn new(debug_name: &str) -> Self {
        Self {
            debug_name: debug_name.to_string(),
            children: Box::new(HashMap::new()),
            classes: Box::new(HashMap::new()),
            interfaces: Box::new(HashMap::new()),
            functions: Box::new(HashMap::new()),
        }
    }

    pub fn find_class(&self, fqn_parts: &[Rc<str>]) -> Option<Rc<RefCell<ClassSymbol>>> {
        match fqn_parts.len() {
            0 => None,
            1 => self.classes.get(&fqn_parts[0]).cloned(),
            _ => self
                .children
                .get(&fqn_parts[0])
                .and_then(|child_tree| child_tree.find_class(&fqn_parts[1..])),
        }
    }

    pub fn find_interface(&self, fqn_parts: &[Rc<str>]) -> Option<Rc<RefCell<InterfaceSymbol>>> {
        match fqn_parts.len() {
            0 => None,
            1 => self.interfaces.get(&fqn_parts[0]).cloned(),
            _ => self
                .children
                .get(&fqn_parts[0])
                .and_then(|child_tree| child_tree.find_interface(&fqn_parts[1..])),
        }
    }

    pub fn find_function(&self, fqn_parts: &[Rc<str>]) -> Option<Rc<RefCell<FunctionSymbol>>> {
        match fqn_parts.len() {
            0 => None,
            1 => self.functions.get(&fqn_parts[0]).cloned(),
            _ => self
                .children
                .get(&fqn_parts[0])
                .and_then(|child_tree| child_tree.find_function(&fqn_parts[1..])),
        }
    }

    pub fn register_module(&mut self, module_symbol: Rc<RefCell<ModuleSymbol>>) {
        self.recurse_register_module(&module_symbol.borrow().fqn_parts());
    }

    pub fn register_module_by_fqn_parts(&mut self, fqn_parts: &[&str]) {
        self.recurse_register_module(
            &fqn_parts
                .iter()
                .map(|part| Rc::from(*part))
                .collect::<Vec<Rc<str>>>(),
        );
    }

    fn recurse_register_module(&mut self, fqn_parts: &[Rc<str>]) {
        if fqn_parts.len() == 0 {
            panic!("Unable to register module fqn with no parts.")
        }
        if fqn_parts.len() == 1 {
            if !self.children.contains_key(&fqn_parts[0]) {
                self.children
                    .insert(fqn_parts[0].clone(), SymbolTree::new(fqn_parts[0].as_ref()));
            }
        } else {
            if self.children.contains_key(fqn_parts[0].as_ref()) {
                let child = self.children.get_mut(fqn_parts[0].as_ref()).unwrap();
                child.recurse_register_module(&fqn_parts[1..]);
            } else {
                let mut child = SymbolTree::new(fqn_parts[0].as_ref());
                child.recurse_register_module(&fqn_parts[1..]);
                self.children.insert(fqn_parts[0].clone(), child);
            }
        }
    }

    pub fn register_function(&mut self, function_symbol: Rc<RefCell<FunctionSymbol>>) {
        let fqn_parts = function_symbol.borrow().fqn_parts_owned();
        self.recurse_register_function(function_symbol, &fqn_parts);
    }

    fn recurse_register_function(
        &mut self,
        function_symbol: Rc<RefCell<FunctionSymbol>>,
        fqn_parts: &[Rc<str>],
    ) {
        if fqn_parts.len() == 0 {
            panic!("Unable to register function fqn with no parts.");
        }
        if fqn_parts.len() == 1 {
            if self.functions.contains_key(fqn_parts[0].as_ref()) {
                panic!(
                    "There is already a function registered in this symbol tree with the name {}",
                    fqn_parts[0]
                );
            }
            self.functions.insert(fqn_parts[0].clone(), function_symbol);
        } else {
            if self.children.contains_key(fqn_parts[0].as_ref()) {
                let child = self.children.get_mut(fqn_parts[0].as_ref()).unwrap();
                child.recurse_register_function(function_symbol, &fqn_parts[1..]);
            } else {
                panic!("No such inner module registered: {}", fqn_parts[0])
            }
        }
    }

    pub fn find_all_by_base_fqn(
        &self,
        fqn_parts: &[Rc<str>],
    ) -> Result<Vec<UsableSymbol>, SymbolLookupError> {
        match fqn_parts.len() {
            0 => {
                let mut all_symbols: Vec<UsableSymbol> = vec![];
                for interface_symbol in self.interfaces.values() {
                    all_symbols.push(UsableSymbol::Interface(interface_symbol.clone()));
                }
                for class_symbol in self.classes.values() {
                    all_symbols.push(UsableSymbol::Class(class_symbol.clone()));
                }
                for function_symbol in self.functions.values() {
                    all_symbols.push(UsableSymbol::Function(function_symbol.clone()));
                }
                Ok(all_symbols)
            }
            _ => {
                if self.children.contains_key(fqn_parts[0].as_ref()) {
                    let child = self.children.get(fqn_parts[0].as_ref()).unwrap();
                    child.find_all_by_base_fqn(&fqn_parts[1..])
                } else {
                    Err(SymbolLookupError::NoSuchNamespace)
                }
            }
        }
    }

    pub fn find_usable_symbol(
        &self,
        fqn_parts: &[Rc<str>],
    ) -> Result<UsableSymbol, SymbolLookupError> {
        match fqn_parts.len() {
            0 => panic!("Cannot resolve a usable symbol with no fqn part."),
            1 => self
                .find_interface(fqn_parts)
                .map(|interface_symbol| UsableSymbol::Interface(interface_symbol.clone()))
                .or_else(|| {
                    self.find_class(fqn_parts)
                        .map(|class_symbol| UsableSymbol::Class(class_symbol.clone()))
                })
                .or_else(|| {
                    self.find_function(fqn_parts)
                        .map(|function_symbol| UsableSymbol::Function(function_symbol))
                })
                .map_or_else(
                    || Err(SymbolLookupError::NoDefinition),
                    |usable_symbol| Ok(usable_symbol),
                ),
            _ => {
                if self.children.contains_key(fqn_parts[0].as_ref()) {
                    let child = self.children.get(fqn_parts[0].as_ref()).unwrap();
                    child.find_usable_symbol(&fqn_parts[1..])
                } else {
                    Err(SymbolLookupError::NoSuchNamespace)
                }
            }
        }
    }
}

enum FormatAction<'a> {
    PrintSymbolTree(&'a SymbolTree),
    IncreaseIndent,
    DecreaseIndent,
}

impl Display for SymbolTree {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let mut acc = String::new();
        let mut indent_writer = IndentWriter::new(0, "  ", &mut acc);
        let mut stack: Vec<FormatAction> = vec![];
        stack.push(FormatAction::PrintSymbolTree(self));

        while let Some(format_action) = stack.pop() {
            match format_action {
                FormatAction::PrintSymbolTree(symbol_tree) => {
                    // reverse order: start with decrease
                    stack.push(FormatAction::DecreaseIndent);
                    // add children
                    for child in symbol_tree.children.values() {
                        stack.push(FormatAction::PrintSymbolTree(child));
                    }
                    // increase for the first child
                    stack.push(FormatAction::IncreaseIndent);

                    indent_writer.writeln_indented(&format!(
                        "SymbolTree(debug_name = {})",
                        symbol_tree.debug_name
                    ))?;

                    indent_writer.increase_indent();
                    for interface in symbol_tree.interfaces.values() {
                        indent_writer.writeln_indented(&format!("Interface({:?}", interface))?;
                    }
                    for class in symbol_tree.classes.values() {
                        indent_writer.writeln_indented(&format!("Class({:?})", class))?;
                    }
                    for function in symbol_tree.functions.values() {
                        indent_writer.writeln_indented(&format!("Function({:?})", function))?;
                    }
                    indent_writer.decrease_indent();
                }
                FormatAction::IncreaseIndent => {
                    indent_writer.increase_indent();
                }
                FormatAction::DecreaseIndent => {
                    indent_writer.decrease_indent();
                }
            }
        }

        writeln!(f, "{}", acc)
    }
}