Lexer, Parser: Add support for parsing for loops

mrunix00 · Jul 9, 2024 · 33bb2e3 · 33bb2e3
1 parent cc1247d
commit 33bb2e3
Show file tree

Hide file tree

Showing 6 changed files with 77 additions and 3 deletions.
diff --git a/include/ast.h b/include/ast.h
@@ -21,6 +21,7 @@ struct AbstractSyntaxTree {
         FunctionCall,
         IfStatement,
         WhileStatement,
+        ForLoop,
     } nodeType{Type::Invalid};
     virtual ~AbstractSyntaxTree() = default;
     virtual bool operator==(const AbstractSyntaxTree &other) const = 0;
@@ -122,5 +123,14 @@ struct WhileStatement : public AbstractSyntaxTree {
     void compile(Program &program, Segment &segment) const override;
 };
 
+struct ForLoop : public AbstractSyntaxTree {
+    AbstractSyntaxTree *initialization;
+    AbstractSyntaxTree *condition;
+    AbstractSyntaxTree *step;
+    AbstractSyntaxTree *body;
+    ForLoop(AbstractSyntaxTree *initialization, AbstractSyntaxTree *condition, AbstractSyntaxTree *step, AbstractSyntaxTree *body);
+    bool operator==(const AbstractSyntaxTree &other) const override;
+};
+
 std::vector<AbstractSyntaxTree *> parse(const char *input);
 Program compile(const char *input);
diff --git a/src/ast.cpp b/src/ast.cpp
@@ -392,6 +392,32 @@ bool UnaryExpression::operator==(const AbstractSyntaxTree &other) const {
            side == otherUnaryExpression.side;
 }
 
+ForLoop::ForLoop(AbstractSyntaxTree *initialization, AbstractSyntaxTree *condition, AbstractSyntaxTree *step, AbstractSyntaxTree *body)
+    : initialization(initialization), condition(condition), step(step), body(body) {
+    nodeType = Type::ForLoop;
+    typeStr = "ForLoop";
+}
+bool ForLoop::operator==(const AbstractSyntaxTree &other) const {
+    if (other.nodeType != nodeType) return false;
+    auto &otherForLoop = dynamic_cast<const ForLoop &>(other);
+    if ((initialization == nullptr || otherForLoop.initialization == nullptr) &&
+        initialization != otherForLoop.initialization)
+        return false;
+    if ((condition == nullptr || otherForLoop.condition == nullptr) &&
+        condition != otherForLoop.condition)
+        return false;
+    if ((step == nullptr || otherForLoop.step == nullptr) &&
+        step != otherForLoop.step)
+        return false;
+    if ((body == nullptr || otherForLoop.body == nullptr) &&
+        body != otherForLoop.body)
+        return false;
+    return *initialization == *otherForLoop.initialization &&
+           *condition == *otherForLoop.condition &&
+           *step == *otherForLoop.step &&
+           *body == *otherForLoop.body;
+}
+
 Program compile(const char *input) {
     Program program;
     auto ast = parse(input);

diff --git a/src/lexer.l b/src/lexer.l
@@ -36,6 +36,7 @@
 "if"                        { return If; }
 "else"                      { return Else; }
 "while"                     { return While; }
+"for"                       { return For; }
 "return"                    { return Return; }
 "u8"                        { return U8; }
 "u16"                       { return U16; }

diff --git a/src/parser.y b/src/parser.y
@@ -22,13 +22,13 @@
 %token String Plus Minus Multiply Divide Modulo Assign
 %token Increment Decrement IncrementAssign DecrementAssign
 %token Equal NotEqual Less Greater LessEqual GreaterEqual And Or Not
-%token Define Function If Else While Return
+%token Define Function If Else While For Return
 %token U8 U16 U32 U64 I8 I16 I32 I64 F32 F64 Bool True False
 %token Colon Comma Semicolon Arrow Newline
 %token LParen RParen LBrace RBrace LBracket RBracket
 
 %token <str> Number Identifier
-%type <ast> Expression Expressions VarType ScopedBody TypeCast FunctionCall IfStatement WhileStatement
+%type <ast> Expression Expressions VarType ScopedBody TypeCast FunctionCall IfStatement WhileStatement ForLoop
 %type <ast> ArgumentDeclaration ArgumentDeclarationsList Arguments FunctionDeclaration UnaryExpression
 
 %left Plus Minus
@@ -59,6 +59,12 @@ WhileStatement:
     }
 ;
 
+ForLoop:
+    For Expression Semicolon Expression Semicolon Expression ScopedBody {
+        $$ = new ForLoop(static_cast<AbstractSyntaxTree*>($2), static_cast<AbstractSyntaxTree*>($4), static_cast<AbstractSyntaxTree*>($6), static_cast<AbstractSyntaxTree*>($7));
+    }
+;
+
 VarType:
     U8  { $$ = new Node({U8, "u8"}); }
     | U16 { $$ = new Node({U16, "u16"}); }
@@ -167,6 +173,7 @@ Expression:
     | FunctionCall { $$ = $1; }
     | IfStatement { $$ = $1; }
     | WhileStatement { $$ = $1; }
+    | ForLoop { $$ = $1; }
     | UnaryExpression { $$ = $1; }
     | Return Expression {
         $$ = new ReturnStatement(static_cast<AbstractSyntaxTree*>($2));

diff --git a/tests/lexer_tests.cpp b/tests/lexer_tests.cpp
@@ -109,7 +109,7 @@ TEST(Lexer, Symbols) {
 
 TEST(Lexer, Keywords) {
     auto actual = lex(
-            "define if else while return "
+            "define if else while for return "
             "u8 u16 u32 u64 i8 i16 i32 i64 "
             "f32 f64 bool true false");
 
@@ -118,6 +118,7 @@ TEST(Lexer, Keywords) {
             {If, "if"},
             {Else, "else"},
             {While, "while"},
+            {For, "for"},
             {Return, "return"},
             {U8, "u8"},
             {U16, "u16"},

diff --git a/tests/parser_tests.cpp b/tests/parser_tests.cpp
@@ -361,6 +361,35 @@ TEST(ParserTests, WhileStatement) {
         ASSERT_EQ(*expectedResult[i], *actualResult[i]);
 }
 
+TEST(ParserTests, ForLoop) {
+    const char *input = "for i = 0; i < 10; i++ {\n"
+                        "\treturn i;\n"
+                        "};";
+    auto expectedResult = std::vector<AbstractSyntaxTree *>{
+            new ForLoop(
+                    new BinaryExpression(
+                            new Node({Identifier, "i"}),
+                            new Node({Number, "0"}),
+                            {Assign, "="}),
+                    new BinaryExpression(
+                            new Node({Identifier, "i"}),
+                            new Node({Number, "10"}),
+                            {Less, "<"}),
+                    new UnaryExpression(
+                            new Node({Identifier, "i"}),
+                            {Increment, "++"},
+                            UnaryExpression::Side::RIGHT),
+                    new ScopedBody({
+                            new ReturnStatement(new Node({Identifier, "i"})),
+                    })),
+    };
+    auto actualResult = parse(input);
+
+    ASSERT_EQ(expectedResult.size(), actualResult.size());
+    for (int i = 0; i < expectedResult.size(); i++)
+        ASSERT_EQ(*expectedResult[i], *actualResult[i]);
+}
+
 TEST(ParserTests, UnaryExpression) {
     const char *input = "x++;"
                         "x--;"