fix(ir): fix bug where iterator was not remembering visited nodes

es-ude · Nov 14, 2024 · 3da0dbc · 3da0dbc
1 parent 1cd2a35
commit 3da0dbc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/elasticai/creator/ir/graph_delegate_test.py b/elasticai/creator/ir/graph_delegate_test.py
@@ -30,10 +30,7 @@ def test_iterating_breadth_first_upwards():
     )
 
     actual = tuple(bfs_iter_up(g.get_predecessors, "5"))
-    assert set(actual[0:2]) == {"3", "6"}
-    assert (set(actual[2:4]) == {"1", "2"} and actual[4] == "4") or (
-        set(actual[3:5]) == {"1", "2"} and actual[2] == "4"
-    )
+    assert actual == ("3", "6", "1", "2", "4", "0")
 
 
 def test_iterating_depth_first_preorder():

diff --git a/elasticai/creator/ir/graph_iterators.py b/elasticai/creator/ir/graph_iterators.py
@@ -27,6 +27,7 @@ def bfs_iter_down(successors: NodeNeighbourFn, start: HashableT) -> Iterator[Has
         for p in successors(current):
             if p not in visited:
                 yield p
+                visited.add(p)
                 visit_next.append(p)
         visited.add(current)