diff --git a/hugr-core/src/hugr.rs b/hugr-core/src/hugr.rs index 0c0edcf5ea..54e4c24b1a 100644 --- a/hugr-core/src/hugr.rs +++ b/hugr-core/src/hugr.rs @@ -135,6 +135,16 @@ impl Hugr { .ok_or(HugrError::UnsupportedEntrypoint { op: op_name }) } + /// Reserves enough capacity to insert at least the given number of + /// additional nodes and links. + /// + /// This method does not take into account already allocated free space left + /// after node removals, and may overallocate capacity. + pub fn reserve(&mut self, nodes: usize, links: usize) { + let ports = links * 2; + self.graph.reserve(nodes, ports); + } + /// Load a Hugr from a json reader. /// /// Validates the Hugr against the provided extension registry, ensuring all diff --git a/hugr-core/src/hugr/hugrmut.rs b/hugr-core/src/hugr/hugrmut.rs index f9df28d868..cada5a5357 100644 --- a/hugr-core/src/hugr/hugrmut.rs +++ b/hugr-core/src/hugr/hugrmut.rs @@ -1,10 +1,8 @@ //! Low-level interface for modifying a HUGR. -use core::panic; -use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; +use std::collections::{BTreeMap, HashMap, VecDeque}; use std::sync::Arc; -use portgraph::view::{NodeFilter, NodeFiltered}; use portgraph::{LinkMut, PortMut, PortView, SecondaryMap}; use crate::core::HugrNode; @@ -399,28 +397,29 @@ impl HugrMut for Hugr { root: Self::Node, mut other: Hugr, ) -> InsertionResult { - let (new_entrypoint, node_map) = insert_hugr_internal(self, root, &other); + let node_map = insert_hugr_internal(self, &other, other.entry_descendants(), |&n| { + if n == other.entrypoint() { + Some(root) + } else { + None + } + }); + // Merge the extension sets. + self.extensions.extend(other.extensions()); // Update the optypes and metadata, taking them from the other graph. // // No need to compute each node's extensions here, as we merge `other.extensions` directly. for (&node, &new_node) in node_map.iter() { - let optype = other.op_types.take(node); - self.op_types.set(new_node, optype); - let meta = other.metadata.take(node); - self.metadata.set(new_node, meta); + let node_pg = node.into_portgraph(); + let new_node_pg = new_node.into_portgraph(); + let optype = other.op_types.take(node_pg); + self.op_types.set(new_node_pg, optype); + let meta = other.metadata.take(node_pg); + self.metadata.set(new_node_pg, meta); } - debug_assert_eq!( - Some(&new_entrypoint.into_portgraph()), - node_map.get(&other.entrypoint().into_portgraph()) - ); InsertionResult { - inserted_entrypoint: new_entrypoint, - node_map: translate_indices( - |n| other.from_portgraph_node(n), - |n| self.from_portgraph_node(n), - node_map, - ) - .collect(), + inserted_entrypoint: node_map[&other.entrypoint()], + node_map, } } @@ -429,31 +428,31 @@ impl HugrMut for Hugr { root: Self::Node, other: &H, ) -> InsertionResult { - let (inserted_entrypoint, node_map) = insert_hugr_internal(self, root, other); + let node_map = insert_hugr_internal(self, other, other.entry_descendants(), |&n| { + if n == other.entrypoint() { + Some(root) + } else { + None + } + }); + // Merge the extension sets. + self.extensions.extend(other.extensions()); // Update the optypes and metadata, copying them from the other graph. // // No need to compute each node's extensions here, as we merge `other.extensions` directly. for (&node, &new_node) in node_map.iter() { - let node = other.from_portgraph_node(node); let nodetype = other.get_optype(node); - self.op_types.set(new_node, nodetype.clone()); + self.op_types + .set(new_node.into_portgraph(), nodetype.clone()); let meta = other.node_metadata_map(node); if !meta.is_empty() { - self.metadata.set(new_node, Some(meta.clone())); + self.metadata + .set(new_node.into_portgraph(), Some(meta.clone())); } } - debug_assert_eq!( - Some(&inserted_entrypoint.into_portgraph()), - node_map.get(&other.to_portgraph_node(other.entrypoint())) - ); InsertionResult { - inserted_entrypoint, - node_map: translate_indices( - |n| other.from_portgraph_node(n), - |n| self.from_portgraph_node(n), - node_map, - ) - .collect(), + inserted_entrypoint: node_map[&other.entrypoint()], + node_map, } } @@ -463,39 +462,25 @@ impl HugrMut for Hugr { other: &H, subgraph: &SiblingSubgraph, ) -> HashMap { - // Create a portgraph view with the explicit list of nodes defined by the subgraph. - let context: HashSet = subgraph - .nodes() - .iter() - .map(|&n| other.to_portgraph_node(n)) - .collect(); - let portgraph: NodeFiltered<_, NodeFilter>, _> = - NodeFiltered::new_node_filtered( - other.portgraph(), - |node, ctx| ctx.contains(&node), - context, - ); - let node_map = insert_subgraph_internal(self, root, other, &portgraph); + let node_map = insert_hugr_internal(self, other, subgraph.nodes().iter().cloned(), |_| { + Some(root) + }); // Update the optypes and metadata, copying them from the other graph. for (&node, &new_node) in node_map.iter() { - let node = other.from_portgraph_node(node); let nodetype = other.get_optype(node); - self.op_types.set(new_node, nodetype.clone()); + self.op_types + .set(new_node.into_portgraph(), nodetype.clone()); let meta = other.node_metadata_map(node); if !meta.is_empty() { - self.metadata.set(new_node, Some(meta.clone())); + self.metadata + .set(new_node.into_portgraph(), Some(meta.clone())); } // Add the required extensions to the registry. if let Ok(exts) = nodetype.used_extensions() { self.use_extensions(exts); } } - translate_indices( - |n| other.from_portgraph_node(n), - |n| self.from_portgraph_node(n), - node_map, - ) - .collect() + node_map } fn copy_descendants( @@ -552,85 +537,77 @@ impl HugrMut for Hugr { } } -/// 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`. +/// Internal implementation of `insert_hugr`, `insert_view`, and +/// `insert_subgraph`. /// -/// 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: &H, -) -> (Node, HashMap) { - let other_entrypoint = other.to_portgraph_node(other.entrypoint()); - // TODO: We may need to replicate `portgraph::LinkMut::insert_graph` at the - // hugr level to avoid calling `hierarchy()` here. - let other_portgraph = - portgraph::view::Region::new(other.portgraph(), other.hierarchy(), other_entrypoint); - let node_map = hugr - .graph - .insert_graph(&other_portgraph) - .unwrap_or_else(|e| panic!("Internal error while inserting a hugr into another: {e}")); - let other_root = node_map[&other_entrypoint]; - - // Update hierarchy and optypes - hugr.hierarchy - .push_child(other_root, root.into_portgraph()) - .expect("Inserting a newly-created node into the hierarchy should never fail."); - for (&node, &new_node) in node_map.iter() { - other - .children(other.from_portgraph_node(node)) - .for_each(|child| { - hugr.hierarchy - .push_child(node_map[&other.to_portgraph_node(child)], new_node) - .expect("Inserting a newly-created node into the hierarchy should never fail."); - }); - } - - // Merge the extension sets. - hugr.extensions.extend(other.extensions()); - - (other_root.into(), node_map) -} - -/// Internal implementation of the `insert_subgraph` method for AsMut. +/// Inserts all the nodes in `other_nodes` into `hugr`, under the given `root` node. /// /// Returns 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. +/// This function does not update the optypes of the inserted nodes, the +/// metadata, nor the hugr extensions, so the caller must do that. /// -/// In contrast to `insert_hugr_internal`, this function does not preserve -/// sibling order in the hierarchy. This is due to the subgraph not necessarily -/// having a single root, so the logic for reconstructing the hierarchy is not -/// able to just do a BFS. -fn insert_subgraph_internal( +/// # Parameters +/// - `hugr`: The hugr to insert into. +/// - `other`: The other graph to insert from. +/// - `other_nodes`: The nodes in the other graph to insert. +/// - `reroot`: A function that returns the new parent for each inserted node. +/// If `None`, the parent is set to the original parent after it has been inserted into `hugr`. +/// If that is the case, the parent must come before the child in the `other_nodes` iterator. +fn insert_hugr_internal( hugr: &mut Hugr, - root: Node, - other: &impl HugrView, - portgraph: &impl portgraph::LinkView, -) -> HashMap { - let node_map = hugr - .graph - .insert_graph(&portgraph) - .expect("Internal error while inserting a subgraph into another"); - - // A map for nodes that we inserted before their parent, so we couldn't - // update the hierarchy with their new id. - for (&node, &new_node) in node_map.iter() { - let new_parent = other - .get_parent(other.from_portgraph_node(node)) - .and_then(|parent| node_map.get(&other.to_portgraph_node(parent)).copied()) - .unwrap_or(root.into_portgraph()); - hugr.hierarchy - .push_child(new_node, new_parent) - .expect("Inserting a newly-created node into the hierarchy should never fail."); + other: &H, + other_nodes: impl Iterator, + reroot: impl Fn(&H::Node) -> Option, +) -> HashMap { + let new_node_count_hint = other_nodes.size_hint().1.unwrap_or_default(); + + // Insert the nodes from the other graph into this one. + let mut node_map = HashMap::with_capacity(new_node_count_hint); + hugr.reserve(new_node_count_hint, 0); + + for old in other_nodes { + // We use a dummy optype here. The callers take care of updating the + // correct optype, avoiding cloning if possible. + let op = OpType::default(); + let new = hugr.add_node(op); + node_map.insert(old, new); + + hugr.set_num_ports(new, other.num_inputs(old), other.num_outputs(old)); + + let new_parent = match reroot(&old) { + Some(new_parent) => new_parent, + None => { + let old_parent = other.get_parent(old).unwrap(); + *node_map + .get(&old_parent) + .expect("Child node came before parent in `other_nodes` iterator") + } + }; + hugr.set_parent(new, new_parent); + + // Reconnect the edges to the new node. + for tgt in other.node_inputs(old) { + for (neigh, src) in other.linked_outputs(old, tgt) { + let Some(&neigh) = node_map.get(&neigh) else { + continue; + }; + hugr.connect(neigh, src, new, tgt); + } + } + for src in other.node_outputs(old) { + for (neigh, tgt) in other.linked_inputs(old, src) { + if neigh == old { + continue; + } + let Some(&neigh) = node_map.get(&neigh) else { + continue; + }; + hugr.connect(new, src, neigh, tgt); + } + } } - node_map }