Added array

4 files changed, 378 insertions, 15 deletions

diff --git a/lib/CodeGen/ArrayHandler.cc b/lib/CodeGen/ArrayHandler.cc
new file mode 100644
index 0000000..470cffa
--- /dev/null
+++ b/lib/CodeGen/ArrayHandler.cc
@@ -0,0 +1,230 @@
+#include "cps/ArrayHandler.h"
+#include "cps/CodeGen.h"
+#include "llvm/IR/Type.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Function.h"
+
+using namespace llvm;
+using namespace cps;
+
+ArrayHandler::ArrayHandler(LLVMContext &C, IRBuilder<> &B, Module &M, std::map<std::string, AllocaInst*> &NV)
+    : TheContext(&C), Builder(&B), TheModule(&M), NamedValues(&NV) {
+    setupExternalFunctions();
+}
+
+void ArrayHandler::setupExternalFunctions() {
+    std::vector<Type*> MallocArgs;
+    MallocArgs.push_back(Type::getInt64Ty(*TheContext));
+    
+    FunctionType *MallocType = FunctionType::get(PointerType::getUnqual(*TheContext), MallocArgs, false);
+    MallocFunc = TheModule->getOrInsertFunction("malloc", MallocType);
+
+    std::vector<Type*> FreeArgs;
+    FreeArgs.push_back(PointerType::getUnqual(*TheContext));
+    FunctionType *FreeType = FunctionType::get(Type::getVoidTy(*TheContext), FreeArgs, false);
+    FreeFunc = TheModule->getOrInsertFunction("free", FreeType);
+}
+
+void ArrayHandler::emitArrayDeclare(ArrayDeclareStmtAST *Stmt, CodeGen &CG) {
+    std::string Name = Stmt->getName();
+    int Rank = Stmt->getBounds().size();
+    
+    std::vector<Value*> Lows, Highs, Dims;
+    Value *TotalElements = ConstantInt::get(*TheContext, APInt(64, 1));
+    
+    for (const auto &Pair : Stmt->getBounds()) {
+        Value *L = CG.emitExpr(Pair.first.get());
+        Value *R = CG.emitExpr(Pair.second.get());
+        
+        if (!L || !R) return;
+        
+        Value *L64 = Builder->CreateSExt(L, Type::getInt64Ty(*TheContext));
+        Value *R64 = Builder->CreateSExt(R, Type::getInt64Ty(*TheContext));
+        
+        Lows.push_back(L);
+        Highs.push_back(R);
+        
+        Value *Diff = Builder->CreateSub(R64, L64);
+        Value *DimSize = Builder->CreateAdd(Diff, ConstantInt::get(*TheContext, APInt(64, 1)));
+        Dims.push_back(DimSize);
+        
+        TotalElements = Builder->CreateMul(TotalElements, DimSize);
+    }
+    
+    std::vector<Value*> Multipliers(Rank);
+    Multipliers[Rank - 1] = ConstantInt::get(*TheContext, APInt(32, 1));
+    
+    for (int i = Rank - 2; i >= 0; --i) {
+        Value *NextMult = Multipliers[i+1];
+        Value *NextDimSize = Dims[i+1];
+        Value *NextDimSize32 = Builder->CreateTrunc(NextDimSize, Type::getInt32Ty(*TheContext));
+        Multipliers[i] = Builder->CreateMul(NextMult, NextDimSize32);
+    }
+
+    ArrayMetadata Meta;
+    Meta.Rank = Rank;
+    Meta.ElementType = Type::getInt32Ty(*TheContext);
+    Meta.LowerBounds = Lows;
+    Meta.UpperBounds = Highs;
+    Meta.Multipliers = Multipliers;
+    ArrayTable[Name] = Meta;
+    
+
+    Value *ElemSize = ConstantInt::get(*TheContext, APInt(64, 4)); 
+    Value *TotalBytes = Builder->CreateMul(TotalElements, ElemSize);
+    
+    CallInst *Ptr = Builder->CreateCall(MallocFunc, TotalBytes);
+    
+    Function *TheFunction = Builder->GetInsertBlock()->getParent();
+    IRBuilder<> TmpB(&TheFunction->getEntryBlock(), TheFunction->getEntryBlock().begin());
+    
+    AllocaInst *Alloca = TmpB.CreateAlloca(PointerType::getUnqual(*TheContext), nullptr, Name);
+    
+    Builder->CreateStore(Ptr, Alloca);
+    (*NamedValues)[Name] = Alloca;
+}
+
+Value* ArrayHandler::computeFlatIndex(const std::string &Name, const std::vector<Value*> &Indices) {
+    if (ArrayTable.find(Name) == ArrayTable.end()) return nullptr;
+    const auto &Meta = ArrayTable[Name];
+    
+    Value *Offset = ConstantInt::get(*TheContext, APInt(32, 0));
+    
+    for (size_t i = 0; i < Indices.size(); ++i) {
+        Value *Idx = Indices[i];
+        Value *Lower = Meta.LowerBounds[i];
+        Value *Diff = Builder->CreateSub(Idx, Lower);
+        Value *Term = Builder->CreateMul(Diff, Meta.Multipliers[i]);
+        Offset = Builder->CreateAdd(Offset, Term);
+    }
+    return Offset;
+}
+
+Value* ArrayHandler::emitArrayAccess(ArrayAccessExprAST *Expr, CodeGen &CG) {
+    std::string Name = Expr->getName();
+    if (ArrayTable.find(Name) == ArrayTable.end()) {
+        fprintf(stderr, "Error: Undeclared array %s\n", Name.c_str());
+        return nullptr;
+    }
+    
+    const auto &Meta = ArrayTable[Name];
+    if (Expr->getIndices().size() != Meta.Rank) {
+        fprintf(stderr, "Error: Incorrect number of indices for %s\n", Name.c_str());
+        return nullptr;
+    }
+
+    std::vector<Value*> Indices;
+    for (const auto &IdxExpr : Expr->getIndices()) {
+        Indices.push_back(CG.emitExpr(IdxExpr.get()));
+    }
+    
+    Value *Offset = computeFlatIndex(Name, Indices);
+    
+    AllocaInst *Alloca = (*NamedValues)[Name];
+    Value *RawPtr = Builder->CreateLoad(PointerType::getUnqual(*TheContext), Alloca);
+    Value *IntPtr = Builder->CreateBitCast(RawPtr, PointerType::getUnqual(*TheContext));
+    Value *ElemPtr = Builder->CreateGEP(Type::getInt32Ty(*TheContext), IntPtr, Offset);
+    
+    return Builder->CreateLoad(Type::getInt32Ty(*TheContext), ElemPtr);
+}
+
+void ArrayHandler::emitArrayAssign(ArrayAssignStmtAST *Stmt, CodeGen &CG) {
+    std::string Name = Stmt->getName();
+    if (ArrayTable.find(Name) == ArrayTable.end()) {
+        fprintf(stderr, "Error: Undeclared array %s\n", Name.c_str());
+        return;
+    }
+    
+    std::vector<Value*> Indices;
+    for (const auto &IdxExpr : Stmt->getIndices()) {
+        Indices.push_back(CG.emitExpr(IdxExpr.get()));
+    }
+    
+    Value *Val = CG.emitExpr(Stmt->getExpr());
+    Value *Offset = computeFlatIndex(Name, Indices);
+    
+    AllocaInst *Alloca = (*NamedValues)[Name];
+    Value *RawPtr = Builder->CreateLoad(PointerType::getUnqual(*TheContext), Alloca);
+    Value *IntPtr = Builder->CreateBitCast(RawPtr, PointerType::getUnqual(*TheContext));
+    Value *ElemPtr = Builder->CreateGEP(Type::getInt32Ty(*TheContext), IntPtr, Offset);
+    
+    Builder->CreateStore(Val, ElemPtr);
+}
+
+bool ArrayHandler::tryEmitArrayOutput(ExprAST *Expr, CodeGen &CG, llvm::FunctionCallee PrintfFunc, llvm::Value* FmtStr) {
+    std::string Name;
+    std::vector<Value*> Indices;
+    
+    if (auto *Var = dynamic_cast<VariableExprAST*>(Expr)) {
+        Name = Var->getName();
+        if (ArrayTable.find(Name) == ArrayTable.end()) return false; 
+    } 
+    else if (auto *Acc = dynamic_cast<ArrayAccessExprAST*>(Expr)) {
+        Name = Acc->getName();
+        if (ArrayTable.find(Name) == ArrayTable.end()) return false;
+        
+        for (const auto &Idx : Acc->getIndices()) {
+            Indices.push_back(CG.emitExpr(Idx.get()));
+        }
+    } else {
+        return false;
+    }
+    
+    const auto &Meta = ArrayTable[Name];
+    
+    if (Indices.size() == Meta.Rank) {
+        return false;
+    }
+    
+    emitPrintLoop(Name, Indices.size(), Indices, PrintfFunc, FmtStr);
+    return true;
+}
+
+void ArrayHandler::emitPrintLoop(const std::string &Name, int CurrentDim, std::vector<Value*> CurrentIndices, llvm::FunctionCallee PrintfFunc, llvm::Value* FmtStr) {
+    const auto &Meta = ArrayTable[Name];
+    
+    if (CurrentDim == Meta.Rank) {
+        Value *Offset = computeFlatIndex(Name, CurrentIndices);
+        
+        AllocaInst *Alloca = (*NamedValues)[Name];
+        Value *RawPtr = Builder->CreateLoad(PointerType::getUnqual(*TheContext), Alloca);
+        Value *IntPtr = Builder->CreateBitCast(RawPtr, PointerType::getUnqual(*TheContext));
+        Value *ElemPtr = Builder->CreateGEP(Type::getInt32Ty(*TheContext), IntPtr, Offset);
+        Value *Val = Builder->CreateLoad(Type::getInt32Ty(*TheContext), ElemPtr);
+        
+        std::vector<Value*> Args;
+        Args.push_back(FmtStr);
+        Args.push_back(Val);
+        Builder->CreateCall(PrintfFunc, Args);
+        return;
+    }
+    
+    Function *TheFunction = Builder->GetInsertBlock()->getParent();
+    BasicBlock *LoopBB = BasicBlock::Create(*TheContext, "arr_loop", TheFunction);
+    BasicBlock *AfterBB = BasicBlock::Create(*TheContext, "arr_after", TheFunction);
+    
+    IRBuilder<> TmpB(&TheFunction->getEntryBlock(), TheFunction->getEntryBlock().begin());
+    AllocaInst *LoopVar = TmpB.CreateAlloca(Type::getInt32Ty(*TheContext), nullptr, "idx");
+    
+    Value *Start = Meta.LowerBounds[CurrentDim];
+    Value *End = Meta.UpperBounds[CurrentDim];
+    
+    Builder->CreateStore(Start, LoopVar);
+    Builder->CreateBr(LoopBB);
+    
+    Builder->SetInsertPoint(LoopBB);
+    Value *CurVal = Builder->CreateLoad(Type::getInt32Ty(*TheContext), LoopVar);
+    
+    std::vector<Value*> NextIndices = CurrentIndices;
+    NextIndices.push_back(CurVal);
+    emitPrintLoop(Name, CurrentDim + 1, NextIndices, PrintfFunc, FmtStr);
+    
+    Value *NextVal = Builder->CreateAdd(CurVal, ConstantInt::get(*TheContext, APInt(32, 1)));
+    Builder->CreateStore(NextVal, LoopVar);
+    
+    Value *Cond = Builder->CreateICmpSLE(NextVal, End);
+    Builder->CreateCondBr(Cond, LoopBB, AfterBB);
+    
+    Builder->SetInsertPoint(AfterBB);
+}
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 31d113a..57d76eb 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -1,15 +1,20 @@
 #include "cps/CodeGen.h"
+#include "cps/ArrayHandler.h"
 #include "llvm/IR/Verifier.h"
 #include "cps/Lexer.h"
 
 using namespace llvm;
 using namespace cps;
 
+CodeGen::~CodeGen() = default;
+
 CodeGen::CodeGen() {
     TheContext = std::make_unique<LLVMContext>();
     TheModule = std::make_unique<Module>("cps_module", *TheContext);
     Builder = std::make_unique<IRBuilder<>>(*TheContext);
     
+    Arrays = std::make_unique<ArrayHandler>(*TheContext, *Builder, *TheModule, NamedValues);
+    
     SetupExternalFunctions();
 }
 
@@ -24,6 +29,8 @@ void CodeGen::SetupExternalFunctions() {
 
     PrintfFormatStr = Builder->CreateGlobalStringPtr("%d\n", "fmt_nl", 0, TheModule.get());
     ScanfFormatStr = Builder->CreateGlobalStringPtr("%d", "fmt_in", 0, TheModule.get());
+    
+    Arrays->setupExternalFunctions();
 }
 
 AllocaInst *CodeGen::CreateEntryBlockAlloca(Function *TheFunction, const std::string &VarName) {
@@ -64,8 +71,14 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
             fprintf(stderr, "Error: Unknown variable name %s\n", Var->getName().c_str());
             return nullptr;
         }
+        
         return Builder->CreateLoad(Type::getInt32Ty(*TheContext), A, Var->getName().c_str());
     }
+    
+
+    if (auto *ArrAcc = dynamic_cast<ArrayAccessExprAST*>(Expr)) {
+        return Arrays->emitArrayAccess(ArrAcc, *this);
+    }
 
     if (auto *Bin = dynamic_cast<BinaryExprAST*>(Expr)) {
         Value *L = emitExpr(Bin->getLHS());
@@ -234,6 +247,16 @@ void CodeGen::emitForStmt(ForStmtAST *Stmt) {
 void CodeGen::emitStmt(StmtAST *Stmt) {
     if (!Stmt) return;
 
+    if (auto *ArrDecl = dynamic_cast<ArrayDeclareStmtAST*>(Stmt)) {
+        Arrays->emitArrayDeclare(ArrDecl, *this);
+        return;
+    }
+    
+    if (auto *ArrAssign = dynamic_cast<ArrayAssignStmtAST*>(Stmt)) {
+        Arrays->emitArrayAssign(ArrAssign, *this);
+        return;
+    }
+
     if (auto *Decl = dynamic_cast<DeclareStmtAST*>(Stmt)) {
         Function *TheFunction = Builder->GetInsertBlock()->getParent();
         AllocaInst *Alloca = CreateEntryBlockAlloca(TheFunction, Decl->getName());
@@ -261,12 +284,16 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
         }
     }
     else if (auto *Out = dynamic_cast<OutputStmtAST*>(Stmt)) {
-        Value *Val = emitExpr(Out->getExpr());
-        if (Val) {
-            std::vector<Value*> Args;
-            Args.push_back(PrintfFormatStr);
-            Args.push_back(Val);
-            Builder->CreateCall(PrintfFunc, Args);
+        bool handled = Arrays->tryEmitArrayOutput(Out->getExpr(), *this, PrintfFunc, PrintfFormatStr);
+        
+        if (!handled) {
+            Value *Val = emitExpr(Out->getExpr());
+            if (Val) {
+                std::vector<Value*> Args;
+                Args.push_back(PrintfFormatStr);
+                Args.push_back(Val);
+                Builder->CreateCall(PrintfFunc, Args);
+            }
         }
     }
     else if (auto *IfStmt = dynamic_cast<IfStmtAST*>(Stmt)) {
diff --git a/lib/Lexer/Lexer.cc b/lib/Lexer/Lexer.cc
index 7f7a167..fdaa3a5 100644
--- a/lib/Lexer/Lexer.cc
+++ b/lib/Lexer/Lexer.cc
@@ -39,6 +39,9 @@ int Lexer::gettok() {
         if (IdentifierStr == "STEP") return tok_step;
         if (IdentifierStr == "NEXT") return tok_next;
 
+        if (IdentifierStr == "ARRAY") return tok_array;
+        if (IdentifierStr == "OF") return tok_of;
+
         return tok_identifier;
     }
 
diff --git a/lib/Parser/Parser.cc b/lib/Parser/Parser.cc
index 7a1af89..a85b151 100644
--- a/lib/Parser/Parser.cc
+++ b/lib/Parser/Parser.cc
@@ -39,6 +39,30 @@ std::unique_ptr<ExprAST> Parser::ParseNumberExpr() {
 std::unique_ptr<ExprAST> Parser::ParseIdentifierExpr() {
     std::string IdName = Lex.IdentifierStr;
     getNextToken();
+    
+    if (CurTok == '[') {
+        getNextToken();
+        std::vector<std::unique_ptr<ExprAST>> Indices;
+        
+        while (true) {
+            auto Exp = ParseExpression();
+            if (!Exp) return nullptr;
+            Indices.push_back(std::move(Exp));
+            
+            if (CurTok == ']') {
+                break;
+            }
+            if (CurTok == ',') {
+                getNextToken();
+                continue;
+            }
+            fprintf(stderr, "Error: Expected ',' or ']' in array index\n");
+            return nullptr;
+        }
+        getNextToken(); 
+        return std::make_unique<ArrayAccessExprAST>(IdName, std::move(Indices));
+    }
+    
     return std::make_unique<VariableExprAST>(IdName);
 }
 
@@ -104,6 +128,72 @@ std::vector<std::unique_ptr<StmtAST>> Parser::Parse() {
     return Statements;
 }
 
+std::unique_ptr<StmtAST> Parser::ParseDeclare() {
+    getNextToken(); 
+    if (CurTok != tok_identifier) return nullptr;
+    std::string Name = Lex.IdentifierStr;
+    getNextToken();
+    
+    if (CurTok != tok_colon) return nullptr;
+    getNextToken();
+    
+    if (CurTok == tok_array) {
+        getNextToken(); 
+        if (CurTok != '[') {
+            fprintf(stderr, "Error: Expected '[' after ARRAY\n");
+            return nullptr;
+        }
+        getNextToken();
+        
+        std::vector<std::pair<std::unique_ptr<ExprAST>, std::unique_ptr<ExprAST>>> Bounds;
+        
+        while (true) {
+            auto Lower = ParseExpression();
+            if (!Lower) return nullptr;
+            
+            if (CurTok != tok_colon) {
+                fprintf(stderr, "Error: Expected ':' in array range (e.g. 1:10)\n");
+                return nullptr;
+            }
+            getNextToken();
+            
+            auto Upper = ParseExpression();
+            if (!Upper) return nullptr;
+            
+            Bounds.push_back({std::move(Lower), std::move(Upper)});
+            
+            if (CurTok == ']') break;
+            if (CurTok == ',') {
+                getNextToken();
+                continue;
+            }
+            fprintf(stderr, "Error: Expected ',' or ']' in array dims\n");
+            return nullptr;
+        }
+        getNextToken();
+        
+        if (CurTok != tok_of) {
+            fprintf(stderr, "Error: Expected OF after array dims\n");
+            return nullptr;
+        }
+        getNextToken();
+        
+        if (CurTok != tok_integer_kw) {
+             fprintf(stderr, "Error: Only INTEGER arrays supported currently\n");
+             return nullptr;
+        }
+        getNextToken();
+        
+        return std::make_unique<ArrayDeclareStmtAST>(Name, std::move(Bounds), "INTEGER");
+    } 
+    else if (CurTok == tok_integer_kw) {
+        getNextToken();
+        return std::make_unique<DeclareStmtAST>(Name);
+    }
+    
+    return nullptr;
+}
+
 std::unique_ptr<StmtAST> Parser::ParseIfStmt() {
     getNextToken();
     auto Cond = ParseExpression();
@@ -248,19 +338,32 @@ std::unique_ptr<StmtAST> Parser::ParseForStmt() {
 
 std::unique_ptr<StmtAST> Parser::ParseStatement() {
     if (CurTok == tok_declare) {
-        getNextToken();
-        if (CurTok != tok_identifier) return nullptr;
-        std::string Name = Lex.IdentifierStr;
-        getNextToken();
-        if (CurTok != tok_colon) return nullptr;
-        getNextToken();
-        if (CurTok != tok_integer_kw) return nullptr;
-        getNextToken();
-        return std::make_unique<DeclareStmtAST>(Name);
+        return ParseDeclare();
     }
     else if (CurTok == tok_identifier) {
         std::string Name = Lex.IdentifierStr;
         getNextToken();
+        
+        if (CurTok == '[') {
+            getNextToken();
+            std::vector<std::unique_ptr<ExprAST>> Indices;
+            while (true) {
+                auto Exp = ParseExpression();
+                Indices.push_back(std::move(Exp));
+                if (CurTok == ']') break;
+                if (CurTok == ',') getNextToken();
+            }
+            getNextToken();
+            
+            if (CurTok != tok_assign) {
+                fprintf(stderr, "Error: Expected '<-' after array access in assignment\n");
+                return nullptr;
+            }
+            getNextToken();
+            auto Expr = ParseExpression();
+            return std::make_unique<ArrayAssignStmtAST>(Name, std::move(Indices), std::move(Expr));
+        }
+        
         if (CurTok != tok_assign) return nullptr;
         getNextToken();
         auto Expr = ParseExpression();