feat: insert_hugr copies Module-level children, _with_defns links by Node

hugr-core/src/builder.rs

@@ -90,6 +90,7 @@ use thiserror::Error;
 use crate::extension::SignatureError;
 use crate::extension::simple_op::OpLoadError;
 use crate::hugr::ValidationError;
+use crate::hugr::hugrmut::InsertDefnError;
 use crate::ops::handle::{BasicBlockID, CfgID, ConditionalID, DfgID, FuncID, TailLoopID};
 use crate::ops::{NamedOp, OpType};
 use crate::types::Type;
@@ -177,6 +178,16 @@ pub enum BuildError {
         node: Node,
     },
 
+    /// From [Dataflow::add_hugr_with_wires_defns]
+    #[error{"In inserting Hugr: {0}"}]
+    HugrInsertionError(#[from] InsertDefnError<Node>),
+
+    /// From [Dataflow::add_hugr_view_with_wires_defns].
+    /// Note that because the type of node in the [InsertDefnError] depends
+    /// upon the view being inserted, we convert the error to a string here.
+    #[error("In inserting HugrView: {0}")]
+    HugrViewInsertionError(String),
+
     /// Wire not found in Hugr
     #[error("Wire not found in Hugr: {0}.")]
     WireNotFound(Wire),
@@ -352,4 +363,33 @@ pub(crate) mod test {
         );
         hugr
     }
+
+    /// A DFG-entrypoint Hugr (no inputs, one bool_t output) containing two calls,
+    /// to a FuncDefn and a FuncDecl each bool_t->bool_t, and their handles.
+    pub(crate) fn dfg_calling_defn_decl() -> (Hugr, FuncID<true>, FuncID<false>) {
+        let mut dfb = DFGBuilder::new(Signature::new(vec![], bool_t())).unwrap();
+        let new_defn = {
+            let mut mb = dfb.module_root_builder();
+            let fb = mb
+                .define_function("helper_id", Signature::new_endo(bool_t()))
+                .unwrap();
+            let [f_inp] = fb.input_wires_arr();
+            fb.finish_with_outputs([f_inp]).unwrap()
+        };
+        let new_decl = dfb
+            .module_root_builder()
+            .declare("helper2", Signature::new_endo(bool_t()).into())
+            .unwrap();
+        let cst = dfb.add_load_value(ops::Value::true_val());
+        let [c1] = dfb
+            .call(new_defn.handle(), &[], [cst])
+            .unwrap()
+            .outputs_arr();
+        let [c2] = dfb.call(&new_decl, &[], [c1]).unwrap().outputs_arr();
+        (
+            dfb.finish_hugr_with_outputs([c2]).unwrap(),
+            *new_defn.handle(),
+            new_decl,
+        )
+    }
 }

hugr-core/src/builder/build_traits.rs

@@ -1,11 +1,12 @@
 use crate::extension::prelude::MakeTuple;
-use crate::hugr::hugrmut::InsertionResult;
+use crate::hugr::hugrmut::{InsertDefnMode, InsertionResult};
 use crate::hugr::views::HugrView;
 use crate::hugr::{NodeMetadata, ValidationError};
 use crate::ops::{self, OpTag, OpTrait, OpType, Tag, TailLoop};
 use crate::utils::collect_array;
 use crate::{Extension, IncomingPort, Node, OutgoingPort};
 
+use std::collections::HashMap;
 use std::iter;
 use std::sync::Arc;
 
@@ -82,13 +83,17 @@ pub trait Container {
         self.add_child_node(constant.into()).into()
     }
 
-    /// Insert a HUGR as a child of the container.
+    /// Insert a HUGR. Its entrypoint will become a child of the container; other
+    /// children of the Module root will be added to the Module root of the Hugr
+    /// being built. (See [Dataflow::add_hugr_with_wires_defns]).
     fn add_hugr(&mut self, child: Hugr) -> InsertionResult {
         let parent = self.container_node();
         self.hugr_mut().insert_hugr(parent, child)
     }
 
     /// Insert a copy of a HUGR as a child of the container.
+    /// (Only the portion below the entrypoint will be inserted;
+    /// see [Dataflow::add_hugr_view_with_wires_defns])
     fn add_hugr_view<H: HugrView>(&mut self, child: &H) -> InsertionResult<H::Node, Node> {
         let parent = self.container_node();
         self.hugr_mut().insert_from_view(parent, child)
@@ -197,8 +202,10 @@ pub trait Dataflow: Container {
         Ok(outs.into())
     }
 
-    /// Insert a hugr-defined op to the sibling graph, wiring up the
-    /// `input_wires` to the incoming ports of the resulting root node.
+    /// Insert a hugr, adding its entrypoint to the sibling graph and wiring up the
+    /// `input_wires` to the incoming ports of the resulting root node. Other children
+    /// of the module root of `hugr` will be added to the module root being built.
+    /// (See [Self::add_hugr_with_wires_defns].)
     ///
     /// # Errors
     ///
@@ -209,20 +216,32 @@ pub trait Dataflow: Container {
         hugr: Hugr,
         input_wires: impl IntoIterator<Item = Wire>,
     ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
-        let optype = hugr.get_optype(hugr.entrypoint()).clone();
-        let num_outputs = optype.value_output_count();
         let node = self.add_hugr(hugr).inserted_entrypoint;
+        wire_ins_return_outs(input_wires, node, self)
+    }
 
-        wire_up_inputs(input_wires, node, self).map_err(|error| BuildError::OperationWiring {
-            op: Box::new(optype),
-            error,
-        })?;
-
-        Ok((node, num_outputs).into())
+    /// Insert a hugr, adding its entrypoint to the sibling graph and wiring up the
+    /// `input_wires` to the incoming ports of the resulting root node. `defns` may
+    /// contain other children of the module root of `hugr`, which will be added to
+    /// the module root being built.
+    fn add_hugr_with_wires_defns(
+        &mut self,
+        hugr: Hugr,
+        input_wires: impl IntoIterator<Item = Wire>,
+        defns: HashMap<Node, InsertDefnMode>,
+    ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
+        let parent = self.container_node();
+        let node = self
+            .hugr_mut()
+            .insert_hugr_with_defns(parent, hugr, defns)?
+            .inserted_entrypoint;
+        wire_ins_return_outs(input_wires, node, self)
     }
 
     /// Copy a hugr-defined op into the sibling graph, wiring up the
-    /// `input_wires` to the incoming ports of the resulting root node.
+    /// `input_wires` to the incoming ports of the node that was the entrypoint.
+    /// (Any part of `hugr` outside the entrypoint is not copied; see
+    /// [Self::add_hugr_view_with_wires_defns])
     ///
     /// # Errors
     ///
@@ -234,15 +253,25 @@ pub trait Dataflow: Container {
         input_wires: impl IntoIterator<Item = Wire>,
     ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
         let node = self.add_hugr_view(hugr).inserted_entrypoint;
-        let optype = hugr.get_optype(hugr.entrypoint()).clone();
-        let num_outputs = optype.value_output_count();
-
-        wire_up_inputs(input_wires, node, self).map_err(|error| BuildError::OperationWiring {
-            op: Box::new(optype),
-            error,
-        })?;
+        wire_ins_return_outs(input_wires, node, self)
+    }
 
-        Ok((node, num_outputs).into())
+    /// Copy a Hugr, adding its entrypoint into the sibling graph and wiring up the
+    /// `input_wires` to the incoming ports. `defns` may contain other children of
+    /// the module root of `hugr`, which will be added to the module root being built.
+    fn add_hugr_view_with_wires_defns<H: HugrView>(
+        &mut self,
+        hugr: &H,
+        input_wires: impl IntoIterator<Item = Wire>,
+        defns: HashMap<H::Node, InsertDefnMode>,
+    ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
+        let parent = self.container_node();
+        let node = self
+            .hugr_mut()
+            .insert_from_view_with_defns(parent, hugr, defns)
+            .map_err(|ins_err| BuildError::HugrViewInsertionError(ins_err.to_string()))?
+            .inserted_entrypoint;
+        wire_ins_return_outs(input_wires, node, self)
     }
 
     /// Wire up the `output_wires` to the input ports of the Output node.
@@ -692,6 +721,20 @@ fn wire_up_inputs<T: Dataflow + ?Sized>(
     Ok(())
 }
 
+fn wire_ins_return_outs<T: Dataflow + ?Sized>(
+    inputs: impl IntoIterator<Item = Wire>,
+    node: Node,
+    data_builder: &mut T,
+) -> Result<BuildHandle<DataflowOpID>, BuildError> {
+    let op = data_builder.hugr().get_optype(node).clone();
+    let num_outputs = op.value_output_count();
+    wire_up_inputs(inputs, node, data_builder).map_err(|error| BuildError::OperationWiring {
+        op: Box::new(op),
+        error,
+    })?;
+    Ok((node, num_outputs).into())
+}
+
 /// Add edge from src to dst.
 ///
 /// # Errors

hugr-core/src/builder/dataflow.rs

@@ -339,20 +339,21 @@ impl<T> HugrBuilder for DFGWrapper<Hugr, T> {
 #[cfg(test)]
 pub(crate) mod test {
     use cool_asserts::assert_matches;
-    use ops::OpParent;
     use rstest::rstest;
     use serde_json::json;
+    use std::collections::HashMap;
 
     use crate::builder::build_traits::DataflowHugr;
+    use crate::builder::test::dfg_calling_defn_decl;
     use crate::builder::{
         BuilderWiringError, DataflowSubContainer, ModuleBuilder, endo_sig, inout_sig,
     };
     use crate::extension::SignatureError;
     use crate::extension::prelude::Noop;
     use crate::extension::prelude::{bool_t, qb_t, usize_t};
+    use crate::hugr::hugrmut::InsertDefnMode;
     use crate::hugr::validate::InterGraphEdgeError;
-    use crate::ops::{OpTag, handle::NodeHandle};
-    use crate::ops::{OpTrait, Value};
+    use crate::ops::{FuncDecl, FuncDefn, OpParent, OpTag, OpTrait, Value, handle::NodeHandle};
 
     use crate::std_extensions::logic::test::and_op;
     use crate::types::type_param::TypeParam;
@@ -545,7 +546,7 @@ pub(crate) mod test {
     }
 
     #[test]
-    fn insert_hugr() -> Result<(), BuildError> {
+    fn add_hugr() -> Result<(), BuildError> {
         // Create a simple DFG
         let mut dfg_builder = DFGBuilder::new(Signature::new(vec![bool_t()], vec![bool_t()]))?;
         let [i1] = dfg_builder.input_wires_arr();
@@ -576,6 +577,62 @@ pub(crate) mod test {
         Ok(())
     }
 
+    #[rstest]
+    fn add_hugr_with_defns(
+        #[values(false, true)] replace: bool,
+        #[values(true, false)] view: bool,
+    ) {
+        let mut fb = FunctionBuilder::new("main", Signature::new_endo(bool_t())).unwrap();
+        let my_decl = fb
+            .module_root_builder()
+            .declare("func1", Signature::new_endo(bool_t()).into())
+            .unwrap();
+        let (insert, ins_defn, ins_decl) = dfg_calling_defn_decl();
+        let ins_defn_name = insert
+            .get_optype(ins_defn.node())
+            .as_func_defn()
+            .unwrap()
+            .func_name()
+            .clone();
+        let ins_decl_name = insert
+            .get_optype(ins_decl.node())
+            .as_func_decl()
+            .unwrap()
+            .func_name()
+            .clone();
+        let decl_mode = if replace {
+            InsertDefnMode::Replace(my_decl.node())
+        } else {
+            InsertDefnMode::Add
+        };
+        let link_spec = HashMap::from([
+            (ins_defn.node(), InsertDefnMode::Add),
+            (ins_decl.node(), decl_mode),
+        ]);
+        let inserted = if view {
+            fb.add_hugr_view_with_wires_defns(&insert, [], link_spec)
+                .unwrap()
+        } else {
+            fb.add_hugr_with_wires_defns(insert, [], link_spec).unwrap()
+        };
+        let h = fb.finish_hugr_with_outputs(inserted.outputs()).unwrap();
+        let defn_names = h
+            .nodes()
+            .filter_map(|n| h.get_optype(n).as_func_defn().map(FuncDefn::func_name))
+            .collect_vec();
+        assert_eq!(defn_names, [&"main".to_string(), &ins_defn_name]);
+        let decl_names = h
+            .nodes()
+            .filter_map(|n| h.get_optype(n).as_func_decl().map(FuncDecl::func_name))
+            .cloned()
+            .collect_vec();
+        let mut expected_decl_names = vec!["func1".to_string()];
+        if !replace {
+            expected_decl_names.push(ins_decl_name)
+        }
+        assert_eq!(decl_names, expected_decl_names);
+    }
+
     #[test]
     fn barrier_node() -> Result<(), BuildError> {
         let mut parent = DFGBuilder::new(endo_sig(bool_t()))?;