feat!: Entrypoints in hugr-py (#2148)

Adds entrypoint definitions to `hugr-py`. Has similar caveats to #2147.

The main issue here was builders that initialized a hugr for a dataflow
op without knowing its output types. The automatic machinery that wraps
the op in a function definition needs to connect the outputs to the new
function's output once we know the types, so I had to add some extra
machinery to `DataflowOp` to signal back to the hugr that it should
connect things.
This feels a bit hacky, perhaps there's a cleaner way to do it...

Blocked by #2147.

BREAKING CHANGE: Hugrs now define an `entrypoint` in addition to a
module root.

Author: aborgna-q

hugr-py/src/hugr/_serialization/serial_hugr.py

@@ -35,6 +35,7 @@ class SerialHugr(ConfiguredBaseModel):
     encoder: str | None = Field(
         default=None, description="The name of the encoder used to generate the Hugr."
     )
+    entrypoint: NodeIdx | None = None
 
     def to_json(self) -> str:
         """Return a JSON representation of the Hugr."""

hugr-py/src/hugr/build/cfg.py

@@ -23,14 +23,14 @@ class Block(DfBase[ops.DataflowBlock]):
     """Builder class for a basic block in a HUGR control flow graph."""
 
     def set_outputs(self, *outputs: Wire) -> None:
-        super().set_outputs(*outputs)
-
         assert len(outputs) > 0
         branching = outputs[0]
         branch_type = self.hugr.port_type(branching.out_port())
         assert isinstance(branch_type, tys.Sum)
         self._set_parent_output_count(len(branch_type.variant_rows))
 
+        super().set_outputs(*outputs)
+
     def set_block_outputs(self, branching: Wire, *other_outputs: Wire) -> None:
         self.set_outputs(branching, *other_outputs)
 
@@ -50,7 +50,7 @@ def _wire_up_port(self, node: Node, offset: PortOffset, p: Wire) -> Type:
             # it does not check for valid dominance between basic blocks
             # that is deferred to full HUGR validation.
             while cfg_node != src_parent:
-                if src_parent is None or src_parent == self.hugr.root:
+                if src_parent is None or src_parent == self.hugr.module_root:
                     raise NotInSameCfg(src.node.idx, node.idx) from e
                 src_parent = self.hugr[src_parent].parent
 
@@ -60,7 +60,7 @@ def _wire_up_port(self, node: Node, offset: PortOffset, p: Wire) -> Type:
 
 @dataclass
 class Cfg(ParentBuilder[ops.CFG], AbstractContextManager):
-    """Builder class for a HUGR control flow graph, with the HUGR root node
+    """Builder class for a HUGR control flow graph, with the HUGR entrypoint node
     being a :class:`CFG <hugr.ops.CFG>`.
 
     Args:
@@ -85,13 +85,17 @@ def __init__(self, *input_types: Type) -> None:
         input_typs = list(input_types)
         root_op = ops.CFG(inputs=input_typs)
         hugr = Hugr(root_op)
-        self._init_impl(hugr, hugr.root, input_typs)
+        self._init_impl(hugr, hugr.entrypoint, input_typs)
 
-    def _init_impl(self: Cfg, hugr: Hugr, root: Node, input_types: TypeRow) -> None:
+    def _init_impl(
+        self: Cfg, hugr: Hugr, entrypoint: Node, input_types: TypeRow
+    ) -> None:
         self.hugr = hugr
-        self.parent_node = root
+        self.parent_node = entrypoint
         # to ensure entry is first child, add a dummy entry at the start
-        self._entry_block = Block.new_nested(ops.DataflowBlock(input_types), hugr, root)
+        self._entry_block = Block.new_nested(
+            ops.DataflowBlock(input_types), hugr, entrypoint
+        )
 
         self.exit = self.hugr.add_node(ops.ExitBlock(), self.parent_node)
 
@@ -128,7 +132,7 @@ def new_nested(
         new = cls.__new__(cls)
         root = hugr.add_node(
             ops.CFG(inputs=input_types),
-            parent or hugr.root,
+            parent or hugr.entrypoint,
         )
         new._init_impl(hugr, root, input_types)
         return new
@@ -261,3 +265,8 @@ def branch_exit(self, src: Wire) -> None:
             self.parent_node = self.hugr._update_node_outs(
                 self.parent_node, len(out_types)
             )
+            if (
+                self.parent_op._entrypoint_requires_wiring
+                and self.hugr.entrypoint == self.parent_node
+            ):
+                self.hugr._connect_df_entrypoint_outputs()

hugr-py/src/hugr/build/cond_loop.py

@@ -112,11 +112,13 @@ class Conditional(ParentBuilder[ops.Conditional], AbstractContextManager):
     def __init__(self, sum_ty: Sum, other_inputs: TypeRow) -> None:
         root_op = ops.Conditional(sum_ty, other_inputs)
         hugr = Hugr(root_op)
-        self._init_impl(hugr, hugr.root, len(sum_ty.variant_rows))
+        self._init_impl(hugr, hugr.entrypoint, len(sum_ty.variant_rows))
 
-    def _init_impl(self: Conditional, hugr: Hugr, root: Node, n_cases: int) -> None:
+    def _init_impl(
+        self: Conditional, hugr: Hugr, entrypoint: Node, n_cases: int
+    ) -> None:
         self.hugr = hugr
-        self.parent_node = root
+        self.parent_node = entrypoint
         self._case_builders = []
 
         for case_id in range(n_cases):
@@ -163,18 +165,18 @@ def new_nested(
             other_inputs: The inputs for the conditional that aren't included in the
                 sum variants. These are passed to all cases.
             hugr: The HUGR instance this Conditional is part of.
-            parent: The parent node for the Conditional: defaults to the root of
+            parent: The parent node for the Conditional: defaults to the entrypoint of
               the HUGR instance.
 
         Returns:
             The new Conditional builder.
         """
         new = cls.__new__(cls)
-        root = hugr.add_node(
+        entrypoint = hugr.add_node(
             ops.Conditional(sum_ty, other_inputs),
-            parent or hugr.root,
+            parent or hugr.entrypoint,
         )
-        new._init_impl(hugr, root, len(sum_ty.variant_rows))
+        new._init_impl(hugr, entrypoint, len(sum_ty.variant_rows))
         return new
 
     def _update_outputs(self, outputs: TypeRow) -> None:
@@ -183,6 +185,11 @@ def _update_outputs(self, outputs: TypeRow) -> None:
             self.parent_node = self.hugr._update_node_outs(
                 self.parent_node, len(outputs)
             )
+            if (
+                self.parent_op._entrypoint_requires_wiring
+                and self.hugr.entrypoint == self.parent_node
+            ):
+                self.hugr._connect_df_entrypoint_outputs()
         else:
             if outputs != self.parent_op._outputs:
                 msg = "Mismatched case outputs."
@@ -239,15 +246,15 @@ def __init__(self, just_inputs: TypeRow, rest: TypeRow) -> None:
         super().__init__(root_op)
 
     def set_outputs(self, *outputs: Wire) -> None:
-        super().set_outputs(*outputs)
-
         assert len(outputs) > 0
         sum_wire = outputs[0]
         sum_type = self.hugr.port_type(sum_wire.out_port())
         assert isinstance(sum_type, Sum)
         assert len(sum_type.variant_rows) == 2
         self._set_parent_output_count(len(sum_type.variant_rows[1]) + len(outputs) - 1)
 
+        super().set_outputs(*outputs)
+
     def set_loop_outputs(self, sum_wire: Wire, *rest: Wire) -> None:
         """Set the outputs of the loop body. The first wire must be the sum type
         that controls loop termination.

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

@@ -30,6 +30,10 @@
 OpVar = TypeVar("OpVar", bound=ops.Op)
 
 
+class DataflowError(Exception):
+    """Error building a :class:`DfBase` dataflow graph."""
+
+
 @dataclass()
 class DefinitionBuilder(Generic[OpVar]):
     """Base class for builders that can define functions, constants, and aliases.
@@ -55,12 +59,12 @@ def define_function(
             output_types: The output types for the function.
                 If not provided, it will be inferred after the function is built.
             type_params: The type parameters for the function, if polymorphic.
-            parent: The parent node of the constant. Defaults to the root node.
+            parent: The parent node of the constant. Defaults to the entrypoint node.
 
         Returns:
             The new function builder.
         """
-        parent_node = parent or self.hugr.root
+        parent_node = parent or self.hugr.entrypoint
         parent_op = ops.FuncDefn(name, input_types, type_params or [])
         func = Function.new_nested(parent_op, self.hugr, parent_node)
         if output_types is not None:
@@ -72,7 +76,7 @@ def add_const(self, value: val.Value, parent: ToNode | None = None) -> Node:
 
         Args:
             value: The constant value to add.
-            parent: The parent node of the constant. Defaults to the root node.
+            parent: The parent node of the constant. Defaults to the entrypoint node.
 
         Returns:
             The node holding the :class:`Const <hugr.ops.Const>` operation.
@@ -83,12 +87,12 @@ def add_const(self, value: val.Value, parent: ToNode | None = None) -> Node:
             >>> dfg.hugr[const_n].op
             Const(TRUE)
         """
-        parent_node = parent or self.hugr.root
+        parent_node = parent or self.hugr.entrypoint
         return self.hugr.add_node(ops.Const(value), parent_node)
 
     def add_alias_defn(self, name: str, ty: Type, parent: ToNode | None = None) -> Node:
         """Add a type alias definition."""
-        parent_node = parent or self.hugr.root
+        parent_node = parent or self.hugr.entrypoint
         return self.hugr.add_node(ops.AliasDefn(name, ty), parent_node)
 
 
@@ -114,7 +118,7 @@ class DfBase(ParentBuilder[DP], DefinitionBuilder, AbstractContextManager):
 
     def __init__(self, parent_op: DP) -> None:
         self.hugr = Hugr(parent_op)
-        self.parent_node = self.hugr.root
+        self.parent_node = self.hugr.entrypoint
         self._init_io_nodes(parent_op)
 
     def _init_io_nodes(self, parent_op: DP):
@@ -141,7 +145,7 @@ def new_nested(
             parent_op: The parent operation of the new dataflow graph.
             hugr: The host HUGR instance to build the dataflow graph in.
             parent: Parent of new dataflow graph's root node: defaults to the
-            host HUGR root.
+            host HUGR entrypoint.
 
         Example:
             >>> hugr = Hugr()
@@ -152,10 +156,44 @@ def new_nested(
         new = cls.__new__(cls)
 
         new.hugr = hugr
-        new.parent_node = hugr.add_node(parent_op, parent or hugr.root)
+        new.parent_node = hugr.add_node(parent_op, parent or hugr.entrypoint)
         new._init_io_nodes(parent_op)
         return new
 
+    @classmethod
+    def _new_existing(cls, hugr: Hugr, root: ToNode | None = None) -> Self:
+        """Start a dataflow graph builder for an existing node.
+
+        Args:
+            hugr: The host HUGR instance to build the dataflow graph in.
+            root: The dataflow graph's root node.
+                Defaults to the host HUGR's entrypoint.
+
+        Example:
+            >>> hugr = Hugr(ops.DFG([]))
+            >>> dfg = Dfg._new_existing(hugr)
+            >>> dfg.parent_node
+            Node(4)
+
+        Raises:
+            :class:`DataflowError` if the `root` operation is not a dataflow
+            parent.
+        """
+        root = root or hugr.entrypoint
+
+        if not ops.is_df_parent_op(hugr[root].op):
+            msg = f"{hugr[root].op} is not a dataflow parent"
+            raise DataflowError(msg)
+
+        new = cls.__new__(cls)
+        new.hugr = hugr
+        new.parent_node = root.to_node()
+        [inp, out] = hugr.children(root)[:2]
+        new.input_node = inp
+        new.output_node = out
+
+        return new
+
     def _input_op(self) -> ops.Input:
         return self.hugr._get_typed_op(self.input_node, ops.Input)
 
@@ -256,7 +294,7 @@ def insert_nested(self, dfg: Dfg, *args: Wire) -> Node:
             args: The input wires to the graph.
 
         Returns:
-            The root node of the inserted graph.
+            The entrypoint node of the inserted graph.
 
         Example:
             >>> dfg = Dfg(tys.Bool)
@@ -482,7 +520,14 @@ def set_outputs(self, *args: Wire) -> None:
             >>> dfg.set_outputs(dfg.inputs()[0]) # connect input to output
         """
         self._wire_up(self.output_node, args)
-        self.parent_op._set_out_types(self._output_op().types)
+        out_types = self._output_op().types
+        self.parent_op._set_out_types(out_types)
+        if (
+            isinstance(self.parent_op, ops.DataflowOp)
+            and self.parent_op._entrypoint_requires_wiring
+            and self.hugr.entrypoint == self.parent_node
+        ):
+            self.hugr._connect_df_entrypoint_outputs()
 
     def _set_parent_output_count(self, count: int) -> None:
         """Set the final number of output ports on the parent operation.
@@ -644,7 +689,7 @@ def _wire_up_port(self, node: Node, offset: PortOffset, p: Wire) -> tys.Type:
 
 
 class Dfg(DfBase[ops.DFG]):
-    """Builder for a simple nested Dataflow graph, with root node of type
+    """Builder for a simple nested Dataflow graph, with entrypoint node of type
     :class:`DFG <hugr.ops.DFG>`.
 
     Args:
@@ -662,8 +707,8 @@ def __init__(self, *input_types: tys.Type) -> None:
         super().__init__(parent_op)
 
     def set_outputs(self, *outputs: Wire) -> None:
-        super().set_outputs(*outputs)
         self._set_parent_output_count(len(outputs))
+        super().set_outputs(*outputs)
 
 
 def _ancestral_sibling(h: Hugr, src: Node, tgt: Node) -> Node | None:

hugr-py/src/hugr/build/function.py

@@ -22,7 +22,7 @@ class Module(DefinitionBuilder[ops.Module]):
 
     Examples:
         >>> m = Module()
-        >>> m.hugr.root_op()
+        >>> m.hugr.entrypoint_op()
         Module()
     """
 
@@ -52,13 +52,13 @@ def declare_function(self, name: str, signature: PolyFuncType) -> Node:
             >>> m.declare_function("f", sig)
             Node(1)
         """
-        return self.hugr.add_node(ops.FuncDecl(name, signature), self.hugr.root)
+        return self.hugr.add_node(ops.FuncDecl(name, signature), self.hugr.entrypoint)
 
     def add_alias_decl(self, name: str, bound: TypeBound) -> Node:
         """Add a type alias declaration."""
-        return self.hugr.add_node(ops.AliasDecl(name, bound), self.hugr.root)
+        return self.hugr.add_node(ops.AliasDecl(name, bound), self.hugr.entrypoint)
 
     @property
     def metadata(self) -> dict[str, object]:
         """Metadata associated with this module."""
-        return self.hugr.root.metadata
+        return self.hugr.entrypoint.metadata