Added binary operations +-*/

6 files changed, 152 insertions, 39 deletions

diff --git a/include/cps/AST.h b/include/cps/AST.h
index 5b28e8a..4d72040 100644
--- a/include/cps/AST.h
+++ b/include/cps/AST.h
@@ -60,4 +60,16 @@ public:
     ExprAST *getExpr() const { return Expr.get(); }
 };
 
+class BinaryExprAST : public ExprAST {
+    char Op; // '+', '-', '*', '/'
+    std::unique_ptr<ExprAST> LHS, RHS;
+public:
+    BinaryExprAST(char Op, std::unique_ptr<ExprAST> LHS, std::unique_ptr<ExprAST> RHS)
+        : Op(Op), LHS(std::move(LHS)), RHS(std::move(RHS)) {}
+
+    char getOp() const { return Op; }
+    ExprAST *getLHS() const { return LHS.get(); }
+    ExprAST *getRHS() const { return RHS.get(); }
+};
+
 } // namespace cps
diff --git a/include/cps/CodeGen.h b/include/cps/CodeGen.h
index 2621795..d6a6013 100644
--- a/include/cps/CodeGen.h
+++ b/include/cps/CodeGen.h
@@ -20,6 +20,7 @@ class CodeGen {
     llvm::FunctionCallee ScanfFunc;
     llvm::Value *PrintfFormatStr;
     llvm::Value *ScanfFormatStr;
+    llvm::Value *emitExpr(ExprAST *Expr);
 
     void SetupExternalFunctions();
 
diff --git a/include/cps/Lexer.h b/include/cps/Lexer.h
index 0ad5609..b5ab37c 100644
--- a/include/cps/Lexer.h
+++ b/include/cps/Lexer.h
@@ -13,7 +13,17 @@ enum Token {
     tok_identifier = -6,
     tok_number = -7,
     tok_assign = -8,     // <-
-    tok_colon = -9       // :
+    tok_colon = -9,      // :
+
+    tok_if = -10,
+    tok_then = -11,
+    tok_else = -12,
+    tok_endif = -13,
+    
+    tok_eq = -20, // =
+    tok_ne = -21, // <>
+    tok_le = -22, // <=
+    tok_ge = -23  // >=
 };
 
 class Lexer {
diff --git a/include/cps/Parser.h b/include/cps/Parser.h
index c24184f..1b27a22 100644
--- a/include/cps/Parser.h
+++ b/include/cps/Parser.h
@@ -3,22 +3,29 @@
 #include "cps/AST.h"
 #include <vector>
 #include <memory>
+#include <map>
 
 namespace cps {
 
 class Parser {
     Lexer &Lex;
     int CurTok;
+    
+    std::map<char, int> BinopPrecedence;
 
     int getNextToken();
+    int GetTokPrecedence();
+
     std::unique_ptr<ExprAST> ParseExpression();
+    std::unique_ptr<ExprAST> ParsePrimary();
+    std::unique_ptr<ExprAST> ParseBinOpRHS(int ExprPrec, std::unique_ptr<ExprAST> LHS); 
+
     std::unique_ptr<ExprAST> ParseNumberExpr();
-    std::unique_ptr<ExprAST> ParseParenExpr(); // 预留
     std::unique_ptr<ExprAST> ParseIdentifierExpr();
+    std::unique_ptr<ExprAST> ParseParenExpr();
 
 public:
-    Parser(Lexer &L) : Lex(L) { getNextToken(); }
-    
+    Parser(Lexer &L);
     std::vector<std::unique_ptr<StmtAST>> Parse();
 };
 
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index a2de939..91339e0 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -49,17 +49,10 @@ void CodeGen::compile(const std::vector<std::unique_ptr<StmtAST>> &Statements) {
         }
         // ASSIGNMENT (x <- expr) 
         else if (auto *Assign = dynamic_cast<AssignStmtAST*>(Stmt.get())) {
-            Value *Val = nullptr;
-            if (auto *Num = dynamic_cast<NumberExprAST*>(Assign->getExpr())) {
-                Val = ConstantInt::get(*TheContext, APInt(32, Num->getVal()));
-            } else if (auto *Var = dynamic_cast<VariableExprAST*>(Assign->getExpr())) {
-                 AllocaInst *A = NamedValues[Var->getName()];
-                 Val = Builder->CreateLoad(Type::getInt32Ty(*TheContext), A, Var->getName());
-            }
-
-            AllocaInst *Alloca = NamedValues[Assign->getName()];
-            if (Alloca && Val) {
-                Builder->CreateStore(Val, Alloca);
+            Value *Val = emitExpr(Assign->getExpr());
+            if (Val) {
+                AllocaInst *Alloca = NamedValues[Assign->getName()];
+                if (Alloca) Builder->CreateStore(Val, Alloca);
             }
         }
         // INPUT 
@@ -74,13 +67,7 @@ void CodeGen::compile(const std::vector<std::unique_ptr<StmtAST>> &Statements) {
         }
         // OUTPUT
         else if (auto *Out = dynamic_cast<OutputStmtAST*>(Stmt.get())) {
-            Value *Val = nullptr;
-            if (auto *Num = dynamic_cast<NumberExprAST*>(Out->getExpr())) {
-                Val = ConstantInt::get(*TheContext, APInt(32, Num->getVal()));
-            } else if (auto *Var = dynamic_cast<VariableExprAST*>(Out->getExpr())) {
-                 AllocaInst *A = NamedValues[Var->getName()];
-                 Val = Builder->CreateLoad(Type::getInt32Ty(*TheContext), A, Var->getName());
-            }
+            Value *Val = emitExpr(Out->getExpr());
 
             if (Val) {
                 // printf("%d\n", val)
@@ -100,3 +87,38 @@ void CodeGen::compile(const std::vector<std::unique_ptr<StmtAST>> &Statements) {
 void CodeGen::print() {
     TheModule->print(errs(), nullptr);
 }
+
+Value *CodeGen::emitExpr(ExprAST *Expr) {
+    if (!Expr) return nullptr;
+
+    if (auto *Num = dynamic_cast<NumberExprAST*>(Expr)) {
+        return ConstantInt::get(*TheContext, APInt(32, Num->getVal(), true)); 
+    }
+    
+    if (auto *Var = dynamic_cast<VariableExprAST*>(Expr)) {
+        AllocaInst *A = NamedValues[Var->getName()];
+        if (!A) {
+            fprintf(stderr, "Error: Unknown variable name %s\n", Var->getName().c_str());
+            return nullptr;
+        }
+        return Builder->CreateLoad(Type::getInt32Ty(*TheContext), A, Var->getName().c_str());
+    }
+
+    if (auto *Bin = dynamic_cast<BinaryExprAST*>(Expr)) {
+        Value *L = emitExpr(Bin->getLHS());
+        Value *R = emitExpr(Bin->getRHS());
+        if (!L || !R) return nullptr;
+
+        switch (Bin->getOp()) {
+        case '+': return Builder->CreateAdd(L, R, "addtmp");
+        case '-': return Builder->CreateSub(L, R, "subtmp");
+        case '*': return Builder->CreateMul(L, R, "multmp");
+        case '/': return Builder->CreateSDiv(L, R, "divtmp");
+        default:  
+            fprintf(stderr, "Error: Invalid binary operator\n");
+            return nullptr;
+        }
+    }
+
+    return nullptr;
+}
diff --git a/lib/Parser/Parser.cc b/lib/Parser/Parser.cc
index 90dc3a0..76b124a 100644
--- a/lib/Parser/Parser.cc
+++ b/lib/Parser/Parser.cc
@@ -2,6 +2,69 @@
 
 using namespace cps;
 
+Parser::Parser(Lexer &L) : Lex(L) {
+    getNextToken();
+
+    BinopPrecedence['+'] = 20;
+    BinopPrecedence['-'] = 20;
+    BinopPrecedence['*'] = 40;
+    BinopPrecedence['/'] = 40; 
+}
+
+int Parser::GetTokPrecedence() {
+    if (!isascii(CurTok)) return -1;
+    int TokPrec = BinopPrecedence[CurTok];
+    if (TokPrec <= 0) return -1;
+    return TokPrec;
+}
+
+std::unique_ptr<ExprAST> Parser::ParseParenExpr() {
+    getNextToken();
+    auto V = ParseExpression();
+    if (!V) return nullptr;
+    if (CurTok != ')') {
+        fprintf(stderr, "Error: expected ')'\n");
+        return nullptr;
+    }
+    getNextToken();
+    return V;
+}
+
+std::unique_ptr<ExprAST> Parser::ParsePrimary() {
+    switch (CurTok) {
+    case tok_identifier: return ParseIdentifierExpr();
+    case tok_number:     return ParseNumberExpr();
+    case '(':            return ParseParenExpr();
+    default:
+        fprintf(stderr, "Error: unknown token when expecting an expression\n");
+        return nullptr;
+    }
+}
+
+std::unique_ptr<ExprAST> Parser::ParseBinOpRHS(int ExprPrec, std::unique_ptr<ExprAST> LHS) {
+    while (true) {
+        int TokPrec = GetTokPrecedence();
+
+        if (TokPrec < ExprPrec)
+            return LHS;
+
+        int BinOp = CurTok;
+        getNextToken();
+
+        auto RHS = ParsePrimary();
+        if (!RHS) return nullptr;
+
+        int NextPrec = GetTokPrecedence();
+        if (TokPrec < NextPrec) {
+            RHS = ParseBinOpRHS(TokPrec + 1, std::move(RHS));
+            if (!RHS) return nullptr;
+        }
+
+        LHS = std::make_unique<BinaryExprAST>(BinOp, std::move(LHS), std::move(RHS));
+    }
+}
+
+
 int Parser::getNextToken() {
     return CurTok = Lex.gettok();
 }
@@ -19,11 +82,9 @@ std::unique_ptr<ExprAST> Parser::ParseIdentifierExpr() {
 }
 
 std::unique_ptr<ExprAST> Parser::ParseExpression() {
-    if (CurTok == tok_number)
-        return ParseNumberExpr();
-    if (CurTok == tok_identifier)
-        return ParseIdentifierExpr();
-    return nullptr;
+    auto LHS = ParsePrimary();
+    if (!LHS) return nullptr;
+    return ParseBinOpRHS(0, std::move(LHS));
 }
 
 std::vector<std::unique_ptr<StmtAST>> Parser::Parse() {
@@ -32,45 +93,45 @@ std::vector<std::unique_ptr<StmtAST>> Parser::Parse() {
     while (CurTok != tok_eof) {
         if (CurTok == tok_declare) {
             // DECLARE <id> : INTEGER
-            getNextToken(); // eat DECLARE
+            getNextToken();
             if (CurTok != tok_identifier) break; // Error
             std::string Name = Lex.IdentifierStr;
-            getNextToken(); // eat Name
+            getNextToken();
             
-            if (CurTok != tok_colon) break; // Error: expected ':'
-            getNextToken(); // eat ':'
+            if (CurTok != tok_colon) break; // Error
+            getNextToken();
 
-            if (CurTok != tok_integer_kw) break; // Error: expected INTEGER
-            getNextToken(); // eat INTEGER
+            if (CurTok != tok_integer_kw) break; // Error
+            getNextToken();
 
             Statements.push_back(std::make_unique<DeclareStmtAST>(Name));
         } 
         else if (CurTok == tok_identifier) {
             // <id> <- <expr>
             std::string Name = Lex.IdentifierStr;
-            getNextToken(); // eat Name
+            getNextToken();
 
-            if (CurTok != tok_assign) break; // Error: expected '<-'
-            getNextToken(); // eat '<-'
+            if (CurTok != tok_assign) break; // Error
+            getNextToken();
 
             auto Expr = ParseExpression();
             Statements.push_back(std::make_unique<AssignStmtAST>(Name, std::move(Expr)));
         }
         else if (CurTok == tok_input) {
             // INPUT <id>
-            getNextToken(); // eat INPUT
+            getNextToken();
             if (CurTok != tok_identifier) break;
             Statements.push_back(std::make_unique<InputStmtAST>(Lex.IdentifierStr));
             getNextToken();
         }
         else if (CurTok == tok_output) {
             // OUTPUT <expr>
-            getNextToken(); // eat OUTPUT
+            getNextToken();
             auto Expr = ParseExpression();
             Statements.push_back(std::make_unique<OutputStmtAST>(std::move(Expr)));
         }
         else {
-            getNextToken(); // Skip unknown
+            getNextToken();
         }
     }
     return Statements;