Export order edges via metadata in Python.

Author: zrho

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

@@ -1,5 +1,6 @@
 """Helpers to export hugr graphs from their python representation to hugr model."""
 
+import json
 from collections.abc import Sequence
 from typing import Generic, TypeVar, cast
 
@@ -59,6 +60,22 @@ def export_node(self, node: Node) -> model.Node | None:
 
         inputs = [self.link_name(InPort(node, i)) for i in range(node_data._num_inps)]
         outputs = [self.link_name(OutPort(node, i)) for i in range(node_data._num_outs)]
+        meta = []
+
+        # Export JSON metadata
+        for meta_name, meta_value in node_data.metadata.items():
+            # TODO: Is this the correct way to convert the metadata as JSON?
+            meta_json = json.dumps(meta_value)
+            meta.append(
+                model.Apply(
+                    "compat.meta_json",
+                    [model.Literal(meta_name), model.Literal(meta_json)],
+                )
+            )
+
+        # Add an order hint key to the node if necessary
+        if _needs_order_key(self.hugr, node):
+            meta.append(model.Apply("order_hint.key", [model.Literal(node.idx)]))
 
         match node_data.op:
             case DFG() as op:
@@ -70,6 +87,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=op.outer_signature().to_model(),
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case Custom() as op:
@@ -82,6 +100,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case AsExtOp() as op:
@@ -96,6 +115,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case Conditional() as op:
@@ -111,6 +131,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case TailLoop() as op:
@@ -122,6 +143,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case FuncDefn() as op:
@@ -132,30 +154,29 @@ def export_node(self, node: Node) -> model.Node | None:
                 region = self.export_region_dfg(node)
 
                 return model.Node(
-                    operation=model.DefineFunc(symbol),
-                    regions=[region],
+                    operation=model.DefineFunc(symbol), regions=[region], meta=meta
                 )
 
             case FuncDecl() as op:
                 name = _mangle_name(node, op.f_name)
                 symbol = self.export_symbol(
                     name, op.signature.params, op.signature.body
                 )
-                return model.Node(
-                    operation=model.DeclareFunc(symbol),
-                )
+                return model.Node(operation=model.DeclareFunc(symbol), meta=meta)
 
             case AliasDecl() as op:
                 symbol = model.Symbol(name=op.alias, signature=model.Apply("core.type"))
 
-                return model.Node(operation=model.DeclareAlias(symbol))
+                return model.Node(operation=model.DeclareAlias(symbol), meta=meta)
 
             case AliasDefn() as op:
                 symbol = model.Symbol(name=op.alias, signature=model.Apply("core.type"))
 
                 alias_value = cast(model.Term, op.definition.to_model())
 
-                return model.Node(operation=model.DefineAlias(symbol, alias_value))
+                return model.Node(
+                    operation=model.DefineAlias(symbol, alias_value), meta=meta
+                )
 
             case Call() as op:
                 input_types = [type.to_model() for type in op.instantiation.input]
@@ -187,6 +208,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case LoadFunc() as op:
@@ -211,6 +233,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case CallIndirect() as op:
@@ -245,6 +268,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case LoadConst() as op:
@@ -264,6 +288,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
+                    meta=meta,
                 )
 
             case Const() as op:
@@ -273,13 +298,13 @@ def export_node(self, node: Node) -> model.Node | None:
                 signature = op.outer_signature().to_model()
                 region = self.export_region_cfg(node)
 
-                # TODO: Export CFGs
                 return model.Node(
                     operation=model.Cfg(),
                     signature=signature,
                     inputs=inputs,
                     outputs=outputs,
                     regions=[region],
+                    meta=meta,
                 )
 
             case DataflowBlock() as op:
@@ -319,6 +344,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     outputs=outputs,
                     regions=[region],
                     signature=signature,
+                    meta=meta,
                 )
 
             case Tag() as op:
@@ -343,6 +369,7 @@ def export_node(self, node: Node) -> model.Node | None:
                     inputs=inputs,
                     outputs=outputs,
                     signature=signature,
+                    meta=meta,
                 )
 
             case op:
@@ -370,6 +397,7 @@ def export_region_dfg(self, node: Node) -> model.Region:
         target_types: model.Term = model.Wildcard()
         sources = []
         targets = []
+        meta = []
 
         for child in node_data.children:
             child_data = self.hugr[child]
@@ -392,8 +420,19 @@ def export_region_dfg(self, node: Node) -> model.Region:
                 case _:
                     child_node = self.export_node(child)
 
-                    if child_node is not None:
-                        children.append(child_node)
+                    if child_node is None:
+                        continue
+
+                    children.append(child_node)
+
+                    meta += [
+                        model.Apply(
+                            "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)
+                    ]
 
         signature = model.Apply("core.fn", [source_types, target_types, model.ExtSet()])
 
@@ -564,3 +603,21 @@ def union(self, a: T, b: T):
 
         self.parents[b] = a
         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
+
+    return False