fix(hugr-py): output and delete node issues (#1971)

Closes #1969 
Closes #1970

drive-by add `add_order_link` to `Hugr`

Author: ss2165

hugr-py/src/hugr/build/dfg.py

@@ -510,10 +510,7 @@ def add_state_order(self, src: Node, dst: Node) -> None:
             [Node(2)]
         """
         # adds edge to the right of all existing edges
-        source = src.out(-1)
-        target = dst.inp(-1)
-        if not self.hugr.has_link(source, target):
-            self.hugr.add_link(source, target)
+        self.hugr.add_order_link(src, dst)
 
     def load(
         self, const: ToNode | val.Value, const_parent: ToNode | None = None

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

@@ -174,6 +174,8 @@ def _add_node(
         node = replace(node, _num_out_ports=num_outs, _metadata=node_data.metadata)
         if parent:
             self[parent].children.append(node)
+
+        self._update_node_outs(node, num_outs)
         return node
 
     def _update_node_outs(self, node: Node, num_outs: int | None) -> Node:
@@ -278,10 +280,10 @@ def delete_node(self, node: ToNode) -> NodeData | None:
         parent = self[node].parent
         if parent:
             self[parent].children.remove(node)
-        for offset in range(self.num_in_ports(node)):
-            self._links.delete_right(_SubPort(node.inp(offset)))
-        for offset in range(self.num_out_ports(node)):
-            self._links.delete_left(_SubPort(node.out(offset)))
+        for inp, _ in self.incoming_links(node):
+            self._links.delete_right(_SubPort(inp))
+        for out, _ in self.outgoing_links(node):
+            self._links.delete_left(_SubPort(out))
 
         weight, self._nodes[node.idx] = self._nodes[node.idx], None
 
@@ -344,6 +346,24 @@ def add_link(self, src: OutPort, dst: InPort) -> None:
         self[src.node]._num_outs = max(self[src.node]._num_outs, src.offset + 1)
         self[dst.node]._num_inps = max(self[dst.node]._num_inps, dst.offset + 1)
 
+    def add_order_link(self, src: ToNode, dst: ToNode) -> None:
+        """Add a state order link between two nodes.
+
+        Args:
+            src: Source node.
+            dst: Destination node.
+
+        Examples:
+            >>> df = dfg.Dfg()
+            >>> df.hugr.add_order_link(df.input_node, df.output_node)
+            >>> list(df.hugr.outgoing_order_links(df.input_node))
+            [Node(2)]
+        """
+        source = src.out(-1)
+        target = dst.inp(-1)
+        if not self.has_link(source, target):
+            self.add_link(source, target)
+
     def delete_link(self, src: OutPort, dst: InPort) -> None:
         """Delete a link (edge) between two nodes from the HUGR.
 

hugr-py/tests/test_hugr_build.py

@@ -310,6 +310,20 @@ def test_higher_order() -> None:
     validate(d.hugr)
 
 
+def test_state_order() -> None:
+    mod = Module()
+    f_id = mod.define_function("id", [tys.Bool])
+    f_id.set_outputs(f_id.input_node[0])
+
+    f_main = mod.define_main([tys.Bool])
+    b = f_main.input_node[0]
+    call1 = f_main.call(f_id, b)
+    f_main.add_state_order(call1, f_main.output_node)
+    # implicit discard of bool to test state order port logic
+    f_main.set_outputs()
+    validate(mod.hugr)
+
+
 def test_alias() -> None:
     mod = Module()
     _dfn = mod.add_alias_defn("my_int", INT_T)