deimos-lang/dmc-lib/src/ast/function.rs

use crate::ast::fqn_context::FqnContext;
use crate::ast::fqn_util::fqn_parts_to_string;
use crate::ast::helpers::{
    collect_diagnostics_into_enumerated_mut, collect_diagnostics_into_mut,
    collect_parameter_symbols_into, insert_declared_types_into, insert_resolved_names_into,
    insert_resolved_types_into, resolve_parameter_names_into,
};
use crate::ast::ir_builder::IrBuilder;
use crate::ast::parameter::Parameter;
use crate::ast::statement::Statement;
use crate::ast::type_use::TypeUse;
use crate::ast::{FunctionReturnTypes, NodeId, NodesToSymbols, NodesToTypes, SymbolsToTypes};
use crate::diagnostic::{Diagnostic, Diagnostics};
use crate::ir::ir_function::IrFunction;
use crate::ir::ir_parameter::IrParameter;
use crate::ir::ir_parameter_or_variable::IrParameterOrVariable;
use crate::source_range::SourceRange;
use crate::symbol::Symbol;
use crate::symbol::class_symbol::ClassSymbol;
use crate::symbol::function_symbol::FunctionSymbol;
use crate::symbol_table::SymbolTable;
use crate::type_info::TypeInfo;
use crate::types_table::TypesTable;
use crate::{diagnostics_result, handle_diagnostics};
use std::ops::Neg;
use std::rc::Rc;

pub struct Function {
    node_id: NodeId,
    declared_name: Rc<str>,
    declared_name_source_range: SourceRange,
    is_public: bool,
    parameters: Vec<Parameter>,
    return_type: Option<TypeUse>,
    statements: Vec<Statement>,
    scope_id: Option<usize>,
}

impl Function {
    pub fn new(
        node_id: NodeId,
        declared_name: &str,
        declared_name_source_range: SourceRange,
        is_public: bool,
        parameters: Vec<Parameter>,
        return_type: Option<TypeUse>,
        statements: Vec<Statement>,
    ) -> Self {
        Self {
            node_id,
            declared_name: declared_name.into(),
            declared_name_source_range,
            is_public,
            parameters,
            return_type,
            statements,
            scope_id: None,
        }
    }

    pub fn node_id(&self) -> NodeId {
        self.node_id
    }

    pub fn declared_name(&self) -> &str {
        &self.declared_name
    }

    pub fn statements(&self) -> Vec<&Statement> {
        self.statements.iter().collect()
    }

    pub fn init_scopes(&mut self, symbol_table: &mut SymbolTable, container_scope: usize) {
        self.scope_id = Some(container_scope);
        let function_scope =
            symbol_table.push_function_scope(&format!("function_scope({})", self.declared_name));

        for parameter in &mut self.parameters {
            parameter.init_scopes(symbol_table, function_scope);
        }
        if let Some(type_use) = &mut self.return_type {
            type_use.init_scopes(symbol_table, function_scope);
        }

        let body_scope =
            symbol_table.push_block_scope(&format!("body_scope({})", self.declared_name));

        for statement in &mut self.statements {
            statement.init_scopes(symbol_table, body_scope);
        }

        symbol_table.pop_scope(); // body
        symbol_table.pop_scope(); // function
    }

    /// Return value contains self FunctionSymbol followed by all symbols (including self symbol).
    pub fn declared_symbols(
        &self,
        fqn_context: &FqnContext,
        is_method: bool,
    ) -> (Rc<FunctionSymbol>, Vec<Symbol>) {
        let mut all_symbols: Vec<Symbol> = vec![];

        let mut parameter_symbols = Vec::new();
        collect_parameter_symbols_into(&self.parameters, &mut all_symbols, &mut parameter_symbols);

        let function_symbol = Rc::new(FunctionSymbol::new(
            &self.declared_name,
            self.declared_name_source_range.clone(),
            fqn_context.resolve(self.declared_name()),
            false,
            is_method,
            self.scope_id.unwrap(),
            parameter_symbols,
        ));
        all_symbols.push(Symbol::Function(function_symbol.clone()));

        (function_symbol, all_symbols)
    }

    fn resolve_names_common(&self, symbol_table: &SymbolTable) -> (NodesToSymbols, Diagnostics) {
        let mut diagnostics = Diagnostics::new();
        let mut nodes_to_symbols = NodesToSymbols::new();

        resolve_parameter_names_into(
            &self.parameters,
            symbol_table,
            &mut nodes_to_symbols,
            &mut diagnostics,
        );

        if let Some(type_use) = &self.return_type {
            let (ns, mut ds) = type_use.resolve_names(symbol_table);
            insert_resolved_names_into(ns, &mut nodes_to_symbols);
            diagnostics.append(&mut ds);
        }

        (nodes_to_symbols, diagnostics)
    }

    pub fn resolve_names_static(
        &self,
        symbol_table: &mut SymbolTable,
    ) -> (NodesToSymbols, Diagnostics) {
        let (mut nodes_to_symbols, mut diagnostics) = self.resolve_names_common(symbol_table);

        for statement in &self.statements {
            let (ns, mut ds) = statement.resolve_names_static(symbol_table);
            insert_resolved_names_into(ns, &mut nodes_to_symbols);
            diagnostics.append(&mut ds);
        }

        (nodes_to_symbols, diagnostics)
    }

    pub fn resolve_names_method(
        &self,
        symbol_table: &mut SymbolTable,
        self_class_symbol: &ClassSymbol,
    ) -> (NodesToSymbols, Diagnostics) {
        let (mut nodes_to_symbols, mut diagnostics) = self.resolve_names_common(symbol_table);

        for statement in &self.statements {
            let (ns, mut ds) = statement.resolve_names_method(symbol_table, self_class_symbol);
            insert_resolved_names_into(ns, &mut nodes_to_symbols);
            diagnostics.append(&mut ds);
        }

        (nodes_to_symbols, diagnostics)
    }

    #[deprecated]
    pub fn check_names(&self, symbol_table: &SymbolTable) -> Vec<Diagnostic> {
        let mut diagnostics = Vec::new();
        for parameter in &self.parameters {
            diagnostics.append(&mut parameter.check_names(symbol_table));
        }
        if let Some(type_use) = &self.return_type {
            diagnostics.append(&mut type_use.check_names(symbol_table));
        }
        diagnostics
    }

    #[deprecated]
    pub fn analyze_method_local_names(
        &self,
        symbol_table: &mut SymbolTable,
        class_symbol: &ClassSymbol,
    ) -> Vec<Diagnostic> {
        self.statements
            .iter()
            .flat_map(|statement| statement.analyze_method_local_names(symbol_table, class_symbol))
            .collect()
    }

    #[deprecated]
    pub fn analyze_static_fn_local_names(&self, symbol_table: &mut SymbolTable) -> Vec<Diagnostic> {
        self.statements
            .iter()
            .flat_map(|statement| statement.analyze_static_fn_local_names(symbol_table))
            .collect()
    }

    pub fn declared_types(&self, names_table: &NodesToSymbols) -> (SymbolsToTypes, Diagnostics) {
        let mut diagnostics = Diagnostics::new();
        let mut declared_types = SymbolsToTypes::new();

        for parameter in &self.parameters {
            let symbol = names_table.get(&parameter.node_id()).unwrap();
            let (type_info, mut ds) = parameter.declared_type(names_table);
            declared_types.insert(symbol.clone(), type_info);
            diagnostics.append(&mut ds);
        }

        (declared_types, diagnostics)
    }

    pub fn resolve_types(
        &self,
        nodes_to_symbols: &NodesToSymbols,
        symbols_to_types: &SymbolsToTypes,
    ) -> (SymbolsToTypes, NodesToTypes, Diagnostics) {
        let mut diagnostics = Diagnostics::new();
        let mut nodes_to_types = NodesToTypes::new();
        let mut symbols_to_types = symbols_to_types.clone();

        for statement in &self.statements {
            let (sts, nts, mut ds) = statement.resolve_types(nodes_to_symbols, &symbols_to_types);
            insert_declared_types_into(sts, &mut symbols_to_types); // merge!
            insert_resolved_types_into(nts, &mut nodes_to_types);
            diagnostics.append(&mut ds);
        }

        // todo: check last statement for return type

        (symbols_to_types, nodes_to_types, diagnostics)
    }

    #[deprecated]
    pub fn gather_types(&self, symbol_table: &SymbolTable, types_table: &mut TypesTable) {
        let function_symbol = symbol_table
            .get_function_symbol_owned(self.scope_id.unwrap(), self.declared_name())
            .unwrap();

        // self type (the signature)
        types_table.function_types_mut().insert(
            function_symbol.clone(),
            TypeInfo::Function(function_symbol.clone()),
        );

        // put return type (temporary, this is deprecated)
        if let Some(type_use) = &self.return_type {
            let resolved_return_type = type_use.type_info(symbol_table, types_table).clone();
            types_table
                .function_return_types_mut()
                .insert(function_symbol, resolved_return_type);
        } else {
            types_table
                .function_return_types_mut()
                .insert(function_symbol, TypeInfo::Void);
        }

        // parameters
        for parameter in &self.parameters {
            parameter.gather_types_into(symbol_table, types_table);
        }
    }

    fn get_return_type_info(
        types_table: &TypesTable,
        function_symbol: &FunctionSymbol,
    ) -> TypeInfo {
        types_table
            .function_return_types()
            .get(function_symbol)
            .cloned()
            .unwrap()
    }

    /// Type checks parameters.
    fn type_check_parameters(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
        diagnostics: &mut Vec<Diagnostic>,
    ) {
        collect_diagnostics_into_mut(
            &mut self.parameters,
            |p| p.type_check(symbol_table, types_table),
            diagnostics,
        )
    }

    /// Type checks return type.
    fn type_check_return_type(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
        diagnostics: &mut Vec<Diagnostic>,
    ) {
        if let Some(type_use) = &mut self.return_type {
            handle_diagnostics!(type_use.type_check(symbol_table, types_table), diagnostics);
        }
    }

    /// Type checks statements, making sure the last statement matches return type, if necessary.
    fn type_check_statements(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
        diagnostics: &mut Vec<Diagnostic>,
        function_symbol: &FunctionSymbol,
    ) {
        let return_type_info = Self::get_return_type_info(types_table, function_symbol);
        let statements_len = self.statements.len();

        collect_diagnostics_into_enumerated_mut(
            &mut self.statements,
            |i, s| {
                let is_last = i == statements_len - 1;
                if is_last {
                    s.type_check(symbol_table, types_table, Some(&return_type_info))
                } else {
                    s.type_check(symbol_table, types_table, None)
                }
            },
            diagnostics,
        );
    }

    #[deprecated]
    pub fn type_check(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &mut TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics = vec![];
        let function_symbol = symbol_table
            .get_function_symbol(self.scope_id.unwrap(), self.declared_name())
            .unwrap();

        // parameters
        self.type_check_parameters(symbol_table, types_table, &mut diagnostics);

        // return type
        self.type_check_return_type(symbol_table, types_table, &mut diagnostics);

        // statements
        self.type_check_statements(symbol_table, types_table, &mut diagnostics, function_symbol);

        diagnostics_result!(diagnostics)
    }

    /// Converts all parameters to ir. Saves the IrParameter to the associated parameter symbol.
    fn parameters_to_ir(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) {
        for (i, parameter) in self.parameters.iter().enumerate() {
            let parameter_symbol = symbol_table
                .get_parameter_symbol_owned(parameter.scope_id(), parameter.declared_name())
                .unwrap();
            let parameter_type_info = types_table
                .parameter_types()
                .get(&parameter_symbol)
                .unwrap();
            let stack_offset = (self.parameters.len() as isize).neg() + (i as isize);
            let ir_parameter = IrParameter::new(
                parameter_symbol.declared_name(),
                parameter_type_info.clone(),
                stack_offset,
            );
            let as_rc = Rc::new(ir_parameter);
            builder.push_parameter(&parameter_symbol, as_rc.clone());
        }
    }

    /// If `class_context.is_some()`, set parameter 0 to the self parameter/variable on the builder.
    fn handle_method_case(&self, builder: &mut IrBuilder, class_context: Option<&ClassSymbol>) {
        // if we are a method, we need to set the self parameter on the builder
        if class_context.is_some() {
            let parameter_0 = builder.parameters()[0].clone();
            // put it in the self parameter
            builder.set_self_parameter_or_variable(IrParameterOrVariable::IrParameter(parameter_0));
        }
    }

    /// Convert all statements to ir.
    fn statements_to_ir(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
        function_symbol: &FunctionSymbol,
    ) {
        let return_type_info = Self::get_return_type_info(types_table, function_symbol);
        let should_return_value = !matches!(return_type_info, TypeInfo::Void);
        for (i, statement) in self.statements.iter().enumerate() {
            let is_last = i == self.statements.len() - 1;
            statement.to_ir(
                builder,
                symbol_table,
                types_table,
                should_return_value && is_last,
            );
        }
    }

    pub fn to_ir(
        &self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
        class_context: Option<&ClassSymbol>,
    ) -> IrFunction {
        let mut builder = IrBuilder::new();
        let function_symbol = symbol_table
            .get_function_symbol(self.scope_id.unwrap(), self.declared_name())
            .unwrap();

        // parameters
        self.parameters_to_ir(&mut builder, symbol_table, types_table);

        let entry_block_id = builder.new_block();

        // preamble
        self.handle_method_case(&mut builder, class_context);

        // body
        self.statements_to_ir(&mut builder, symbol_table, types_table, function_symbol);

        builder.finish_block();

        let entry_block = builder.get_block(entry_block_id).clone();
        IrFunction::new(
            fqn_parts_to_string(function_symbol.fqn_parts()),
            builder.parameters().iter().map(|p| (*p).clone()).collect(),
            &Self::get_return_type_info(types_table, function_symbol),
            entry_block,
        )
    }

    pub fn lower_static(
        &self,
        nodes_to_symbols: &NodesToSymbols,
        symbols_to_types: &SymbolsToTypes,
        nodes_to_types: &NodesToTypes,
        function_return_types: &FunctionReturnTypes,
    ) -> IrFunction {
        let mut builder = IrBuilder::new();

        let function_symbol = nodes_to_symbols
            .get(&self.node_id)
            .unwrap()
            .unwrap_function_symbol();

        // put parameters in builder
        for (i, parameter_symbol) in function_symbol.parameters().iter().enumerate() {
            let parameter_type_info = symbols_to_types
                .get(&Symbol::Parameter(parameter_symbol.clone()))
                .unwrap();
            let stack_offset = (function_symbol.parameters().len() as isize).neg() + (i as isize);
            let ir_parameter = Rc::new(IrParameter::new(
                parameter_symbol.declared_name(),
                parameter_type_info.clone(),
                stack_offset,
            ));
            builder.push_parameter(parameter_symbol, ir_parameter);
        }

        let entry_block_id = builder.new_block();

        // lower statements
        let return_type_info = function_return_types.get(function_symbol).unwrap();
        let should_return_value = !matches!(return_type_info, TypeInfo::Void);
        for (i, statement) in self.statements.iter().enumerate() {
            let is_last = i == self.statements.len() - 1;
            statement.lower(
                &mut builder,
                nodes_to_symbols,
                symbols_to_types,
                nodes_to_types,
                should_return_value && is_last,
            );
        }

        builder.finish_block();

        let entry_block = builder.get_block(entry_block_id).clone();
        IrFunction::new(
            fqn_parts_to_string(function_symbol.fqn_parts()),
            builder.parameters().iter().map(|p| (*p).clone()).collect(),
            return_type_info,
            entry_block,
        )
    }
}