CodeGenerator.cpp

Go to the documentation of this file.

 
#include "CodeGenerator.hpp"
#include "Builtins.hpp"
#include "PostfixBuilder.hpp"
#include "PostfixHelper.hpp"
#include "types.hpp"
 
#include <format>
#include <functional>
#include <utility>
 
namespace infix2postfix {
 
CodeGenerator::CodeGenerator(Mode mode, int num_inputs,
                             int num_intermediate_inputs)
    : mode(mode), num_inputs(num_inputs),
      num_intermediate_inputs(num_intermediate_inputs) {}
 
std::string CodeGenerator::generate(const Program* program) {
    PostfixBuilder main_builder;
    for (const auto& stmt : program->statements) {
        generate(stmt.get(), main_builder);
    }
 
    if (mode == Mode::Expr || mode == Mode::VkExpr) {
        main_builder.addVariableLoad("RESULT");
    }
 
    return main_builder.getExpression();
}
 
CodeGenerator::ExprResult CodeGenerator::generateExpr(Expr* expr) {
    PostfixBuilder b;
    Type t = generate(expr, b);
    return {.postfix = b, .type = t};
}
 
std::string CodeGenerator::generateExprToString(Expr* expr) {
    if (expr == nullptr) {
        return "";
    }
    PostfixBuilder temp_builder;
    generate(expr, temp_builder);
    return temp_builder.getExpression();
}
 
Type CodeGenerator::generate(Expr* expr, PostfixBuilder& builder) {
    if (expr == nullptr) {
        return Type::Value;
    }
    return std::visit(
        [this, &builder](auto& e) { return this->handle(e, builder); },
        expr->value);
}
 
void CodeGenerator::generate(Stmt* stmt, PostfixBuilder& builder) {
    if (stmt == nullptr) {
        return;
    }
    std::visit([this, &builder](auto& s) { this->handle(s, builder); },
               stmt->value);
}
 
Type CodeGenerator::handle(const NumberExpr& expr, PostfixBuilder& builder) {
    builder.addNumber(expr.value.value);
    return Type::Literal;
}
 
Type CodeGenerator::handle(const VariableExpr& expr, PostfixBuilder& builder) {
    std::string name = expr.name.value;
 
    if (param_substitutions.contains(name)) {
        return generate(param_substitutions.at(name), builder);
    }
 
    if (name.starts_with("$")) {
        std::string base_name = name.substr(1);
        builder.addConstant(base_name);
 
        if (auto buf_idx = get_buffer_index(base_name)) {
            if (*buf_idx >= 0 && *buf_idx < num_intermediate_inputs) {
                return Type::Clip;
            }
        }
 
        if (is_clip_name(base_name)) {
            return Type::Clip;
        }
        return Type::Value;
    }
 
    std::string var_name = name;
    if (var_rename_map.contains(name)) {
        var_name = var_rename_map.at(name);
    } else if (expr.symbol) {
        var_name = expr.symbol->name;
    }
 
    builder.addVariableLoad(var_name);
 
    if (expr.symbol) {
        return expr.symbol->type;
    }
 
    return Type::Value;
}
 
Type CodeGenerator::handle(const UnaryExpr& expr, PostfixBuilder& builder) {
    if (expr.op.type == TokenType::Minus) {
        if (auto* num = get_if<NumberExpr>(expr.right.get())) {
            std::string val = num->value.value;
            if (!val.starts_with("-")) {
                val = std::format("-{}", val);
            }
            builder.addNumber(val);
            return Type::Literal;
        }
    }
 
    generate(expr.right.get(), builder);
 
    if (expr.op.type == TokenType::Not) {
        builder.addNumber("0");
        builder.addOp(TokenType::Eq);
        return Type::Value;
    }
 
    builder.addUnaryOp(expr.op.type);
    return Type::Value;
}
 
Type CodeGenerator::handle(const BinaryExpr& expr, PostfixBuilder& builder) {
    generate(expr.left.get(), builder);
 
    if (expr.op.type == TokenType::LogicalAnd ||
        expr.op.type == TokenType::LogicalOr) {
        builder.addNumber("0");
        builder.addOp(TokenType::Eq);
        builder.addUnaryOp(TokenType::Not);
 
        generate(expr.right.get(), builder);
        builder.addNumber("0");
        builder.addOp(TokenType::Eq);
        builder.addUnaryOp(TokenType::Not);
 
        builder.addOp(expr.op.type);
        return Type::Value;
    }
 
    generate(expr.right.get(), builder);
    builder.addOp(expr.op.type);
    return Type::Value;
}
 
Type CodeGenerator::handle(const TernaryExpr& expr, PostfixBuilder& builder) {
    generate(expr.cond.get(), builder);
    builder.addNumber("0");
    builder.addOp(TokenType::Eq);
    builder.addUnaryOp(TokenType::Not);
 
    PostfixBuilder true_branch_builder;
    generate(expr.true_expr.get(), true_branch_builder);
 
    PostfixBuilder false_branch_builder;
    generate(expr.false_expr.get(), false_branch_builder);
 
    builder.append(true_branch_builder);
    builder.append(false_branch_builder);
    builder.addTernaryOp();
    return Type::Value;
}
 
Type CodeGenerator::handle(const CallExpr& expr, PostfixBuilder& builder) {
    // User-defined functions are inlined
    if (expr.resolved_signature != nullptr && expr.resolved_def != nullptr) {
        const auto& sig = *expr.resolved_signature;
        FunctionDef* func_def = expr.resolved_def;
 
        inlineFunctionCall(sig, func_def, expr.args, expr.range, builder);
        return Type::Value;
    }
 
    // Built-in functions
    if (expr.resolved_builtin != nullptr) {
        const BuiltinFunction* builtin = expr.resolved_builtin;
 
        if (builtin->special_handler) {
            builder.append(builtin->special_handler(this, expr));
            return Type::Value;
        }
 
        for (size_t i = 0; i < expr.args.size(); ++i) {
            if (builtin->param_types[i] != Type::LiteralString) {
                generate(expr.args[i].get(), builder);
            }
        }
        builder.addFunctionCall(expr.callee);
        return Type::Value;
    }
 
    // nth_N functions
    if (expr.callee.starts_with("nth_")) {
        std::string n_str = expr.callee.substr(4);
        int n = std::stoi(n_str);
        int arg_count = static_cast<int>(expr.args.size());
 
        for (const auto& arg : expr.args) {
            generate(arg.get(), builder);
        }
        builder.addSortN(arg_count);
        builder.addDropN(n - 1);
        builder.addSwapN(arg_count - n);
        builder.addDropN(arg_count - n);
 
        return Type::Value;
    }
 
    std::unreachable();
}
 
Type CodeGenerator::handle(const PropAccessExpr& expr,
                           PostfixBuilder& builder) {
    std::string clip_name = expr.clip.value;
    if (param_substitutions.contains(clip_name)) {
        auto* subst_expr = param_substitutions.at(clip_name);
        clip_name = generateExprToString(subst_expr);
    } else {
        clip_name = clip_name.substr(1);
    }
 
    builder.addPropAccess(clip_name, expr.prop.value);
    return Type::Value;
}
 
Type CodeGenerator::handle(const StaticRelPixelAccessExpr& expr,
                           PostfixBuilder& builder) {
    std::string clip_name = expr.clip.value;
    if (param_substitutions.contains(clip_name)) {
        auto* subst_expr = param_substitutions.at(clip_name);
        clip_name = generateExprToString(subst_expr);
    } else {
        clip_name = clip_name.substr(1);
    }
 
    builder.addStaticPixelAccess(clip_name, expr.offset_x.value,
                                 expr.offset_y.value, expr.boundary_suffix);
    return Type::Value;
}
 
Type CodeGenerator::handle(const FrameDimensionExpr& expr,
                           PostfixBuilder& builder) {
    std::string plane_idx_str =
        generateExprToString(expr.plane_index_expr.get());
    builder.addFrameDimension(expr.dimension_name, plane_idx_str);
    return Type::Value;
}
 
Type CodeGenerator::handle(const ArrayAccessExpr& expr,
                           PostfixBuilder& builder) {
    // Get the array variable name
    std::string array_name;
    auto* var_expr = get_if<VariableExpr>(expr.array.get());
    if (var_expr != nullptr) {
        array_name = var_expr->name.value;
 
        if (var_rename_map.contains(array_name)) {
            array_name = var_rename_map.at(array_name);
        } else if (expr.array_symbol) {
            array_name = expr.array_symbol->name;
        }
    } else if (expr.array_symbol) {
        array_name = expr.array_symbol->name;
    }
 
    generate(expr.index.get(), builder);
    builder.addArrayLoad(array_name);
 
    return Type::Value;
}
 
void CodeGenerator::handle(const ExprStmt& stmt, PostfixBuilder& builder) {
#ifndef NDEBUG
    PostfixBuilder temp_builder;
    generate(stmt.expr.get(), temp_builder);
    checkStackEffect(temp_builder.getExpression(), 0, stmt.range);
    builder.append(temp_builder);
#else
    generate(stmt.expr.get(), builder);
#endif
}
 
void CodeGenerator::handle(const AssignStmt& stmt, PostfixBuilder& builder) {
    std::string name = stmt.name.value;
 
    std::string var_name = name;
    if (var_rename_map.contains(name)) {
        var_name = var_rename_map.at(name);
    } else if (stmt.symbol) {
        var_name = stmt.symbol->name;
    }
 
    // Check if this is a new() or resize() call for array allocation
    if (auto* call_expr = get_if<CallExpr>(stmt.value.get())) {
        if (call_expr->callee == "new" || call_expr->callee == "resize") {
#ifndef NDEBUG
            PostfixBuilder b;
#else
            PostfixBuilder& b = builder;
#endif
            if (mode == Mode::Expr || mode == Mode::VkExpr) {
                auto* arg_expr = call_expr->args[0].get();
 
                std::string size_value;
 
                if (auto* num_expr = get_if<NumberExpr>(arg_expr)) {
                    size_value = num_expr->value.value;
                } else if (auto* var_expr = get_if<VariableExpr>(arg_expr)) {
                    std::string var_name_arg = var_expr->name.value;
                    if (param_substitutions.contains(var_name_arg)) {
                        auto* subst_expr = param_substitutions.at(var_name_arg);
                        if (auto* subst_num = get_if<NumberExpr>(subst_expr)) {
                            size_value = subst_num->value.value;
                        }
                    }
                }
 
                b.addArrayAllocStatic(var_name, size_value);
            } else {
                generate(call_expr->args[0].get(), b);
                b.addArrayAllocDynamic(var_name);
            }
#ifndef NDEBUG
            checkStackEffect(b.getExpression(), 0, stmt.range);
            builder.append(b);
#endif
            return;
        }
    }
 
#ifndef NDEBUG
    PostfixBuilder b;
    generate(stmt.value.get(), b);
    b.addVariableStore(var_name);
    checkStackEffect(b.getExpression(), 0, stmt.range);
    builder.append(b);
#else
    generate(stmt.value.get(), builder);
    builder.addVariableStore(var_name);
#endif
}
 
void CodeGenerator::handle(const ArrayAssignStmt& stmt,
                           PostfixBuilder& builder) {
    auto* array_access = get_if<ArrayAccessExpr>(stmt.target.get());
 
    std::string array_name;
    auto* var_expr = get_if<VariableExpr>(array_access->array.get());
    if (var_expr != nullptr) {
        array_name = var_expr->name.value;
 
        if (var_rename_map.contains(array_name)) {
            array_name = var_rename_map.at(array_name);
        } else if (array_access->array_symbol) {
            array_name = array_access->array_symbol->name;
        }
    } else if (array_access->array_symbol) {
        array_name = array_access->array_symbol->name;
    }
 
#ifndef NDEBUG
    PostfixBuilder b;
    generate(stmt.value.get(), b);
    generate(array_access->index.get(), b);
    b.addArrayStore(array_name);
    checkStackEffect(b.getExpression(), 0, stmt.range);
    builder.append(b);
#else
    generate(stmt.value.get(), builder);
    generate(array_access->index.get(), builder);
    builder.addArrayStore(array_name);
#endif
}
 
void CodeGenerator::handle(const BlockStmt& stmt, PostfixBuilder& builder) {
    for (const auto& s : stmt.statements) {
        generate(s.get(), builder);
    }
}
 
void CodeGenerator::handle(const IfStmt& stmt, PostfixBuilder& builder) {
    std::string else_label = std::format("__internal_else_{}", label_counter++);
    std::string endif_label =
        std::format("__internal_endif_{}", label_counter++);
 
    generate(stmt.condition.get(), builder);
    builder.addNumber("0");
    builder.addOp(TokenType::Eq);
    builder.addConditionalJump(else_label);
 
    generate(stmt.then_branch.get(), builder);
 
    if (stmt.else_branch) {
        builder.addUnconditionalJump(endif_label);
        builder.addLabel(else_label);
        generate(stmt.else_branch.get(), builder);
        builder.addLabel(endif_label);
    } else {
        builder.addLabel(else_label);
    }
}
 
void CodeGenerator::handle(const WhileStmt& stmt, PostfixBuilder& builder) {
    std::string start_label =
        std::format("__internal_while_start_{}", label_counter++);
    std::string end_label =
        std::format("__internal_while_end_{}", label_counter++);
 
    builder.addLabel(start_label);
    generate(stmt.condition.get(), builder);
    builder.addNumber("0");
    builder.addOp(TokenType::Eq);
    builder.addConditionalJump(end_label);
    generate(stmt.body.get(), builder);
    builder.addUnconditionalJump(start_label);
    builder.addLabel(end_label);
}
 
void CodeGenerator::handle(const ReturnStmt& stmt, PostfixBuilder& builder) {
    if (stmt.value) {
        generate(stmt.value.get(), builder);
        const int call_id = call_site_id_stack.back();
        std::string ret_var_name = std::format("__internal_ret_{}_{}",
                                               current_function->name, call_id);
        builder.addVariableStore(ret_var_name);
    }
 
    const int call_id = call_site_id_stack.back();
    std::string ret_label = std::format("__internal_ret_label_{}_{}",
                                        current_function->name, call_id);
    builder.addUnconditionalJump(ret_label);
}
 
void CodeGenerator::handle(const LabelStmt& stmt, PostfixBuilder& builder) {
    std::string label_name = stmt.symbol ? stmt.symbol->name : stmt.name.value;
    builder.addLabel(label_name);
}
 
void CodeGenerator::handle(const GotoStmt& stmt, PostfixBuilder& builder) {
    std::string label_name = stmt.target_label_symbol
                                 ? stmt.target_label_symbol->name
                                 : stmt.label.value;
    if (stmt.condition) {
        generate(stmt.condition.get(), builder);
        builder.addNumber("0");
        builder.addOp(TokenType::Eq);
        builder.addUnaryOp(TokenType::Not);
        builder.addConditionalJump(label_name);
    } else {
        builder.addUnconditionalJump(label_name);
    }
}
 
void CodeGenerator::handle([[maybe_unused]] const FunctionDef& stmt,
                           [[maybe_unused]] PostfixBuilder& builder) {
    // Inlined at call sites
}
 
void CodeGenerator::handle([[maybe_unused]] const GlobalDecl& stmt,
                           [[maybe_unused]] PostfixBuilder& builder) {
    // Global declarations don't generate code
}
 
void CodeGenerator::checkStackEffect([[maybe_unused]] const std::string& s,
                                     [[maybe_unused]] int expected,
                                     [[maybe_unused]] const Range& range) {
#ifndef NDEBUG
    int effect = computeStackEffect(s, range);
    if (effect != expected) {
        throw CodeGenError(std::format("Unbalanced stack. "
                                       "Expected {}, got {}",
                                       expected, effect),
                           range);
    }
#endif
}
 
int CodeGenerator::computeStackEffect(const std::string& s,
                                      const Range& range) {
    PostfixMode postfix_mode = PostfixMode::SingleExpr;
    if (mode == Mode::Expr || mode == Mode::VkExpr) {
        postfix_mode = PostfixMode::Expr;
    }
 
    try {
        return compute_postfix_stack_effect(s, postfix_mode, range.start.line,
                                            num_inputs,
                                            num_intermediate_inputs);
    } catch (const std::exception& e) {
        throw CodeGenError(e.what(), range);
    }
}
 
Expr* CodeGenerator::resolveParamSubstitution(Expr* expr) const {
    Expr* resolved = expr;
    std::set<const Expr*> seen;
 
    while (resolved != nullptr) {
        auto* var_expr = get_if<VariableExpr>(resolved);
        if (var_expr == nullptr) {
            break;
        }
 
        auto it = param_substitutions.find(var_expr->name.value);
        if (it == param_substitutions.end() || it->second == nullptr) {
            break;
        }
 
        if (!seen.insert(resolved).second || it->second == resolved) {
            break;
        }
 
        resolved = it->second;
    }
 
    return resolved;
}
 
void CodeGenerator::inlineFunctionCall(
    const FunctionSignature& sig, FunctionDef* func_def,
    const std::vector<std::unique_ptr<Expr>>& args,
    [[maybe_unused]] const Range& call_range, PostfixBuilder& builder) {
 
    const auto& func_name = sig.name;
    const int call_id = call_site_counter++;
    call_site_id_stack.push_back(call_id);
 
    // Save current state
    auto saved_param_substitutions = param_substitutions;
    auto saved_var_rename_map = var_rename_map;
    const auto* saved_current_function = current_function;
 
    std::map<std::string, Expr*> next_param_substitutions;
    std::map<std::string, std::string> param_map;
 
    PostfixBuilder param_assignments;
 
    for (size_t i = 0; i < sig.params.size(); ++i) {
        const auto& param_info = sig.params[i];
        const std::string& param_name = param_info.name;
        const Type param_type = param_info.type;
 
        auto* arg_expr = args[i].get();
 
        if (param_type == Type::Literal || param_type == Type::Clip) {
            next_param_substitutions[param_name] =
                resolveParamSubstitution(arg_expr);
        } else if (param_type == Type::Array) {
            auto* var_expr = get_if<VariableExpr>(arg_expr);
            if (var_expr != nullptr) {
                std::string arg_var_name = var_expr->name.value;
                if (saved_var_rename_map.contains(arg_var_name)) {
                    param_map[param_name] =
                        saved_var_rename_map.at(arg_var_name);
                } else if (var_expr->symbol) {
                    param_map[param_name] = var_expr->symbol->name;
                } else {
                    param_map[param_name] = arg_var_name;
                }
            }
        } else {
            std::string renamed_param =
                std::format("__internal_func_{}_"
                            "{}_{}",
                            func_name, call_id, param_name);
 
            generate(args[i].get(), param_assignments);
            param_assignments.addVariableStore(renamed_param);
            param_map[param_name] = renamed_param;
        }
    }
 
    // Collect local variables from function body
    std::function<void(Stmt*)> collect_locals = [&](Stmt* stmt) {
        if (!stmt) {
            return;
        }
        if (auto* assign = get_if<AssignStmt>(stmt)) {
            std::string var_name = assign->name.value;
            if (!param_map.contains(var_name)) {
                std::string renamed_var =
                    std::format("__"
                                "interna"
                                "l_func_"
                                "{}_{}_{"
                                "}",
                                func_name, call_id, var_name);
                param_map[var_name] = renamed_var;
            }
        } else if (auto* block = get_if<BlockStmt>(stmt)) {
            for (const auto& s : block->statements) {
                collect_locals(s.get());
            }
        } else if (auto* if_stmt = get_if<IfStmt>(stmt)) {
            collect_locals(if_stmt->then_branch.get());
            if (if_stmt->else_branch) {
                collect_locals(if_stmt->else_branch.get());
            }
        } else if (auto* while_stmt = get_if<WhileStmt>(stmt)) {
            collect_locals(while_stmt->body.get());
        }
    };
 
    for (const auto& s : func_def->body->statements) {
        collect_locals(s.get());
    }
 
    // Update context for function body generation
    param_substitutions = std::move(next_param_substitutions);
    var_rename_map = param_map;
    current_function = &sig;
 
    std::string label_prefix =
        std::format("__internal_{}_{}_", func_name, call_id);
 
    PostfixBuilder body_builder;
    handle(*func_def->body, body_builder);
    body_builder.prefixLabels(label_prefix);
 
    std::string ret_label =
        std::format("__internal_ret_label_{}_{}", func_name, call_id);
    body_builder.addLabel(ret_label);
 
    if (sig.returns_value) {
        std::string ret_var =
            std::format("__internal_ret_{}_{}", func_name, call_id);
        body_builder.addVariableLoad(ret_var);
    }
 
    // Restore state
    param_substitutions = saved_param_substitutions;
    var_rename_map = saved_var_rename_map;
    current_function = saved_current_function;
    call_site_id_stack.pop_back();
 
    PostfixBuilder inlined_builder;
    inlined_builder.append(param_assignments);
    inlined_builder.append(body_builder);
 
#ifndef NDEBUG
    int expected_effect = sig.returns_value ? 1 : 0;
    try {
        int actual_effect =
            computeStackEffect(inlined_builder.getExpression(), call_range);
        if (actual_effect != expected_effect) {
            throw CodeGenError(std::format("Function '{}' has "
                                           "unbalanced stack. "
                                           "Expected "
                                           "effect: {}, "
                                           "actual: {}",
                                           func_name, expected_effect,
                                           actual_effect),
                               sig.range);
        }
    } catch (const CodeGenError& e) {
        throw CodeGenError(
            std::format("In function '{}': {}", func_name, e.what()),
            sig.range);
    }
#endif
    builder.append(inlined_builder);
}
 
} // namespace infix2postfix