Add unstructured control flow for for-loop, break and continue

src/mlir/cxx/mlir/codegen.h

@@ -263,6 +263,16 @@ class Codegen {
   void branch(mlir::Location loc, mlir::Block* block,
               mlir::ValueRange operands = {});
 
+  struct Loop {
+    mlir::Block* continueBlock = nullptr;
+    mlir::Block* breakBlock = nullptr;
+
+    Loop() = default;
+
+    Loop(mlir::Block* continueBlock, mlir::Block* breakBlock)
+        : continueBlock(continueBlock), breakBlock(breakBlock) {}
+  };
+
   struct UnitVisitor;
   struct DeclarationVisitor;
   struct StatementVisitor;
@@ -295,6 +305,7 @@ class Codegen {
   mlir::cxx::AllocaOp exitValue_;
   std::unordered_map<ClassSymbol*, mlir::Type> classNames_;
   std::unordered_map<Symbol*, mlir::Value> locals_;
+  Loop loop_;
   int count_ = 0;
 };
 

src/mlir/cxx/mlir/codegen_statements.cc

@@ -129,7 +129,9 @@ void Codegen::StatementVisitor::operator()(IfStatementAST* ast) {
   gen.branch(gen.getLocation(ast->statement->lastSourceLocation()), mergeBlock);
   gen.builder_.setInsertionPointToEnd(falseBlock);
   gen.statement(ast->elseStatement);
-  gen.branch(gen.getLocation(ast->elseStatement->lastSourceLocation()),
+  gen.branch(gen.getLocation(ast->elseStatement
+                                 ? ast->elseStatement->lastSourceLocation()
+                                 : ast->elseLoc),
              mergeBlock);
   gen.builder_.setInsertionPointToEnd(mergeBlock);
 }
@@ -158,6 +160,10 @@ void Codegen::StatementVisitor::operator()(WhileStatementAST* ast) {
   auto bodyLoopBlock = gen.newBlock();
   auto endLoopBlock = gen.newBlock();
 
+  Loop loop{beginLoopBlock, endLoopBlock};
+
+  std::swap(gen.loop_, loop);
+
   gen.branch(gen.getLocation(ast->condition->firstSourceLocation()),
              beginLoopBlock);
 
@@ -170,19 +176,26 @@ void Codegen::StatementVisitor::operator()(WhileStatementAST* ast) {
   gen.branch(gen.getLocation(ast->statement->lastSourceLocation()),
              beginLoopBlock);
   gen.builder_.setInsertionPointToEnd(endLoopBlock);
+
+  std::swap(gen.loop_, loop);
 }
 
 void Codegen::StatementVisitor::operator()(DoStatementAST* ast) {
   auto loopBlock = gen.newBlock();
   auto endLoopBlock = gen.newBlock();
 
+  Loop loop{loopBlock, endLoopBlock};
+  std::swap(gen.loop_, loop);
+
   gen.branch(gen.getLocation(ast->statement->firstSourceLocation()), loopBlock);
 
   gen.builder_.setInsertionPointToEnd(loopBlock);
   gen.statement(ast->statement);
   gen.condition(ast->expression, loopBlock, endLoopBlock);
 
   gen.builder_.setInsertionPointToEnd(endLoopBlock);
+
+  std::swap(gen.loop_, loop);
 }
 
 void Codegen::StatementVisitor::operator()(ForRangeStatementAST* ast) {
@@ -197,21 +210,61 @@ void Codegen::StatementVisitor::operator()(ForRangeStatementAST* ast) {
 }
 
 void Codegen::StatementVisitor::operator()(ForStatementAST* ast) {
-  (void)gen.emitTodoStmt(ast->firstSourceLocation(), to_string(ast->kind()));
-
-#if false
   gen.statement(ast->initializer);
-  auto conditionResult = gen.expression(ast->condition);
-  auto expressionResult = gen.expression(ast->expression);
+
+  auto beginLoopBlock = gen.newBlock();
+  auto loopBodyBlock = gen.newBlock();
+  auto stepLoopBlock = gen.newBlock();
+  auto endLoopBlock = gen.newBlock();
+
+  Loop loop{stepLoopBlock, endLoopBlock};
+  std::swap(gen.loop_, loop);
+
+  gen.branch(
+      gen.getLocation(ast->condition ? ast->condition->firstSourceLocation()
+                                     : ast->semicolonLoc),
+      beginLoopBlock);
+
+  gen.builder_.setInsertionPointToEnd(beginLoopBlock);
+  gen.condition(ast->condition, loopBodyBlock, endLoopBlock);
+
+  gen.builder_.setInsertionPointToEnd(loopBodyBlock);
   gen.statement(ast->statement);
-#endif
+
+  gen.branch(gen.getLocation(ast->statement->lastSourceLocation()),
+             stepLoopBlock);
+
+  gen.builder_.setInsertionPointToEnd(stepLoopBlock);
+
+  (void)gen.expression(ast->expression, ExpressionFormat::kSideEffect);
+
+  gen.branch(
+      gen.getLocation(ast->expression ? ast->expression->lastSourceLocation()
+                                      : ast->rparenLoc),
+      beginLoopBlock);
+
+  gen.builder_.setInsertionPointToEnd(endLoopBlock);
+
+  std::swap(gen.loop_, loop);
 }
 
 void Codegen::StatementVisitor::operator()(BreakStatementAST* ast) {
+  if (auto target = gen.loop_.breakBlock) {
+    gen.builder_.create<mlir::cf::BranchOp>(
+        gen.getLocation(ast->firstSourceLocation()), target);
+    return;
+  }
+
   (void)gen.emitTodoStmt(ast->firstSourceLocation(), to_string(ast->kind()));
 }
 
 void Codegen::StatementVisitor::operator()(ContinueStatementAST* ast) {
+  if (auto target = gen.loop_.continueBlock) {
+    gen.builder_.create<mlir::cf::BranchOp>(
+        gen.getLocation(ast->firstSourceLocation()), target);
+    return;
+  }
+
   (void)gen.emitTodoStmt(ast->firstSourceLocation(), to_string(ast->kind()));
 }