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

use crate::ast::assign_statement::AssignStatement;
use crate::ast::binary_expression::BinaryExpression;
use crate::ast::call::Call;
use crate::ast::compilation_unit::CompilationUnit;
use crate::ast::expression::Expression;
use crate::ast::expression_statement::ExpressionStatement;
use crate::ast::extern_function::ExternFunction;
use crate::ast::function::Function;
use crate::ast::identifier::Identifier;
use crate::ast::let_statement::LetStatement;
use crate::ast::negative_expression::NegativeExpression;
use crate::ast::statement::Statement;
use crate::semantic_analysis::analysis_context::AnalysisContext;

/// N.b.: the parent_scope_id refers to the **parent** of the compilation unit scope, and must
/// be a valid (non-panicking) index of `scopes`.
pub fn collect_scopes_compilation_unit(
    compilation_unit: &CompilationUnit,
    ctx: &mut AnalysisContext,
) {
    ctx.push_scope("compilation_unit");
    for function in compilation_unit.functions() {
        collect_scopes_function(function, ctx);
    }
    for extern_function in compilation_unit.extern_functions() {
        collect_scopes_extern_function(extern_function, ctx);
    }
    ctx.pop_scope(); // compilation_unit
}

fn collect_scopes_function(function: &Function, ctx: &mut AnalysisContext) {
    // associate function with current scope
    ctx.associate_node_to_current_scope(function.node_id());

    // push function params scope
    ctx.push_scope("function_parameters");

    // save return-type and parameters' scope id
    for parameter in function.parameters() {
        ctx.associate_node_to_current_scope(parameter.node_id());
    }
    match function.return_type() {
        None => {
            // no-op
        }
        Some(type_use) => {
            ctx.associate_node_to_current_scope(type_use.node_id());
        }
    }

    // push block scope for body
    ctx.push_scope("function_body");

    for statement in function.statements() {
        collect_scopes_statement(statement, ctx);
    }

    ctx.pop_scope(); // body
    ctx.pop_scope(); // parameters
}

fn collect_scopes_extern_function(extern_function: &ExternFunction, ctx: &mut AnalysisContext) {
    // associate extern_function with current scope
    ctx.associate_node_to_current_scope(extern_function.node_id());

    // parameters scope
    ctx.push_scope("extern_function_parameters");
    for parameter in extern_function.parameters() {
        ctx.associate_node_to_current_scope(parameter.node_id());
    }
    ctx.associate_node_to_current_scope(extern_function.return_type().node_id());
    ctx.pop_scope(); // parameters
}

pub fn collect_scopes_statement(statement: &Statement, ctx: &mut AnalysisContext) {
    match statement {
        Statement::Let(let_statement) => collect_scopes_let_statement(let_statement, ctx),
        Statement::Expression(expression_statement) => {
            collect_scopes_expression_statement(expression_statement, ctx)
        }
        Statement::Assign(assign_statement) => {
            collect_scopes_assign_statement(assign_statement, ctx)
        }
    }
}

fn collect_scopes_let_statement(let_statement: &LetStatement, ctx: &mut AnalysisContext) {
    collect_scopes_expression(let_statement.initializer(), ctx);
    ctx.associate_node_to_current_scope(let_statement.node_id());
}

fn collect_scopes_expression_statement(
    expression_statement: &ExpressionStatement,
    ctx: &mut AnalysisContext,
) {
    collect_scopes_expression(expression_statement.expression(), ctx);
}

fn collect_scopes_assign_statement(assign_statement: &AssignStatement, ctx: &mut AnalysisContext) {
    collect_scopes_expression(assign_statement.value(), ctx);
    collect_scopes_expression(assign_statement.destination(), ctx);
}

fn collect_scopes_expression(expression: &Expression, ctx: &mut AnalysisContext) {
    match expression {
        Expression::Binary(binary_expression) => {
            collect_scopes_binary_expression(binary_expression, ctx);
        }
        Expression::Negative(negative_expression) => {
            collect_scopes_negative_expression(negative_expression, ctx);
        }
        Expression::Call(call) => {
            collect_scopes_call(call, ctx);
        }
        Expression::Identifier(identifier) => {
            collect_scopes_identifier(identifier, ctx);
        }
        Expression::Integer(_) => {}
        Expression::Double(_) => {}
        Expression::String(_) => {}
    }
}

fn collect_scopes_binary_expression(
    binary_expression: &BinaryExpression,
    ctx: &mut AnalysisContext,
) {
    collect_scopes_expression(binary_expression.lhs(), ctx);
    collect_scopes_expression(binary_expression.rhs(), ctx);
}

fn collect_scopes_negative_expression(
    negative_expression: &NegativeExpression,
    ctx: &mut AnalysisContext,
) {
    collect_scopes_expression(negative_expression.operand(), ctx);
}

fn collect_scopes_call(call: &Call, ctx: &mut AnalysisContext) {
    for argument in call.arguments() {
        collect_scopes_expression(argument, ctx);
    }
    collect_scopes_expression(call.callee(), ctx);
}

fn collect_scopes_identifier(identifier: &Identifier, ctx: &mut AnalysisContext) {
    ctx.associate_node_to_current_scope(identifier.node_id());
}