main.cpp

#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Verifier.h>
#include <llvm/Support/raw_ostream.h>
#include <vector>
#include <string>
#include <memory>
#include <iostream>
#include <map>
#include <fstream>
#include <sstream>
#include <cctype>

using namespace llvm;

// ============ Lexer/Tokenizer ============
enum TokenType {
    tok_eof = 0,
    tok_fn = 1,
    tok_let = 2,
    tok_return = 3,
    tok_identifier = 4,
    tok_number = 5,
    tok_lparen = 6,
    tok_rparen = 7,
    tok_lbrace = 8,
    tok_rbrace = 9,
    tok_semicolon = 10,
    tok_equal = 11,
    tok_plus = 12,
    tok_minus = 13,
    tok_star = 14,
    tok_slash = 15,
    tok_arrow = 16,  // ->
    tok_i64 = 17,
};

struct Token {
    TokenType type;
    std::string value;
};

class Lexer {
public:
    explicit Lexer(const std::string &input) : input(input), pos(0) {}
    
    Token nextToken() {
        skipWhitespace();
        if (pos >= input.size()) return {tok_eof, ""};
        
        char c = input[pos];
        
        // Single-char tokens
        if (c == '(') { pos++; return {tok_lparen, "("}; }
        if (c == ')') { pos++; return {tok_rparen, ")"}; }
        if (c == '{') { pos++; return {tok_lbrace, "{"}; }
        if (c == '}') { pos++; return {tok_rbrace, "}"}; }
        if (c == ';') { pos++; return {tok_semicolon, ";"}; }
        if (c == '+') { pos++; return {tok_plus, "+"}; }
        if (c == '-') { 
            pos++; 
            if (pos < input.size() && input[pos] == '>') {
                pos++;
                return {tok_arrow, "->"};
            }
            return {tok_minus, "-"};
        }
        if (c == '*') { pos++; return {tok_star, "*"}; }
        if (c == '/') { pos++; return {tok_slash, "/"}; }
        if (c == '=') { pos++; return {tok_equal, "="}; }
        
        // Keywords and identifiers
        if (std::isalpha(c) || c == '_') {
            return readIdentifierOrKeyword();
        }
        
        // Numbers
        if (std::isdigit(c)) {
            return readNumber();
        }
        
        pos++; // Skip unknown chars
        return nextToken();
    }

private:
    std::string input;
    size_t pos;
    
    void skipWhitespace() {
        while (pos < input.size() && std::isspace(input[pos])) pos++;
    }
    
    Token readIdentifierOrKeyword() {
        size_t start = pos;
        while (pos < input.size() && (std::isalnum(input[pos]) || input[pos] == '_')) pos++;
        std::string value = input.substr(start, pos - start);
        
        if (value == "fn") return {tok_fn, value};
        if (value == "let") return {tok_let, value};
        if (value == "return") return {tok_return, value};
        if (value == "i64") return {tok_i64, value};
        return {tok_identifier, value};
    }
    
    Token readNumber() {
        size_t start = pos;
        while (pos < input.size() && std::isdigit(input[pos])) pos++;
        return {tok_number, input.substr(start, pos - start)};
    }
};

enum class MiniRustType {
    I64,
    RefI64 // Corresponds to &i64
};

// ============ Ownership System ============
enum class OwnershipState {
    Uninitialized,  // Variable not yet assigned
    Owned,          // Variable is owned by current scope
    Moved,          // Variable has been moved (no longer usable)
    BorrowedImmut,  // Variable has immutable borrow(s) active
    BorrowedMut     // Variable has mutable borrow active
};

struct VariableInfo {
    std::string name;
    OwnershipState state;
    int immutBorrowCount;  // Count of active immutable borrows
    bool hasMutBorrow;     // Does it have an active mutable borrow?
};

class OwnershipChecker {
private:
    std::map<std::string, VariableInfo> variables;
    std::vector<std::string> errors;
    
public:
    // Initialize a new variable
    void declareVariable(const std::string &name) {
        if (variables.find(name) != variables.end()) {
            errors.push_back("Error: Variable '" + name + "' already declared in this scope");
            return;
        }
        variables[name] = {name, OwnershipState::Uninitialized, 0, false};
    }
    
    // Assign a value to a variable (takes ownership)
    void assignVariable(const std::string &name) {
        if (variables.find(name) == variables.end()) {
            errors.push_back("Error: Variable '" + name + "' not declared");
            return;
        }
        variables[name].state = OwnershipState::Owned;
    }
    
    // Use a variable (requires ownership or immutable borrow)
    void useVariable(const std::string &name) {
        if (variables.find(name) == variables.end()) {
            errors.push_back("Error: Variable '" + name + "' not declared");
            return;
        }
        
        VariableInfo &var = variables[name];
        if (var.state == OwnershipState::Moved) {
            errors.push_back("Error: Variable '" + name + "' was moved and cannot be used");
            return;
        }
        if (var.state == OwnershipState::Uninitialized) {
            errors.push_back("Error: Variable '" + name + "' is not initialized");
            return;
        }
        // Using a variable moves it if it's owned
        if (var.state == OwnershipState::Owned) {
            var.state = OwnershipState::Moved;
        }
    }
    
    // Borrow a variable immutably
    void borrowImmutable(const std::string &name) {
        if (variables.find(name) == variables.end()) {
            errors.push_back("Error: Variable '" + name + "' not declared");
            return;
        }
        
        VariableInfo &var = variables[name];
        if (var.state == OwnershipState::Moved) {
            errors.push_back("Error: Cannot borrow '" + name + "' after it was moved");
            return;
        }
        if (var.hasMutBorrow) {
            errors.push_back("Error: Cannot take immutable borrow of '" + name + "' while mutable borrow is active");
            return;
        }
        var.immutBorrowCount++;
    }
    
    // Borrow a variable mutably
    void borrowMutable(const std::string &name) {
        if (variables.find(name) == variables.end()) {
            errors.push_back("Error: Variable '" + name + "' not declared");
            return;
        }
        
        VariableInfo &var = variables[name];
        if (var.state == OwnershipState::Moved) {
            errors.push_back("Error: Cannot borrow '" + name + "' after it was moved");
            return;
        }
        if (var.hasMutBorrow) {
            errors.push_back("Error: Cannot take mutable borrow of '" + name + "' - already has mutable borrow");
            return;
        }
        if (var.immutBorrowCount > 0) {
            errors.push_back("Error: Cannot take mutable borrow of '" + name + "' while immutable borrows are active");
            return;
        }
        var.hasMutBorrow = true;
    }
    
    // Return immutable borrow
    void returnBorrowImmutable(const std::string &name) {
        if (variables.find(name) != variables.end() && variables[name].immutBorrowCount > 0) {
            variables[name].immutBorrowCount--;
        }
    }
    
    // Return mutable borrow
    void returnBorrowMutable(const std::string &name) {
        if (variables.find(name) != variables.end()) {
            variables[name].hasMutBorrow = false;
        }
    }
    
    // Get all errors collected
    const std::vector<std::string> &getErrors() const {
        return errors;
    }
    
    // Check if there are any errors
    bool hasErrors() const {
        return !errors.empty();
    }
    
    // Print current state of all variables
    void printState() {
        std::cout << "\n--- Ownership State ---" << std::endl;
        for (const auto &pair : variables) {
            const VariableInfo &var = pair.second;
            std::string state;
            switch (var.state) {
                case OwnershipState::Uninitialized: state = "Uninitialized"; break;
                case OwnershipState::Owned: state = "Owned"; break;
                case OwnershipState::Moved: state = "Moved"; break;
                case OwnershipState::BorrowedImmut: state = "BorrowedImmut"; break;
                case OwnershipState::BorrowedMut: state = "BorrowedMut"; break;
            }
            std::cout << "  " << var.name << ": " << state;
            if (var.immutBorrowCount > 0) {
                std::cout << " (immut borrows: " << var.immutBorrowCount << ")";
            }
            if (var.hasMutBorrow) {
                std::cout << " (mut borrow: active)";
            }
            std::cout << std::endl;
        }
    }
};

// Global LLVM state
static std::unique_ptr<llvm::LLVMContext> TheContext;
static std::unique_ptr<llvm::Module> TheModule;
static std::unique_ptr<llvm::IRBuilder<>> Builder;

// Global ownership checker
static OwnershipChecker OwnershipMgr;

class ExprAST {
public:
    virtual ~ExprAST() = default;
    virtual llvm::Value *codegen() = 0;
};

// Represents a variable access: 'x'
class VariableExprAST : public ExprAST {
    std::string Name;
public:
    VariableExprAST(const std::string &Name) : Name(Name) {}
    llvm::Value *codegen() override;
};

// Represents a borrow: '&x'
class BorrowExprAST : public ExprAST {
    std::string Name;
public:
    BorrowExprAST(const std::string &Name) : Name(Name) {}
    llvm::Value *codegen() override;
};

// Stores the memory address (AllocaInst) for each variable name
static std::map<std::string, llvm::AllocaInst *> NamedValues;

// Helper to create an alloca in the entry block of the function
static llvm::AllocaInst *CreateEntryBlockAlloca(llvm::Function *TheFunction, 
                                               const std::string &VarName) {
    llvm::IRBuilder<> TmpB(&TheFunction->getEntryBlock(), 
                           TheFunction->getEntryBlock().begin());
    return TmpB.CreateAlloca(llvm::Type::getInt64Ty(*TheContext), nullptr, VarName);
}

// Generating code for 'x'
llvm::Value *VariableExprAST::codegen() {
    llvm::AllocaInst *V = NamedValues[Name];
    if (!V) return nullptr;
    // For a normal variable access, we LOAD the value from the stack
    return Builder->CreateLoad(V->getAllocatedType(), V, Name.c_str());
}

// Generating code for '&x'
llvm::Value *BorrowExprAST::codegen() {
    llvm::AllocaInst *V = NamedValues[Name];
    if (!V) return nullptr;
    // A reference IS the address, so we return the pointer itself (no Load!)
    return V;
}

// Minimal stub to allow compilation; real parser not implemented here.
std::unique_ptr<ExprAST> ParseExpression() {
    return nullptr;
}
// --- Initialization ---
void InitializeModule() {
    TheContext = std::make_unique<LLVMContext>();
    TheModule = std::make_unique<Module>("MiniRustCompiler", *TheContext);
    Builder = std::make_unique<IRBuilder<>>(*TheContext);
}

// --- A Simple "Let" Statement Mockup ---
// In a real parser, this would be part of your ParseLetStatement()
void CreateLetStatement(std::string VarName, uint64_t InitValue) {
    Function *TheFunction = Builder->GetInsertBlock()->getParent();
    
    // 1. Create alloca (allocate memory on the stack)
    AllocaInst *Alloca = Builder->CreateAlloca(Type::getInt64Ty(*TheContext), nullptr, VarName);
    
    // 2. Store initial value
    llvm::Value *InitConst = llvm::ConstantInt::get(*TheContext, llvm::APInt(64, InitValue, true));
    Builder->CreateStore(InitConst, Alloca);
    
    // 3. Add to symbol table
    NamedValues[VarName] = Alloca;
}

int main(int argc, char* argv[]) {
    if (argc != 2) {
        std::cerr << "Usage: minirust <input.rs>" << std::endl;
        return 1;
    }
    
    // Read input file
    std::ifstream file(argv[1]);
    if (!file.is_open()) {
        std::cerr << "Error: Could not open file '" << argv[1] << "'" << std::endl;
        return 1;
    }
    
    std::stringstream buffer;
    buffer << file.rdbuf();
    std::string source = buffer.str();
    file.close();
    
    std::cout << "Read " << source.size() << " bytes from '" << argv[1] << "'" << std::endl;
    std::cout << "\n--- Source Code ---\n" << source << "\n" << std::endl;
    
    // Tokenize
    Lexer lexer(source);
    std::cout << "--- Tokens ---" << std::endl;
    Token tok;
    std::vector<Token> tokens;
    do {
        tok = lexer.nextToken();
        tokens.push_back(tok);
        if (tok.type != tok_eof) {
            std::cout << "Token: type=" << tok.type << " value='" << tok.value << "'" << std::endl;
        }
    } while (tok.type != tok_eof);
    std::cout << std::endl;
    
    // ============ Semantic Analysis: Ownership Checking ============
    std::cout << "--- Ownership Analysis ---" << std::endl;
    
    // Simulate analyzing: let x = 42;
    OwnershipMgr.declareVariable("x");
    OwnershipMgr.assignVariable("x");
    std::cout << "After 'let x = 42;' - x is owned" << std::endl;
    
    // Simulate analyzing: let y = &x;
    OwnershipMgr.declareVariable("y");
    OwnershipMgr.borrowImmutable("x");
    std::cout << "After 'let y = &x;' - y borrows x immutably" << std::endl;
    
    // Simulate analyzing: return x;
    OwnershipMgr.useVariable("x");
    std::cout << "After 'return x;' - x is moved" << std::endl;
    
    // Print ownership state
    OwnershipMgr.printState();
    
    // Check for errors
    if (OwnershipMgr.hasErrors()) {
        std::cout << "\n--- Ownership Errors ---" << std::endl;
        for (const auto &error : OwnershipMgr.getErrors()) {
            std::cout << error << std::endl;
        }
        std::cout << std::endl;
    } else {
        std::cout << "\n--- No ownership errors detected ---\n" << std::endl;
    }
    
    // 1. Setup LLVM
    InitializeModule();

    // 2. Create a "main" function to hold our code
    // fn main() -> i64
    FunctionType *FT = FunctionType::get(Type::getInt64Ty(*TheContext), false);
    Function *MainFn = Function::Create(FT, Function::ExternalLinkage, "main", TheModule.get());
    BasicBlock *BB = BasicBlock::Create(*TheContext, "entry", MainFn);
    Builder->SetInsertPoint(BB);

    // 3. Simulate parsing: "let x = 42;"
    std::cout << "; Generating IR for 'let x = 42;'..." << std::endl;
    CreateLetStatement("x", 42);

    // 4. Simulate parsing: "return x;"
    // We load the value from the memory address stored in our map
    AllocaInst *V = NamedValues["x"];
    Value *LoadedVal = Builder->CreateLoad(V->getAllocatedType(), V, "x_val");
    Builder->CreateRet(LoadedVal);

    // 5. Verify the code is valid
    verifyFunction(*MainFn);

    // 6. Print the result!
    std::cout << "\n--- Generated LLVM IR ---\n" << std::endl;
    TheModule->print(errs(), nullptr);

    return 0;
}