fix: use -1 indexing to add state order edges

CQCL · Jun 6, 2024 · 10e8440 · 10e8440
1 parent 8cf18ef
commit 10e8440
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 15 deletions.
diff --git 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
@@ -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
@@ -228,18 +228,18 @@ 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))
     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)
+    h.set_outputs(nested.root, b)
 
-    _validate(h.hugr)
+    _validate(h.hugr, True)
 
 
 def test_ancestral_sibling():