Implementation of bfs_nodes now uses bfs_edges

dace/sdfg/graph.py

@@ -6,7 +6,7 @@
 import networkx as nx
 from dace.dtypes import deduplicate
 import dace.serialize
-from typing import Any, Callable, Generic, Iterable, List, Sequence, TypeVar, Union
+from typing import Any, Callable, Generic, Iterable, Iterator, List, Optional, Sequence, TypeVar, Union
 
 
 class NodeNotFoundError(Exception):
@@ -364,28 +364,21 @@ def sink_nodes(self) -> List[NodeT]:
         """Returns nodes with no outgoing edges."""
         return [n for n in self.nodes() if self.out_degree(n) == 0]
 
-    def bfs_nodes(self, source: NodeT = None) -> Sequence[NodeT]:
+    def bfs_nodes(self, source: Optional[NodeT] = None) -> Iterator[NodeT]:
         """Returns nodes in topological order iff the graph contains exactly
         one node with no incoming edges."""
-        if source is not None:
-            sources = [source]
-        else:
+        if source is None:
             sources = self.source_nodes()
-            if len(sources) == 0:
-                sources = [self.nodes()[0]]
-            if len(sources) > 1:
-                sources = [self.nodes()[0]]
-        seen = OrderedDict()  # No OrderedSet in Python
-        queue = deque(sources)
-        while len(queue) > 0:
-            node = queue.popleft()
-            seen[node] = None
-            for e in self.out_edges(node):
-                succ = e.dst
-                if succ not in seen:
-                    seen[succ] = None
-                    queue.append(succ)
-        return seen.keys()
+            if len(sources) == 1:
+                source = sources[0]
+            else:
+                source = next(iter(self.nodes()), None)
+                if source is None:
+                    return # graph has no nodes
+
+        yield source # needed if `source` has no outgoing edges
+        for edge in self.bfs_edges(source):
+            yield edge.dst
 
     def all_simple_paths(self,
                          source_node: NodeT,