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

use crate::ast::expression::Expression;
use crate::ast::fqn_util::fqn_parts_to_string;
use crate::ast::ir_builder::IrBuilder;
use crate::diagnostic::Diagnostic;
use crate::ir::ir_call::IrCall;
use crate::ir::ir_expression::IrExpression;
use crate::source_range::SourceRange;
use crate::symbol::Symbol;
use crate::symbol::callable_symbol::CallableSymbol;
use crate::symbol::expressible_symbol::ExpressibleSymbol;
use crate::symbol_table::SymbolTable;
use crate::type_info::TypeInfo;

pub struct Call {
    callee: Box<Expression>,
    arguments: Vec<Expression>,
    source_range: SourceRange,
    return_type_info: Option<TypeInfo>,
}

impl Call {
    pub fn new(callee: Expression, arguments: Vec<Expression>, source_range: SourceRange) -> Self {
        Self {
            callee: callee.into(),
            arguments,
            source_range,
            return_type_info: None,
        }
    }

    pub fn callee(&self) -> &Expression {
        &self.callee
    }

    pub fn arguments(&self) -> Vec<&Expression> {
        self.arguments.iter().collect()
    }

    pub fn gather_declared_names(
        &mut self,
        symbol_table: &mut SymbolTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut to_gather = vec![];
        to_gather.push(self.callee.as_mut());
        to_gather.extend(&mut self.arguments);

        let diagnostics: Vec<Diagnostic> = to_gather
            .iter_mut()
            .map(|expression| expression.gather_declared_names(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect();
        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn check_name_usages(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        let mut to_check = vec![];
        to_check.push(self.callee.as_mut());
        to_check.extend(&mut self.arguments);
        let diagnostics: Vec<Diagnostic> = to_check
            .iter_mut()
            .map(|expression| expression.check_name_usages(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect();
        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn type_check(&mut self, symbol_table: &SymbolTable) -> Result<(), Vec<Diagnostic>> {
        self.callee.as_mut().type_check(symbol_table)?;

        let mut diagnostics: Vec<Diagnostic> = self
            .arguments
            .iter_mut()
            .map(|argument| argument.type_check(symbol_table))
            .filter_map(Result::err)
            .flatten()
            .collect();

        // check that callee is callable
        let callable_symbol = match self.callee.type_info() {
            TypeInfo::Function(function_symbol) => {
                CallableSymbol::Function(function_symbol.clone())
            }
            TypeInfo::ClassInstance(class_symbol) => CallableSymbol::Class(class_symbol.clone()),
            _ => {
                diagnostics.push(Diagnostic::new(
                    &format!(
                        "Receiver of type {} is not callable.",
                        self.callee.type_info()
                    ),
                    self.callee.source_range().start(),
                    self.callee.source_range().end(),
                ));
                return Err(diagnostics);
            }
        };

        // set return type
        self.return_type_info = Some(callable_symbol.return_type_info());

        // check arguments length
        let parameters = callable_symbol.parameters();
        if parameters.len() != self.arguments.len() {
            diagnostics.push(Diagnostic::new(
                &format!(
                    "Wrong number of arguments; expected {} but found {}",
                    parameters.len(),
                    self.arguments.len()
                ),
                self.source_range.start(),
                self.source_range.end(),
            ));
        }

        if !diagnostics.is_empty() {
            return Err(diagnostics);
        }

        // check argument types
        for i in 0..parameters.len() {
            let parameter = &parameters[i].borrow();
            let argument = &self.arguments[i];
            let parameter_type_info = parameter.type_info();
            let argument_type_info = argument.type_info();
            if !parameter_type_info.is_assignable_from(argument_type_info) {
                diagnostics.push(Diagnostic::new(
                    &format!(
                        "Mismatched types: expected {} but found {}",
                        parameter_type_info, argument_type_info
                    ),
                    argument.source_range().start(),
                    argument.source_range().end(),
                ))
            }
        }

        if diagnostics.is_empty() {
            Ok(())
        } else {
            Err(diagnostics)
        }
    }

    pub fn return_type_info(&self) -> &TypeInfo {
        self.return_type_info.as_ref().unwrap()
    }

    fn get_callee_symbol(&self) -> CallableSymbol {
        match self.callee() {
            Expression::Identifier(identifier) => {
                let expressible_symbol = identifier.expressible_symbol();
                match expressible_symbol {
                    ExpressibleSymbol::Function(function_symbol) => {
                        CallableSymbol::Function(function_symbol.clone())
                    }
                    ExpressibleSymbol::Class(class_symbol) => {
                        CallableSymbol::Class(class_symbol.clone())
                    }
                    _ => panic!("Calling things other than functions not yet supported."),
                }
            }
            _ => panic!("Calling things other than identifiers not yet supported."),
        }
    }

    pub fn to_ir(&self, builder: &mut IrBuilder, symbol_table: &SymbolTable) -> IrCall {
        let arguments: Vec<IrExpression> = self
            .arguments
            .iter()
            .map(|argument| argument.to_ir(builder, symbol_table))
            .inspect(|expression| {
                if expression.is_none() {
                    panic!("Attempt to pass non-expression")
                }
            })
            .map(Option::unwrap)
            .collect();
        let callable_symbol = self.get_callee_symbol();
        match callable_symbol {
            CallableSymbol::Function(function_symbol) => IrCall::new(
                fqn_parts_to_string(function_symbol.borrow().fqn_parts()),
                arguments,
                function_symbol.borrow().is_extern(),
            ),
            CallableSymbol::Class(class_symbol) => {
                let constructor_symbol = class_symbol
                    .borrow()
                    .constructor_symbol()
                    .cloned()
                    .expect("Default constructors not supported yet.");
                IrCall::new(
                    fqn_parts_to_string(constructor_symbol.borrow().fqn_parts()),
                    arguments,
                    false,
                )
            }
        }
    }

    pub fn source_range(&self) -> &SourceRange {
        &self.source_range
    }
}