Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions hugr-core/src/hugr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
231 changes: 104 additions & 127 deletions hugr-core/src/hugr/hugrmut.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -399,28 +397,29 @@ impl HugrMut for Hugr {
root: Self::Node,
mut other: Hugr,
) -> InsertionResult<Node, Self::Node> {
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,
}
}

Expand All @@ -429,31 +428,31 @@ impl HugrMut for Hugr {
root: Self::Node,
other: &H,
) -> InsertionResult<H::Node, Self::Node> {
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,
}
}

Expand All @@ -463,39 +462,25 @@ impl HugrMut for Hugr {
other: &H,
subgraph: &SiblingSubgraph<H::Node>,
) -> HashMap<H::Node, Self::Node> {
// Create a portgraph view with the explicit list of nodes defined by the subgraph.
let context: HashSet<portgraph::NodeIndex> = subgraph
.nodes()
.iter()
.map(|&n| other.to_portgraph_node(n))
.collect();
let portgraph: NodeFiltered<_, NodeFilter<HashSet<portgraph::NodeIndex>>, _> =
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(
Expand Down Expand Up @@ -552,85 +537,77 @@ impl HugrMut for Hugr {
}
}

/// Internal implementation of `insert_hugr` and `insert_view` methods for
/// AsMut<Hugr>.
///
/// 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<H: HugrView>(
hugr: &mut Hugr,
root: Node,
other: &H,
) -> (Node, HashMap<portgraph::NodeIndex, portgraph::NodeIndex>) {
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<Hugr>.
/// 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<N: HugrNode>(
/// # 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<H: HugrView>(
hugr: &mut Hugr,
root: Node,
other: &impl HugrView<Node = N>,
portgraph: &impl portgraph::LinkView,
) -> HashMap<portgraph::NodeIndex, portgraph::NodeIndex> {
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<Item = H::Node>,
reroot: impl Fn(&H::Node) -> Option<Node>,
) -> HashMap<H::Node, Node> {
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
}

Expand Down
Loading