semantic.cpp

/**
 * semantic.cpp
 * author: Bao Le
 */

#include <iostream>
#include <optional>
#include <vector>
#include <set>
#include <string>
#include "ast.hpp"
#include "semantic.hpp"
#include "ast.hpp"
#include "lexer.hpp"
#include "parser.hpp"

// Main public entry point
void SemanticAnalyzer::analyze() {
    register_libc_builtins();
    for (auto& child : root->children) {
        visit_top_level(child.get());
    }
}

void SemanticAnalyzer::register_libc_builtins() {
    declare("printf", Symbol{ConstantType::Int, std::nullopt, true, true});
}

void SemanticAnalyzer::visit_top_level(ASTNode* node) {
    if (auto* fn = dynamic_cast<Function*>(node)) {
        visit_function(fn);
        return;
    }
    if (auto* decl = dynamic_cast<Declaration*>(node)) {
        visit_declaration(decl);
        return;
    }
     if (auto* assign = dynamic_cast<Assignment*>(node)) {
        visit_assignment(assign);
        return;
    }
    throw std::runtime_error("[SEMANTIC] unexpected top-level node\n");
}

// Register a variable into the current scope
void SemanticAnalyzer::declare(const std::string& name, Symbol symbol) {
    // Variable is not found
    if (scopes.back().find(name) == scopes.back().end()) {
        scopes.back()[name] = symbol;
        return;
    }
    throw std::runtime_error("[DECLARE] Redeclaration of variable '" + name + "'");
}

// Check if variable exists from innermost to outermost scope
Symbol SemanticAnalyzer::lookup(const std::string& name) {
    for (int i = scopes.size() - 1; i >= 0; --i) {
        auto it = scopes[i].find(name);
        if (it != scopes[i].end()) return it->second;
    }
    throw std::runtime_error("[LOOKUP] variable " + name + "doesn't exist in scope\n");
}

void SemanticAnalyzer::push_scope() {
    scopes.push_back({});
}

void SemanticAnalyzer::pop_scope() {
    scopes.pop_back();
}

void SemanticAnalyzer::visit_function(Function* fn) {
    std::vector<ConstantType> param_types;
    for (auto& param : fn->params) {
        param_types.push_back(param.type);
    }
    declare(fn->name, Symbol{fn->type, param_types});

    push_scope();
    for (auto& param : fn->params) {
        declare(param.name, Symbol{param.type, std::nullopt});
    }
    visit_statements_no_push(fn->body->statements, fn->type);
    pop_scope();
}

void SemanticAnalyzer::visit_statements_no_push(std::vector<std::unique_ptr<ASTNode>>& statements, ConstantType expected_return_type) {
    for (auto& stmt : statements) {
        visit_statement(stmt.get(), expected_return_type);    
    }
}
        
// For if-else statement body
void SemanticAnalyzer::visit_statement_body(StatementBody* body, ConstantType expected_return_value) {
    push_scope();
    visit_statements_no_push(body->statements, expected_return_value);
    pop_scope();
}
        
// Route to the correct visit function based on the statement's type
void SemanticAnalyzer::visit_statement(ASTNode* stmt, ConstantType expected_return_type) {
    if (auto* if_stmt = dynamic_cast<If*>(stmt)) {
        visit_if(if_stmt, expected_return_type);
        return;
    }
    if (auto* decl = dynamic_cast<Declaration*>(stmt)) {
        visit_declaration(decl);
        return;
    }
    if (auto* ret = dynamic_cast<Return*>(stmt)) {
        visit_return(ret, expected_return_type);
        return;
    }
    if (auto* assignment = dynamic_cast<Assignment*>(stmt)) {
        visit_assignment(assignment);
        return;
    }
    if (auto* fn_call = dynamic_cast<FunctionCall*>(stmt)) {
        visit_function_call_statement(fn_call);
        return;
    }
    if (auto* body = dynamic_cast<StatementBody*>(stmt)) {
        visit_statement_body(body, expected_return_type);
        return;
    }
    if (auto* while_stmt = dynamic_cast<While*>(stmt)) {
        visit_while(while_stmt);
        return;
    }
    throw std::runtime_error("[STATEMENT] Mismatched statement found");
}

void SemanticAnalyzer::visit_while(While* while_stmt) {
    auto cond_type = visit_expression(while_stmt->cond.get());
    if (cond_type != ConstantType::Bool) throw std::runtime_error("[WHILE] Didn't find Boolean type in condition\n");
    push_scope();
    visit_statements_no_push(while_stmt->body->statements, ConstantType::Void);
    pop_scope();
}

void SemanticAnalyzer::visit_declaration(Declaration* decl) {
    declare(decl->var.name, Symbol{decl->type, std::nullopt});
    if (decl->init_val) {
        auto actual_type = visit_expression(decl->init_val->get());
        if (actual_type != decl->type) throw std::runtime_error("[DECLARATION] Type mismatch between declared and actual value\n");
    }
}

void SemanticAnalyzer::visit_assignment(Assignment* assign) {
    // get() returns ASTNode* from unique_ptr<ASTNode>
    auto target_var = dynamic_cast<Variable*>(assign->target.get());
    auto existing_symbol = lookup(target_var->name);
    auto actual_type = visit_expression(assign->rhs.get());
    if (actual_type != existing_symbol.type) throw std::runtime_error("[ASSIGNMENT] Type mismatch between declared and actual rhs value\n");
}

void SemanticAnalyzer::visit_if(If* if_stmt, ConstantType expected_return_type) {
    auto cond_type = visit_expression(if_stmt->cond.get());
    if (cond_type != ConstantType::Bool) throw std::runtime_error("[IF] Didn't find Boolean type in condition\n");
    visit_statement_body(if_stmt->if_body.get(), expected_return_type);
    if (if_stmt->else_body) visit_statement_body(if_stmt->else_body->get(), expected_return_type);
}

void SemanticAnalyzer::visit_return(Return* ret, ConstantType expected_return_type) {
    // Void return
    if (ret->val.has_value()) {
        auto actual_type = visit_expression(ret->val.value().get());
        if (actual_type != expected_return_type) throw std::runtime_error("[RETURN] Type mistmatch between actual type and expected return type\n");
    }
    else {
        if (expected_return_type != ConstantType::Void) throw std::runtime_error("[RETURN] Type mistmatch between return and expected return\n");
    }
}

void SemanticAnalyzer::visit_function_call_statement(FunctionCall* call) {
    Symbol fn_symbol = lookup(call->name);

    if (fn_symbol.is_external && fn_symbol.is_variadic) {
        if (call->args.empty()) {
            throw std::runtime_error("[FUNCTION CALL] printf requires at least a format string\n");
        }
        auto fmt_type = visit_expression(call->args[0].get());
        if (fmt_type != ConstantType::String) {
            throw std::runtime_error("[FUNCTION CALL] printf format must be a string literal\n");
        }
        for (size_t i = 1; i < call->args.size(); ++i) {
            auto arg_type = visit_expression(call->args[i].get());
            if (arg_type != ConstantType::Int) {
                throw std::runtime_error("[FUNCTION CALL] printf extra args must be int\n");
            }
        }
        return;
    }

    if (!fn_symbol.params.has_value()) throw std::runtime_error("[FUNCTION CALL] Missing parameters inside function\n");
    if (call->args.size() != fn_symbol.params->size()) throw std::runtime_error("[FUNCTION CALL] Parameters count mismatch\n");
    for (int i{}; i < call->args.size(); ++i){
        ConstantType arg_type = visit_expression(call->args[i].get());
        // unwrap std::optional with *
        if (arg_type != (*fn_symbol.params)[i]) throw std::runtime_error("[FUNCTION CALL] Argument type mismatch\n");
    }
}

// Dispatches based on Constant, Variable, BinaryOperation, or FunctionCall node
ConstantType SemanticAnalyzer::visit_expression(ASTNode* expr) {
    if (auto constant = dynamic_cast<Constant*>(expr)) return constant->type;
    // check a variable's type based on the lookup table
    if (auto var = dynamic_cast<Variable*>(expr)) return lookup(var->name).type;
    if (auto bin_op = dynamic_cast<BinaryOperation*>(expr)) {
        // Check the left and right operands
        auto left_type = visit_expression(bin_op->left.get());
        auto right_type = visit_expression(bin_op->right.get());
        if (left_type != right_type) throw std::runtime_error("[EXPRESSION] Type mistmatch between left an right operands\n");
        // Check the operator for the exact return type
        std::set<BinaryOperator> rel_ops {BinaryOperator::LessThan, BinaryOperator::GreaterThan, BinaryOperator::Equal, BinaryOperator::NotEqual};
        std::set<BinaryOperator> math_ops {BinaryOperator::Plus, BinaryOperator::Minus, BinaryOperator::Star, BinaryOperator::Slash};
        if (rel_ops.count(bin_op->op)) return ConstantType::Bool;
        if (math_ops.count(bin_op->op)) return ConstantType::Int;
        throw std::runtime_error("[EXPRESSION] Binary operator's type cannot be deduced\n");
    }
    // determine function call's type
    if (auto fn_call = dynamic_cast<FunctionCall*>(expr)) {
        // check for valid params
        visit_function_call_statement(fn_call);
        // return the function's type through lookup
        return lookup(fn_call->name).type;
    }
    throw std::runtime_error("[EXPRESSION] unknown expression node type");
}