Added half way of IF statement

3 files changed, 230 insertions, 78 deletions

diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 91339e0..0a946bf 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -1,5 +1,6 @@
 #include "cps/CodeGen.h"
 #include "llvm/IR/Verifier.h"
+#include "cps/Lexer.h"
 
 using namespace llvm;
 using namespace cps;
@@ -40,47 +41,13 @@ void CodeGen::compile(const std::vector<std::unique_ptr<StmtAST>> &Statements) {
     Builder->SetInsertPoint(BB);
 
     for (const auto &Stmt : Statements) {
-        // DECLARE 
-        if (auto *Decl = dynamic_cast<DeclareStmtAST*>(Stmt.get())) {
-            AllocaInst *Alloca = CreateEntryBlockAlloca(F, Decl->getName());
-            NamedValues[Decl->getName()] = Alloca;
-            // 0
-            Builder->CreateStore(ConstantInt::get(*TheContext, APInt(32, 0)), Alloca);
-        }
-        // ASSIGNMENT (x <- expr) 
-        else if (auto *Assign = dynamic_cast<AssignStmtAST*>(Stmt.get())) {
-            Value *Val = emitExpr(Assign->getExpr());
-            if (Val) {
-                AllocaInst *Alloca = NamedValues[Assign->getName()];
-                if (Alloca) Builder->CreateStore(Val, Alloca);
-            }
-        }
-        // INPUT 
-        else if (auto *In = dynamic_cast<InputStmtAST*>(Stmt.get())) {
-            AllocaInst *Alloca = NamedValues[In->getName()];
-            if (Alloca) {
-                std::vector<Value*> Args;
-                Args.push_back(ScanfFormatStr);
-                Args.push_back(Alloca);
-                Builder->CreateCall(ScanfFunc, Args);
-            }
-        }
-        // OUTPUT
-        else if (auto *Out = dynamic_cast<OutputStmtAST*>(Stmt.get())) {
-            Value *Val = emitExpr(Out->getExpr());
-
-            if (Val) {
-                // printf("%d\n", val)
-                std::vector<Value*> Args;
-                Args.push_back(PrintfFormatStr);
-                Args.push_back(Val);
-                Builder->CreateCall(PrintfFunc, Args);
-            }
-        }
+        emitStmt(Stmt.get());
     }
 
     // main return 0
-    Builder->CreateRet(ConstantInt::get(*TheContext, APInt(32, 0)));
+    if (!Builder->GetInsertBlock()->getTerminator())
+        Builder->CreateRet(ConstantInt::get(*TheContext, APInt(32, 0)));
+        
     verifyFunction(*F);
 }
 
@@ -114,6 +81,12 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
         case '-': return Builder->CreateSub(L, R, "subtmp");
         case '*': return Builder->CreateMul(L, R, "multmp");
         case '/': return Builder->CreateSDiv(L, R, "divtmp");
+        case tok_eq: return Builder->CreateICmpEQ(L, R, "eqtmp");
+        case tok_ne: return Builder->CreateICmpNE(L, R, "netmp");
+        case '<':    return Builder->CreateICmpSLT(L, R, "slttmp");
+        case '>':    return Builder->CreateICmpSGT(L, R, "sgttmp");
+        case tok_le: return Builder->CreateICmpSLE(L, R, "sletmp");
+        case tok_ge: return Builder->CreateICmpSGE(L, R, "sgetmp");
         default:  
             fprintf(stderr, "Error: Invalid binary operator\n");
             return nullptr;
@@ -122,3 +95,86 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
 
     return nullptr;
 }
+
+void CodeGen::emitIfStmt(IfStmtAST *Stmt) {
+    Value *CondV = emitExpr(Stmt->getCond());
+    if (!CondV) return;
+
+    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);
+
+    Builder->CreateCondBr(CondV, ThenBB, ElseBB);
+
+    Builder->SetInsertPoint(ThenBB);
+    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::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) {
+            AllocaInst *Alloca = NamedValues[Assign->getName()];
+            if (!Alloca) {
+                fprintf(stderr, "Error: Unknown variable name %s\n", Assign->getName().c_str());
+                return;
+            }
+            Builder->CreateStore(Val, Alloca);
+        }
+    }
+    // INPUT
+    else if (auto *In = dynamic_cast<InputStmtAST*>(Stmt)) {
+        AllocaInst *Alloca = NamedValues[In->getName()];
+        if (Alloca) {
+            std::vector<Value*> Args;
+            Args.push_back(ScanfFormatStr);
+            Args.push_back(Alloca); 
+            Builder->CreateCall(ScanfFunc, Args);
+        }
+    }
+    // OUTPUT
+    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);
+        }
+    }
+    // IF
+    else if (auto *IfStmt = dynamic_cast<IfStmtAST*>(Stmt)) {
+        emitIfStmt(IfStmt);
+    }
+}
diff --git a/lib/Lexer/Lexer.cc b/lib/Lexer/Lexer.cc
index b04438b..9ac01d9 100644
--- a/lib/Lexer/Lexer.cc
+++ b/lib/Lexer/Lexer.cc
@@ -19,6 +19,10 @@ 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;
         return tok_identifier;
     }
 
@@ -34,13 +38,35 @@ int Lexer::gettok() {
 
     if (LastChar == '<') {
         LastChar = getchar();
-        if (LastChar == '-') {
+        if (LastChar == '-') { // <-
             LastChar = getchar();
             return tok_assign;
         }
+        if (LastChar == '=') { // <=
+            LastChar = getchar();
+            return tok_le;
+        }
+        if (LastChar == '>') { // <>
+            LastChar = getchar();
+            return tok_ne;
+        }
         return '<';
     }
 
+    if (LastChar == '>') {
+        LastChar = getchar();
+        if (LastChar == '=') { // >=
+            LastChar = getchar();
+            return tok_ge;
+        }
+        return '>';
+    }
+
+    if (LastChar == '=') {
+        LastChar = getchar();
+        return tok_eq;
+    }
+
     // for comments
     
     if (LastChar == ':') {
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;
 }