feat: automatically add state order edges for inter-graph edges

hugr-py/src/hugr/_dfg.py

@@ -4,6 +4,7 @@
 from ._hugr import Hugr, Node, Wire, OutPort
 
 from ._ops import Op, Command, Input, Output, DFG
+from ._exceptions import NoSiblingAncestor
 from hugr.serialization.tys import FunctionType, Type
 
 
@@ -72,6 +73,11 @@ def add_state_order(self, src: Node, dst: Node) -> None:
     def _wire_up(self, node: Node, ports: Iterable[Wire]):
         for i, p in enumerate(ports):
             src = p.out_port()
+            node_ancestor = _ancestral_sibling(self.hugr, src.node, node)
+            if node_ancestor is None:
+                raise NoSiblingAncestor(src.node.idx, node.idx)
+            if node_ancestor != src.node:
+                self.add_state_order(src.node, node_ancestor)
             self.hugr.add_link(src, node.inp(i))
 
 

hugr-py/src/hugr/_exceptions.py

@@ -4,3 +4,13 @@
 @dataclass
 class ParentBeforeChild(Exception):
     msg: str = "Parent node must be added before child node."
+
+
+@dataclass
+class NoSiblingAncestor(Exception):
+    src: int
+    tgt: int
+
+    @property
+    def msg(self):
+        return f"Source {self.src} has no sibling ancestor of target {self.tgt}, so cannot wire up."

hugr-py/src/hugr/_hugr.py

@@ -292,6 +292,7 @@ def _node_links(
             return
         # iterate over known offsets
         for offset in range(self.num_ports(node, direction)):
+            # TODO should this also look for -1 state order edges?
             port = cast(P, node.port(offset, direction))
             yield port, list(self._linked_ports(port, links))
 

hugr-py/tests/test_hugr_build.py

@@ -3,7 +3,7 @@
 import subprocess
 import os
 import pathlib
-from hugr._hugr import Hugr, Node, Wire
+from hugr._hugr import Hugr, Node, Wire, _SubPort
 from hugr._dfg import Dfg, _ancestral_sibling
 from hugr._ops import Custom, Command
 import hugr._ops as ops
@@ -234,11 +234,11 @@ def test_build_inter_graph():
 
     nt = nested.add(Not(a))
     nested.set_outputs(nt)
-    # TODO a context manager could add this state order edge on
-    # exit by tracking parents of source nodes
-    h.add_state_order(h.input_node, nested.root)
+
     h.set_outputs(nested.root, b)
 
+    assert _SubPort(h.input_node.out(-1)) in h.hugr._links
+    assert h.hugr.num_outgoing(h.input_node) == 2  # doesn't count state order
     _validate(h.hugr, True)