From c21a999ccbd6a8a2659cf3ffdccbd9be800d3ef8 Mon Sep 17 00:00:00 2001
From: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>
Date: Mon, 22 Jul 2024 15:21:59 +0100
Subject: [PATCH] fix: dfg wrapper build handles incorrect output wire numbers
 (#1332)

DFGWrapper handles incorrectly were using internal dfg output number as
the outputs from the container node. Fixed for function, case and
tailloop. Also fixed a test that relied on this incorrect behaviour.

Closes #1257
---
 hugr-core/src/builder/handle.rs     | 16 +++-------------
 hugr-core/src/builder/tail_loop.rs  | 17 +++++++++--------
 hugr-core/src/hugr/validate/test.rs | 11 ++++-------
 3 files changed, 16 insertions(+), 28 deletions(-)

diff --git a/hugr-core/src/builder/handle.rs b/hugr-core/src/builder/handle.rs
index 1530392a0..f6058e2d8 100644
--- a/hugr-core/src/builder/handle.rs
+++ b/hugr-core/src/builder/handle.rs
@@ -1,5 +1,5 @@
 //! Handles to nodes in HUGR used during the building phase.
-use crate::ops::handle::{BasicBlockID, CaseID, DfgID, FuncID, NodeHandle, TailLoopID};
+use crate::ops::handle::{BasicBlockID, CaseID, DfgID, FuncID, NodeHandle};
 use crate::ops::OpTag;
 use crate::utils::collect_array;
 use crate::{Node, OutgoingPort, Wire};
@@ -70,7 +70,7 @@ impl From<BuildHandle<DfgID>> for BuildHandle<FuncID<true>> {
     fn from(value: BuildHandle<DfgID>) -> Self {
         Self {
             node_handle: value.node().into(),
-            num_value_outputs: value.num_value_outputs,
+            num_value_outputs: 0,
         }
     }
 }
@@ -87,17 +87,7 @@ impl From<BuildHandle<DfgID>> for BuildHandle<CaseID> {
     fn from(value: BuildHandle<DfgID>) -> Self {
         Self {
             node_handle: value.node().into(),
-            num_value_outputs: value.num_value_outputs,
-        }
-    }
-}
-
-impl From<BuildHandle<DfgID>> for BuildHandle<TailLoopID> {
-    #[inline]
-    fn from(value: BuildHandle<DfgID>) -> Self {
-        Self {
-            node_handle: value.node().into(),
-            num_value_outputs: value.num_value_outputs,
+            num_value_outputs: 0,
         }
     }
 }
diff --git a/hugr-core/src/builder/tail_loop.rs b/hugr-core/src/builder/tail_loop.rs
index c7d5544e1..7d713cc2c 100644
--- a/hugr-core/src/builder/tail_loop.rs
+++ b/hugr-core/src/builder/tail_loop.rs
@@ -1,11 +1,10 @@
 use crate::extension::ExtensionSet;
-use crate::ops;
+use crate::ops::{self, DataflowOpTrait};
 
 use crate::hugr::views::HugrView;
 use crate::types::{Signature, TypeRow};
 use crate::{Hugr, Node};
 
-use super::build_traits::SubContainer;
 use super::handle::BuildHandle;
 use super::{
     dataflow::{DFGBuilder, DFGWrapper},
@@ -52,20 +51,22 @@ impl<B: AsMut<Hugr> + AsRef<Hugr>> TailLoopBuilder<B> {
     pub fn internal_output_row(&self) -> Result<TypeRow, BuildError> {
         self.loop_signature().map(ops::TailLoop::body_output_row)
     }
-}
 
-impl<H: AsMut<Hugr> + AsRef<Hugr>> TailLoopBuilder<H> {
     /// Set outputs and finish, see [`TailLoopBuilder::set_outputs`]
     pub fn finish_with_outputs(
         mut self,
         out_variant: Wire,
         rest: impl IntoIterator<Item = Wire>,
-    ) -> Result<<Self as SubContainer>::ContainerHandle, BuildError>
+    ) -> Result<BuildHandle<TailLoopID>, BuildError>
     where
         Self: Sized,
     {
         self.set_outputs(out_variant, rest)?;
-        self.finish_sub_container()
+        Ok((
+            self.container_node(),
+            self.loop_signature()?.signature().output_count(),
+        )
+            .into())
     }
 }
 
@@ -96,7 +97,7 @@ mod test {
     use crate::{
         builder::{
             test::{BIT, NAT},
-            DataflowSubContainer, HugrBuilder, ModuleBuilder,
+            DataflowSubContainer, HugrBuilder, ModuleBuilder, SubContainer,
         },
         extension::prelude::{ConstUsize, PRELUDE_ID, USIZE_T},
         hugr::ValidationError,
@@ -196,7 +197,7 @@ mod test {
     }
 
     #[test]
-    #[should_panic] // issue 1257: When building a TailLoop, calling outputs_arr, you are given an OrderEdge "output wire"
+    // fixed: issue 1257: When building a TailLoop, calling outputs_arr, you are given an OrderEdge "output wire"
     fn tailloop_output_arr() {
         let mut builder =
             TailLoopBuilder::new(type_row![], type_row![], type_row![], ExtensionSet::new())
diff --git a/hugr-core/src/hugr/validate/test.rs b/hugr-core/src/hugr/validate/test.rs
index fbc847f1c..4fc752419 100644
--- a/hugr-core/src/hugr/validate/test.rs
+++ b/hugr-core/src/hugr/validate/test.rs
@@ -13,7 +13,7 @@ use crate::extension::prelude::{BOOL_T, PRELUDE, PRELUDE_ID, QB_T, USIZE_T};
 use crate::extension::{Extension, ExtensionSet, TypeDefBound, EMPTY_REG, PRELUDE_REGISTRY};
 use crate::hugr::internal::HugrMutInternals;
 use crate::hugr::HugrMut;
-use crate::ops::dataflow::{IOTrait, LoadFunction};
+use crate::ops::dataflow::IOTrait;
 use crate::ops::handle::NodeHandle;
 use crate::ops::leaf::MakeTuple;
 use crate::ops::{self, Noop, OpType, Value};
@@ -650,14 +650,11 @@ fn row_variables() -> Result<(), Box<dyn std::error::Error>> {
     )?;
     // All the wires here are carrying higher-order Function values
     let [func_arg] = fb.input_wires_arr();
-    let [id_usz] = {
+    let id_usz = {
         let bldr = fb.define_function("id_usz", Signature::new_endo(USIZE_T))?;
         let vals = bldr.input_wires();
-        let [inner_def] = bldr.finish_with_outputs(vals)?.outputs_arr();
-        let loadf =
-            LoadFunction::try_new(Signature::new_endo(USIZE_T).into(), [], &PRELUDE_REGISTRY)
-                .unwrap();
-        fb.add_dataflow_op(loadf, [inner_def])?.outputs_arr()
+        let inner_def = bldr.finish_with_outputs(vals)?;
+        fb.load_func(inner_def.handle(), &[], &PRELUDE_REGISTRY)?
     };
     let par = e.instantiate_extension_op(
         "parallel",