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

use crate::ast::NodeId;
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::scope::{Scope, ScopeId};
use std::collections::HashMap;

pub struct ScopeCollectionResult {
    pub scopes: Vec<Scope>,
    pub nodes_to_scopes: HashMap<NodeId, ScopeId>,
}

struct ScopeCollectionContext {
    scopes: Vec<Scope>,
    current_scope_id: Option<ScopeId>,
    nodes_to_scopes: HashMap<NodeId, ScopeId>,
}

impl ScopeCollectionContext {
    pub fn new() -> Self {
        Self {
            scopes: Vec::new(),
            current_scope_id: None,
            nodes_to_scopes: HashMap::new(),
        }
    }

    pub fn push_scope(&mut self, scope: Scope) -> ScopeId {
        self.scopes.push(scope);
        self.current_scope_id = Some(self.scopes.len() - 1);
        self.current_scope_id.unwrap() // guaranteed because we just set it
    }

    pub fn pop_scope(&mut self) {
        self.current_scope_id = self.scopes[self.current_scope_id.unwrap()].parent_id();
    }

    pub fn current_scope_id(&self) -> Option<ScopeId> {
        self.current_scope_id
    }

    pub fn nodes_to_scopes_mut(&mut self) -> &mut HashMap<NodeId, ScopeId> {
        &mut self.nodes_to_scopes
    }
}

pub fn collect_scopes(compilation_unit: &CompilationUnit) -> ScopeCollectionResult {
    let mut ctx = ScopeCollectionContext::new();
    ctx.push_scope(Scope::new(None));
    for function in compilation_unit.functions() {
        collect_scopes_function(function, &mut ctx);
    }
    for extern_function in compilation_unit.extern_functions() {
        collect_scopes_extern_function(extern_function, &mut ctx);
    }
    ctx.pop_scope();
    ScopeCollectionResult {
        scopes: ctx.scopes,
        nodes_to_scopes: ctx.nodes_to_scopes,
    }
}

fn collect_scopes_function(function: &Function, ctx: &mut ScopeCollectionContext) {
    // get containing scope id
    let containing_scope_id = ctx.current_scope_id().unwrap(); // guaranteed because functions are in modules

    // save function's containing scope id
    ctx.nodes_to_scopes_mut()
        .insert(function.node_id(), containing_scope_id);

    // push function scope
    let function_scope = Scope::new(Some(containing_scope_id));
    let function_scope_id = ctx.push_scope(function_scope);

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

    // push block scope for body
    let block_scope = Scope::new(Some(function_scope_id));
    ctx.push_scope(block_scope);

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

    ctx.pop_scope(); // block
    ctx.pop_scope(); // function
}

fn collect_scopes_extern_function(
    extern_function: &ExternFunction,
    ctx: &mut ScopeCollectionContext,
) {
    let containing_scope_id = ctx.current_scope_id().unwrap();

    ctx.nodes_to_scopes_mut()
        .insert(extern_function.node_id(), containing_scope_id);

    let function_scope = Scope::new(Some(containing_scope_id));
    let function_scope_id = ctx.push_scope(function_scope);

    for parameter in extern_function.parameters() {
        ctx.nodes_to_scopes_mut()
            .insert(parameter.node_id(), function_scope_id);
    }
    ctx.nodes_to_scopes_mut()
        .insert(extern_function.return_type().node_id(), function_scope_id);

    ctx.pop_scope();
}

fn collect_scopes_statement(statement: &Statement, ctx: &mut ScopeCollectionContext) {
    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 ScopeCollectionContext) {
    collect_scopes_expression(let_statement.initializer(), ctx);
}

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

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

fn collect_scopes_expression(expression: &Expression, ctx: &mut ScopeCollectionContext) {
    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 ScopeCollectionContext,
) {
    collect_scopes_expression(binary_expression.lhs(), ctx);
    collect_scopes_expression(binary_expression.rhs(), ctx);
}

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

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

fn collect_scopes_identifier(identifier: &Identifier, ctx: &mut ScopeCollectionContext) {
    let current_scope_id = ctx.current_scope_id().unwrap(); // if this fails, we tried to do this without any context
    ctx.nodes_to_scopes_mut()
        .insert(identifier.node_id(), current_scope_id);
}