deimos-lang/src/ir/lower_ast.rs

use crate::ast::node::{
    CompilationUnit, Expression, Function, FunctionBlockBody, FunctionBody, Literal,
    ModuleLevelDeclaration, Statement, VariableDeclaration,
};
use crate::ir::{
    Ir, IrAssign, IrConst, IrExpression, IrFunction, IrKind, IrLiteral, IrPrimitiveKind,
    IrStatement, IrVariable,
};
use crate::name_analysis::symbol::{ParameterSymbol, PrimitiveTypeSymbol, TypeSymbol};
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

pub fn lower_compilation_unit(compilation_unit: &CompilationUnit) -> Vec<Ir> {
    let mut result: Vec<Ir> = vec![];
    for declaration in compilation_unit.module_level_declarations() {
        result.append(&mut lower_module_level_declaration(declaration));
    }
    result
}

fn lower_module_level_declaration(declaration: &ModuleLevelDeclaration) -> Vec<Ir> {
    match declaration {
        ModuleLevelDeclaration::Module(module) => {
            todo!()
        }
        ModuleLevelDeclaration::Interface(interface) => {
            todo!()
        }
        ModuleLevelDeclaration::Class(class) => {
            todo!()
        }
        ModuleLevelDeclaration::Function(function) => lower_function(function),
        ModuleLevelDeclaration::PlatformFunction(platform_function) => {
            todo!()
        }
    }
}

pub struct LocalsTable {
    t_count: usize,
    stack_slot_count: usize,
    stack_slots: HashMap<IrVariable, usize>,
}

impl LocalsTable {
    pub fn new() -> Self {
        Self {
            t_count: 0,
            stack_slot_count: 0,
            stack_slots: HashMap::new(),
        }
    }

    pub fn next_t_var(&mut self) -> IrVariable {
        let t_num = self.t_count;
        self.t_count += 1;
        let t_name = format!("__t{}", t_num);
        let t_var = IrVariable::new(&t_name);
        self.register_stack_local(&t_var);
        t_var
    }

    pub fn register_stack_local(&mut self, ir_variable: &IrVariable) {
        let stack_slot = self.stack_slot_count;
        self.stack_slot_count += 1;
        self.stack_slots.insert(ir_variable.clone(), stack_slot);
    }
    
    pub fn get_stack_slot(&self, ir_variable: &IrVariable) -> usize {
        *self.stack_slots.get(ir_variable).unwrap()
    }
}

fn lower_function(function: &Function) -> Vec<Ir> {
    let function_symbol = function.function_symbol().unwrap().borrow();
    let name = function_symbol.fqn_parts_owned().join("::");
    let parameters = lower_parameter_symbols(function_symbol.parameter_symbols());
    let return_type = lower_type_symbol(function_symbol.return_type().unwrap());

    let mut locals_table = LocalsTable::new();
    let (statements, consts) = lower_function_body(function.function_body(), &mut locals_table);

    let mut result: Vec<Ir> = vec![];
    result.push(Ir::Function(Box::new(IrFunction::new(
        name,
        parameters,
        Box::new(return_type),
        statements,
        Box::new(locals_table),
    ))));
    result.append(
        &mut consts
            .into_iter()
            .map(|constant| Ir::Const(Box::new(constant)))
            .collect(),
    );
    result
}

fn lower_parameter_symbols(parameter_symbols: &[Rc<RefCell<ParameterSymbol>>]) -> Vec<IrKind> {
    parameter_symbols
        .iter()
        .map(|parameter_symbol_rc| parameter_symbol_rc.borrow())
        .map(|parameter_symbol| lower_type_symbol(parameter_symbol.type_symbol().unwrap()))
        .collect()
}

fn lower_type_symbol(type_symbol: &TypeSymbol) -> IrKind {
    match type_symbol {
        TypeSymbol::Primitive(primitive_type_symbol) => {
            lower_primitive_type_symbol(primitive_type_symbol)
        }
        TypeSymbol::Class(class_type_symbol) => {
            todo!()
        }
        TypeSymbol::Interface(interface_type_symbol) => {
            todo!()
        }
        TypeSymbol::Generic(generic_type_symbol) => {
            todo!()
        }
    }
}

fn lower_primitive_type_symbol(primitive_type_symbol: &PrimitiveTypeSymbol) -> IrKind {
    IrKind::Primitive(Box::new(match primitive_type_symbol {
        PrimitiveTypeSymbol::Byte => IrPrimitiveKind::I8,
        PrimitiveTypeSymbol::Char => IrPrimitiveKind::I8,
        PrimitiveTypeSymbol::Short => IrPrimitiveKind::I16,
        PrimitiveTypeSymbol::Int => IrPrimitiveKind::I32,
        PrimitiveTypeSymbol::Long => IrPrimitiveKind::I64,
        PrimitiveTypeSymbol::Float => IrPrimitiveKind::Float,
        PrimitiveTypeSymbol::Double => IrPrimitiveKind::Double,
        PrimitiveTypeSymbol::Boolean => IrPrimitiveKind::Boolean,
        PrimitiveTypeSymbol::String => IrPrimitiveKind::String,
        PrimitiveTypeSymbol::TypedArray { inner_type } => {
            IrPrimitiveKind::Array(Box::new(lower_type_symbol(inner_type.as_ref().unwrap())))
        }
        PrimitiveTypeSymbol::Any => {
            todo!()
        }
        PrimitiveTypeSymbol::Void => IrPrimitiveKind::Void,
    }))
}

fn lower_function_body(
    function_body: &FunctionBody,
    locals_table: &mut LocalsTable
) -> (Vec<IrStatement>, Vec<IrConst>) {
    match function_body {
        FunctionBody::FunctionAliasBody(alias_body) => {
            todo!()
        }
        FunctionBody::FunctionEqualsBody(equals_body) => {
            todo!()
        }
        FunctionBody::FunctionBlockBody(block_body) => {
            lower_function_block_body(block_body, locals_table)
        }
    }
}

fn lower_function_block_body(
    body: &FunctionBlockBody,
    locals_table: &mut LocalsTable
) -> (Vec<IrStatement>, Vec<IrConst>) {
    let mut statements: Vec<IrStatement> = vec![];
    let mut consts: Vec<IrConst> = vec![];
    for statement in body.statements() {
        statements.append(&mut lower_statement(statement, &mut consts, locals_table));
    }
    (statements, consts)
}

fn lower_statement(
    statement: &Statement,
    consts_pool: &mut Vec<IrConst>,
    locals_table: &mut LocalsTable
) -> Vec<IrStatement> {
    match statement {
        Statement::VariableDeclaration(variable_declaration) => {
            lower_variable_declaration(variable_declaration, consts_pool, locals_table)
        }
        Statement::AssignmentStatement(assignment_statement) => {
            todo!()
        }
        Statement::ExpressionStatement(expression_statement) => {
            todo!()
        }
        Statement::UseStatement(use_statement) => {
            todo!()
        }
        Statement::IfStatement(if_statement) => {
            todo!()
        }
        Statement::WhileStatement(while_statement) => {
            todo!()
        }
        Statement::ForStatement(for_statement) => {
            todo!()
        }
    }
}

fn lower_variable_declaration(
    variable_declaration: &VariableDeclaration,
    consts_pool: &mut Vec<IrConst>,
    locals_table: &mut LocalsTable
) -> Vec<IrStatement> {
    let mut result: Vec<IrStatement> = vec![];
    let expression_t_var = lower_expression(
        variable_declaration.expression().unwrap(),
        &mut result,
        consts_pool,
        locals_table,
    );
    let target = IrVariable::new(variable_declaration.identifier().name());
    locals_table.register_stack_local(&target);
    let assign_stmt = IrAssign::new(
        Box::new(target),
        Box::new(IrExpression::Variable(Box::new(expression_t_var))),
    );
    result.push(IrStatement::Assign(assign_stmt));
    result
}

fn lower_expression(
    expression: &Expression,
    target: &mut Vec<IrStatement>,
    consts_pool: &mut Vec<IrConst>,
    locals_table: &mut LocalsTable
) -> IrVariable {
    match expression {
        Expression::Ternary(ternary) => {
            todo!()
        }
        Expression::Or(or) => {
            todo!()
        }
        Expression::And(and) => {
            todo!()
        }
        Expression::Comparison(comparison) => {
            todo!()
        }
        Expression::Shift(shift) => {
            todo!()
        }
        Expression::Additive(additive) => {
            todo!()
        }
        Expression::Multiplicative(multiplicative) => {
            todo!()
        }
        Expression::Prefix(prefix) => {
            todo!()
        }
        Expression::Suffix(suffix) => {
            todo!()
        }
        Expression::Literal(literal) => lower_literal(literal, target, consts_pool, locals_table),
        Expression::Identifier(identifier) => {
            todo!()
        }
        Expression::Fqn(fqn) => {
            todo!()
        }
        Expression::Closure(closure) => {
            todo!()
        }
        Expression::List(list) => {
            todo!()
        }
    }
}

fn lower_literal(
    literal: &Literal,
    statements: &mut Vec<IrStatement>,
    consts_pool: &mut Vec<IrConst>,
    locals_table: &mut LocalsTable
) -> IrVariable {
    match literal {
        Literal::IntLiteral(int_literal) => {
            let t_var = locals_table.next_t_var();
            let assign_stmt = IrAssign::new(
                Box::new(t_var.clone()),
                Box::new(IrExpression::Literal(Box::new(IrLiteral::I32(
                    *int_literal,
                )))),
            );
            statements.push(IrStatement::Assign(assign_stmt));
            t_var
        }
        Literal::LongLiteral(long_literal) => {
            todo!()
        }
        Literal::DoubleLiteral(double_literal) => {
            todo!()
        }
        Literal::SingleQuoteString(sq_string) => {
            todo!()
        }
        Literal::DString(d_string) => {
            todo!()
        }
        Literal::BacktickString(b_string) => {
            todo!()
        }
        Literal::BooleanLiteral(boolean_literal) => {
            todo!()
        }
    }
}