Loops (#48)

* add grammar * add implementation * add typecheck and tests * add proper scoping * update readme
jackparsonss · Jul 5, 2024 · 8d4ab2e · 8d4ab2e
1 parent e14d21f
commit 8d4ab2e
Show file tree

Hide file tree

Showing 20 changed files with 152 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -118,3 +118,15 @@ if(x == 4){
 }
 ```
 Conditions **must** be booleans
+
+#### Loops
+Similar to c-style for loops, they are composed of 4 parts
+- Variable declaration
+- Loop condition
+- Variable assignment
+- Body
+```
+for(let i: i32 = 0; i < 5; i = i + 1){
+    println(i);
+}
+```
diff --git a/grammar/Fusion.g4 b/grammar/Fusion.g4
@@ -6,14 +6,15 @@ statement
     : function
     | call SEMI
     | if
-    | declaration
-    | assignment
+    | loop
+    | declaration SEMI
+    | assignment SEMI
     | block
     | return
     ;
 
 declaration:
-    variable EQUAL expr SEMI
+    variable EQUAL expr
     ;
 
 block: L_CURLY statement* R_CURLY;
@@ -22,12 +23,14 @@ function: FUNCTION ID L_PAREN variable? (COMMA variable)* R_PAREN COLON type blo
 
 variable: qualifier ID COLON type;
 
-assignment: ID EQUAL expr SEMI;
+assignment: ID EQUAL expr;
 
 call: ID L_PAREN expr? (COMMA expr)* R_PAREN;
 
 if: IF L_PAREN expr R_PAREN block else?;
 
+loop: FOR L_PAREN declaration SEMI expr SEMI assignment R_PAREN block;
+
 else: ELSE (block | if);
 
 return: RETURN expr SEMI;
@@ -63,6 +66,7 @@ CONST: 'const';
 LET: 'let';
 IF: 'if';
 ELSE: 'else';
+FOR: 'for';
 
 // symbols
 SEMI: ';';

diff --git a/include/ast/ast.h b/include/ast/ast.h
@@ -247,4 +247,19 @@ class Conditional : public Node {
 
     void xml(int level) override;
 };
+
+class Loop : public Node {
+   public:
+    shared_ptr<Declaration> variable;
+    shared_ptr<Expression> condition;
+    shared_ptr<Assignment> assignment;
+    shared_ptr<Block> body;
+    Loop(shared_ptr<Declaration> variable,
+         shared_ptr<Expression> condition,
+         shared_ptr<Assignment> assignment,
+         shared_ptr<Block> body,
+         Token* token);
+
+    void xml(int level) override;
+};
 }  // namespace ast
diff --git a/include/ast/builder.h b/include/ast/builder.h
@@ -47,4 +47,5 @@ class Builder : public FusionBaseVisitor {
     std::any visitAssignment(FusionParser::AssignmentContext* ctx) override;
     std::any visitIf(FusionParser::IfContext* ctx) override;
     std::any visitElse(FusionParser::ElseContext* ctx) override;
+    std::any visitLoop(FusionParser::LoopContext* ctx) override;
 };
diff --git a/include/ast/passes/def_ref.h b/include/ast/passes/def_ref.h
@@ -19,4 +19,5 @@ class DefRef : public Pass {
     void visit_function(shared_ptr<ast::Function>) override;
     void visit_variable(shared_ptr<ast::Variable>) override;
     void visit_call(shared_ptr<ast::Call>) override;
+    void visit_loop(shared_ptr<ast::Loop>) override;
 };
diff --git a/include/ast/passes/pass.h b/include/ast/passes/pass.h
@@ -30,6 +30,7 @@ class Pass {
     virtual void visit_binary_operator(shared_ptr<ast::BinaryOperator>);
     virtual void visit_unary_operator(shared_ptr<ast::UnaryOperator>);
     virtual void visit_conditional(shared_ptr<ast::Conditional>);
+    virtual void visit_loop(shared_ptr<ast::Loop>);
 
     static void run_passes(shared_ptr<ast::Block> ast, shared_ptr<SymbolTable>);
 };
diff --git a/include/ast/passes/type_check.h b/include/ast/passes/type_check.h
@@ -21,4 +21,5 @@ class TypeCheck : public Pass {
     void visit_binary_operator(shared_ptr<ast::BinaryOperator>) override;
     void visit_unary_operator(shared_ptr<ast::UnaryOperator>) override;
     void visit_conditional(shared_ptr<ast::Conditional>) override;
+    void visit_loop(shared_ptr<ast::Loop>) override;
 };
diff --git a/include/backend/backend.h b/include/backend/backend.h
@@ -37,4 +37,5 @@ class Backend {
     mlir::Value visit_binary_operator(shared_ptr<ast::BinaryOperator>);
     mlir::Value visit_unary_operator(shared_ptr<ast::UnaryOperator>);
     mlir::Value visit_conditional(shared_ptr<ast::Conditional>);
+    mlir::Value visit_loop(shared_ptr<ast::Loop>);
 };
diff --git a/src/ast/ast.cpp b/src/ast/ast.cpp
@@ -435,3 +435,24 @@ void ast::Conditional::xml(int level) {
 
     std::cout << std::string(level * 4, ' ') << "</conditional>\n";
 }
+
+ast::Loop::Loop(shared_ptr<Declaration> variable,
+                shared_ptr<Expression> condition,
+                shared_ptr<Assignment> assignment,
+                shared_ptr<Block> body,
+                Token* token)
+    : Node(token) {
+    this->variable = variable;
+    this->condition = condition;
+    this->assignment = assignment;
+    this->body = body;
+}
+
+void ast::Loop::xml(int level) {
+    std::cout << std::string(level * 4, ' ') << "<loop>\n";
+    variable->xml(level + 1);
+    condition->xml(level + 1);
+    assignment->xml(level + 1);
+    body->xml(level + 1);
+    std::cout << std::string(level * 4, ' ') << "</loop>\n";
+}
diff --git a/src/ast/builder.cpp b/src/ast/builder.cpp
@@ -68,6 +68,10 @@ std::any Builder::visitStatement(FusionParser::StatementContext* ctx) {
         return visit(ctx->if_());
     }
 
+    if (ctx->loop() != nullptr) {
+        return visit(ctx->loop());
+    }
+
     throw std::runtime_error("found an invalid statement");
 }
 
@@ -410,3 +414,16 @@ std::any Builder::visitElse(FusionParser::ElseContext* ctx) {
     auto node = make_shared<ast::Conditional>(condition, block, token);
     return to_node(node);
 }
+
+std::any Builder::visitLoop(FusionParser::LoopContext* ctx) {
+    Token* token = ctx->FOR()->getSymbol();
+
+    auto variable = cast_node(ast::Declaration, visit(ctx->declaration()));
+    auto condition = cast_node(ast::Expression, visit(ctx->expr()));
+    auto assignment = cast_node(ast::Assignment, visit(ctx->assignment()));
+    auto body = cast_node(ast::Block, visit(ctx->block()));
+
+    auto node =
+        make_shared<ast::Loop>(variable, condition, assignment, body, token);
+    return to_node(node);
+}
diff --git a/src/ast/passes/def_ref.cpp b/src/ast/passes/def_ref.cpp
@@ -118,6 +118,15 @@ void DefRef::visit_call(shared_ptr<ast::Call> node) {
     node->set_function(vs->function);
 }
 
+void DefRef::visit_loop(shared_ptr<ast::Loop> node) {
+    symbol_table->push();
+    visit(node->variable);
+    visit(node->condition);
+    visit(node->assignment);
+    visit(node->body);
+    symbol_table->pop();
+}
+
 bool DefRef::is_builtin(std::string name) {
     if (name == "print" || name == "println") {
         return true;

diff --git a/src/ast/passes/pass.cpp b/src/ast/passes/pass.cpp
@@ -52,6 +52,7 @@ void Pass::visit(shared_ptr<ast::Node> node) {
     try_visit(node, ast::BinaryOperator, this->visit_binary_operator);
     try_visit(node, ast::UnaryOperator, this->visit_unary_operator);
     try_visit(node, ast::Conditional, this->visit_conditional);
+    try_visit(node, ast::Loop, this->visit_loop);
 
     throw std::runtime_error("node not added to pass manager");
 }
@@ -113,3 +114,10 @@ void Pass::visit_conditional(shared_ptr<ast::Conditional> node) {
         visit(node->else_if.value());
     }
 }
+
+void Pass::visit_loop(shared_ptr<ast::Loop> node) {
+    visit(node->variable);
+    visit(node->condition);
+    visit(node->assignment);
+    visit(node->body);
+}
diff --git a/src/ast/passes/type_check.cpp b/src/ast/passes/type_check.cpp
@@ -155,6 +155,15 @@ void TypeCheck::visit_conditional(shared_ptr<ast::Conditional> node) {
     }
 }
 
+void TypeCheck::visit_loop(shared_ptr<ast::Loop> node) {
+    visit(node->variable);
+    visit(node->condition);
+    check_bool(node->condition->get_type(), node->token->getLine());
+
+    visit(node->assignment);
+    visit(node->body);
+}
+
 void TypeCheck::check_numeric(TypePtr type, size_t line) {
     if (!type->is_numeric()) {
         throw TypeError(line, "type(" + type->get_name() + ") is not numeric");

diff --git a/src/backend/visitor.cpp b/src/backend/visitor.cpp
@@ -34,6 +34,7 @@ mlir::Value Backend::visit(shared_ptr<ast::Node> node) {
     try_visit(node, ast::BinaryOperator, this->visit_binary_operator);
     try_visit(node, ast::UnaryOperator, this->visit_unary_operator);
     try_visit(node, ast::Conditional, this->visit_conditional);
+    try_visit(node, ast::Loop, this->visit_loop);
 
     throw std::runtime_error("node not added to backend visit function");
 }
@@ -183,3 +184,25 @@ mlir::Value Backend::visit_conditional(shared_ptr<ast::Conditional> node) {
 
     return nullptr;
 }
+
+mlir::Value Backend::visit_loop(shared_ptr<ast::Loop> node) {
+    mlir::Block* b_cond = ctx::current_function().addBlock();
+    mlir::Block* b_loop = ctx::current_function().addBlock();
+    mlir::Block* b_exit = ctx::current_function().addBlock();
+
+    visit(node->variable);
+
+    flow::jump(b_cond);
+    ctx::builder->setInsertionPointToStart(b_cond);
+
+    mlir::Value condition = visit(node->condition);
+    flow::branch(condition, b_loop, b_exit);
+
+    ctx::builder->setInsertionPointToStart(b_loop);
+    visit(node->body);
+    visit(node->assignment);
+    flow::jump(b_cond);
+
+    ctx::builder->setInsertionPointToStart(b_exit);
+    return nullptr;
+}
diff --git a/tests/input/errors/symbol/loop.in b/tests/input/errors/symbol/loop.in
@@ -0,0 +1,6 @@
+fn main(): i32 {
+    for(let i: i32 = 0; i < 5; i = i + 1) {
+    }
+    println(i);
+    return 0;
+}
diff --git a/tests/input/errors/type/loop.in b/tests/input/errors/type/loop.in
@@ -0,0 +1,6 @@
+fn main(): i32 {
+    for(let i: i32 = 0; i; i = i + 1) {
+        println(5);
+    }
+    return 0;
+}
diff --git a/tests/input/loop/basic.in b/tests/input/loop/basic.in
@@ -0,0 +1,6 @@
+fn main(): i32 {
+    for(let i: i32 = 0; i < 5; i = i + 1) {
+        println(5);
+    }
+    return 0;
+}
diff --git a/tests/output/errors/symbol/loop.out b/tests/output/errors/symbol/loop.out
@@ -0,0 +1 @@
+SymbolError on Line 4
diff --git a/tests/output/errors/type/loop.out b/tests/output/errors/type/loop.out
@@ -0,0 +1 @@
+TypeError on Line 2
diff --git a/tests/output/loop/basic.out b/tests/output/loop/basic.out
@@ -0,0 +1,5 @@
+5
+5
+5
+5
+5