Fixed various bugs in python -> model conversion.

zrho · zrho · commit 49fd26634b4c · 2025-03-19T18:26:24.000Z
diff --git a/hugr-model/src/v0/ast/python.rs b/hugr-model/src/v0/ast/python.rs
@@ -39,7 +39,12 @@ impl<'py> pyo3::FromPyObject<'py> for Term {
                 let region = term.getattr("region")?.extract()?;
                 Self::Func(Arc::new(region))
             }
-            _ => return Err(PyTypeError::new_err("Unknown Term type.")),
+            _ => {
+                return Err(PyTypeError::new_err(format!(
+                    "Unknown Term type: {}.",
+                    name.to_str()?
+                )))
+            }
         })
     }
 }
diff --git a/hugr-py/src/hugr/_hugr.cpython-312-darwin.so b/hugr-py/src/hugr/_hugr.cpython-312-darwin.so
diff --git a/hugr-py/src/hugr/hugr/base.py b/hugr-py/src/hugr/hugr/base.py
@@ -22,7 +22,6 @@
 from hugr.utils import BiMap
 from hugr.val import Value
 import hugr.model as model
-import hugr.model.export as model_export
 
 from .node_port import (
     Direction,
@@ -732,7 +731,8 @@ def to_model(self) -> model.Module:
         return model.Module(self.to_model_region())
 
     def to_model_region(self) -> model.Region:
-        export = model_export.Export(self)
+        from hugr.model.export import ModelExport
+        export = ModelExport(self)
         return export.export_region_module(self.root)
 
     @classmethod
diff --git a/hugr-py/src/hugr/model/export.py b/hugr-py/src/hugr/model/export.py
@@ -1,36 +1,43 @@
 from typing_extensions import Optional
 import hugr.model as model
 from hugr.hugr.base import Hugr, Node
-from hugr.ops import DFG, Input, Output, Custom, ExtOp, Conditional, TailLoop, FuncDefn, FuncDecl, Call, CallIndirect, LoadFunc, AliasDecl, AliasDefn, LoadConst, Const, CFG, ExitBlock, DataflowBlock
+from hugr.ops import DFG, Input, Output, Custom, AsExtOp, Conditional, TailLoop, FuncDefn, FuncDecl, Call, CallIndirect, LoadFunc, AliasDecl, AliasDefn, LoadConst, Const, CFG, ExitBlock, DataflowBlock, Tag
 from hugr.hugr.node_port import InPort, OutPort
 from hugr.tys import FunctionKind, ConstKind, TypeParam, Type, TypeTypeParam, TypeBound
 from typing import cast, Sequence
 
-class Export:
+class ModelExport:
     def __init__(self, hugr: Hugr):
         self.hugr = hugr
         self.link_ports: UnionFind[InPort | OutPort] = UnionFind()
         self.link_names: dict[InPort | OutPort, str] = {}
 
         for (a, b) in self.hugr.links():
+            print(f"Link {a} and {b}")
             self.link_ports.union(a, b)
 
     def link_name(self, port):
         root = self.link_ports[port]
+        # print(f"Port {port} -> {root}")
 
         if root in self.link_names:
+            # print(f"Known name {self.link_names[root]} for {root}")
+            return self.link_names[root]
+        else:
             index = str(len(self.link_names))
+            # print(f"New name {index} for {root}")
             self.link_names[root] = index
             return index
-        else:
-            return self.link_names[root]
 
     def export_node(self, node: Node) -> Optional[model.Node]:
         node_data = self.hugr[node]
 
         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)]
 
+        print(f"node {node} {inputs} {outputs}")
+        print(node_data.op)
+
         match node_data.op:
             case DFG() as op:
                 region = self.export_region_dfg(node)
@@ -55,9 +62,9 @@ def export_node(self, node: Node) -> Optional[model.Node]:
                     outputs = outputs
                 )
 
-            case ExtOp() as op:
+            case AsExtOp() as op:
                 name = op.op_def().qualified_name()
-                args = cast(list[model.Term], [arg.to_model() for arg in op.args])
+                args = cast(list[model.Term], [arg.to_model() for arg in op.type_args()])
                 signature = op.outer_signature().to_model()
 
                 return model.Node(
@@ -241,6 +248,62 @@ def export_node(self, node: Node) -> Optional[model.Node]:
                     regions = [region]
                 )
 
+            case DataflowBlock() as op:
+                region = self.export_region_dfg(node)
+
+                input_types = model.List([
+                    model.Apply("core.ctrl", [
+                        model.List([ type.to_model() for type in op.inputs ])
+                    ])
+                ])
+
+                other_output_types = [ type.to_model() for type in op.other_outputs ]
+                output_types = model.List([
+                    model.Apply("core.ctrl", [model.List([
+                        *[type.to_model() for type in row],
+                        *other_output_types
+                    ])])
+                    for row in op.sum_ty.variant_rows
+                ])
+
+                signature = model.Apply("core.fn", [input_types, output_types, model.ExtSet()])
+
+                return model.Node(
+                    operation = model.Block(),
+                    inputs = inputs,
+                    outputs = outputs,
+                    regions = [region],
+                    signature = signature
+                )
+
+            case Tag() as op:
+                variants = model.List([
+                    model.List([type.to_model() for type in row])
+                    for row in op.sum_ty.variant_rows
+                ])
+
+                types = model.List([
+                    type.to_model()
+                    for type in op.sum_ty.variant_rows[op.tag]
+                ])
+
+                tag = model.Literal(op.tag)
+                signature = op.outer_signature().to_model()
+
+                return model.Node(
+                    operation = model.CustomOp(model.Apply("core.make_adt", [
+                        variants,
+                        types,
+                        tag
+                    ])),
+                    inputs = inputs,
+                    outputs = outputs,
+                    signature = signature
+                )
+
+            case op:
+                raise ValueError(f"Unknown operation: {op}")
+
     def export_region_module(self, node: Node) -> model.Region:
         node_data = self.hugr[node]
         children = []
@@ -368,13 +431,13 @@ def find_func_input(self, node: Node) -> Optional[str]:
             func_node = next(
                 out_port.node
                 for (in_port, out_ports) in self.hugr.incoming_links(node)
-                if isinstance(self.hugr.port_kind, FunctionKind)
+                if isinstance(self.hugr.port_kind(in_port), FunctionKind)
                 for out_port in out_ports
             )
         except StopIteration:
             return None
 
-        match self.hugr[func_node]:
+        match self.hugr[func_node].op:
             case FuncDecl() as func_op:
                 return func_op.f_name
             case FuncDefn() as func_op:
@@ -387,16 +450,16 @@ def find_const_input(self, node: Node) -> Optional[model.Term]:
             const_node = next(
                 out_port.node
                 for (in_port, out_ports) in self.hugr.incoming_links(node)
-                if isinstance(self.hugr.port_kind, ConstKind)
+                if isinstance(self.hugr.port_kind(in_port), ConstKind)
                 for out_port in out_ports
             )
         except StopIteration:
             return None
 
-        match self.hugr[const_node]:
+        match self.hugr[const_node].op:
             case Const() as op:
                 return op.val.to_model()
-            case _:
+            case op:
                 return None
 
 class UnionFind[T]:
diff --git a/hugr-py/src/hugr/std/int.py b/hugr-py/src/hugr/std/int.py
@@ -10,6 +10,7 @@
 from hugr import ext, tys, val
 from hugr.ops import AsExtOp, DataflowOp, ExtOp, RegisteredOp
 from hugr.std import _load_extension
+import hugr.model as model
 
 if TYPE_CHECKING:
     from hugr.ops import Command, ComWire
@@ -71,6 +72,11 @@ def to_value(self) -> val.Extension:
     def __str__(self) -> str:
         return f"{self.v}"
 
+    def to_model(self) -> model.Term:
+        return model.Apply("arithmetic.int.const", [
+            model.Literal(self.width),
+            model.Literal(self.v)
+        ])
 
 INT_OPS_EXTENSION = _load_extension("arithmetic.int")
 
diff --git a/hugr-py/src/hugr/tys.py b/hugr-py/src/hugr/tys.py
@@ -274,8 +274,8 @@ def __str__(self) -> str:
         return f"({comma_sep_str(self.elems)})"
 
     def to_model(self) -> model.Term:
-        # TODO: Should this be a tuple or a list?
-        ...
+        # TODO: We should separate lists and tuples. For now we assume that this is a list.
+        return model.List([elem.to_model() for elem in self.elems])
 
 
 @dataclass(frozen=True)
@@ -686,7 +686,11 @@ def to_model(self) -> model.Term:
         # actual type or a row variable.
         args = cast(list[model.Term], args)
 
-        return model.Apply(self.type_def.name, args)
+        extension_name = self.type_def.get_extension().name
+        type_name = self.type_def.name
+        name = f"{extension_name}.{type_name}"
+
+        return model.Apply(name, args)
 
 
 def _type_str(name: str, args: Sequence[TypeArg]) -> str:
diff --git a/hugr-py/src/hugr/val.py b/hugr-py/src/hugr/val.py
@@ -342,3 +342,7 @@ def type_(self) -> tys.Type:
 
     def _to_serial(self) -> sops.CustomValue:
         return self.to_value()._to_serial()
+
+    def to_model(self) -> model.Term:
+        # Fallback
+        return self.to_value().to_model()

Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,12 @@ impl<'py> pyo3::FromPyObject<'py> for Term {`
`39`	`39`	`let region = term.getattr("region")?.extract()?;`
`40`	`40`	`Self::Func(Arc::new(region))`
`41`	`41`	`}`
`42`		`- _ => return Err(PyTypeError::new_err("Unknown Term type.")),`
	`42`	`+ _ => {`
	`43`	`+ return Err(PyTypeError::new_err(format!(`
	`44`	`+ "Unknown Term type: {}.",`
	`45`	`+ name.to_str()?`
	`46`	`+ )))`
	`47`	`+ }`
`43`	`48`	`})`
`44`	`49`	`}`
`45`	`50`	`}`