feat: add trait+funcs for linking Hugrs explicitly by Node (linking 2/4)

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::linking::NodeLinkingError;
 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_link_hugr_by_node_with_wires]
+    #[error{"In inserting Hugr: {0}"}]
+    HugrInsertionError(#[from] NodeLinkingError<Node, Node>),
+
+    /// From [Dataflow::add_link_view_by_node_with_wires].
+    /// Note that because the type of node in the [NodeLinkingError] 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,35 @@ pub(crate) mod test {
         );
         hugr
     }
+
+    /// Builds a DFG-entrypoint Hugr (no inputs, one bool_t output) containing two calls,
+    /// to a FuncDefn and a FuncDecl each bool_t->bool_t.
+    /// Returns the Hugr and both function handles.
+    #[fixture]
+    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,13 @@
 use crate::extension::prelude::MakeTuple;
 use crate::hugr::hugrmut::InsertionResult;
+use crate::hugr::linking::{HugrLinking, NodeLinkingDirective};
 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;
 
@@ -99,9 +101,8 @@ pub trait Container {
     }
 
     /// Insert a copy of a HUGR as a child of the container.
-    ///
-    /// Only the portion below the entrypoint will be inserted, with any incoming
-    /// edges broken.
+    /// (Only the portion below the entrypoint will be inserted, with any incoming
+    /// edges broken; see [Dataflow::add_link_view_by_node_with_wires])
     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)
@@ -246,22 +247,34 @@ pub trait Dataflow: Container {
         region: Node,
         input_wires: impl IntoIterator<Item = Wire>,
     ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
-        let optype = hugr.get_optype(region).clone();
-        let num_outputs = optype.value_output_count();
         let node = self.add_hugr_region(hugr, region).inserted_entrypoint;
 
-        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())
+    /// 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_link_hugr_by_node_with_wires(
+        &mut self,
+        hugr: Hugr,
+        input_wires: impl IntoIterator<Item = Wire>,
+        defns: HashMap<Node, NodeLinkingDirective>,
+    ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
+        let parent = Some(self.container_node());
+        let ep = hugr.entrypoint();
+        let node = self
+            .hugr_mut()
+            .insert_link_hugr_by_node(parent, hugr, defns)?
+            .node_map[&ep];
+        wire_ins_return_outs(input_wires, node, self)
     }
 
-    /// Copy a hugr-defined op into the sibling graph, wiring up the
+    /// Copy a hugr's entrypoint-subtree (only) into the sibling graph, wiring up the
     /// `input_wires` to the incoming ports of the node that was the entrypoint.
-    /// (Note, any part of `hugr` outside the entrypoint is not copied;
-    /// this may lead to new input ports being disconnected.)
+    /// (Note that any wires from outside the entrypoint-subtree are disconnected in the copy;
+    /// see [Self::add_link_view_by_node_with_wires] for an alternative.)
     ///
     /// # Errors
     ///
@@ -273,15 +286,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_link_view_by_node_with_wires<H: HugrView>(
+        &mut self,
+        hugr: &H,
+        input_wires: impl IntoIterator<Item = Wire>,
+        defns: HashMap<H::Node, NodeLinkingDirective>,
+    ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
+        let parent = Some(self.container_node());
+        let node = self
+            .hugr_mut()
+            .insert_link_view_by_node(parent, hugr, defns)
+            .map_err(|ins_err| BuildError::HugrViewInsertionError(ins_err.to_string()))?
+            .node_map[&hugr.entrypoint()];
+        wire_ins_return_outs(input_wires, node, self)
     }
 
     /// Wire up the `output_wires` to the input ports of the Output node.
@@ -731,6 +754,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::linking::NodeLinkingDirective;
     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,63 @@ pub(crate) mod test {
         Ok(())
     }
 
+    #[rstest]
+    fn add_hugr_link_nodes(
+        #[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 {
+            NodeLinkingDirective::UseExisting(my_decl.node())
+        } else {
+            NodeLinkingDirective::add()
+        };
+        let link_spec = HashMap::from([
+            (ins_defn.node(), NodeLinkingDirective::add()),
+            (ins_decl.node(), decl_mode),
+        ]);
+        let inserted = if view {
+            fb.add_link_view_by_node_with_wires(&insert, [], link_spec)
+                .unwrap()
+        } else {
+            fb.add_link_hugr_by_node_with_wires(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()))?;