1 files changed, 109 insertions, 39 deletions

diff --git a/lib/Parser/Parser.cc b/lib/Parser/Parser.cc
index 76b124a..ee94b36 100644
--- a/lib/Parser/Parser.cc
+++ b/lib/Parser/Parser.cc
@@ -5,10 +5,12 @@ using namespace cps;
 Parser::Parser(Lexer &L) : Lex(L) {
     getNextToken();
 
-    BinopPrecedence['+'] = 20;
-    BinopPrecedence['-'] = 20;
-    BinopPrecedence['*'] = 40;
-    BinopPrecedence['/'] = 40; 
+    BinopPrecedence[tok_eq] = 10;
+    BinopPrecedence[tok_ne] = 10;
+    BinopPrecedence[tok_le] = 10;
+    BinopPrecedence[tok_ge] = 10;
+    BinopPrecedence['<'] = 10;
+    BinopPrecedence['>'] = 10;
 }
 
 int Parser::GetTokPrecedence() {
@@ -91,48 +93,116 @@ std::vector<std::unique_ptr<StmtAST>> Parser::Parse() {
     std::vector<std::unique_ptr<StmtAST>> Statements;
 
     while (CurTok != tok_eof) {
-        if (CurTok == tok_declare) {
-            // DECLARE <id> : INTEGER
-            getNextToken();
-            if (CurTok != tok_identifier) break; // Error
-            std::string Name = Lex.IdentifierStr;
-            getNextToken();
-            
-            if (CurTok != tok_colon) break; // Error
-            getNextToken();
+        if (auto Stmt = ParseStatement()) {
+            Statements.push_back(std::move(Stmt));
+        } else {
+            if (CurTok != tok_eof) getNextToken();
+        }
+    }
+    return Statements;
+}
+
+
+std::unique_ptr<StmtAST> Parser::ParseIfStmt() {
+    getNextToken();
+
+    auto Cond = ParseExpression();
+    if (!Cond) return nullptr;
+
+    if (CurTok != tok_then) {
+        fprintf(stderr, "Error: expected THEN\n");
+        return nullptr;
+    }
+    getNextToken();
+
+
+    std::vector<std::unique_ptr<StmtAST>> ThenStmts;
+    while (CurTok != tok_else && CurTok != tok_endif && CurTok != tok_eof) {
+        auto Stmt = ParseStatement();
+        if (Stmt) ThenStmts.push_back(std::move(Stmt));
+    }
 
-            if (CurTok != tok_integer_kw) break; // Error
-            getNextToken();
+    std::vector<std::unique_ptr<StmtAST>> ElseStmts;
+    if (CurTok == tok_else) {
+        getNextToken();
+        while (CurTok != tok_endif && CurTok != tok_eof) {
+            auto Stmt = ParseStatement();
+            if (Stmt) ElseStmts.push_back(std::move(Stmt));
+        }
+    }
 
-            Statements.push_back(std::make_unique<DeclareStmtAST>(Name));
-        } 
-        else if (CurTok == tok_identifier) {
-            // <id> <- <expr>
-            std::string Name = Lex.IdentifierStr;
-            getNextToken();
+    if (CurTok != tok_endif) {
+        fprintf(stderr, "Error: expected ENDIF\n");
+        return nullptr;
+    }
+    getNextToken();
 
-            if (CurTok != tok_assign) break; // Error
-            getNextToken();
+    return std::make_unique<IfStmtAST>(std::move(Cond), std::move(ThenStmts), std::move(ElseStmts));
+}
 
-            auto Expr = ParseExpression();
-            Statements.push_back(std::make_unique<AssignStmtAST>(Name, std::move(Expr)));
+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;
+        }
+        std::string Name = Lex.IdentifierStr;
+        getNextToken();
+        
+        if (CurTok != tok_colon) {
+            fprintf(stderr, "Error: expected ':' after identifier\n");
+            return nullptr;
         }
-        else if (CurTok == tok_input) {
-            // INPUT <id>
-            getNextToken();
-            if (CurTok != tok_identifier) break;
-            Statements.push_back(std::make_unique<InputStmtAST>(Lex.IdentifierStr));
-            getNextToken();
+        getNextToken();
+
+        if (CurTok != tok_integer_kw) {
+            fprintf(stderr, "Error: expected INTEGER type\n");
+            return nullptr;
         }
-        else if (CurTok == tok_output) {
-            // OUTPUT <expr>
-            getNextToken();
-            auto Expr = ParseExpression();
-            Statements.push_back(std::make_unique<OutputStmtAST>(std::move(Expr)));
+        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;
         }
-        else {
-            getNextToken();
+        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;
         }
+        std::string Name = Lex.IdentifierStr;
+        getNextToken();
+        return std::make_unique<InputStmtAST>(Name);
     }
-    return Statements;
+    
+    // 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();
+    }
+
+    return nullptr;
 }