generator.hpp

#ifndef __GENERATOR_H
#define __GENERATOR_H

#include "parser.hpp"
#include "fstream"
#include <sstream>
#include <map>
#include <stack>

class Generator{
  public:
    Generator(AST_NODE* root, std::string filename){
        AST_ROOT=root;
        assemblyFile.open(filename + ".asm");
        labelCounter = 0;
        variableOffset = 0;
        currentScope = 0;
    }
    
    void writeSections(){
        assemblyFile << sectionData.str();
        assemblyFile << sectionText.str();
        assemblyFile.close();
    }
    
    std::string generateLabel(const std::string& prefix) {
        return prefix + std::to_string(labelCounter++);
    }
    
    std::string generateExpression(AST_NODE* node) {
        if(!node) return "";
        
        switch(node->TYPE) {
            case NODE_INT: {
                return "mov rax, " + *node->VALUE + "\n";
            }
            case NODE_IDENTIFIER: {
                std::string varName = *node->VALUE;
                if(variables.find(varName) != variables.end()) {
                    return "mov rax, [rbp - " + std::to_string(variables[varName]) + "]\n";
                } else {
                    std::cout << "Error: Undefined variable '" << varName << "'" << std::endl;
                    exit(1);
                }
            }
            case NODE_BINARY_OP: {
                std::string leftCode = generateExpression(node->LEFT);
                std::string rightCode = generateExpression(node->RIGHT);
                
                // Push left operand
                sectionText << leftCode;
                sectionText << "push rax\n";
                
                // Generate right operand
                sectionText << rightCode;
                sectionText << "mov rbx, rax\n";
                sectionText << "pop rax\n";
                
                // Perform operation
                if(*node->VALUE == "+") {
                    sectionText << "add rax, rbx\n";
                } else if(*node->VALUE == "-") {
                    sectionText << "sub rax, rbx\n";
                } else if(*node->VALUE == "*") {
                    sectionText << "mul rbx\n";
                } else if(*node->VALUE == "/") {
                    sectionText << "xor rdx, rdx\n";
                    sectionText << "div rbx\n";
                } else {
                    std::cout << "Error: Unsupported operator '" << *node->VALUE << "'" << std::endl;
                    exit(1);
                }
                return "";
            }
            default: {
                std::cout << "Error: Unsupported expression type" << std::endl;
                exit(1);
            }
        }
    }
    
    std::string generateReturn(AST_NODE* STATEMENT){
        if(!STATEMENT->CHILD){
            std::cout<<"Generation error: return statement has no value"<<std::endl;
            exit(1);
        }
        
        generateExpression(STATEMENT->CHILD);
        sectionText<<"mov rdi, rax\n";
        sectionText<<"mov rax, 60\n";
        sectionText<<"syscall\n";
        
        return "";
    }
    
    void generateVariable(AST_NODE* node) {
        std::string varName = *node->VALUE;
        variableOffset += 8; // 64-bit variable
        variables[varName] = variableOffset;
        
        generateExpression(node->CHILD);
        sectionText << "mov [rbp - " << variableOffset << "], rax\n";
    }
    
    void generatePrint(AST_NODE* node) {
        if(!node->CHILD || node->CHILD->TYPE != NODE_STRING) {
            std::cout << "Error: Print statement requires a string literal" << std::endl;
            exit(1);
        }
        
        std::string message = node->CHILD->VALUE ? *node->CHILD->VALUE : "";
        std::string label = generateLabel("msg_");
        
        // Add string to data section
        sectionData << label << " db \"" << message << "\", 0\n";
        sectionData << label << "_len equ $ - " << label << "\n";
        
        // Generate print code
        sectionText << "mov rax, 1\n";
        sectionText << "mov rdi, 1\n";
        sectionText << "mov rsi, " << label << "\n";
        sectionText << "mov rdx, " << label << "_len\n";
        sectionText << "syscall\n";
    }
    
    void generateIf(AST_NODE* node) {
        std::string elseLabel = generateLabel("else_");
        std::string endLabel = generateLabel("endif_");
        
        // Generate condition
        generateExpression(node->LEFT);
        sectionText << "cmp rax, 0\n";
        sectionText << "je " << elseLabel << "\n";
        
        // Generate if block
        for(AST_NODE* stmt : node->SUB_STATEMENTS) {
            generateStatement(stmt);
        }
        
        sectionText << "jmp " << endLabel << "\n";
        sectionText << elseLabel << ":\n";
        
        // Generate else block (if any)
        // Note: This is simplified - in a full implementation, 
        // we'd need to track which statements belong to else
        
        sectionText << endLabel << ":\n";
    }
    
    void generateWhile(AST_NODE* node) {
        std::string loopLabel = generateLabel("while_");
        std::string endLabel = generateLabel("endwhile_");
        
        sectionText << loopLabel << ":\n";
        
        // Generate condition
        generateExpression(node->LEFT);
        sectionText << "cmp rax, 0\n";
        sectionText << "je " << endLabel << "\n";
        
        // Generate loop body
        for(AST_NODE* stmt : node->SUB_STATEMENTS) {
            generateStatement(stmt);
        }
        
        sectionText << "jmp " << loopLabel << "\n";
        sectionText << endLabel << ":\n";
    }
    
    void generateFunction(AST_NODE* node) {
        std::string funcName = node->VALUE ? *node->VALUE : "unknown";
        
        sectionText << "global " << funcName << "\n";
        sectionText << funcName << ":\n";
        sectionText << "push rbp\n";
        sectionText << "mov rbp, rsp\n";
        
        // Generate function body
        for(AST_NODE* stmt : node->SUB_STATEMENTS) {
            generateStatement(stmt);
        }
        
        sectionText << "pop rbp\n";
        sectionText << "ret\n";
    }
    
    void generateFunctionCall(AST_NODE* node) {
        std::string funcName = node->VALUE ? *node->VALUE : "unknown";
        
        // Generate arguments (simplified - just push them)
        for(AST_NODE* arg : node->ARGUMENTS) {
            generateExpression(arg);
            sectionText << "push rax\n";
        }
        
        sectionText << "call " << funcName << "\n";
        
        // Clean up arguments
        if(!node->ARGUMENTS.empty()) {
            sectionText << "add rsp, " << (node->ARGUMENTS.size() * 8) << "\n";
        }
    }
    
    void generateStatement(AST_NODE* node) {
        if(!node) return;
        
        switch(node->TYPE) {
            case NODE_VARIABLE:
                generateVariable(node);
                break;
            case NODE_RETURN:
                generateReturn(node);
                break;
            case NODE_PRINT:
                generatePrint(node);
                break;
            case NODE_IF:
                generateIf(node);
                break;
            case NODE_WHILE:
                generateWhile(node);
                break;
            case NODE_FUNCTION:
                generateFunction(node);
                break;
            case NODE_FUNCTION_CALL:
                generateFunctionCall(node);
                break;
            default:
                std::cout << "Error: Unsupported statement type" << std::endl;
                exit(1);
        }
    }
    
    void generate(){
        sectionData << "section .data\n\n";
        sectionText << "section .text\n";
        sectionText << "global _start\n";
        sectionText << "_start:\n";
        sectionText << "push rbp\n";
        sectionText << "mov rbp, rsp\n\n";
        
        if(AST_ROOT && !AST_ROOT->SUB_STATEMENTS.empty()) {
            for(AST_NODE* current : AST_ROOT->SUB_STATEMENTS){
                generateStatement(current);
            }
        }
        
        // If no explicit return, add default exit
        sectionText << "mov rax, 60\n";
        sectionText << "mov rdi, 0\n";
        sectionText << "syscall\n";
        
        writeSections();
    }
    
  private:
    AST_NODE* AST_ROOT;
    std::ofstream assemblyFile;
    std::stringstream sectionData;
    std::stringstream sectionText;
    std::map<std::string, int> variables;
    int labelCounter;
    int variableOffset;
    int currentScope;
};

#endif