diff --git a/hugr-py/src/hugr/dfg.py b/hugr-py/src/hugr/dfg.py index 141e2e15f..a06df03c0 100644 --- a/hugr-py/src/hugr/dfg.py +++ b/hugr-py/src/hugr/dfg.py @@ -42,6 +42,7 @@ def define_function( name: str, input_types: TypeRow, type_params: list[TypeParam] | None = None, + parent: ToNode | None = None, ) -> Function: """Start building a function definition in the graph. @@ -49,18 +50,21 @@ def define_function( name: The name of the function. input_types: The input types for the function. type_params: The type parameters for the function, if polymorphic. + parent: The parent node of the constant. Defaults to the root node. Returns: The new function builder. """ + parent_node = parent or self.hugr.root parent_op = ops.FuncDefn(name, input_types, type_params or []) - return Function.new_nested(parent_op, self.hugr) + return Function.new_nested(parent_op, self.hugr, parent_node) - def add_const(self, value: val.Value) -> Node: + def add_const(self, value: val.Value, parent: ToNode | None = None) -> Node: """Add a static constant to the graph. Args: value: The constant value to add. + parent: The parent node of the constant. Defaults to the root node. Returns: The node holding the :class:`Const ` operation. @@ -71,11 +75,13 @@ def add_const(self, value: val.Value) -> Node: >>> dfg.hugr[const_n].op Const(TRUE) """ - return self.hugr.add_node(ops.Const(value), self.hugr.root) + parent_node = parent or self.hugr.root + return self.hugr.add_node(ops.Const(value), parent_node) - def add_alias_defn(self, name: str, ty: Type) -> Node: + def add_alias_defn(self, name: str, ty: Type, parent: ToNode | None = None) -> Node: """Add a type alias definition.""" - return self.hugr.add_node(ops.AliasDefn(name, ty), self.hugr.root) + parent_node = parent or self.hugr.root + return self.hugr.add_node(ops.AliasDefn(name, ty), parent_node) DP = TypeVar("DP", bound=ops.DfParentOp) @@ -482,13 +488,17 @@ def add_state_order(self, src: Node, dst: Node) -> None: # adds edge to the right of all existing edges self.hugr.add_link(src.out(-1), dst.inp(-1)) - def load(self, const: ToNode | val.Value) -> Node: + def load( + self, const: ToNode | val.Value, const_parent: ToNode | None = None + ) -> Node: """Load a constant into the graph as a dataflow value. Args: const: The constant to load, either a Value that will be added as a - child Const node then loaded, or a node corresponding to an existing - Const. + child Const node then loaded, or a node corresponding to an + existing Const. + const_parent: If `const` is a Value, the parent node for the new + constant definition. Defaults to the current dataflow container. Returns: The node holding the :class:`LoadConst ` @@ -503,7 +513,8 @@ def load(self, const: ToNode | val.Value) -> Node: LoadConst(Bool) """ if isinstance(const, val.Value): - const = self.add_const(const) + const_parent = const_parent or self.parent_node + const = self.add_const(const, parent=const_parent) const_op = self.hugr._get_typed_op(const, ops.Const) load_op = ops.LoadConst(const_op.val.type_())