feat: add toposort to HUGR-py

hugr-py/src/hugr/hugr/base.py

@@ -227,6 +227,65 @@ def nodes(self) -> Iterable[tuple[Node, NodeData]]:
         """
         return self.items()
 
+    def sorted_region_nodes(self, parent: Node) -> Iterator[Node]:
+        """Iterator over a topological ordering of all the hugr nodes.
+
+        Note that the sort is performed within a hugr region and non-local
+        edges are ignored.
+
+        Args:
+            parent: The parent node of the region to sort.
+
+        Raises:
+            ValueError: If the region contains a cycle.
+
+        Examples:
+            >>> from hugr.build.tracked_dfg import TrackedDfg
+            >>> from hugr.std.logic import Not
+            >>> dfg = TrackedDfg(tys.Bool)
+            >>> [b] = dfg.track_inputs()
+            >>> for _ in range(6):
+            ...     _= dfg.add(Not(b));
+            >>> dfg.set_tracked_outputs()
+            >>> nodes = list(dfg.hugr)
+            >>> list(dfg.hugr.sorted_region_nodes(nodes[4]))
+            [Node(5), Node(7), Node(8), Node(9), Node(10), Node(11), Node(12), Node(6)]
+        """
+        # A dict to keep track of how many times we see a node.
+        # Store the Nodes with the input degrees as values.
+        # Implementation uses Kahn's algorithm
+        # https://en.wikipedia.org/wiki/Topological_sorting#Kahn's_algorithm
+        visit_dict: dict[Node, int] = {}
+        queue: Queue[Node] = Queue()
+        for node in self.children(parent):
+            incoming = 0
+            for n in self.input_neighbours(node):
+                same_region = self[n].parent == parent
+                # Only update the degree of the node if edge is within the same region.
+                # We do not count non-local edges.
+                if same_region:
+                    incoming += 1
+            if incoming:
+                visit_dict[node] = incoming
+            # If a Node has no dependencies, add it to the queue.
+            else:
+                queue.put(node)
+
+        while not queue.empty():
+            new_node = queue.get()
+            yield new_node
+
+            for neigh in self.output_neighbours(new_node):
+                visit_dict[neigh] -= 1
+                if visit_dict[neigh] == 0:
+                    del visit_dict[neigh]
+                    queue.put(neigh)
+
+        # If our dict is non-empty here then our graph contains a cycle
+        if visit_dict:
+            err = "Graph contains a cycle. No topological ordering exists."
+            raise ValueError(err)
+
     def links(self) -> Iterator[tuple[OutPort, InPort]]:
         """Iterator over all the links in the HUGR.
 

hugr-py/tests/test_hugr_build.py

@@ -10,10 +10,11 @@
 from hugr.hugr import Hugr
 from hugr.hugr.node_port import Node, _SubPort
 from hugr.ops import NoConcreteFunc
+from hugr.package import Package
 from hugr.std.int import INT_T, DivMod, IntVal
 from hugr.std.logic import Not
 
-from .conftest import validate
+from .conftest import QUANTUM_EXT, H, validate
 
 
 def test_stable_indices():
@@ -406,6 +407,46 @@ def test_option() -> None:
     validate(dfg.hugr)
 
 
+# a helper for the toposort tests
+@pytest.fixture
+def simple_fn() -> Function:
+    f = Function("prepare_qubit", [tys.Bool, tys.Qubit])
+    [b, q] = f.inputs()
+
+    h = f.add_op(H, q)
+    q = h.out(0)
+
+    nnot = f.add_op(Not, b)
+
+    f.set_outputs(q, nnot, b)
+    validate(Package([f.hugr], [QUANTUM_EXT]))
+    return f
+
+
+# https://github.com/CQCL/hugr/issues/2350
+def test_toposort(simple_fn: Function) -> None:
+    nodes = list(simple_fn.hugr)
+    func_node = nodes[1]
+
+    sorted_nodes = list(simple_fn.hugr.sorted_region_nodes(func_node))
+    assert set(sorted_nodes) == set(simple_fn.hugr.children(simple_fn))
+    assert sorted_nodes[0] == simple_fn.input_node
+    assert sorted_nodes[-1] == simple_fn.output_node
+
+
+def test_toposort_error(simple_fn: Function) -> None:
+    # Test that we get an error if we toposort an invalid hugr containing a cycle
+    nodes = list(simple_fn.hugr)
+    func_node = nodes[1]
+
+    # Add a loop, invalidating the HUGR
+    simple_fn.hugr.add_link(nodes[4].out_port(), nodes[4].inp(0))
+    with pytest.raises(
+        ValueError, match="Graph contains a cycle. No topological ordering exists."
+    ):
+        list(simple_fn.hugr.sorted_region_nodes(func_node))
+
+
 def test_html_labels(snapshot) -> None:
     """Ensures that HTML-like labels can be processed correctly by both the builder and
     the renderer.