feat: insert_hugr/insert_view return node map

src/builder/build_traits.rs

@@ -3,6 +3,7 @@ use crate::hugr::views::HugrView;
 use crate::hugr::{Node, NodeMetadata, Port, ValidationError};
 use crate::ops::{self, LeafOp, OpTrait, OpType};
 
+use std::collections::HashMap;
 use std::iter;
 
 use super::FunctionBuilder;
@@ -108,13 +109,16 @@ pub trait Container {
     }
 
     /// Insert a HUGR as a child of the container.
-    fn add_hugr(&mut self, child: Hugr) -> Result<Node, BuildError> {
+    fn add_hugr(&mut self, child: Hugr) -> Result<(Node, HashMap<Node, Node>), BuildError> {
         let parent = self.container_node();
         Ok(self.hugr_mut().insert_hugr(parent, child)?)
     }
 
     /// Insert a copy of a HUGR as a child of the container.
-    fn add_hugr_view(&mut self, child: &impl HugrView) -> Result<Node, BuildError> {
+    fn add_hugr_view(
+        &mut self,
+        child: &impl HugrView,
+    ) -> Result<(Node, HashMap<Node, Node>), BuildError> {
         let parent = self.container_node();
         Ok(self.hugr_mut().insert_from_view(parent, child)?)
     }
@@ -230,7 +234,7 @@ pub trait Dataflow: Container {
         input_wires: impl IntoIterator<Item = Wire>,
     ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
         let num_outputs = hugr.get_optype(hugr.root()).signature().output_count();
-        let node = self.add_hugr(hugr)?;
+        let node = self.add_hugr(hugr)?.0;
 
         let inputs = input_wires.into_iter().collect();
         wire_up_inputs(inputs, node, self)?;
@@ -251,7 +255,7 @@ pub trait Dataflow: Container {
         input_wires: impl IntoIterator<Item = Wire>,
     ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
         let num_outputs = hugr.get_optype(hugr.root()).signature().output_count();
-        let node = self.add_hugr_view(hugr)?;
+        let node = self.add_hugr_view(hugr)?.0;
 
         let inputs = input_wires.into_iter().collect();
         wire_up_inputs(inputs, node, self)?;

src/hugr/hugrmut.rs

@@ -142,18 +142,30 @@ pub trait HugrMut: HugrView + HugrMutInternals {
 
     /// Insert another hugr into this one, under a given root node.
     ///
-    /// Returns the root node of the inserted hugr.
+    /// Returns a tuple of
+    /// * the root node of the inserted hugr.
+    /// * a map from each Node in `other` to the equivalent new Node in `self`
     #[inline]
-    fn insert_hugr(&mut self, root: Node, other: Hugr) -> Result<Node, HugrError> {
+    fn insert_hugr(
+        &mut self,
+        root: Node,
+        other: Hugr,
+    ) -> Result<(Node, HashMap<Node, Node>), HugrError> {
         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.
+    /// Returns a tuple of
+    /// * the root node of the inserted hugr.
+    /// * a map from each Node in `other` to the equivalent new Node in `self`
     #[inline]
-    fn insert_from_view(&mut self, root: Node, other: &impl HugrView) -> Result<Node, HugrError> {
+    fn insert_from_view(
+        &mut self,
+        root: Node,
+        other: &impl HugrView,
+    ) -> Result<(Node, HashMap<Node, Node>), HugrError> {
         self.valid_node(root)?;
         self.hugr_mut().insert_from_view(root, other)
     }
@@ -258,7 +270,11 @@ where
         Ok((src_port, dst_port))
     }
 
-    fn insert_hugr(&mut self, root: Node, mut other: Hugr) -> Result<Node, HugrError> {
+    fn insert_hugr(
+        &mut self,
+        root: Node,
+        mut other: Hugr,
+    ) -> Result<(Node, HashMap<Node, Node>), HugrError> {
         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() {
@@ -267,10 +283,17 @@ where
             let meta = other.metadata.take(node);
             self.as_mut().set_metadata(node.into(), meta).unwrap();
         }
-        Ok(other_root)
+        Ok((
+            other_root,
+            HashMap::from_iter(node_map.into_iter().map(|(k, v)| (k.into(), v.into()))),
+        ))
     }
 
-    fn insert_from_view(&mut self, root: Node, other: &impl HugrView) -> Result<Node, HugrError> {
+    fn insert_from_view(
+        &mut self,
+        root: Node,
+        other: &impl HugrView,
+    ) -> Result<(Node, HashMap<Node, Node>), HugrError> {
         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() {
@@ -281,7 +304,10 @@ where
                 .set_metadata(node.into(), meta.clone())
                 .unwrap();
         }
-        Ok(other_root)
+        Ok((
+            other_root,
+            HashMap::from_iter(node_map.into_iter().map(|(k, v)| (k.into(), v.into()))),
+        ))
     }
 }
 

src/hugr/rewrite/outline_cfg.rs

@@ -9,7 +9,8 @@ use crate::extension::ExtensionSet;
 use crate::hugr::rewrite::Rewrite;
 use crate::hugr::{HugrMut, HugrView};
 use crate::ops;
-use crate::ops::{BasicBlock, OpTag, OpTrait, OpType};
+use crate::ops::handle::NodeHandle;
+use crate::ops::{BasicBlock, OpTrait, OpType};
 use crate::{type_row, Node};
 
 /// Moves part of a Control-flow Sibling Graph into a new CFG-node
@@ -114,7 +115,7 @@ impl Rewrite for OutlineCfg {
         let outer_entry = h.children(outer_cfg).next().unwrap();
 
         // 2. new_block contains input node, sub-cfg, exit node all connected
-        let new_block = {
+        let (new_block, cfg_node) = {
             let mut new_block_bldr = BlockBuilder::new(
                 inputs.clone(),
                 vec![type_row![]],
@@ -128,26 +129,21 @@ impl Rewrite for OutlineCfg {
             let cfg = new_block_bldr
                 .cfg_builder(wires_in, outputs, extension_delta)
                 .unwrap();
-            let cfg_outputs = cfg.finish_sub_container().unwrap().outputs();
+            let cfg = cfg.finish_sub_container().unwrap();
             let predicate = new_block_bldr
                 .add_constant(ops::Const::simple_unary_predicate(), ExtensionSet::new())
                 .unwrap();
             let pred_wire = new_block_bldr.load_const(&predicate).unwrap();
-            new_block_bldr.set_outputs(pred_wire, cfg_outputs).unwrap();
-            h.insert_hugr(outer_cfg, new_block_bldr.hugr().clone())
-                .unwrap()
+            new_block_bldr
+                .set_outputs(pred_wire, cfg.outputs())
+                .unwrap();
+            let (new_block, node_map) = h
+                .insert_hugr(outer_cfg, new_block_bldr.hugr().clone())
+                .unwrap();
+            (new_block, *node_map.get(&cfg.node()).unwrap())
         };
 
-        // 3. Extract Cfg node created above (it moved when we called insert_hugr)
-        let cfg_node = h
-            .children(new_block)
-            .filter(|n| h.get_optype(*n).tag() == OpTag::Cfg)
-            .exactly_one()
-            .ok() // HugrMut::Children is not Debug
-            .unwrap();
-        let inner_exit = h.children(cfg_node).exactly_one().ok().unwrap();
-
-        // 4. Entry edges. Change any edges into entry_block from outside, to target new_block
+        // 3. Entry edges. Change any edges into entry_block from outside, to target new_block
         let preds: Vec<_> = h
             .linked_ports(entry, h.node_inputs(entry).exactly_one().ok().unwrap())
             .collect();
@@ -163,7 +159,8 @@ impl Rewrite for OutlineCfg {
             h.move_before_sibling(new_block, outer_entry).unwrap();
         }
 
-        // 5. Children of new CFG.
+        // 4. Children of new CFG.
+        let inner_exit = h.children(cfg_node).exactly_one().ok().unwrap();
         // Entry node must be first
         h.move_before_sibling(entry, inner_exit).unwrap();
         // And remaining nodes
@@ -174,7 +171,7 @@ impl Rewrite for OutlineCfg {
             }
         }
 
-        // 6. Exit edges.
+        // 5. Exit edges.
         // Retarget edge from exit_node (that used to target outside) to inner_exit
         let exit_port = h
             .node_outputs(exit)