Added loop statements, FOR, WHILE and REPEAT

4 files changed, 67 insertions, 12 deletions

diff --git a/include/cps/AST.h b/include/cps/AST.h
index 3c58642..1473c45 100644
--- a/include/cps/AST.h
+++ b/include/cps/AST.h
@@ -24,6 +24,17 @@ public:
     const std::string &getName() const { return Name; }
 };
 
+class BinaryExprAST : public ExprAST {
+    int Op;
+    std::unique_ptr<ExprAST> LHS, RHS;
+public:
+    BinaryExprAST(int Op, std::unique_ptr<ExprAST> LHS, std::unique_ptr<ExprAST> RHS)
+        : Op(Op), LHS(std::move(LHS)), RHS(std::move(RHS)) {}
+    int getOp() const { return Op; }
+    ExprAST *getLHS() const { return LHS.get(); }
+    ExprAST *getRHS() const { return RHS.get(); }
+};
+
 class StmtAST {
 public:
     virtual ~StmtAST() = default;
@@ -60,17 +71,6 @@ public:
     ExprAST *getExpr() const { return Expr.get(); }
 };
 
-class BinaryExprAST : public ExprAST {
-    int Op;
-    std::unique_ptr<ExprAST> LHS, RHS;
-public:
-    BinaryExprAST(int Op, std::unique_ptr<ExprAST> LHS, std::unique_ptr<ExprAST> RHS)
-        : Op(Op), LHS(std::move(LHS)), RHS(std::move(RHS)) {}
-    int getOp() const { return Op; }
-    ExprAST *getLHS() const { return LHS.get(); }
-    ExprAST *getRHS() const { return RHS.get(); }
-};
-
 class IfStmtAST : public StmtAST {
     std::unique_ptr<ExprAST> Cond;
     std::vector<std::unique_ptr<StmtAST>> ThenStmts;
@@ -86,5 +86,43 @@ public:
     const std::vector<std::unique_ptr<StmtAST>> &getElseStmts() const { return ElseStmts; }
 };
 
+class WhileStmtAST : public StmtAST {
+    std::unique_ptr<ExprAST> Cond;
+    std::vector<std::unique_ptr<StmtAST>> Body;
+public:
+    WhileStmtAST(std::unique_ptr<ExprAST> Cond, std::vector<std::unique_ptr<StmtAST>> Body)
+        : Cond(std::move(Cond)), Body(std::move(Body)) {}
+    ExprAST *getCond() const { return Cond.get(); }
+    const std::vector<std::unique_ptr<StmtAST>> &getBody() const { return Body; }
+};
+
+class RepeatStmtAST : public StmtAST {
+    std::vector<std::unique_ptr<StmtAST>> Body;
+    std::unique_ptr<ExprAST> Cond;
+public:
+    RepeatStmtAST(std::vector<std::unique_ptr<StmtAST>> Body, std::unique_ptr<ExprAST> Cond)
+        : Body(std::move(Body)), Cond(std::move(Cond)) {}
+    ExprAST *getCond() const { return Cond.get(); }
+    const std::vector<std::unique_ptr<StmtAST>> &getBody() const { return Body; }
+};
+
+class ForStmtAST : public StmtAST {
+    std::string VarName;
+    std::unique_ptr<ExprAST> Start;
+    std::unique_ptr<ExprAST> End;
+    std::unique_ptr<ExprAST> Step;
+    std::vector<std::unique_ptr<StmtAST>> Body;
+public:
+    ForStmtAST(const std::string &VarName, std::unique_ptr<ExprAST> Start, 
+               std::unique_ptr<ExprAST> End, std::unique_ptr<ExprAST> Step,
+               std::vector<std::unique_ptr<StmtAST>> Body)
+        : VarName(VarName), Start(std::move(Start)), End(std::move(End)), 
+          Step(std::move(Step)), Body(std::move(Body)) {}
+    const std::string &getVarName() const { return VarName; }
+    ExprAST *getStart() const { return Start.get(); }
+    ExprAST *getEnd() const { return End.get(); }
+    ExprAST *getStep() const { return Step.get(); }
+    const std::vector<std::unique_ptr<StmtAST>> &getBody() const { return Body; }
+};
 
 } // namespace cps
diff --git a/include/cps/CodeGen.h b/include/cps/CodeGen.h
index 88e4755..9be6cc7 100644
--- a/include/cps/CodeGen.h
+++ b/include/cps/CodeGen.h
@@ -20,10 +20,14 @@ class CodeGen {
     llvm::FunctionCallee ScanfFunc;
     llvm::Value *PrintfFormatStr;
     llvm::Value *ScanfFormatStr;
+    
     llvm::Value *emitExpr(ExprAST *Expr);
     void emitStmt(StmtAST *Stmt);
-    void emitIfStmt(IfStmtAST *Stmt);
     
+    void emitIfStmt(IfStmtAST *Stmt);
+    void emitWhileStmt(WhileStmtAST *Stmt);
+    void emitRepeatStmt(RepeatStmtAST *Stmt);
+    void emitForStmt(ForStmtAST *Stmt);
 
     void SetupExternalFunctions();
 
diff --git a/include/cps/Lexer.h b/include/cps/Lexer.h
index b5ab37c..a03c3f6 100644
--- a/include/cps/Lexer.h
+++ b/include/cps/Lexer.h
@@ -20,6 +20,16 @@ enum Token {
     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_eq = -20, // =
     tok_ne = -21, // <>
     tok_le = -22, // <=
diff --git a/include/cps/Parser.h b/include/cps/Parser.h
index 58b6366..b685fb2 100644
--- a/include/cps/Parser.h
+++ b/include/cps/Parser.h
@@ -26,6 +26,9 @@ class Parser {
 
     std::unique_ptr<StmtAST> ParseStatement();
     std::unique_ptr<StmtAST> ParseIfStmt();
+    std::unique_ptr<StmtAST> ParseWhileStmt();
+    std::unique_ptr<StmtAST> ParseRepeatStmt();
+    std::unique_ptr<StmtAST> ParseForStmt();
 
 public:
     Parser(Lexer &L);