Control Flow (#51)

* implementation

* fix looping with control flow

* update readme

README.md

@@ -130,3 +130,7 @@ for(let i: i32 = 0; i < 5; i = i + 1){
     println(i);
 }
 ```
+
+Loops also allow for control flow:
+- `continue`: goes to the next iteration of the loop
+- `break`: exits the loop early

grammar/Fusion.g4

@@ -11,6 +11,8 @@ statement
     | assignment SEMI
     | block
     | return
+    | CONTINUE SEMI
+    | BREAK SEMI
     ;
 
 declaration:
@@ -67,6 +69,8 @@ LET: 'let';
 IF: 'if';
 ELSE: 'else';
 FOR: 'for';
+BREAK: 'break';
+CONTINUE: 'continue';
 
 // symbols
 SEMI: ';';

include/ast/ast.h

@@ -262,4 +262,16 @@ class Loop : public Node {
 
     void xml(int level) override;
 };
+
+class Continue : public Node {
+   public:
+    Continue(Token* token) : Node(token) {}
+    void xml(int level);
+};
+
+class Break : public Node {
+   public:
+    Break(Token* token) : Node(token) {}
+    void xml(int level);
+};
 }  // namespace ast

include/ast/passes/pass.h

@@ -31,6 +31,8 @@ class Pass {
     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>);
+    virtual void visit_continue(shared_ptr<ast::Continue>);
+    virtual void visit_break(shared_ptr<ast::Break>);
 
     static void run_passes(shared_ptr<ast::Block> ast, shared_ptr<SymbolTable>);
 };

include/backend/backend.h

@@ -2,6 +2,7 @@
 
 #include <memory>
 #include <ostream>
+#include <stack>
 #include <string>
 #include <unordered_map>
 
@@ -14,6 +15,11 @@ using std::shared_ptr;
 class Backend {
    private:
     std::unordered_map<std::string, mlir::Value> variables;
+
+    // used by continue/break
+    std::stack<mlir::Block*> loop_conditions;
+    std::stack<mlir::Block*> loop_exits;
+
     mlir::Value visit(shared_ptr<ast::Node>);
 
    public:
@@ -38,4 +44,6 @@ class Backend {
     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>);
+    mlir::Value visit_continue(shared_ptr<ast::Continue>);
+    mlir::Value visit_break(shared_ptr<ast::Break>);
 };

src/ast/ast.cpp

@@ -456,3 +456,11 @@ void ast::Loop::xml(int level) {
     body->xml(level + 1);
     std::cout << std::string(level * 4, ' ') << "</loop>\n";
 }
+
+void ast::Continue::xml(int level) {
+    std::cout << std::string(level * 4, ' ') << "</continue>\n";
+}
+
+void ast::Break::xml(int level) {
+    std::cout << std::string(level * 4, ' ') << "</break>\n";
+}

src/ast/builder.cpp

@@ -72,6 +72,16 @@ std::any Builder::visitStatement(FusionParser::StatementContext* ctx) {
         return visit(ctx->loop());
     }
 
+    if (ctx->CONTINUE() != nullptr) {
+        auto node = make_shared<ast::Continue>(ctx->CONTINUE()->getSymbol());
+        return to_node(node);
+    }
+
+    if (ctx->BREAK() != nullptr) {
+        auto node = make_shared<ast::Break>(ctx->BREAK()->getSymbol());
+        return to_node(node);
+    }
+
     throw std::runtime_error("found an invalid statement");
 }
 

src/ast/passes/pass.cpp

@@ -53,6 +53,8 @@ void Pass::visit(shared_ptr<ast::Node> node) {
     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);
+    try_visit(node, ast::Continue, this->visit_continue);
+    try_visit(node, ast::Break, this->visit_break);
 
     throw std::runtime_error("node not added to pass manager");
 }
@@ -68,6 +70,8 @@ void Pass::visit_character_literal(shared_ptr<ast::CharacterLiteral> node) {}
 void Pass::visit_boolean_literal(shared_ptr<ast::BooleanLiteral> node) {}
 void Pass::visit_variable(shared_ptr<ast::Variable> node) {}
 void Pass::visit_parameter(shared_ptr<ast::Parameter> node) {}
+void Pass::visit_continue(shared_ptr<ast::Continue> node) {}
+void Pass::visit_break(shared_ptr<ast::Break> node) {}
 
 void Pass::visit_declaration(shared_ptr<ast::Declaration> node) {
     visit(node->var);

src/backend/visitor.cpp

@@ -35,6 +35,8 @@ mlir::Value Backend::visit(shared_ptr<ast::Node> node) {
     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);
+    try_visit(node, ast::Continue, this->visit_continue);
+    try_visit(node, ast::Break, this->visit_break);
 
     throw std::runtime_error("node not added to backend visit function");
 }
@@ -188,8 +190,12 @@ mlir::Value Backend::visit_conditional(shared_ptr<ast::Conditional> node) {
 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_assn = ctx::current_function().addBlock();
     mlir::Block* b_exit = ctx::current_function().addBlock();
 
+    loop_conditions.push(b_assn);
+    loop_exits.push(b_exit);
+
     visit(node->variable);
 
     flow::jump(b_cond);
@@ -200,9 +206,41 @@ mlir::Value Backend::visit_loop(shared_ptr<ast::Loop> node) {
 
     ctx::builder->setInsertionPointToStart(b_loop);
     visit(node->body);
+    flow::jump(b_assn);
+
+    ctx::builder->setInsertionPointToStart(b_assn);
     visit(node->assignment);
     flow::jump(b_cond);
 
     ctx::builder->setInsertionPointToStart(b_exit);
+    loop_conditions.pop();
+    loop_exits.pop();
+
+    return nullptr;
+}
+
+mlir::Value Backend::visit_continue(shared_ptr<ast::Continue> node) {
+    if (loop_conditions.empty()) {
+        throw std::runtime_error("backend found continue outside of a loop");
+    }
+
+    mlir::Block* b_body = ctx::current_function().addBlock();
+    flow::jump(loop_conditions.top());
+
+    ctx::builder->setInsertionPointToStart(b_body);
+
+    return nullptr;
+}
+
+mlir::Value Backend::visit_break(shared_ptr<ast::Break> node) {
+    if (loop_exits.empty()) {
+        throw std::runtime_error("backend found break outside of a loop");
+    }
+
+    mlir::Block* b_body = ctx::current_function().addBlock();
+    flow::jump(loop_exits.top());
+
+    ctx::builder->setInsertionPointToStart(b_body);
+
     return nullptr;
 }

tests/input/loop/flow.in

@@ -0,0 +1,13 @@
+fn main(): i32 {
+    for(let i: i32 = 0; i < 5; i = i + 1) {
+      if(i == 0) {
+          continue;
+      }
+
+      if(i > 2) {
+        break;
+      }
+      println(i);
+    }
+    return 0;
+}

tests/output/loop/flow.out

@@ -0,0 +1,2 @@
+1
+2