Added loop statements, FOR, WHILE and REPEAT

7 files changed, 354 insertions, 98 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);
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 0a946bf..31d113a 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -14,17 +14,14 @@ CodeGen::CodeGen() {
 }
 
 void CodeGen::SetupExternalFunctions() {
-    // printf(i8*, ...)
     std::vector<Type*> PrintfArgs;
     PrintfArgs.push_back(PointerType::getUnqual(*TheContext));
     FunctionType *PrintfType = FunctionType::get(Type::getInt32Ty(*TheContext), PrintfArgs, true);
     PrintfFunc = TheModule->getOrInsertFunction("printf", PrintfType);
 
-    // scanf(i8*, ...)
     FunctionType *ScanfType = FunctionType::get(Type::getInt32Ty(*TheContext), PrintfArgs, true);
     ScanfFunc = TheModule->getOrInsertFunction("scanf", ScanfType);
 
-    // format %d\n and %d
     PrintfFormatStr = Builder->CreateGlobalStringPtr("%d\n", "fmt_nl", 0, TheModule.get());
     ScanfFormatStr = Builder->CreateGlobalStringPtr("%d", "fmt_in", 0, TheModule.get());
 }
@@ -44,7 +41,6 @@ void CodeGen::compile(const std::vector<std::unique_ptr<StmtAST>> &Statements) {
         emitStmt(Stmt.get());
     }
 
-    // main return 0
     if (!Builder->GetInsertBlock()->getTerminator())
         Builder->CreateRet(ConstantInt::get(*TheContext, APInt(32, 0)));
         
@@ -92,7 +88,6 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
             return nullptr;
         }
     }
-
     return nullptr;
 }
 
@@ -100,12 +95,10 @@ void CodeGen::emitIfStmt(IfStmtAST *Stmt) {
     Value *CondV = emitExpr(Stmt->getCond());
     if (!CondV) return;
 
-    if (CondV->getType()->isIntegerTy(32)) {
+    if (CondV->getType()->isIntegerTy(32))
         CondV = Builder->CreateICmpNE(CondV, ConstantInt::get(*TheContext, APInt(32, 0)), "ifcond");
-    }
 
     Function *TheFunction = Builder->GetInsertBlock()->getParent();
-
     BasicBlock *ThenBB = BasicBlock::Create(*TheContext, "then", TheFunction);
     BasicBlock *ElseBB = BasicBlock::Create(*TheContext, "else", TheFunction);
     BasicBlock *MergeBB = BasicBlock::Create(*TheContext, "ifcont", TheFunction);
@@ -113,35 +106,140 @@ void CodeGen::emitIfStmt(IfStmtAST *Stmt) {
     Builder->CreateCondBr(CondV, ThenBB, ElseBB);
 
     Builder->SetInsertPoint(ThenBB);
-    for (const auto &S : Stmt->getThenStmts()) {
-        emitStmt(S.get());
-    }
-    if (!Builder->GetInsertBlock()->getTerminator())
-        Builder->CreateBr(MergeBB);
+    for (const auto &S : Stmt->getThenStmts()) emitStmt(S.get());
+    if (!Builder->GetInsertBlock()->getTerminator()) Builder->CreateBr(MergeBB);
 
     Builder->SetInsertPoint(ElseBB);
+    for (const auto &S : Stmt->getElseStmts()) emitStmt(S.get());
+    if (!Builder->GetInsertBlock()->getTerminator()) Builder->CreateBr(MergeBB);
+
+    Builder->SetInsertPoint(MergeBB);
+}
+
+void CodeGen::emitWhileStmt(WhileStmtAST *Stmt) {
+    Function *TheFunction = Builder->GetInsertBlock()->getParent();
+    
+    BasicBlock *CondBB = BasicBlock::Create(*TheContext, "whilecond", TheFunction);
+    BasicBlock *LoopBB = BasicBlock::Create(*TheContext, "whileloop", TheFunction);
+    BasicBlock *AfterBB = BasicBlock::Create(*TheContext, "whilecont", TheFunction);
+
+    Builder->CreateBr(CondBB);
+
+    Builder->SetInsertPoint(CondBB);
+    Value *CondV = emitExpr(Stmt->getCond());
+    if (!CondV) return;
+    if (CondV->getType()->isIntegerTy(32))
+        CondV = Builder->CreateICmpNE(CondV, ConstantInt::get(*TheContext, APInt(32, 0)), "loopcond");
+    
+    Builder->CreateCondBr(CondV, LoopBB, AfterBB);
+
+    Builder->SetInsertPoint(LoopBB);
+    for (const auto &S : Stmt->getBody()) emitStmt(S.get());
     
-    for (const auto &S : Stmt->getElseStmts()) {
-        emitStmt(S.get());
-    }
     if (!Builder->GetInsertBlock()->getTerminator())
-        Builder->CreateBr(MergeBB);
+        Builder->CreateBr(CondBB);
 
-    Builder->SetInsertPoint(MergeBB);
+    Builder->SetInsertPoint(AfterBB);
 }
 
+void CodeGen::emitRepeatStmt(RepeatStmtAST *Stmt) {
+    Function *TheFunction = Builder->GetInsertBlock()->getParent();
+
+    BasicBlock *LoopBB = BasicBlock::Create(*TheContext, "repeatloop", TheFunction);
+    BasicBlock *CondBB = BasicBlock::Create(*TheContext, "repeatcond", TheFunction);
+    BasicBlock *AfterBB = BasicBlock::Create(*TheContext, "repeatcont", TheFunction);
+
+    Builder->CreateBr(LoopBB);
+
+    Builder->SetInsertPoint(LoopBB);
+    for (const auto &S : Stmt->getBody()) emitStmt(S.get());
+    if (!Builder->GetInsertBlock()->getTerminator())
+        Builder->CreateBr(CondBB);
+
+    Builder->SetInsertPoint(CondBB);
+    Value *CondV = emitExpr(Stmt->getCond());
+    if (!CondV) return;
+    
+    if (CondV->getType()->isIntegerTy(32))
+        CondV = Builder->CreateICmpNE(CondV, ConstantInt::get(*TheContext, APInt(32, 0)), "untilcond");
+
+    Builder->CreateCondBr(CondV, AfterBB, LoopBB);
+
+    Builder->SetInsertPoint(AfterBB);
+}
+
+void CodeGen::emitForStmt(ForStmtAST *Stmt) {
+    Function *TheFunction = Builder->GetInsertBlock()->getParent();
+    std::string VarName = Stmt->getVarName();
+    
+    Value *StartVal = emitExpr(Stmt->getStart());
+    if (!StartVal) return;
+
+    AllocaInst *Alloca = NamedValues[VarName];
+    if (!Alloca) {
+        fprintf(stderr, "Error: Unknown variable in FOR loop %s\n", VarName.c_str());
+        return;
+    }
+    Builder->CreateStore(StartVal, Alloca);
+
+    BasicBlock *CondBB = BasicBlock::Create(*TheContext, "forcond", TheFunction);
+    BasicBlock *LoopBB = BasicBlock::Create(*TheContext, "forloop", TheFunction);
+    BasicBlock *IncBB = BasicBlock::Create(*TheContext, "forinc", TheFunction);
+    BasicBlock *AfterBB = BasicBlock::Create(*TheContext, "forcont", TheFunction);
+
+    Builder->CreateBr(CondBB);
+    Builder->SetInsertPoint(CondBB);
+
+    Value *CurVar = Builder->CreateLoad(Type::getInt32Ty(*TheContext), Alloca, VarName.c_str());
+    Value *EndVal = emitExpr(Stmt->getEnd());
+    if (!EndVal) return;
+
+    bool isNegativeStep = false;
+    if (Stmt->getStep()) {
+        if (auto *Num = dynamic_cast<NumberExprAST*>(Stmt->getStep())) {
+            if (Num->getVal() < 0) isNegativeStep = true;
+        }
+    }
+
+    Value *CondV;
+    if (isNegativeStep) {
+        CondV = Builder->CreateICmpSGE(CurVar, EndVal, "forcond_ge");
+    } else {
+        CondV = Builder->CreateICmpSLE(CurVar, EndVal, "forcond_le");
+    }
+
+    Builder->CreateCondBr(CondV, LoopBB, AfterBB);
+
+    Builder->SetInsertPoint(LoopBB);
+    for (const auto &S : Stmt->getBody()) emitStmt(S.get());
+    if (!Builder->GetInsertBlock()->getTerminator())
+        Builder->CreateBr(IncBB);
+
+    Builder->SetInsertPoint(IncBB);
+    Value *StepVal;
+    if (Stmt->getStep()) {
+        StepVal = emitExpr(Stmt->getStep());
+    } else {
+        StepVal = ConstantInt::get(*TheContext, APInt(32, 1));
+    }
+    
+    Value *CurValForInc = Builder->CreateLoad(Type::getInt32Ty(*TheContext), Alloca, VarName.c_str());
+    Value *NextVal = Builder->CreateAdd(CurValForInc, StepVal, "nextval");
+    Builder->CreateStore(NextVal, Alloca);
+    Builder->CreateBr(CondBB);
+
+    Builder->SetInsertPoint(AfterBB);
+}
 
 void CodeGen::emitStmt(StmtAST *Stmt) {
     if (!Stmt) return;
 
-    // DECLARE
     if (auto *Decl = dynamic_cast<DeclareStmtAST*>(Stmt)) {
         Function *TheFunction = Builder->GetInsertBlock()->getParent();
         AllocaInst *Alloca = CreateEntryBlockAlloca(TheFunction, Decl->getName());
         NamedValues[Decl->getName()] = Alloca;
         Builder->CreateStore(ConstantInt::get(*TheContext, APInt(32, 0)), Alloca);
     }
-    // ASSIGNMENT
     else if (auto *Assign = dynamic_cast<AssignStmtAST*>(Stmt)) {
         Value *Val = emitExpr(Assign->getExpr());
         if (Val) {
@@ -153,7 +251,6 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
             Builder->CreateStore(Val, Alloca);
         }
     }
-    // INPUT
     else if (auto *In = dynamic_cast<InputStmtAST*>(Stmt)) {
         AllocaInst *Alloca = NamedValues[In->getName()];
         if (Alloca) {
@@ -163,7 +260,6 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
             Builder->CreateCall(ScanfFunc, Args);
         }
     }
-    // OUTPUT
     else if (auto *Out = dynamic_cast<OutputStmtAST*>(Stmt)) {
         Value *Val = emitExpr(Out->getExpr());
         if (Val) {
@@ -173,8 +269,16 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
             Builder->CreateCall(PrintfFunc, Args);
         }
     }
-    // IF
     else if (auto *IfStmt = dynamic_cast<IfStmtAST*>(Stmt)) {
         emitIfStmt(IfStmt);
     }
+    else if (auto *WhileStmt = dynamic_cast<WhileStmtAST*>(Stmt)) {
+        emitWhileStmt(WhileStmt);
+    }
+    else if (auto *RepeatStmt = dynamic_cast<RepeatStmtAST*>(Stmt)) {
+        emitRepeatStmt(RepeatStmt);
+    }
+    else if (auto *ForStmt = dynamic_cast<ForStmtAST*>(Stmt)) {
+        emitForStmt(ForStmt);
+    }
 }
diff --git a/lib/Lexer/Lexer.cc b/lib/Lexer/Lexer.cc
index c50897b..7f7a167 100644
--- a/lib/Lexer/Lexer.cc
+++ b/lib/Lexer/Lexer.cc
@@ -1,6 +1,8 @@
 #include "cps/Lexer.h"
 #include <cctype>
 #include <cstdio>
+#include <string>
+#include <algorithm>
 
 using namespace cps;
 
@@ -19,10 +21,24 @@ int Lexer::gettok() {
         if (IdentifierStr == "INTEGER") return tok_integer_kw;
         if (IdentifierStr == "INPUT") return tok_input;
         if (IdentifierStr == "OUTPUT") return tok_output;
+        
         if (IdentifierStr == "IF") return tok_if;
         if (IdentifierStr == "THEN") return tok_then;
         if (IdentifierStr == "ELSE") return tok_else;
         if (IdentifierStr == "ENDIF") return tok_endif;
+
+        if (IdentifierStr == "WHILE") return tok_while;
+        if (IdentifierStr == "DO") return tok_do;
+        if (IdentifierStr == "ENDWHILE") return tok_endwhile;
+        
+        if (IdentifierStr == "REPEAT") return tok_repeat;
+        if (IdentifierStr == "UNTIL") return tok_until;
+        
+        if (IdentifierStr == "FOR") return tok_for;
+        if (IdentifierStr == "TO") return tok_to;
+        if (IdentifierStr == "STEP") return tok_step;
+        if (IdentifierStr == "NEXT") return tok_next;
+
         return tok_identifier;
     }
 
@@ -67,6 +83,11 @@ int Lexer::gettok() {
         return tok_eq;
     }
 
+    if (LastChar == ':') {
+        LastChar = getchar();
+        return tok_colon;
+    }
+
     if (LastChar == '/') {
         int NextChar = getchar();
         if (NextChar == '/') {
@@ -81,11 +102,6 @@ int Lexer::gettok() {
             return '/';
         }
     }
-    
-    if (LastChar == ':') {
-        LastChar = getchar();
-        return tok_colon;
-    }
 
     if (LastChar == EOF)
         return tok_eof;
diff --git a/lib/Parser/Parser.cc b/lib/Parser/Parser.cc
index 998d50b..7a1af89 100644
--- a/lib/Parser/Parser.cc
+++ b/lib/Parser/Parser.cc
@@ -26,6 +26,22 @@ int Parser::GetTokPrecedence() {
     return TokPrec;
 }
 
+int Parser::getNextToken() {
+    return CurTok = Lex.gettok();
+}
+
+std::unique_ptr<ExprAST> Parser::ParseNumberExpr() {
+    auto Result = std::make_unique<NumberExprAST>(Lex.NumVal);
+    getNextToken();
+    return Result;
+}
+
+std::unique_ptr<ExprAST> Parser::ParseIdentifierExpr() {
+    std::string IdName = Lex.IdentifierStr;
+    getNextToken();
+    return std::make_unique<VariableExprAST>(IdName);
+}
+
 std::unique_ptr<ExprAST> Parser::ParseParenExpr() {
     getNextToken();
     auto V = ParseExpression();
@@ -52,9 +68,7 @@ std::unique_ptr<ExprAST> Parser::ParsePrimary() {
 std::unique_ptr<ExprAST> Parser::ParseBinOpRHS(int ExprPrec, std::unique_ptr<ExprAST> LHS) {
     while (true) {
         int TokPrec = GetTokPrecedence();
-
-        if (TokPrec < ExprPrec)
-            return LHS;
+        if (TokPrec < ExprPrec) return LHS;
 
         int BinOp = CurTok;
         getNextToken();
@@ -72,23 +86,6 @@ std::unique_ptr<ExprAST> Parser::ParseBinOpRHS(int ExprPrec, std::unique_ptr<Exp
     }
 }
 
-
-int Parser::getNextToken() {
-    return CurTok = Lex.gettok();
-}
-
-std::unique_ptr<ExprAST> Parser::ParseNumberExpr() {
-    auto Result = std::make_unique<NumberExprAST>(Lex.NumVal);
-    getNextToken();
-    return Result;
-}
-
-std::unique_ptr<ExprAST> Parser::ParseIdentifierExpr() {
-    std::string IdName = Lex.IdentifierStr;
-    getNextToken();
-    return std::make_unique<VariableExprAST>(IdName);
-}
-
 std::unique_ptr<ExprAST> Parser::ParseExpression() {
     auto LHS = ParsePrimary();
     if (!LHS) return nullptr;
@@ -97,7 +94,6 @@ std::unique_ptr<ExprAST> Parser::ParseExpression() {
 
 std::vector<std::unique_ptr<StmtAST>> Parser::Parse() {
     std::vector<std::unique_ptr<StmtAST>> Statements;
-
     while (CurTok != tok_eof) {
         if (auto Stmt = ParseStatement()) {
             Statements.push_back(std::move(Stmt));
@@ -108,10 +104,8 @@ std::vector<std::unique_ptr<StmtAST>> Parser::Parse() {
     return Statements;
 }
 
-
 std::unique_ptr<StmtAST> Parser::ParseIfStmt() {
     getNextToken();
-
     auto Cond = ParseExpression();
     if (!Cond) return nullptr;
 
@@ -121,7 +115,6 @@ std::unique_ptr<StmtAST> Parser::ParseIfStmt() {
     }
     getNextToken();
 
-
     std::vector<std::unique_ptr<StmtAST>> ThenStmts;
     while (CurTok != tok_else && CurTok != tok_endif && CurTok != tok_eof) {
         auto Stmt = ParseStatement();
@@ -146,69 +139,157 @@ std::unique_ptr<StmtAST> Parser::ParseIfStmt() {
     return std::make_unique<IfStmtAST>(std::move(Cond), std::move(ThenStmts), std::move(ElseStmts));
 }
 
+std::unique_ptr<StmtAST> Parser::ParseWhileStmt() {
+    getNextToken();
+    auto Cond = ParseExpression();
+    if (!Cond) return nullptr;
+
+    if (CurTok != tok_do) {
+        fprintf(stderr, "Error: expected DO after WHILE condition\n");
+        return nullptr;
+    }
+    getNextToken();
+
+    std::vector<std::unique_ptr<StmtAST>> Body;
+    while (CurTok != tok_endwhile && CurTok != tok_eof) {
+        auto Stmt = ParseStatement();
+        if (Stmt) Body.push_back(std::move(Stmt));
+    }
+
+    if (CurTok != tok_endwhile) {
+        fprintf(stderr, "Error: expected ENDWHILE\n");
+        return nullptr;
+    }
+    getNextToken();
+
+    return std::make_unique<WhileStmtAST>(std::move(Cond), std::move(Body));
+}
+
+std::unique_ptr<StmtAST> Parser::ParseRepeatStmt() {
+    getNextToken();
+
+    std::vector<std::unique_ptr<StmtAST>> Body;
+    while (CurTok != tok_until && CurTok != tok_eof) {
+        auto Stmt = ParseStatement();
+        if (Stmt) Body.push_back(std::move(Stmt));
+    }
+
+    if (CurTok != tok_until) {
+        fprintf(stderr, "Error: expected UNTIL\n");
+        return nullptr;
+    }
+    getNextToken();
+
+    auto Cond = ParseExpression();
+    if (!Cond) return nullptr;
+
+    return std::make_unique<RepeatStmtAST>(std::move(Body), std::move(Cond));
+}
+
+std::unique_ptr<StmtAST> Parser::ParseForStmt() {
+    getNextToken();
+    if (CurTok != tok_identifier) {
+        fprintf(stderr, "Error: expected identifier after FOR\n");
+        return nullptr;
+    }
+    std::string VarName = Lex.IdentifierStr;
+    getNextToken();
+
+    if (CurTok != tok_assign) {
+        fprintf(stderr, "Error: expected '<-' in FOR loop\n");
+        return nullptr;
+    }
+    getNextToken();
+
+    auto Start = ParseExpression();
+    if (!Start) return nullptr;
+
+    if (CurTok != tok_to) {
+        fprintf(stderr, "Error: expected TO in FOR loop\n");
+        return nullptr;
+    }
+    getNextToken();
+
+    auto End = ParseExpression();
+    if (!End) return nullptr;
+
+    std::unique_ptr<ExprAST> Step = nullptr;
+    if (CurTok == tok_step) {
+        getNextToken();
+        Step = ParseExpression();
+    }
+
+    std::vector<std::unique_ptr<StmtAST>> Body;
+    while (CurTok != tok_next && CurTok != tok_eof) {
+        auto Stmt = ParseStatement();
+        if (Stmt) Body.push_back(std::move(Stmt));
+    }
+
+    if (CurTok != tok_next) {
+        fprintf(stderr, "Error: expected NEXT\n");
+        return nullptr;
+    }
+    getNextToken();
+    
+    if (CurTok != tok_identifier) {
+        fprintf(stderr, "Error: expected identifier after NEXT (e.g., NEXT %s)\n", VarName.c_str());
+        return nullptr;
+    }
+
+    if (Lex.IdentifierStr != VarName) {
+        fprintf(stderr, "Error: NEXT identifier '%s' does not match FOR variable '%s'\n", 
+                Lex.IdentifierStr.c_str(), VarName.c_str());
+        return nullptr;
+    }
+    getNextToken();
+
+    return std::make_unique<ForStmtAST>(VarName, std::move(Start), std::move(End), std::move(Step), std::move(Body));
+}
+
 std::unique_ptr<StmtAST> Parser::ParseStatement() {
-    // DECLARE
     if (CurTok == tok_declare) {
         getNextToken();
-        if (CurTok != tok_identifier) {
-            fprintf(stderr, "Error: expected identifier after DECLARE\n");
-            return nullptr;
-        }
+        if (CurTok != tok_identifier) return nullptr;
         std::string Name = Lex.IdentifierStr;
         getNextToken();
-        
-        if (CurTok != tok_colon) {
-            fprintf(stderr, "Error: expected ':' after identifier\n");
-            return nullptr;
-        }
+        if (CurTok != tok_colon) return nullptr;
         getNextToken();
-
-        if (CurTok != tok_integer_kw) {
-            fprintf(stderr, "Error: expected INTEGER type\n");
-            return nullptr;
-        }
+        if (CurTok != tok_integer_kw) return nullptr;
         getNextToken();
         return std::make_unique<DeclareStmtAST>(Name);
     }
-    
-    // ASSIGNMENT (x <- expr)
     else if (CurTok == tok_identifier) {
         std::string Name = Lex.IdentifierStr;
-        getNextToken(); // eat Name
-
-        if (CurTok != tok_assign) {
-            fprintf(stderr, "Error: expected '<-' for assignment\n");
-            return nullptr;
-        }
         getNextToken();
-
+        if (CurTok != tok_assign) return nullptr;
+        getNextToken();
         auto Expr = ParseExpression();
         return std::make_unique<AssignStmtAST>(Name, std::move(Expr));
     }
-    
-    // INPUT
     else if (CurTok == tok_input) {
         getNextToken();
-        if (CurTok != tok_identifier) {
-            fprintf(stderr, "Error: expected identifier after INPUT\n");
-            return nullptr;
-        }
+        if (CurTok != tok_identifier) return nullptr;
         std::string Name = Lex.IdentifierStr;
         getNextToken();
         return std::make_unique<InputStmtAST>(Name);
     }
-    
-    // OUTPUT
     else if (CurTok == tok_output) {
         getNextToken();
         auto Expr = ParseExpression();
         return std::make_unique<OutputStmtAST>(std::move(Expr));
     }
-    
-    // IF
     else if (CurTok == tok_if) {
         return ParseIfStmt();
     }
+    else if (CurTok == tok_while) {
+        return ParseWhileStmt();
+    }
+    else if (CurTok == tok_repeat) {
+        return ParseRepeatStmt();
+    }
+    else if (CurTok == tok_for) {
+        return ParseForStmt();
+    }
 
     return nullptr;
 }