Export order hints in Python.

Author: zrho

hugr-core/src/export.rs

@@ -633,7 +633,7 @@ impl<'a> Context<'a> {
                     sources = self.make_ports(child, Direction::Outgoing, input.types.len());
                     input_types = Some(&input.types);
 
-                    if has_order_edges(&self.hugr, child) {
+                    if has_order_edges(self.hugr, child) {
                         let key = self.make_term(model::Literal::Nat(child.index() as u64).into());
                         meta.push(self.make_term_apply(model::ORDER_HINT_INPUT_KEY, &[key]));
                     }
@@ -642,7 +642,7 @@ impl<'a> Context<'a> {
                     targets = self.make_ports(child, Direction::Incoming, output.types.len());
                     output_types = Some(&output.types);
 
-                    if has_order_edges(&self.hugr, child) {
+                    if has_order_edges(self.hugr, child) {
                         let key = self.make_term(model::Literal::Nat(child.index() as u64).into());
                         meta.push(self.make_term_apply(model::ORDER_HINT_OUTPUT_KEY, &[key]));
                     }
@@ -1106,7 +1106,7 @@ impl<'a> Context<'a> {
     }
 
     fn export_node_order_metadata(&mut self, node: Node, meta: &mut Vec<table::TermId>) {
-        if has_order_edges(&self.hugr, node) {
+        if has_order_edges(self.hugr, node) {
             let key = self.make_term(model::Literal::Nat(node.index() as u64).into());
             meta.push(self.make_term_apply(model::ORDER_HINT_KEY, &[key]));
         }

hugr-core/src/import.rs

@@ -654,7 +654,7 @@ impl<'a> Context<'a> {
 
                 // Check that the node has order ports.
                 // NOTE: This assumes that a node has an input order port iff it has an output one.
-                if !child_optype.other_output_port().is_some() {
+                if child_optype.other_output_port().is_none() {
                     return Err(OrderHintError::NoOrderPort(*child_id).into());
                 }
 

hugr-py/src/hugr/model/export.py

@@ -71,7 +71,7 @@ def export_node(
         meta = self.export_json_meta(node)
 
         # Add an order hint key to the node if necessary
-        if _needs_order_key(self.hugr, node):
+        if _has_order_links(self.hugr, node):
             meta.append(model.Apply("core.order_hint.key", [model.Literal(node.idx)]))
 
         match node_data.op:
@@ -411,13 +411,27 @@ def export_region_dfg(self, node: Node) -> model.Region:
                         for i in range(child_data._num_outs)
                     ]
 
+                    if _has_order_links(self.hugr, child):
+                        meta.append(
+                            model.Apply(
+                                "core.order_hint.input_key", [model.Literal(child.idx)]
+                            )
+                        )
+
                 case Output() as op:
                     target_types = model.List([type.to_model() for type in op.types])
                     targets = [
                         self.link_name(InPort(child, i))
                         for i in range(child_data._num_inps)
                     ]
 
+                    if _has_order_links(self.hugr, child):
+                        meta.append(
+                            model.Apply(
+                                "core.order_hint.output_key", [model.Literal(child.idx)]
+                            )
+                        )
+
                 case _:
                     child_node = self.export_node(child)
 
@@ -426,14 +440,13 @@ def export_region_dfg(self, node: Node) -> model.Region:
 
                     children.append(child_node)
 
-                    meta += [
-                        model.Apply(
-                            "core.order_hint.order",
-                            [model.Literal(child.idx), model.Literal(successor.idx)],
-                        )
-                        for successor in self.hugr.outgoing_order_links(child)
-                        if not isinstance(self.hugr[successor].op, Output)
-                    ]
+            meta += [
+                model.Apply(
+                    "core.order_hint.order",
+                    [model.Literal(child.idx), model.Literal(successor.idx)],
+                )
+                for successor in self.hugr.outgoing_order_links(child)
+            ]
 
         signature = model.Apply("core.fn", [source_types, target_types])
 
@@ -618,19 +631,12 @@ def union(self, a: T, b: T):
         self.sizes[a] += self.sizes[b]
 
 
-def _needs_order_key(hugr: Hugr, node: Node) -> bool:
-    """Checks whether the node has any order links for the purposes of
-    exporting order hint metadata. Order links to `Input` or `Output`
-    operations are ignored, since they are not present in the model format.
-    """
-    for succ in hugr.outgoing_order_links(node):
-        succ_op = hugr[succ].op
-        if not isinstance(succ_op, Output):
-            return True
-
-    for pred in hugr.incoming_order_links(node):
-        pred_op = hugr[pred].op
-        if not isinstance(pred_op, Input):
-            return True
+def _has_order_links(hugr: Hugr, node: Node) -> bool:
+    """Checks whether the node has any order links."""
+    for _succ in hugr.outgoing_order_links(node):
+        return True
+
+    for _pred in hugr.incoming_order_links(node):
+        return True
 
     return False