From be7c998d5a1ef254bd242fee3dac8ad6b4999632 Mon Sep 17 00:00:00 2001
From: Kartik Singhal <kartik.singhal@quantinuum.com>
Date: Thu, 19 Dec 2024 09:52:02 -0600
Subject: [PATCH] fix: allow add_input/add_output on all DFGs

Closes: #1819
---
 hugr-core/src/builder/dataflow.rs | 168 +++++++++++++++---------------
 1 file changed, 84 insertions(+), 84 deletions(-)

diff --git a/hugr-core/src/builder/dataflow.rs b/hugr-core/src/builder/dataflow.rs
index 2176d0ae6..914fca092 100644
--- a/hugr-core/src/builder/dataflow.rs
+++ b/hugr-core/src/builder/dataflow.rs
@@ -79,86 +79,6 @@ impl DFGBuilder<Hugr> {
         let root = base.root();
         DFGBuilder::create_with_io(base, root, signature)
     }
-}
-
-impl HugrBuilder for DFGBuilder<Hugr> {
-    fn finish_hugr(mut self) -> Result<Hugr, ValidationError> {
-        if cfg!(feature = "extension_inference") {
-            self.base.infer_extensions(false)?;
-        }
-        self.base.validate()?;
-        Ok(self.base)
-    }
-}
-
-impl<T: AsMut<Hugr> + AsRef<Hugr>> Container for DFGBuilder<T> {
-    #[inline]
-    fn container_node(&self) -> Node {
-        self.dfg_node
-    }
-
-    #[inline]
-    fn hugr_mut(&mut self) -> &mut Hugr {
-        self.base.as_mut()
-    }
-
-    #[inline]
-    fn hugr(&self) -> &Hugr {
-        self.base.as_ref()
-    }
-}
-
-impl<T: AsMut<Hugr> + AsRef<Hugr>> SubContainer for DFGBuilder<T> {
-    type ContainerHandle = BuildHandle<DfgID>;
-    #[inline]
-    fn finish_sub_container(self) -> Result<Self::ContainerHandle, BuildError> {
-        Ok((self.dfg_node, self.num_out_wires).into())
-    }
-}
-
-impl<T: AsMut<Hugr> + AsRef<Hugr>> Dataflow for DFGBuilder<T> {
-    #[inline]
-    fn num_inputs(&self) -> usize {
-        self.num_in_wires
-    }
-}
-
-/// Wrapper around [`DFGBuilder`] used to build other dataflow regions.
-// Stores option of DFGBuilder so it can be taken out without moving.
-#[derive(Debug, Clone, PartialEq)]
-pub struct DFGWrapper<B, T>(DFGBuilder<B>, PhantomData<T>);
-
-impl<B, T> DFGWrapper<B, T> {
-    pub(super) fn from_dfg_builder(db: DFGBuilder<B>) -> Self {
-        Self(db, PhantomData)
-    }
-}
-
-/// Builder for a [`ops::FuncDefn`] node
-pub type FunctionBuilder<B> = DFGWrapper<B, BuildHandle<FuncID<true>>>;
-
-impl FunctionBuilder<Hugr> {
-    /// Initialize a builder for a FuncDefn rooted HUGR
-    /// # Errors
-    ///
-    /// Error in adding DFG child nodes.
-    pub fn new(
-        name: impl Into<String>,
-        signature: impl Into<PolyFuncType>,
-    ) -> Result<Self, BuildError> {
-        let signature = signature.into();
-        let body = signature.body().clone();
-        let op = ops::FuncDefn {
-            signature,
-            name: name.into(),
-        };
-
-        let base = Hugr::new(op);
-        let root = base.root();
-
-        let db = DFGBuilder::create_with_io(base, root, body)?;
-        Ok(Self::from_dfg_builder(db))
-    }
 
     /// Add a new input to the function being constructed.
     ///
@@ -194,7 +114,7 @@ impl FunctionBuilder<Hugr> {
         }
 
         // Update the builder metadata
-        self.0.num_in_wires += 1;
+        self.num_in_wires += 1;
 
         self.input_wires().last().unwrap()
     }
@@ -231,7 +151,7 @@ impl FunctionBuilder<Hugr> {
         }
 
         // Update the builder metadata
-        self.0.num_out_wires += 1;
+        self.num_out_wires += 1;
     }
 
     /// Update the function builder's parent signature.
@@ -264,6 +184,86 @@ impl FunctionBuilder<Hugr> {
     }
 }
 
+impl HugrBuilder for DFGBuilder<Hugr> {
+    fn finish_hugr(mut self) -> Result<Hugr, ValidationError> {
+        if cfg!(feature = "extension_inference") {
+            self.base.infer_extensions(false)?;
+        }
+        self.base.validate()?;
+        Ok(self.base)
+    }
+}
+
+impl<T: AsMut<Hugr> + AsRef<Hugr>> Container for DFGBuilder<T> {
+    #[inline]
+    fn container_node(&self) -> Node {
+        self.dfg_node
+    }
+
+    #[inline]
+    fn hugr_mut(&mut self) -> &mut Hugr {
+        self.base.as_mut()
+    }
+
+    #[inline]
+    fn hugr(&self) -> &Hugr {
+        self.base.as_ref()
+    }
+}
+
+impl<T: AsMut<Hugr> + AsRef<Hugr>> SubContainer for DFGBuilder<T> {
+    type ContainerHandle = BuildHandle<DfgID>;
+    #[inline]
+    fn finish_sub_container(self) -> Result<Self::ContainerHandle, BuildError> {
+        Ok((self.dfg_node, self.num_out_wires).into())
+    }
+}
+
+impl<T: AsMut<Hugr> + AsRef<Hugr>> Dataflow for DFGBuilder<T> {
+    #[inline]
+    fn num_inputs(&self) -> usize {
+        self.num_in_wires
+    }
+}
+
+/// Wrapper around [`DFGBuilder`] used to build other dataflow regions.
+// Stores option of DFGBuilder so it can be taken out without moving.
+#[derive(Debug, Clone, PartialEq)]
+pub struct DFGWrapper<B, T>(DFGBuilder<B>, PhantomData<T>);
+
+impl<B, T> DFGWrapper<B, T> {
+    pub(super) fn from_dfg_builder(db: DFGBuilder<B>) -> Self {
+        Self(db, PhantomData)
+    }
+}
+
+/// Builder for a [`ops::FuncDefn`] node
+pub type FunctionBuilder<B> = DFGWrapper<B, BuildHandle<FuncID<true>>>;
+
+impl FunctionBuilder<Hugr> {
+    /// Initialize a builder for a FuncDefn rooted HUGR
+    /// # Errors
+    ///
+    /// Error in adding DFG child nodes.
+    pub fn new(
+        name: impl Into<String>,
+        signature: impl Into<PolyFuncType>,
+    ) -> Result<Self, BuildError> {
+        let signature = signature.into();
+        let body = signature.body().clone();
+        let op = ops::FuncDefn {
+            signature,
+            name: name.into(),
+        };
+
+        let base = Hugr::new(op);
+        let root = base.root();
+
+        let db = DFGBuilder::create_with_io(base, root, body)?;
+        Ok(Self::from_dfg_builder(db))
+    }
+}
+
 impl<B: AsMut<Hugr> + AsRef<Hugr>, T> Container for DFGWrapper<B, T> {
     #[inline]
     fn container_node(&self) -> Node {
@@ -465,8 +465,8 @@ pub(crate) mod test {
             f_build.set_order(&noop0.node(), &f_build.io()[1]);
 
             // Add a new input and output, and connect them with a noop in between
-            f_build.add_output(qb_t());
-            let i1 = f_build.add_input(qb_t());
+            f_build.0.add_output(qb_t());
+            let i1 = f_build.0.add_input(qb_t());
             let noop1 = f_build.add_dataflow_op(Noop(qb_t()), [i1])?;
 
             let hugr = f_build.finish_hugr_with_outputs([noop0.out_wire(0), noop1.out_wire(0)])?;