diff --git a/src/builder/build_traits.rs b/src/builder/build_traits.rs index e2077e477..50cc91c3c 100644 --- a/src/builder/build_traits.rs +++ b/src/builder/build_traits.rs @@ -29,7 +29,7 @@ use super::{ use crate::Hugr; -use crate::hugr::HugrInternalsMut; +use crate::hugr::HugrMut; /// Trait for HUGR container builders. /// Containers are nodes that are parents of sibling graphs. diff --git a/src/builder/cfg.rs b/src/builder/cfg.rs index 914f53c50..9be7bc0b0 100644 --- a/src/builder/cfg.rs +++ b/src/builder/cfg.rs @@ -12,7 +12,7 @@ use crate::{ops::handle::NodeHandle, types::Type}; use crate::Node; use crate::{ - hugr::{HugrInternalsMut, NodeType}, + hugr::{HugrMut, NodeType}, type_row, Hugr, }; diff --git a/src/builder/conditional.rs b/src/builder/conditional.rs index 2312a2046..e894e56bc 100644 --- a/src/builder/conditional.rs +++ b/src/builder/conditional.rs @@ -15,7 +15,7 @@ use super::{ use crate::Node; use crate::{ - hugr::{HugrInternalsMut, NodeType}, + hugr::{HugrMut, NodeType}, Hugr, }; diff --git a/src/builder/dataflow.rs b/src/builder/dataflow.rs index 24e580a2a..56fe0b755 100644 --- a/src/builder/dataflow.rs +++ b/src/builder/dataflow.rs @@ -11,7 +11,7 @@ use crate::types::{FunctionType, Signature}; use crate::extension::ExtensionSet; use crate::Node; -use crate::{hugr::HugrInternalsMut, Hugr}; +use crate::{hugr::HugrMut, Hugr}; /// Builder for a [`ops::DFG`] node. #[derive(Debug, Clone, PartialEq)] diff --git a/src/builder/module.rs b/src/builder/module.rs index d16634039..09e914092 100644 --- a/src/builder/module.rs +++ b/src/builder/module.rs @@ -4,6 +4,7 @@ use super::{ BuildError, Container, }; +use crate::hugr::hugrmut::sealed::HugrMutInternals; use crate::{ hugr::{views::HugrView, ValidationError}, ops, @@ -18,10 +19,7 @@ use crate::types::Signature; use crate::Node; use smol_str::SmolStr; -use crate::{ - hugr::{HugrInternalsMut, NodeType}, - Hugr, -}; +use crate::{hugr::NodeType, Hugr}; /// Builder for a HUGR module. #[derive(Debug, Clone, PartialEq)] diff --git a/src/extension/infer.rs b/src/extension/infer.rs index 9024c66c9..bae062066 100644 --- a/src/extension/infer.rs +++ b/src/extension/infer.rs @@ -626,7 +626,7 @@ mod test { use super::*; use crate::builder::{BuildError, DFGBuilder, Dataflow, DataflowHugr}; use crate::extension::ExtensionSet; - use crate::hugr::HugrInternalsMut; + use crate::hugr::HugrMut; use crate::hugr::{validate::ValidationError, Hugr, HugrView, NodeType}; use crate::ops::{self, dataflow::IOTrait, handle::NodeHandle}; use crate::type_row; diff --git a/src/hugr.rs b/src/hugr.rs index 4e1c52edf..87698e02f 100644 --- a/src/hugr.rs +++ b/src/hugr.rs @@ -1,6 +1,6 @@ //! The Hugr data structure, and its basic component handles. -mod hugrmut; +pub mod hugrmut; pub mod rewrite; pub mod serialize; @@ -10,14 +10,14 @@ pub mod views; use std::collections::{HashMap, VecDeque}; use std::iter; -pub(crate) use self::hugrmut::HugrInternalsMut; +pub(crate) use self::hugrmut::HugrMut; pub use self::validate::ValidationError; use derive_more::From; pub use rewrite::{Rewrite, SimpleReplacement, SimpleReplacementError}; use portgraph::multiportgraph::MultiPortGraph; -use portgraph::{Hierarchy, PortMut, UnmanagedDenseMap}; +use portgraph::{Hierarchy, NodeIndex, PortMut, UnmanagedDenseMap}; use thiserror::Error; #[cfg(feature = "pyo3")] @@ -238,12 +238,31 @@ impl Hugr { } } - /// Produce a canonical ordering of the nodes. + /// Add a node to the graph, with the default conversion from OpType to NodeType + pub(crate) fn add_op(&mut self, op: impl Into) -> Node { + // TODO: Default to `NodeType::open_extensions` once we can infer extensions + self.add_node(NodeType::pure(op)) + } + + /// Add a node to the graph. + pub(crate) fn add_node(&mut self, nodetype: NodeType) -> Node { + let node = self + .graph + .add_node(nodetype.input_count(), nodetype.output_count()); + self.op_types[node] = nodetype; + node.into() + } + + /// Produce a canonical ordering of the descendant nodes of a root, + /// following the graph hierarchy. + /// + /// This starts with the root, and then proceeds in BFS order through the + /// contained regions. /// /// Used by [`HugrMut::canonicalize_nodes`] and the serialization code. - fn canonical_order(&self) -> impl Iterator + '_ { + fn canonical_order(&self, root: Node) -> impl Iterator + '_ { // Generate a BFS-ordered list of nodes based on the hierarchy - let mut queue = VecDeque::from([self.root.into()]); + let mut queue = VecDeque::from([root]); iter::from_fn(move || { let node = queue.pop_front()?; for child in self.children(node) { @@ -252,6 +271,46 @@ impl Hugr { Some(node) }) } + + /// Compact the nodes indices of the hugr to be contiguous, and order them as a breadth-first + /// traversal of the hierarchy. + /// + /// The rekey function is called for each moved node with the old and new indices. + /// + /// After this operation, a serialization and deserialization of the Hugr is guaranteed to + /// preserve the indices. + pub fn canonicalize_nodes(&mut self, mut rekey: impl FnMut(Node, Node)) { + // Generate the ordered list of nodes + let mut ordered = Vec::with_capacity(self.node_count()); + let root = self.root(); + ordered.extend(self.as_mut().canonical_order(root)); + + // Permute the nodes in the graph to match the order. + // + // Invariant: All the elements before `position` are in the correct place. + for position in 0..ordered.len() { + // Find the element's location. If it originally came from a previous position + // then it has been swapped somewhere else, so we follow the permutation chain. + let mut source: Node = ordered[position]; + while position > source.index.index() { + source = ordered[source.index.index()]; + } + + let target: Node = NodeIndex::new(position).into(); + if target != source { + self.graph.swap_nodes(target.index, source.index); + self.op_types.swap(target.index, source.index); + self.hierarchy.swap_nodes(target.index, source.index); + rekey(source, target); + } + } + self.root = NodeIndex::new(0); + + // Finish by compacting the copy nodes. + // The operation nodes will be left in place. + // This step is not strictly necessary. + self.graph.compact_nodes(|_, _| {}); + } } impl Port { @@ -354,6 +413,9 @@ pub enum HugrError { /// An error occurred while manipulating the hierarchy. #[error("An error occurred while manipulating the hierarchy.")] HierarchyError(#[from] portgraph::hierarchy::AttachError), + /// The node doesn't exist. + #[error("Invalid node {0:?}.")] + InvalidNode(Node), } #[cfg(feature = "pyo3")] diff --git a/src/hugr/hugrmut.rs b/src/hugr/hugrmut.rs index 4892b672e..d0250f7aa 100644 --- a/src/hugr/hugrmut.rs +++ b/src/hugr/hugrmut.rs @@ -1,4 +1,4 @@ -//! Base HUGR builder providing low-level building blocks. +//! Low-level interface for modifying a HUGR. use std::collections::HashMap; use std::ops::Range; @@ -8,110 +8,368 @@ use portgraph::{LinkMut, NodeIndex, PortMut, PortView, SecondaryMap}; use crate::hugr::{Direction, HugrError, HugrView, Node, NodeType}; use crate::ops::OpType; -use crate::ops::handle::NodeHandle; use crate::{Hugr, Port}; +use self::sealed::HugrMutInternals; + use super::NodeMetadata; -pub(crate) use sealed::HugrInternalsMut; +/// Functions for low-level building of a HUGR. +pub trait HugrMut: HugrView + HugrMutInternals { + /// Returns the metadata associated with a node. + fn get_metadata_mut(&mut self, node: Node) -> &mut NodeMetadata; + + /// Sets the metadata associated with a node. + fn set_metadata(&mut self, node: Node, metadata: NodeMetadata) { + *self.get_metadata_mut(node) = metadata; + } -pub trait HugrMut: sealed::HugrInternalsMut { - /// The kind of handle that can be used to refer to the root node. + /// Add a node to the graph with a parent in the hierarchy. /// - /// The handle is guaranteed to be able to contain the operation returned by - /// [`HugrView::root_type`]. - type RootHandle: NodeHandle; + /// The node becomes the parent's last child. + #[inline] + fn add_op_with_parent( + &mut self, + parent: Node, + op: impl Into, + ) -> Result { + self.valid_node(parent)?; + self.hugr_mut().add_op_with_parent(parent, op) + } - /// The view interface for this mutable reference to a Hugr. - type View<'h>: 'h + HugrView - where - Self: 'h; + /// Add a node to the graph with a parent in the hierarchy. + /// + /// The node becomes the parent's last child. + #[inline] + fn add_node_with_parent(&mut self, parent: Node, op: NodeType) -> Result { + self.valid_node(parent)?; + self.hugr_mut().add_node_with_parent(parent, op) + } + + /// Add a node to the graph as the previous sibling of another node. + /// + /// The sibling node's parent becomes the new node's parent. + /// + /// # Errors + /// + /// - If the sibling node does not have a parent. + /// - If the attachment would introduce a cycle. + #[inline] + fn add_op_before(&mut self, sibling: Node, op: impl Into) -> Result { + self.valid_non_root(sibling)?; + self.hugr_mut().add_op_before(sibling, op) + } - /// Returns a shared view of the Hugr. - fn view(&self) -> Self::View<'_>; + /// Add a node to the graph as the next sibling of another node. + /// + /// The sibling node's parent becomes the new node's parent. + /// + /// # Errors + /// + /// - If the sibling node does not have a parent. + /// - If the attachment would introduce a cycle. + #[inline] + fn add_op_after(&mut self, sibling: Node, op: impl Into) -> Result { + self.valid_non_root(sibling)?; + self.hugr_mut().add_op_after(sibling, op) + } + + /// Remove a node from the graph. + /// + /// # Panics + /// + /// Panics if the node is the root node. + #[inline] + fn remove_node(&mut self, node: Node) -> Result<(), HugrError> { + self.valid_non_root(node)?; + self.hugr_mut().remove_node(node) + } + + /// Connect two nodes at the given ports. + /// + /// The port must have already been created. See [`add_ports`] and [`set_num_ports`]. + /// + /// [`add_ports`]: #method.add_ports + /// [`set_num_ports`]: #method.set_num_ports. + #[inline] + fn connect( + &mut self, + src: Node, + src_port: usize, + dst: Node, + dst_port: usize, + ) -> Result<(), HugrError> { + self.valid_node(src)?; + self.valid_node(dst)?; + self.hugr_mut().connect(src, src_port, dst, dst_port) + } + + /// Disconnects all edges from the given port. + /// + /// The port is left in place. + #[inline] + fn disconnect(&mut self, node: Node, port: Port) -> Result<(), HugrError> { + self.valid_node(node)?; + self.hugr_mut().disconnect(node, port) + } + + /// Adds a non-dataflow edge between two nodes. The kind is given by the + /// operation's [`OpTrait::other_input`] or [`OpTrait::other_output`]. + /// + /// Returns the offsets of the new input and output ports, or an error if + /// the connection failed. + /// + /// [`OpTrait::other_input`]: crate::ops::OpTrait::other_input + /// [`OpTrait::other_output`]: crate::ops::OpTrait::other_output + fn add_other_edge(&mut self, src: Node, dst: Node) -> Result<(Port, Port), HugrError> { + self.valid_node(src)?; + self.valid_node(dst)?; + self.hugr_mut().add_other_edge(src, dst) + } + + /// Insert another hugr into this one, under a given root node. + /// + /// Returns the root node of the inserted hugr. + #[inline] + fn insert_hugr(&mut self, root: Node, other: Hugr) -> Result { + self.valid_node(root)?; + self.hugr_mut().insert_hugr(root, other) + } + + /// Copy another hugr into this one, under a given root node. + /// + /// Returns the root node of the inserted hugr. + #[inline] + fn insert_from_view(&mut self, root: Node, other: &impl HugrView) -> Result { + self.valid_node(root)?; + self.hugr_mut().insert_from_view(root, other) + } } +/// Impl for non-wrapped Hugrs. Overwrites the recursive default-impls to directly use the hugr. impl HugrMut for T where - T: AsRef + AsMut, + T: HugrView + AsMut, { - type RootHandle = Node; + fn get_metadata_mut(&mut self, node: Node) -> &mut NodeMetadata { + self.as_mut().metadata.get_mut(node.index) + } + + fn add_op_with_parent( + &mut self, + parent: Node, + op: impl Into, + ) -> Result { + // TODO: Default to `NodeType::open_extensions` once we can infer extensions + self.add_node_with_parent(parent, NodeType::pure(op)) + } + + fn add_node_with_parent(&mut self, parent: Node, node: NodeType) -> Result { + let node = self.add_node(node); + self.as_mut() + .hierarchy + .push_child(node.index, parent.index)?; + Ok(node) + } + + fn add_op_before(&mut self, sibling: Node, op: impl Into) -> Result { + let node = self.add_op(op); + self.as_mut() + .hierarchy + .insert_before(node.index, sibling.index)?; + Ok(node) + } + + fn add_op_after(&mut self, sibling: Node, op: impl Into) -> Result { + let node = self.add_op(op); + self.as_mut() + .hierarchy + .insert_after(node.index, sibling.index)?; + Ok(node) + } + + fn remove_node(&mut self, node: Node) -> Result<(), HugrError> { + if node == self.root() { + // TODO: Add a HugrMutError ? + panic!("cannot remove root node"); + } + self.as_mut().hierarchy.remove(node.index); + self.as_mut().graph.remove_node(node.index); + self.as_mut().op_types.remove(node.index); + Ok(()) + } + + fn connect( + &mut self, + src: Node, + src_port: usize, + dst: Node, + dst_port: usize, + ) -> Result<(), HugrError> { + self.as_mut() + .graph + .link_nodes(src.index, src_port, dst.index, dst_port)?; + Ok(()) + } - type View<'h> = &'h Hugr where Self: 'h; + fn disconnect(&mut self, node: Node, port: Port) -> Result<(), HugrError> { + let offset = port.offset; + let port = self.as_mut().graph.port_index(node.index, offset).ok_or( + portgraph::LinkError::UnknownOffset { + node: node.index, + offset, + }, + )?; + self.as_mut().graph.unlink_port(port); + Ok(()) + } + + fn add_other_edge(&mut self, src: Node, dst: Node) -> Result<(Port, Port), HugrError> { + let src_port: Port = self + .get_optype(src) + .other_port_index(Direction::Outgoing) + .expect("Source operation has no non-dataflow outgoing edges"); + let dst_port: Port = self + .get_optype(dst) + .other_port_index(Direction::Incoming) + .expect("Destination operation has no non-dataflow incoming edges"); + self.connect(src, src_port.index(), dst, dst_port.index())?; + Ok((src_port, dst_port)) + } - fn view(&self) -> Self::View<'_> { - self.as_ref() + fn insert_hugr(&mut self, root: Node, mut other: Hugr) -> Result { + let (other_root, node_map) = insert_hugr_internal(self.as_mut(), root, &other)?; + // Update the optypes and metadata, taking them from the other graph. + for (&node, &new_node) in node_map.iter() { + let optype = other.op_types.take(node); + self.as_mut().op_types.set(new_node, optype); + let meta = other.metadata.take(node); + self.as_mut().set_metadata(node.into(), meta); + } + Ok(other_root) + } + + fn insert_from_view(&mut self, root: Node, other: &impl HugrView) -> Result { + let (other_root, node_map) = insert_hugr_internal(self.as_mut(), root, other)?; + // Update the optypes and metadata, copying them from the other graph. + for (&node, &new_node) in node_map.iter() { + let nodetype = other.get_nodetype(node.into()); + self.as_mut().op_types.set(new_node, nodetype.clone()); + let meta = other.get_metadata(node.into()); + self.as_mut().set_metadata(node.into(), meta.clone()); + } + Ok(other_root) } } +/// Internal implementation of `insert_hugr` and `insert_view` methods for +/// AsMut. +/// +/// Returns the root node of the inserted hierarchy and a mapping from the nodes +/// in the inserted graph to their new indices in `hugr`. +/// +/// This function does not update the optypes of the inserted nodes, so the +/// caller must do that. +fn insert_hugr_internal( + hugr: &mut Hugr, + root: Node, + other: &impl HugrView, +) -> Result<(Node, HashMap), HugrError> { + let node_map = hugr.graph.insert_graph(&other.portgraph())?; + let other_root = node_map[&other.root().index]; + + // Update hierarchy and optypes + hugr.hierarchy.push_child(other_root, root.index)?; + for (&node, &new_node) in node_map.iter() { + other + .children(node.into()) + .try_for_each(|child| -> Result<(), HugrError> { + hugr.hierarchy + .push_child(node_map[&child.index], new_node)?; + Ok(()) + })?; + } + + // The root node didn't have any ports. + let root_optype = other.get_optype(other.root()); + hugr.set_num_ports( + other_root.into(), + root_optype.input_count(), + root_optype.output_count(), + ); + + Ok((other_root.into(), node_map)) +} + pub(crate) mod sealed { use super::*; - /// Functions for low-level building of a HUGR. (Or, in the future, a subregion thereof) - pub trait HugrInternalsMut { - /// Add a node to the graph, with the default conversion from OpType to NodeType - fn add_op(&mut self, op: impl Into) -> Node; - /// Add a node to the graph. - fn add_node(&mut self, node: NodeType) -> Node; + /// Trait for accessing the mutable internals of a Hugr(Mut). + /// + /// Specifically, this trait lets you apply arbitrary modifications that may + /// invalidate the HUGR. + pub trait HugrMutInternals: HugrView { + /// Returns the Hugr at the base of a chain of views. + fn hugr_mut(&mut self) -> &mut Hugr; - /// Remove a node from the graph. + /// Validates that a node is valid in the graph. /// - /// # Panics - /// - /// Panics if the node is the root node. - fn remove_node(&mut self, node: Node) -> Result<(), HugrError>; - - /// Returns the metadata associated with a node. - fn get_metadata_mut(&mut self, node: Node) -> &mut NodeMetadata; - - /// Sets the metadata associated with a node. - fn set_metadata(&mut self, node: Node, metadata: NodeMetadata) { - *self.get_metadata_mut(node) = metadata; + /// Returns a [`HugrError::InvalidNode`] otherwise. + #[inline] + fn valid_node(&self, node: Node) -> Result<(), HugrError> { + match self.contains_node(node) { + true => Ok(()), + false => Err(HugrError::InvalidNode(node)), + } } - /// Connect two nodes at the given ports. - /// - /// The port must have already been created. See [`add_ports`] and [`set_num_ports`]. + /// Validates that a node is a valid root descendant in the graph. /// - /// [`add_ports`]: #method.add_ports - /// [`set_num_ports`]: #method.set_num_ports. - fn connect( - &mut self, - src: Node, - src_port: usize, - dst: Node, - dst_port: usize, - ) -> Result<(), HugrError>; - - /// Disconnects all edges from the given port. + /// To include the root node use [`HugrMutInternals::valid_node`] instead. /// - /// The port is left in place. - fn disconnect(&mut self, node: Node, port: Port) -> Result<(), HugrError>; + /// Returns a [`HugrError::InvalidNode`] otherwise. + #[inline] + fn valid_non_root(&self, node: Node) -> Result<(), HugrError> { + match self.root() == node { + true => Err(HugrError::InvalidNode(node)), + false => self.valid_node(node), + } + } - /// Adds a non-dataflow edge between two nodes. The kind is given by the - /// operation's [`OpType::other_input`] or [`OpType::other_output`]. - /// - /// Returns the offsets of the new input and output ports, or an error if - /// the connection failed. - /// - /// [`OpType::other_input`]: crate::ops::OpType::other_input - /// [`OpType::other_output`]: crate::ops::OpType::other_output. - fn add_other_edge(&mut self, src: Node, dst: Node) -> Result<(Port, Port), HugrError>; + /// Add a node to the graph, with the default conversion from OpType to NodeType + fn add_op(&mut self, op: impl Into) -> Node { + self.hugr_mut().add_op(op) + } + + /// Add a node to the graph. + fn add_node(&mut self, nodetype: NodeType) -> Node { + self.hugr_mut().add_node(nodetype) + } /// Set the number of ports on a node. This may invalidate the node's `PortIndex`. - fn set_num_ports(&mut self, node: Node, incoming: usize, outgoing: usize); + fn set_num_ports(&mut self, node: Node, incoming: usize, outgoing: usize) { + self.valid_node(node).unwrap_or_else(|e| panic!("{}", e)); + self.hugr_mut().set_num_ports(node, incoming, outgoing) + } /// Alter the number of ports on a node and returns a range with the new /// port offsets, if any. This may invalidate the node's `PortIndex`. /// /// The `direction` parameter specifies whether to add ports to the incoming /// or outgoing list. - fn add_ports(&mut self, node: Node, direction: Direction, amount: isize) -> Range; + fn add_ports(&mut self, node: Node, direction: Direction, amount: isize) -> Range { + self.valid_node(node).unwrap_or_else(|e| panic!("{}", e)); + self.hugr_mut().add_ports(node, direction, amount) + } /// Sets the parent of a node. /// /// The node becomes the parent's last child. - fn set_parent(&mut self, node: Node, parent: Node) -> Result<(), HugrError>; + fn set_parent(&mut self, node: Node, parent: Node) -> Result<(), HugrError> { + self.valid_node(parent)?; + self.valid_non_root(node)?; + self.hugr_mut().set_parent(node, parent) + } /// Move a node in the hierarchy to be the subsequent sibling of another /// node. @@ -119,166 +377,52 @@ pub(crate) mod sealed { /// The sibling node's parent becomes the new node's parent. /// /// The node becomes the parent's last child. - fn move_after_sibling(&mut self, node: Node, after: Node) -> Result<(), HugrError>; + fn move_after_sibling(&mut self, node: Node, after: Node) -> Result<(), HugrError> { + self.valid_non_root(node)?; + self.valid_non_root(after)?; + self.hugr_mut().move_after_sibling(node, after) + } /// Move a node in the hierarchy to be the prior sibling of another node. /// /// The sibling node's parent becomes the new node's parent. /// /// The node becomes the parent's last child. - fn move_before_sibling(&mut self, node: Node, before: Node) -> Result<(), HugrError>; - - /// Add a node to the graph with a parent in the hierarchy. - /// - /// The node becomes the parent's last child. - fn add_op_with_parent( - &mut self, - parent: Node, - op: impl Into, - ) -> Result; - - /// Add a node to the graph with a parent in the hierarchy. - /// - /// The node becomes the parent's last child. - fn add_node_with_parent(&mut self, parent: Node, op: NodeType) -> Result; - - /// Add a node to the graph as the previous sibling of another node. - /// - /// The sibling node's parent becomes the new node's parent. - /// - /// # Errors - /// - /// - If the sibling node does not have a parent. - /// - If the attachment would introduce a cycle. - fn add_op_before( - &mut self, - sibling: Node, - op: impl Into, - ) -> Result; - - /// Add a node to the graph as the next sibling of another node. - /// - /// The sibling node's parent becomes the new node's parent. - /// - /// # Errors - /// - /// - If the sibling node does not have a parent. - /// - If the attachment would introduce a cycle. - fn add_op_after(&mut self, sibling: Node, op: impl Into) - -> Result; + fn move_before_sibling(&mut self, node: Node, before: Node) -> Result<(), HugrError> { + self.valid_non_root(node)?; + self.valid_non_root(before)?; + self.hugr_mut().move_before_sibling(node, before) + } /// Replace the OpType at node and return the old OpType. /// In general this invalidates the ports, which may need to be resized to /// match the OpType signature. /// TODO: Add a version which ignores input extensions - fn replace_op(&mut self, node: Node, op: NodeType) -> NodeType; - - /// Insert another hugr into this one, under a given root node. - /// - /// Returns the root node of the inserted hugr. - fn insert_hugr(&mut self, root: Node, other: Hugr) -> Result; - - /// Copy another hugr into this one, under a given root node. - /// - /// Returns the root node of the inserted hugr. - fn insert_from_view( - &mut self, - root: Node, - other: &impl HugrView, - ) -> Result; - - /// Compact the nodes indices of the hugr to be contiguous, and order them as a breadth-first - /// traversal of the hierarchy. - /// - /// The rekey function is called for each moved node with the old and new indices. - /// - /// After this operation, a serialization and deserialization of the Hugr is guaranteed to - /// preserve the indices. - fn canonicalize_nodes(&mut self, rekey: impl FnMut(Node, Node)); + fn replace_op(&mut self, node: Node, op: NodeType) -> NodeType { + self.valid_node(node).unwrap_or_else(|e| panic!("{}", e)); + self.hugr_mut().replace_op(node, op) + } } - impl HugrInternalsMut for T + /// Impl for non-wrapped Hugrs. Overwrites the recursive default-impls to directly use the hugr. + impl HugrMutInternals for T where - T: AsRef + AsMut, + T: HugrView + AsMut, { - fn add_node(&mut self, nodetype: NodeType) -> Node { - let node = self - .as_mut() - .graph - .add_node(nodetype.input_count(), nodetype.output_count()); - self.as_mut().op_types[node] = nodetype; - node.into() - } - - fn add_op(&mut self, op: impl Into) -> Node { - // TODO: Default to `NodeType::open_extensions` once we can infer extensions - self.add_node(NodeType::pure(op)) - } - - fn remove_node(&mut self, node: Node) -> Result<(), HugrError> { - if node.index == self.as_ref().root { - // TODO: Add a HugrMutError ? - panic!("cannot remove root node"); - } - self.as_mut().hierarchy.remove(node.index); - self.as_mut().graph.remove_node(node.index); - self.as_mut().op_types.remove(node.index); - Ok(()) - } - - fn get_metadata_mut(&mut self, node: Node) -> &mut NodeMetadata { - self.as_mut().metadata.get_mut(node.index) - } - - fn connect( - &mut self, - src: Node, - src_port: usize, - dst: Node, - dst_port: usize, - ) -> Result<(), HugrError> { + fn hugr_mut(&mut self) -> &mut Hugr { self.as_mut() - .graph - .link_nodes(src.index, src_port, dst.index, dst_port)?; - Ok(()) - } - - fn disconnect(&mut self, node: Node, port: Port) -> Result<(), HugrError> { - let offset = port.offset; - let port = self.as_mut().graph.port_index(node.index, offset).ok_or( - portgraph::LinkError::UnknownOffset { - node: node.index, - offset, - }, - )?; - self.as_mut().graph.unlink_port(port); - Ok(()) - } - - fn add_other_edge(&mut self, src: Node, dst: Node) -> Result<(Port, Port), HugrError> { - let src_port: Port = self - .get_optype(src) - .other_port_index(Direction::Outgoing) - .expect("Source operation has no non-dataflow outgoing edges"); - let dst_port: Port = self - .get_optype(dst) - .other_port_index(Direction::Incoming) - .expect("Destination operation has no non-dataflow incoming edges"); - self.connect(src, src_port.index(), dst, dst_port.index())?; - Ok((src_port, dst_port)) } #[inline] fn set_num_ports(&mut self, node: Node, incoming: usize, outgoing: usize) { - self.as_mut() + self.hugr_mut() .graph .set_num_ports(node.index, incoming, outgoing, |_, _| {}) } - #[inline] fn add_ports(&mut self, node: Node, direction: Direction, amount: isize) -> Range { - let mut incoming = self.as_mut().graph.num_inputs(node.index); - let mut outgoing = self.as_mut().graph.num_outputs(node.index); + let mut incoming = self.hugr_mut().graph.num_inputs(node.index); + let mut outgoing = self.hugr_mut().graph.num_outputs(node.index); let increment = |num: &mut usize| { let new = num.saturating_add_signed(amount); let range = *num..new; @@ -289,185 +433,40 @@ pub(crate) mod sealed { Direction::Incoming => increment(&mut incoming), Direction::Outgoing => increment(&mut outgoing), }; - self.as_mut() + self.hugr_mut() .graph .set_num_ports(node.index, incoming, outgoing, |_, _| {}); range } fn set_parent(&mut self, node: Node, parent: Node) -> Result<(), HugrError> { - self.as_mut().hierarchy.detach(node.index); - self.as_mut() + self.hugr_mut().hierarchy.detach(node.index); + self.hugr_mut() .hierarchy .push_child(node.index, parent.index)?; Ok(()) } fn move_after_sibling(&mut self, node: Node, after: Node) -> Result<(), HugrError> { - self.as_mut().hierarchy.detach(node.index); - self.as_mut() + self.hugr_mut().hierarchy.detach(node.index); + self.hugr_mut() .hierarchy .insert_after(node.index, after.index)?; Ok(()) } fn move_before_sibling(&mut self, node: Node, before: Node) -> Result<(), HugrError> { - self.as_mut().hierarchy.detach(node.index); - self.as_mut() + self.hugr_mut().hierarchy.detach(node.index); + self.hugr_mut() .hierarchy .insert_before(node.index, before.index)?; Ok(()) } - fn add_op_with_parent( - &mut self, - parent: Node, - op: impl Into, - ) -> Result { - // TODO: Default to `NodeType::open_extensions` once we can infer extensions - self.add_node_with_parent(parent, NodeType::pure(op)) - } - - fn add_node_with_parent( - &mut self, - parent: Node, - node: NodeType, - ) -> Result { - let node = self.add_node(node); - self.as_mut() - .hierarchy - .push_child(node.index, parent.index)?; - Ok(node) - } - - fn add_op_before( - &mut self, - sibling: Node, - op: impl Into, - ) -> Result { - let node = self.add_op(op); - self.as_mut() - .hierarchy - .insert_before(node.index, sibling.index)?; - Ok(node) - } - - fn add_op_after( - &mut self, - sibling: Node, - op: impl Into, - ) -> Result { - let node = self.add_op(op); - self.as_mut() - .hierarchy - .insert_after(node.index, sibling.index)?; - Ok(node) - } - fn replace_op(&mut self, node: Node, op: NodeType) -> NodeType { - let cur = self.as_mut().op_types.get_mut(node.index); + let cur = self.hugr_mut().op_types.get_mut(node.index); std::mem::replace(cur, op) } - - fn insert_hugr(&mut self, root: Node, mut other: Hugr) -> Result { - let (other_root, node_map) = insert_hugr_internal(self.as_mut(), root, &other)?; - // Update the optypes and metadata, taking them from the other graph. - for (&node, &new_node) in node_map.iter() { - let optype = other.op_types.take(node); - self.as_mut().op_types.set(new_node, optype); - let meta = other.metadata.take(node); - self.as_mut().set_metadata(node.into(), meta); - } - Ok(other_root) - } - - fn insert_from_view( - &mut self, - root: Node, - other: &impl HugrView, - ) -> Result { - let (other_root, node_map) = insert_hugr_internal(self.as_mut(), root, other)?; - // Update the optypes and metadata, copying them from the other graph. - for (&node, &new_node) in node_map.iter() { - let nodetype = other.get_nodetype(node.into()); - self.as_mut().op_types.set(new_node, nodetype.clone()); - let meta = other.get_metadata(node.into()); - self.as_mut().set_metadata(node.into(), meta.clone()); - } - Ok(other_root) - } - - fn canonicalize_nodes(&mut self, mut rekey: impl FnMut(Node, Node)) { - // Generate the ordered list of nodes - let mut ordered = Vec::with_capacity(self.node_count()); - ordered.extend(self.as_ref().canonical_order()); - - // Permute the nodes in the graph to match the order. - // - // Invariant: All the elements before `position` are in the correct place. - for position in 0..ordered.len() { - // Find the element's location. If it originally came from a previous position - // then it has been swapped somewhere else, so we follow the permutation chain. - let mut source: Node = ordered[position]; - while position > source.index.index() { - source = ordered[source.index.index()]; - } - - let target: Node = NodeIndex::new(position).into(); - if target != source { - let hugr = self.as_mut(); - hugr.graph.swap_nodes(target.index, source.index); - hugr.op_types.swap(target.index, source.index); - hugr.hierarchy.swap_nodes(target.index, source.index); - rekey(source, target); - } - } - self.as_mut().root = NodeIndex::new(0); - - // Finish by compacting the copy nodes. - // The operation nodes will be left in place. - // This step is not strictly necessary. - self.as_mut().graph.compact_nodes(|_, _| {}); - } - } - - /// Internal implementation of `insert_hugr` and `insert_view` methods for - /// AsMut. - /// - /// Returns the root node of the inserted hierarchy and a mapping from the nodes - /// in the inserted graph to their new indices in `hugr`. - /// - /// This function does not update the optypes of the inserted nodes, so the - /// caller must do that. - fn insert_hugr_internal( - hugr: &mut Hugr, - root: Node, - other: &impl HugrView, - ) -> Result<(Node, HashMap), HugrError> { - let node_map = hugr.graph.insert_graph(other.portgraph())?; - let other_root = node_map[&other.root().index]; - - // Update hierarchy and optypes - hugr.hierarchy.push_child(other_root, root.index)?; - for (&node, &new_node) in node_map.iter() { - other - .children(node.into()) - .try_for_each(|child| -> Result<(), HugrError> { - hugr.hierarchy - .push_child(node_map[&child.index], new_node)?; - Ok(()) - })?; - } - - // The root node didn't have any ports. - let root_optype = other.get_optype(other.root()); - hugr.set_num_ports( - other_root.into(), - root_optype.input_count(), - root_optype.output_count(), - ); - - Ok((other_root.into(), node_map)) } } @@ -480,7 +479,6 @@ mod test { types::{test::COPYABLE_T, FunctionType, Type}, }; - use super::sealed::HugrInternalsMut; use super::*; const NAT: Type = COPYABLE_T; diff --git a/src/hugr/rewrite/outline_cfg.rs b/src/hugr/rewrite/outline_cfg.rs index 3349b4dd6..e6f45b77e 100644 --- a/src/hugr/rewrite/outline_cfg.rs +++ b/src/hugr/rewrite/outline_cfg.rs @@ -5,8 +5,9 @@ use itertools::Itertools; use thiserror::Error; use crate::builder::{BlockBuilder, Container, Dataflow, SubContainer}; +use crate::hugr::hugrmut::sealed::HugrMutInternals; use crate::hugr::rewrite::Rewrite; -use crate::hugr::{HugrInternalsMut, HugrView}; +use crate::hugr::{HugrMut, HugrView}; use crate::ops; use crate::ops::{BasicBlock, OpTag, OpTrait, OpType}; use crate::{type_row, Hugr, Node}; diff --git a/src/hugr/rewrite/simple_replace.rs b/src/hugr/rewrite/simple_replace.rs index e6b07ec76..3f323c430 100644 --- a/src/hugr/rewrite/simple_replace.rs +++ b/src/hugr/rewrite/simple_replace.rs @@ -4,7 +4,7 @@ use std::collections::{HashMap, HashSet}; use itertools::Itertools; -use crate::hugr::{HugrInternalsMut, HugrView, NodeMetadata}; +use crate::hugr::{HugrMut, HugrView, NodeMetadata}; use crate::{ hugr::{Node, Rewrite}, ops::{OpTag, OpTrait, OpType}, diff --git a/src/hugr/serialize.rs b/src/hugr/serialize.rs index d0e8d2968..97b0815eb 100644 --- a/src/hugr/serialize.rs +++ b/src/hugr/serialize.rs @@ -9,7 +9,7 @@ use thiserror::Error; use pyo3::prelude::*; use crate::extension::ExtensionSet; -use crate::hugr::{Hugr, HugrInternalsMut, NodeType}; +use crate::hugr::{Hugr, NodeType}; use crate::ops::OpTrait; use crate::ops::OpType; use crate::Node; @@ -18,7 +18,7 @@ use portgraph::{Direction, LinkError, NodeIndex, PortView}; use serde::{Deserialize, Deserializer, Serialize}; -use super::{HugrError, HugrView}; +use super::{HugrError, HugrMut, HugrView}; /// A wrapper over the available HUGR serialization formats. /// @@ -129,7 +129,7 @@ impl TryFrom<&Hugr> for SerHugrV0 { // We compact the operation nodes during the serialization process, // and ignore the copy nodes. let mut node_rekey: HashMap = HashMap::with_capacity(hugr.node_count()); - for (order, node) in hugr.canonical_order().enumerate() { + for (order, node) in hugr.canonical_order(hugr.root()).enumerate() { node_rekey.insert(node, NodeIndex::new(order).into()); } @@ -267,6 +267,7 @@ impl TryFrom for Hugr { pub mod test { use super::*; + use crate::hugr::hugrmut::sealed::HugrMutInternals; use crate::{ builder::{ Container, DFGBuilder, Dataflow, DataflowHugr, DataflowSubContainer, HugrBuilder, diff --git a/src/hugr/validate.rs b/src/hugr/validate.rs index 334632db2..c0e2226c1 100644 --- a/src/hugr/validate.rs +++ b/src/hugr/validate.rs @@ -694,7 +694,8 @@ mod test { }; use crate::extension::prelude::BOOL_T; use crate::extension::ExtensionSet; - use crate::hugr::{HugrError, HugrInternalsMut, NodeType}; + use crate::hugr::hugrmut::sealed::HugrMutInternals; + use crate::hugr::{HugrError, HugrMut, NodeType}; use crate::ops::dataflow::IOTrait; use crate::ops::{self, LeafOp, OpType}; use crate::std_extensions::logic; diff --git a/src/hugr/views.rs b/src/hugr/views.rs index 275a1aabb..bead9168b 100644 --- a/src/hugr/views.rs +++ b/src/hugr/views.rs @@ -52,14 +52,23 @@ pub trait HugrView: sealed::HugrInternals { where Self: 'a; - /// Return index of HUGR root node. - fn root(&self) -> Node; + /// Return the root node of this view. + #[inline] + fn root(&self) -> Node { + self.root_node() + } /// Return the type of the HUGR root node. + #[inline] fn root_type(&self) -> &NodeType { - self.get_nodetype(self.root()) + let node_type = self.get_nodetype(self.root()); + debug_assert!(Self::RootHandle::can_hold(node_type.tag())); + node_type } + /// Returns whether the node exists. + fn contains_node(&self, node: Node) -> bool; + /// Returns the parent of a node. fn get_parent(&self, node: Node) -> Option; @@ -232,8 +241,8 @@ where Self: 'a; #[inline] - fn root(&self) -> Node { - self.as_ref().root.into() + fn contains_node(&self, node: Node) -> bool { + self.as_ref().graph.contains_node(node.index) } #[inline] @@ -352,6 +361,9 @@ pub(crate) mod sealed { /// Returns the Hugr at the base of a chain of views. fn base_hugr(&self) -> &Hugr; + + /// Return the root node of this view. + fn root_node(&self) -> Node; } impl HugrInternals for T @@ -369,5 +381,10 @@ pub(crate) mod sealed { fn base_hugr(&self) -> &Hugr { self.as_ref() } + + #[inline] + fn root_node(&self) -> Node { + self.as_ref().root.into() + } } } diff --git a/src/hugr/views/hierarchy.rs b/src/hugr/views/hierarchy.rs index 9514a841c..98eedd5b4 100644 --- a/src/hugr/views/hierarchy.rs +++ b/src/hugr/views/hierarchy.rs @@ -21,7 +21,7 @@ use std::iter; use ::petgraph::visit as pv; use context_iterators::{ContextIterator, IntoContextIterator, MapWithCtx}; use itertools::{Itertools, MapInto}; -use portgraph::{LinkView, PortIndex, PortView}; +use portgraph::{LinkView, MultiPortGraph, PortIndex, PortView}; use crate::ops::handle::NodeHandle; use crate::ops::OpTrait; @@ -29,8 +29,7 @@ use crate::{hugr::NodeType, hugr::OpType, Direction, Hugr, Node, Port}; use super::{sealed::HugrInternals, HugrView, NodeMetadata}; -type FlatRegionGraph<'g, Base> = - portgraph::view::FlatRegion<'g, ::Portgraph>; +type FlatRegionGraph<'g> = portgraph::view::FlatRegion<'g, MultiPortGraph>; /// View of a HUGR sibling graph. /// @@ -45,7 +44,7 @@ where root: Node, /// The filtered portgraph encoding the adjacency structure of the HUGR. - graph: FlatRegionGraph<'g, Base>, + graph: FlatRegionGraph<'g>, /// The rest of the HUGR. hugr: &'g Base, @@ -75,7 +74,7 @@ where where Self: 'a; - type NodePorts<'a> = MapInto< as PortView>::NodePortOffsets<'a>, Port> + type NodePorts<'a> = MapInto< as PortView>::NodePortOffsets<'a>, Port> where Self: 'a; @@ -83,20 +82,20 @@ where where Self: 'a; - type Neighbours<'a> = MapInto< as LinkView>::Neighbours<'a>, Node> + type Neighbours<'a> = MapInto< as LinkView>::Neighbours<'a>, Node> where Self: 'a; type PortLinks<'a> = MapWithCtx< - as LinkView>::PortLinks<'a>, + as LinkView>::PortLinks<'a>, &'a Self, (Node, Port), > where Self: 'a; #[inline] - fn root(&self) -> Node { - self.root + fn contains_node(&self, node: Node) -> bool { + self.graph.contains_node(node.index) } #[inline] @@ -203,7 +202,7 @@ where } } -type RegionGraph<'g, Base> = portgraph::view::Region<'g, ::Portgraph>; +type RegionGraph<'g> = portgraph::view::Region<'g, MultiPortGraph>; /// View of a HUGR descendants graph. /// @@ -220,7 +219,7 @@ where root: Node, /// The graph encoding the adjacency structure of the HUGR. - graph: RegionGraph<'g, Base>, + graph: RegionGraph<'g>, /// The node hierarchy. hugr: &'g Base, @@ -246,11 +245,11 @@ where { type RootHandle = Root; - type Nodes<'a> = MapInto< as PortView>::Nodes<'a>, Node> + type Nodes<'a> = MapInto< as PortView>::Nodes<'a>, Node> where Self: 'a; - type NodePorts<'a> = MapInto< as PortView>::NodePortOffsets<'a>, Port> + type NodePorts<'a> = MapInto< as PortView>::NodePortOffsets<'a>, Port> where Self: 'a; @@ -258,20 +257,20 @@ where where Self: 'a; - type Neighbours<'a> = MapInto< as LinkView>::Neighbours<'a>, Node> + type Neighbours<'a> = MapInto< as LinkView>::Neighbours<'a>, Node> where Self: 'a; type PortLinks<'a> = MapWithCtx< - as LinkView>::PortLinks<'a>, + as LinkView>::PortLinks<'a>, &'a Self, (Node, Port), > where Self: 'a; #[inline] - fn root(&self) -> Node { - self.root + fn contains_node(&self, node: Node) -> bool { + self.graph.contains_node(node.index) } #[inline] @@ -404,8 +403,8 @@ where } Self { root, - graph: FlatRegionGraph::::new_flat_region( - hugr.portgraph(), + graph: FlatRegionGraph::new_flat_region( + &hugr.base_hugr().graph, &hugr.base_hugr().hierarchy, root.index, ), @@ -430,8 +429,8 @@ where } Self { root, - graph: RegionGraph::::new_region( - hugr.portgraph(), + graph: RegionGraph::new_region( + &hugr.base_hugr().graph, &hugr.base_hugr().hierarchy, root.index, ), @@ -446,7 +445,7 @@ where Root: NodeHandle, Base: HugrInternals, { - type Portgraph = FlatRegionGraph<'g, Base>; + type Portgraph = FlatRegionGraph<'g>; #[inline] fn portgraph(&self) -> &Self::Portgraph { @@ -457,6 +456,11 @@ where fn base_hugr(&self) -> &Hugr { self.hugr.base_hugr() } + + #[inline] + fn root_node(&self) -> Node { + self.root + } } impl<'g, Root, Base> super::sealed::HugrInternals for DescendantsGraph<'g, Root, Base> @@ -464,7 +468,7 @@ where Root: NodeHandle, Base: HugrInternals, { - type Portgraph = RegionGraph<'g, Base>; + type Portgraph = RegionGraph<'g>; #[inline] fn portgraph(&self) -> &Self::Portgraph { @@ -475,6 +479,11 @@ where fn base_hugr(&self) -> &Hugr { self.hugr.base_hugr() } + + #[inline] + fn root_node(&self) -> Node { + self.root + } } #[cfg(test)] diff --git a/src/ops/custom.rs b/src/ops/custom.rs index 9302ee110..2a7cefbd2 100644 --- a/src/ops/custom.rs +++ b/src/ops/custom.rs @@ -6,7 +6,8 @@ use std::sync::Arc; use thiserror::Error; use crate::extension::{ExtensionId, OpDef, SignatureError}; -use crate::hugr::{HugrInternalsMut, HugrView, NodeType}; +use crate::hugr::hugrmut::sealed::HugrMutInternals; +use crate::hugr::{HugrView, NodeType}; use crate::types::{type_param::TypeArg, FunctionType, SignatureDescription}; use crate::{Extension, Hugr, Node}; diff --git a/src/ops/handle.rs b/src/ops/handle.rs index 5a752f5a3..beaff0a91 100644 --- a/src/ops/handle.rs +++ b/src/ops/handle.rs @@ -27,6 +27,11 @@ pub trait NodeHandle: Clone { fn try_cast>(&self) -> Option { T::TAG.is_superset(Self::TAG).then(|| self.node().into()) } + + /// Checks whether the handle can hold an operation with the given tag. + fn can_hold(tag: OpTag) -> bool { + Self::TAG.is_superset(tag) + } } /// Trait for handles that contain children.