diff --git a/pystencils/transformations.py b/pystencils/transformations.py index 06005cbf709fdfeb74a14dc141f0256aee6039cb..09c2e0976ddf14e2745a796e060a79e011d8ca67 100644 --- a/pystencils/transformations.py +++ b/pystencils/transformations.py @@ -638,6 +638,7 @@ def move_constants_before_loop(ast_node): new_symbol = TypedSymbol(sp.Dummy().name, child.lhs.dtype) target.insert_before(ast.SympyAssignment(new_symbol, child.rhs, is_const=child.is_const), child_to_insert_before) + block.append(ast.SympyAssignment(child.lhs, new_symbol, is_const=child.is_const)) def split_inner_loop(ast_node: ast.Node, symbol_groups): diff --git a/pystencils_tests/test_move_constant_before_loop.py b/pystencils_tests/test_move_constant_before_loop.py index bec180deea5a82c3123f9bf14cd6045dbccfe91b..ea736dd183459a896ccf5d86662386a1a396c85c 100644 --- a/pystencils_tests/test_move_constant_before_loop.py +++ b/pystencils_tests/test_move_constant_before_loop.py @@ -25,7 +25,9 @@ def test_symbol_renaming(): loops = block.atoms(LoopOverCoordinate) assert len(loops) == 2 + assert len(block.args[2].body.args) == 1 + assert len(block.args[3].body.args) == 2 + for loop in loops: - assert len(loop.body.args) == 1 assert len(loop.parent.args) == 4 # 2 loops + 2 subexpressions assert loop.parent.args[0].lhs.name != loop.parent.args[1].lhs.name