ast.h

#ifndef AST_H
#define AST_H

#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <memory>
#include <llvm/IR/Value.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Type.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/BasicBlock.h>
#include <llvm/IR/Verifier.h>
#include <llvm/ExecutionEngine/ExecutionEngine.h>
#include <llvm/ExecutionEngine/GenericValue.h>
#include <llvm/Support/TargetSelect.h>

// Forward declarations
class CodeGenContext;

// Base class for all expressions
class Expr {
public:
    virtual ~Expr() {}
    virtual llvm::Value* codegen(CodeGenContext& context) = 0;
};

// Integer literal expression
class IntLiteral : public Expr {
public:
    int value;
    
    IntLiteral(int value) : value(value) {}
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Boolean literal expression
class BoolLiteral : public Expr {
public:
    bool value;
    
    BoolLiteral(bool value) : value(value) {}
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// String literal expression
class StringLiteral : public Expr {
public:
    std::string value;
    
    StringLiteral(const std::string& value) : value(value) {}
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Variable expression
class VarExpr : public Expr {
public:
    std::string name;
    
    VarExpr(const std::string& name) : name(name) {}
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Binary operation types
enum class BinaryOp {
    ADD, SUB, MUL, DIV, EQ, NEQ, LT, GT, AND, OR
};

// Binary expression
class BinaryExpr : public Expr {
public:
    Expr* lhs;
    BinaryOp op;
    Expr* rhs;
    
    BinaryExpr(Expr* lhs, BinaryOp op, Expr* rhs) : lhs(lhs), op(op), rhs(rhs) {}
    virtual ~BinaryExpr() {
        delete lhs;
        delete rhs;
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Base class for all statements
class Stmt {
public:
    virtual ~Stmt() {}
    virtual llvm::Value* codegen(CodeGenContext& context) = 0;
};

// Declaration statement
class DeclStmt : public Stmt {
public:
    std::string name;
    Expr* initializer;
    
    DeclStmt(const std::string& name, Expr* initializer) : name(name), initializer(initializer) {}
    virtual ~DeclStmt() {
        delete initializer;
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Assignment statement
class AssignStmt : public Stmt {
public:
    std::string name;
    Expr* value;
    
    AssignStmt(const std::string& name, Expr* value) : name(name), value(value) {}
    virtual ~AssignStmt() {
        delete value;
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// If statement
class IfStmt : public Stmt {
public:
    Expr* condition;
    std::vector<Stmt*>* thenBlock;
    std::vector<Stmt*>* elseBlock;
    
    IfStmt(Expr* condition, std::vector<Stmt*>* thenBlock, std::vector<Stmt*>* elseBlock)
        : condition(condition), thenBlock(thenBlock), elseBlock(elseBlock) {}
    virtual ~IfStmt() {
        delete condition;
        if (thenBlock) {
            for (auto stmt : *thenBlock) {
                delete stmt;
            }
            delete thenBlock;
        }
        if (elseBlock) {
            for (auto stmt : *elseBlock) {
                delete stmt;
            }
            delete elseBlock;
        }
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Loop statement
class LoopStmt : public Stmt {
public:
    Expr* count;
    std::vector<Stmt*>* body;
    
    LoopStmt(Expr* count, std::vector<Stmt*>* body) : count(count), body(body) {}
    virtual ~LoopStmt() {
        delete count;
        if (body) {
            for (auto stmt : *body) {
                delete stmt;
            }
            delete body;
        }
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Print statement
class PrintStmt : public Stmt {
public:
    Expr* value;
    
    PrintStmt(Expr* value) : value(value) {}
    virtual ~PrintStmt() {
        delete value;
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Return statement
class ReturnStmt : public Stmt {
public:
    Expr* value;
    
    ReturnStmt(Expr* value) : value(value) {}
    virtual ~ReturnStmt() {
        delete value;
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Function definition
class FunctionDef : public Stmt {
public:
    std::string name;
    std::vector<std::string> params; // New: list of parameter names
    std::vector<Stmt*>* body;
    
    // Updated constructor with parameters
    FunctionDef(const std::string& name, const std::vector<std::string>& params, std::vector<Stmt*>* body) 
        : name(name), params(params), body(body) {}
    
    virtual ~FunctionDef() {
        if (body) {
            for (auto stmt : *body) {
                delete stmt;
            }
            delete body;
        }
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Function call
class FunctionCall : public Stmt, public Expr {
public:
    std::string name;
    std::vector<Expr*>* args;
    
    FunctionCall(const std::string& name, std::vector<Expr*>* args) : name(name), args(args) {}
    virtual ~FunctionCall() {
        if (args) {
            for (auto arg : *args) {
                delete arg;
            }
            delete args;
        }
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
};

// Program - the root of our AST
class Program {
public:
    std::vector<Stmt*>* statements;
    std::vector<FunctionDef*>* functions;
    Program(std::vector<Stmt*>* statements) : statements(statements) {}
    virtual ~Program() {
        if (statements) {
            for (auto stmt : *statements) {
                delete stmt;
            }
            delete statements;
        }
    }
    virtual llvm::Value* codegen(CodeGenContext& context);
    
    void printTokens() const {
        if (!statements || statements->empty()) {
            std::cout << "Program is empty" << std::endl;
            return;
        }
        
        std::cout << "Program with " << statements->size() << " statements:" << std::endl;
        for (size_t i = 0; i < statements->size(); i++) {
            std::cout << "  Statement " << (i+1) << ": ";
            // Print statement type information and content
            if (auto decl = dynamic_cast<DeclStmt*>(statements->at(i)))
                std::cout << "Declaration of '" << decl->name << "'" << std::endl;
            else if (auto assign = dynamic_cast<AssignStmt*>(statements->at(i)))
                std::cout << "Assignment to '" << assign->name << "'" << std::endl;
            else if (auto ifStmt = dynamic_cast<IfStmt*>(statements->at(i)))
                std::cout << "If statement with " << ifStmt->thenBlock->size() << " statements in then block"
                          << (ifStmt->elseBlock ? " and " + std::to_string(ifStmt->elseBlock->size()) + " in else block" : "") << std::endl;
            else if (auto loop = dynamic_cast<LoopStmt*>(statements->at(i)))
                std::cout << "Loop statement with " << loop->body->size() << " statements in body" << std::endl;
            else if (dynamic_cast<PrintStmt*>(statements->at(i)))
                std::cout << "Print statement" << std::endl;
            else if (dynamic_cast<ReturnStmt*>(statements->at(i)))
                std::cout << "Return statement" << std::endl;
            else if (auto func = dynamic_cast<FunctionDef*>(statements->at(i)))
                std::cout << "Function definition '" << func->name << "' with " 
                          << (func->params.size() > 0 ? std::to_string(func->params.size()) + " parameters and " : "")
                          << func->body->size() << " statements" << std::endl;
            else if (auto call = dynamic_cast<FunctionCall*>(statements->at(i)))
                std::cout << "Function call to '" << call->name << "' with " 
                          << (call->args ? call->args->size() : 0) << " arguments" << std::endl;
            else
                std::cout << "Unknown statement type" << std::endl;
        }
    }
};

// Code generation context
class CodeGenContext {
public:
    llvm::LLVMContext context;
    std::unique_ptr<llvm::Module> module;
    llvm::IRBuilder<> builder;
    std::map<std::string, llvm::AllocaInst*> namedValues;
    std::map<std::string, llvm::Function*> functions;
    
    CodeGenContext() : 
        module(std::make_unique<llvm::Module>("linguo", context)),
        builder(context) {}
    
    llvm::AllocaInst* createEntryBlockAlloca(llvm::Function* function, 
                                           const std::string& varName,
                                           llvm::Type* type);
    
    void generateCode(Program* program);
    void printIR();
    void writeIR(const std::string& filename);
    void optimize();
    void runJIT();
};

#endif