diff --git a/hugr-py/src/hugr/_dfg.py b/hugr-py/src/hugr/_dfg.py index 803427162..bf19c3856 100644 --- a/hugr-py/src/hugr/_dfg.py +++ b/hugr-py/src/hugr/_dfg.py @@ -67,10 +67,7 @@ def set_outputs(self, *args: Wire) -> None: def add_state_order(self, src: Node, dst: Node) -> None: # adds edge to the right of all existing edges - # breaks if further edges are added - self.hugr.add_link( - src.out(self.hugr.num_outgoing(src)), dst.inp(self.hugr.num_incoming(dst)) - ) + self.hugr.add_link(src.out(-1), dst.inp(-1)) def _wire_up(self, node: Node, ports: Iterable[Wire]): for i, p in enumerate(ports): diff --git a/hugr-py/src/hugr/_hugr.py b/hugr-py/src/hugr/_hugr.py index d84dee5d4..e2443a22a 100644 --- a/hugr-py/src/hugr/_hugr.py +++ b/hugr-py/src/hugr/_hugr.py @@ -345,19 +345,35 @@ def add_dfg(self, input_types: Sequence[Type], output_types: Sequence[Type]) -> def to_serial(self) -> SerialHugr: node_it = (node for node in self._nodes if node is not None) + + def _serialise_link( + link: tuple[_SO, _SI], + ) -> tuple[tuple[int, int], tuple[int, int]]: + src, dst = link + s, d = self._constrain_offset(src.port), self._constrain_offset(dst.port) + return (src.port.node.idx, s), (dst.port.node.idx, d) + return SerialHugr( version="v1", # non contiguous indices will be erased nodes=[node.to_serial(Node(idx), self) for idx, node in enumerate(node_it)], - edges=[ - ( - (src.port.node.idx, src.port.offset), - (dst.port.node.idx, dst.port.offset), - ) - for src, dst in self._links.items() - ], + edges=[_serialise_link(link) for link in self._links.items()], ) + def _constrain_offset(self, p: P) -> int: + # negative offsets are used to refer to the last port + if p.offset < 0: + match p.direction: + case Direction.INCOMING: + current = self.num_incoming(p.node) + case Direction.OUTGOING: + current = self.num_outgoing(p.node) + offset = current + p.offset + 1 + else: + offset = p.offset + + return offset + @classmethod def from_serial(cls, serial: SerialHugr) -> Hugr: assert serial.nodes, "Empty Hugr is invalid" diff --git a/hugr-py/tests/test_hugr_build.py b/hugr-py/tests/test_hugr_build.py index f8209a266..1e4213a6c 100644 --- a/hugr-py/tests/test_hugr_build.py +++ b/hugr-py/tests/test_hugr_build.py @@ -228,8 +228,8 @@ def _nested_nop(dfg: Dfg): def test_build_inter_graph(): - h = Dfg.endo([BOOL_T]) - (a,) = h.inputs() + h = Dfg.endo([BOOL_T, BOOL_T]) + (a, b) = h.inputs() nested = h.add_nested([], [BOOL_T]) nt = nested.add(Not(a)) @@ -237,9 +237,9 @@ def test_build_inter_graph(): # 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) + h.set_outputs(nested.root, b) - _validate(h.hugr) + _validate(h.hugr, True) def test_ancestral_sibling():