Added array

5 files changed, 139 insertions, 37 deletions

diff --git a/include/cps/AST.h b/include/cps/AST.h
index 1473c45..aaf81cb 100644
--- a/include/cps/AST.h
+++ b/include/cps/AST.h
@@ -24,6 +24,16 @@ public:
     const std::string &getName() const { return Name; }
 };
 
+class ArrayAccessExprAST : public ExprAST {
+    std::string Name;
+    std::vector<std::unique_ptr<ExprAST>> Indices;
+public:
+    ArrayAccessExprAST(const std::string &Name, std::vector<std::unique_ptr<ExprAST>> Indices)
+        : Name(Name), Indices(std::move(Indices)) {}
+    const std::string &getName() const { return Name; }
+    const std::vector<std::unique_ptr<ExprAST>> &getIndices() const { return Indices; }
+};
+
 class BinaryExprAST : public ExprAST {
     int Op;
     std::unique_ptr<ExprAST> LHS, RHS;
@@ -47,6 +57,21 @@ public:
     const std::string &getName() const { return Name; }
 };
 
+class ArrayDeclareStmtAST : public StmtAST {
+    std::string Name;
+    std::vector<std::pair<std::unique_ptr<ExprAST>, std::unique_ptr<ExprAST>>> Bounds; 
+    std::string Type;
+public:
+    ArrayDeclareStmtAST(const std::string &Name, 
+        std::vector<std::pair<std::unique_ptr<ExprAST>, std::unique_ptr<ExprAST>>> Bounds,
+        const std::string &Type)
+        : Name(Name), Bounds(std::move(Bounds)), Type(Type) {}
+    
+    const std::string &getName() const { return Name; }
+    const std::vector<std::pair<std::unique_ptr<ExprAST>, std::unique_ptr<ExprAST>>> &getBounds() const { return Bounds; }
+    const std::string &getType() const { return Type; }
+};
+
 class AssignStmtAST : public StmtAST {
     std::string Name;
     std::unique_ptr<ExprAST> Expr;
@@ -57,6 +82,18 @@ public:
     ExprAST *getExpr() const { return Expr.get(); }
 };
 
+class ArrayAssignStmtAST : public StmtAST {
+    std::string Name;
+    std::vector<std::unique_ptr<ExprAST>> Indices;
+    std::unique_ptr<ExprAST> Expr;
+public:
+    ArrayAssignStmtAST(const std::string &Name, std::vector<std::unique_ptr<ExprAST>> Indices, std::unique_ptr<ExprAST> Expr)
+        : Name(Name), Indices(std::move(Indices)), Expr(std::move(Expr)) {}
+    const std::string &getName() const { return Name; }
+    const std::vector<std::unique_ptr<ExprAST>> &getIndices() const { return Indices; }
+    ExprAST *getExpr() const { return Expr.get(); }
+};
+
 class InputStmtAST : public StmtAST {
     std::string Name;
 public:
diff --git a/include/cps/ArrayHandler.h b/include/cps/ArrayHandler.h
new file mode 100644
index 0000000..5052de4
--- /dev/null
+++ b/include/cps/ArrayHandler.h
@@ -0,0 +1,50 @@
+#pragma once
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "cps/AST.h"
+#include <map>
+#include <string>
+#include <vector>
+
+namespace cps {
+
+struct ArrayMetadata {
+    int Rank;
+    llvm::Type* ElementType;
+    std::vector<llvm::Value*> LowerBounds;
+    std::vector<llvm::Value*> UpperBounds;
+    std::vector<llvm::Value*> Multipliers; 
+};
+
+class CodeGen;
+
+class ArrayHandler {
+    llvm::LLVMContext *TheContext;
+    llvm::IRBuilder<> *Builder;
+    llvm::Module *TheModule;
+    std::map<std::string, llvm::AllocaInst*> *NamedValues;
+
+    std::map<std::string, ArrayMetadata> ArrayTable;
+
+    llvm::FunctionCallee MallocFunc;
+    llvm::FunctionCallee FreeFunc;
+
+    llvm::Value* computeFlatIndex(const std::string &Name, const std::vector<llvm::Value*> &Indices);
+    
+    void emitPrintLoop(const std::string &Name, int CurrentDim, std::vector<llvm::Value*> CurrentIndices, llvm::FunctionCallee PrintfFunc, llvm::Value* FmtStr);
+
+public:
+    ArrayHandler(llvm::LLVMContext &C, llvm::IRBuilder<> &B, llvm::Module &M, std::map<std::string, llvm::AllocaInst*> &NV);
+
+    void setupExternalFunctions();
+    
+    void emitArrayDeclare(ArrayDeclareStmtAST *Stmt, CodeGen &CG);
+    void emitArrayAssign(ArrayAssignStmtAST *Stmt, CodeGen &CG);
+    
+    llvm::Value* emitArrayAccess(ArrayAccessExprAST *Expr, CodeGen &CG);
+
+    bool tryEmitArrayOutput(ExprAST *Expr, CodeGen &CG, llvm::FunctionCallee PrintfFunc, llvm::Value* FmtStr);
+};
+
+}
diff --git a/include/cps/CodeGen.h b/include/cps/CodeGen.h
index 9be6cc7..8bf5819 100644
--- a/include/cps/CodeGen.h
+++ b/include/cps/CodeGen.h
@@ -1,40 +1,46 @@
 #pragma once
-#include "cps/AST.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "cps/AST.h"
 #include <map>
+#include <memory>
 
 namespace cps {
 
+class ArrayHandler;
+
 class CodeGen {
     std::unique_ptr<llvm::LLVMContext> TheContext;
     std::unique_ptr<llvm::Module> TheModule;
     std::unique_ptr<llvm::IRBuilder<>> Builder;
     std::map<std::string, llvm::AllocaInst*> NamedValues;
-
-    llvm::AllocaInst *CreateEntryBlockAlloca(llvm::Function *TheFunction, 
-                                            const std::string &VarName);
+    
+    std::unique_ptr<ArrayHandler> Arrays;
 
     llvm::FunctionCallee PrintfFunc;
     llvm::FunctionCallee ScanfFunc;
     llvm::Value *PrintfFormatStr;
     llvm::Value *ScanfFormatStr;
-    
-    llvm::Value *emitExpr(ExprAST *Expr);
-    void emitStmt(StmtAST *Stmt);
-    
+
+    void SetupExternalFunctions();
+    llvm::AllocaInst *CreateEntryBlockAlloca(llvm::Function *TheFunction, const std::string &VarName);
+
     void emitIfStmt(IfStmtAST *Stmt);
     void emitWhileStmt(WhileStmtAST *Stmt);
     void emitRepeatStmt(RepeatStmtAST *Stmt);
     void emitForStmt(ForStmtAST *Stmt);
 
-    void SetupExternalFunctions();
-
 public:
     CodeGen();
+    ~CodeGen();
     void compile(const std::vector<std::unique_ptr<StmtAST>> &Statements);
     void print();
+    
+    llvm::Value *emitExpr(ExprAST *Expr);
+    void emitStmt(StmtAST *Stmt);
+    
+    friend class ArrayHandler;
 };
 
-} // namespace cps
+}
diff --git a/include/cps/Lexer.h b/include/cps/Lexer.h
index a03c3f6..a7a6218 100644
--- a/include/cps/Lexer.h
+++ b/include/cps/Lexer.h
@@ -1,39 +1,46 @@
 #pragma once
 #include <string>
-#include <iostream>
 
 namespace cps {
 
 enum Token {
     tok_eof = -1,
+
     tok_declare = -2,
-    tok_integer_kw = -3, // INTEGER
+    tok_integer_kw = -3,
     tok_input = -4,
     tok_output = -5,
-    tok_identifier = -6,
-    tok_number = -7,
-    tok_assign = -8,     // <-
-    tok_colon = -9,      // :
-
-    tok_if = -10,
-    tok_then = -11,
-    tok_else = -12,
-    tok_endif = -13,
     
-    tok_while = -14,
-    tok_do = -15,
-    tok_endwhile = -16,
-    tok_repeat = -17,
-    tok_until = -18,
-    tok_for = -19,
-    tok_to = -24,
-    tok_step = -25,
-    tok_next = -26,
+    tok_if = -6,
+    tok_then = -7,
+    tok_else = -8,
+    tok_endif = -9,
+    
+    tok_while = -10,
+    tok_do = -11,
+    tok_endwhile = -12,
+    
+    tok_repeat = -13,
+    tok_until = -14,
+    
+    tok_for = -15,
+    tok_to = -16,
+    tok_step = -17,
+    tok_next = -18,
+
+    tok_array = -19,
+    tok_of = -20,
 
-    tok_eq = -20, // =
-    tok_ne = -21, // <>
-    tok_le = -22, // <=
-    tok_ge = -23  // >=
+    tok_identifier = -100,
+    tok_number = -101,
+    
+    tok_assign = -200, // <-
+    tok_eq = -201, // =
+    tok_ne = -202, // <>
+    tok_le = -203, // <=
+    tok_ge = -204, // >=
+    
+    tok_colon = -300 // :
 };
 
 class Lexer {
@@ -44,4 +51,4 @@ public:
     int gettok();
 };
 
-} // namespace cps
+}
diff --git a/include/cps/Parser.h b/include/cps/Parser.h
index b685fb2..f5be655 100644
--- a/include/cps/Parser.h
+++ b/include/cps/Parser.h
@@ -29,7 +29,9 @@ class Parser {
     std::unique_ptr<StmtAST> ParseWhileStmt();
     std::unique_ptr<StmtAST> ParseRepeatStmt();
     std::unique_ptr<StmtAST> ParseForStmt();
-
+    
+    std::unique_ptr<StmtAST> ParseDeclare();
+    
 public:
     Parser(Lexer &L);
     std::vector<std::unique_ptr<StmtAST>> Parse();