From b2d9c35d517399f4fbedd858900de478cd360ab4 Mon Sep 17 00:00:00 2001 From: Roberto Raggi Date: Sat, 9 Aug 2025 12:56:10 +0200 Subject: [PATCH] Fix codegen of for statements Signed-off-by: Roberto Raggi --- src/mlir/cxx/mlir/codegen_statements.cc | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/mlir/cxx/mlir/codegen_statements.cc b/src/mlir/cxx/mlir/codegen_statements.cc index f40f2f28..4ffe6738 100644 --- a/src/mlir/cxx/mlir/codegen_statements.cc +++ b/src/mlir/cxx/mlir/codegen_statements.cc @@ -188,15 +188,21 @@ void Codegen::StatementVisitor::operator()(WhileStatementAST* ast) { void Codegen::StatementVisitor::operator()(DoStatementAST* ast) { auto loopBlock = gen.newBlock(); + auto conditionBlock = gen.newBlock(); auto endLoopBlock = gen.newBlock(); - Loop loop{loopBlock, endLoopBlock}; + Loop loop{conditionBlock, endLoopBlock}; std::swap(gen.loop_, loop); gen.branch(gen.getLocation(ast->statement->firstSourceLocation()), loopBlock); gen.builder_.setInsertionPointToEnd(loopBlock); gen.statement(ast->statement); + + gen.branch(gen.getLocation(ast->statement->lastSourceLocation()), + conditionBlock); + + gen.builder_.setInsertionPointToEnd(conditionBlock); gen.condition(ast->expression, loopBlock, endLoopBlock); gen.builder_.setInsertionPointToEnd(endLoopBlock); @@ -226,13 +232,14 @@ void Codegen::StatementVisitor::operator()(ForStatementAST* ast) { Loop loop{stepLoopBlock, endLoopBlock}; std::swap(gen.loop_, loop); - gen.branch( - gen.getLocation(ast->condition ? ast->condition->firstSourceLocation() - : ast->semicolonLoc), - beginLoopBlock); - + gen.branch(gen.getLocation(ast->firstSourceLocation()), beginLoopBlock); gen.builder_.setInsertionPointToEnd(beginLoopBlock); - gen.condition(ast->condition, loopBodyBlock, endLoopBlock); + + if (ast->condition) { + gen.condition(ast->condition, loopBodyBlock, endLoopBlock); + } else { + gen.branch(gen.getLocation(ast->semicolonLoc), loopBodyBlock); + } gen.builder_.setInsertionPointToEnd(loopBodyBlock); gen.statement(ast->statement);