deimos-lang/dmc-lib/src/semantic_analysis/collect_symbols.rs

use crate::ast::NodeId;
use crate::ast::compilation_unit::CompilationUnit;
use crate::ast::extern_function::ExternFunction;
use crate::ast::function::Function;
use crate::ast::parameter::Parameter;
use crate::ast::statement::Statement;
use crate::diagnostic::Diagnostics;
use crate::semantic_analysis::diagnostic_helpers::symbol_already_declared;
use crate::semantic_analysis::scope::{Scope, ScopeId};
use crate::semantic_analysis::symbol::{FunctionSymbol, ParameterSymbol, Symbol};
use std::collections::HashMap;
use std::rc::Rc;

pub struct SymbolCollectionResult(pub Diagnostics);

struct SymbolCollectionContext<'a> {
    scopes: &'a mut Vec<Scope>,
    nodes_to_scopes: &'a HashMap<NodeId, ScopeId>,
    symbols: &'a mut Vec<Symbol>,
    fqn_context: Vec<Rc<str>>,
    diagnostics: Diagnostics,
}

impl<'a> SymbolCollectionContext<'a> {
    pub fn new(
        scopes: &'a mut Vec<Scope>,
        nodes_to_scopes: &'a HashMap<NodeId, ScopeId>,
        symbols: &'a mut Vec<Symbol>,
    ) -> Self {
        Self {
            scopes,
            nodes_to_scopes,
            symbols,
            fqn_context: Vec::new(),
            diagnostics: Diagnostics::new(),
        }
    }

    pub fn find_symbol_in_scope_for(&self, name: &str, node_id: NodeId) -> Option<&Symbol> {
        let scope_id = self.nodes_to_scopes[&node_id];
        let scope = &self.scopes[scope_id];
        if let Some(symbol_id) = scope.symbols().get(name) {
            Some(&self.symbols[*symbol_id])
        } else {
            None
        }
    }

    pub fn insert_symbol(&mut self, symbol: Symbol, node_id: NodeId) {
        let declared_name = symbol.declared_name_owned();
        self.symbols.push(symbol);
        let symbol_id = self.symbols.len() - 1;
        let scope_id = self.nodes_to_scopes[&node_id];
        let scope = &mut self.scopes[scope_id];
        scope.symbols_mut().insert(declared_name, symbol_id);
    }

    pub fn symbols_mut(&mut self) -> &mut Vec<Symbol> {
        &mut self.symbols
    }

    pub fn diagnostics_mut(&mut self) -> &mut Diagnostics {
        &mut self.diagnostics
    }

    fn join_fqn_parts(parts: &[Rc<str>]) -> String {
        parts.join("::")
    }

    pub fn get_fqn_base(&self) -> String {
        Self::join_fqn_parts(&self.fqn_context)
    }

    pub fn resolve_fqn(&self, suffix: &Rc<str>) -> String {
        let base = self.get_fqn_base();
        if base.is_empty() {
            suffix.to_string()
        } else {
            Self::join_fqn_parts(&[base.into(), suffix.clone()])
        }
    }
}

pub fn collect_symbols_in_compilation_unit(
    compilation_unit: &CompilationUnit,
    scopes: &mut Vec<Scope>,
    nodes_to_scopes: &HashMap<NodeId, ScopeId>,
    symbols: &mut Vec<Symbol>,
) -> SymbolCollectionResult {
    let mut ctx = SymbolCollectionContext::new(scopes, nodes_to_scopes, symbols);

    for function in compilation_unit.functions() {
        collect_symbols_function(function, &mut ctx);
    }

    for extern_function in compilation_unit.extern_functions() {
        collect_symbols_extern_function(extern_function, &mut ctx);
    }

    SymbolCollectionResult(ctx.diagnostics)
}

fn collect_symbols_function(function: &Function, ctx: &mut SymbolCollectionContext) {
    let fqn = ctx.resolve_fqn(&function.declared_name_owned()).into();

    // function itself
    let function_symbol = Symbol::Function(FunctionSymbol::new(
        function.declared_name_owned(),
        Some(function.declared_name_source_range()),
        fqn,
        false,
    ));

    // insert
    if let Some(already_declared) =
        ctx.find_symbol_in_scope_for(function.declared_name(), function.node_id())
    {
        let diagnostic = symbol_already_declared(already_declared, &function_symbol);
        ctx.diagnostics_mut().push(diagnostic);
    } else {
        ctx.insert_symbol(function_symbol, function.node_id());
    }

    // parameters
    for parameter in function.parameters() {
        collect_symbols_parameter(parameter, ctx);
    }

    // n.b. do not do statements yet, because variables are declared and resolved in the resolution pass
}

fn collect_symbols_extern_function(
    extern_function: &ExternFunction,
    ctx: &mut SymbolCollectionContext,
) {
    let fqn = ctx
        .resolve_fqn(&extern_function.declared_name_owned())
        .into();

    // function itself
    let function_symbol = Symbol::Function(FunctionSymbol::new(
        extern_function.declared_name_owned(),
        Some(extern_function.declared_name_source_range()),
        fqn,
        true,
    ));

    // insert function symbol
    if let Some(already_declared) =
        ctx.find_symbol_in_scope_for(extern_function.declared_name(), extern_function.node_id())
    {
        let diagnostic = symbol_already_declared(already_declared, &function_symbol);
        ctx.diagnostics_mut().push(diagnostic);
    } else {
        ctx.insert_symbol(function_symbol, extern_function.node_id());
    }

    // parameters
    for parameter in extern_function.parameters() {
        collect_symbols_parameter(parameter, ctx);
    }
}

fn collect_symbols_parameter(parameter: &Parameter, ctx: &mut SymbolCollectionContext) {
    let parameter_symbol = ParameterSymbol::new(
        parameter.declared_name_owned(),
        Some(parameter.declared_name_source_range()),
    );
    ctx.insert_symbol(Symbol::Parameter(parameter_symbol), parameter.node_id());
}