diff --git a/src/flowrep/parsers/for_parser.py b/src/flowrep/parsers/for_parser.py index ccb2e78b..6362f1b8 100644 --- a/src/flowrep/parsers/for_parser.py +++ b/src/flowrep/parsers/for_parser.py @@ -58,7 +58,10 @@ def parse_for_node( # axis has no parent symbol to remap, so its variable is added fresh instead. body_symbol_map = walker.symbol_map.fork( {axis.port: axis.variable for axis in all_iters if axis.binding is None}, - added_symbols=[axis.variable for axis in all_iters if axis.binding is not None], + added_symbols=[ + axis.variable if axis.binding is not None else axis.port + for axis in all_iters + ], available_accumulators=walker.symbol_map.declared_accumulators.copy(), ) @@ -161,7 +164,7 @@ def _wire_inputs( and s not in {axis.variable for axis in all_iters} ] # Need to keep it consistently ordered, so don't use a simple set op scattered_symbols = [axis.port for axis in all_iters] - inputs = broadcast_symbols + scattered_symbols + inputs = list(dict.fromkeys(broadcast_symbols + scattered_symbols)) broadcast_inputs = { edge_models.TargetHandle( node=FOR_BODY_LABEL, port=port diff --git a/tests/integration/parsers/test_parsing_for_nodes.py b/tests/integration/parsers/test_parsing_for_nodes.py index ece4a652..f1431b60 100644 --- a/tests/integration/parsers/test_parsing_for_nodes.py +++ b/tests/integration/parsers/test_parsing_for_nodes.py @@ -446,5 +446,28 @@ def test_against_static_recipes(self): ) +def _nested_loops_with_inner_collection(xs, offset): + outer = [] + for x in xs: + inner = [] + for y in xs: + d = std.add(y, offset) + inner.append(d) + total = std.add(x, offset) + outer.append(total) + return outer + + +class TestNestedLoops(unittest.TestCase): + def test_loop_in_loop_over_same_collection(self): + node = workflow_parser.parse_workflow(_nested_loops_with_inner_collection) + outer = node.nodes["for_each_0"] + self.assertEqual(outer.inputs, ["offset", "xs"]) + self.assertEqual(outer.nested_ports, ["x"]) + inner = outer.body_node.recipe.nodes["for_each_0"] + self.assertEqual(inner.nested_ports, ["y"]) + self.assertEqual(inner.inputs, ["offset", "xs"]) + + if __name__ == "__main__": unittest.main()