From a79b7663d4fdf18ec77ac5c25cd88f79b14fae66 Mon Sep 17 00:00:00 2001
From: Craig Roy <craig.roy@cambridgequantum.com>
Date: Tue, 14 Nov 2023 15:39:59 +0000
Subject: [PATCH] fix(Inference): Don't add I/O constraints for `FuncDefn`s

---
 src/extension/infer.rs      |  8 +++-
 src/extension/infer/test.rs | 80 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/src/extension/infer.rs b/src/extension/infer.rs
index d468017a7..831dc6fdd 100644
--- a/src/extension/infer.rs
+++ b/src/extension/infer.rs
@@ -279,7 +279,13 @@ impl UnificationContext {
                     let m_input_node = self.make_or_get_meta(input, dir);
                     self.add_constraint(m_input_node, Constraint::Equal(m_input));
                     let m_output_node = self.make_or_get_meta(output, dir);
-                    self.add_constraint(m_output_node, Constraint::Equal(m_output));
+                    // If the parent node is a FuncDefn, it's input_resources
+                    // should always be empty, meaning if the function adds
+                    // resources, they shouldn't be equal to the parent's output
+                    // (which will also be empty.)
+                    if node_type.tag() != OpTag::FuncDefn {
+                        self.add_constraint(m_output_node, Constraint::Equal(m_output));
+                    }
                 }
             }
 
diff --git a/src/extension/infer/test.rs b/src/extension/infer/test.rs
index 480a1d5bc..6f038fb26 100644
--- a/src/extension/infer/test.rs
+++ b/src/extension/infer/test.rs
@@ -1,11 +1,17 @@
+
 use std::error::Error;
 
 use super::*;
 use crate::builder::test::closed_dfg_root_hugr;
-use crate::builder::{DFGBuilder, Dataflow, DataflowHugr};
+use crate::builder::{
+    Container, DFGBuilder, Dataflow, DataflowHugr, DataflowSubContainer, HugrBuilder, ModuleBuilder,
+};
 use crate::extension::prelude::QB_T;
 use crate::extension::ExtensionId;
-use crate::extension::{prelude::PRELUDE_REGISTRY, ExtensionSet};
+use crate::extension::{
+    prelude::{ConstUsize, PRELUDE_REGISTRY, USIZE_T},
+    ExtensionSet,
+};
 use crate::hugr::{validate::ValidationError, Hugr, HugrMut, HugrView, NodeType};
 use crate::macros::const_extension_ids;
 use crate::ops::custom::{ExternalOp, OpaqueOp};
@@ -962,3 +968,73 @@ fn sccs() {
         Some(&ExtensionSet::from_iter([A, B, C, UNKNOWN_EXTENSION]))
     );
 }
+
+#[test]
+fn simple_funcdefn() -> Result<(), Box<dyn Error>> {
+    let mut builder = ModuleBuilder::new();
+    let mut func_builder = builder.define_function(
+        "F",
+        FunctionType::new(vec![NAT], vec![NAT])
+            .with_extension_delta(&ExtensionSet::singleton(&A))
+            .pure(),
+    )?;
+
+    let [w] = func_builder.input_wires_arr();
+    let lift = func_builder.add_dataflow_op(
+        ops::LeafOp::Lift {
+            type_row: type_row![NAT],
+            new_extension: A,
+        },
+        [w],
+    )?;
+    let [w] = lift.outputs_arr();
+    func_builder.finish_with_outputs([w])?;
+    builder.finish_hugr(&PRELUDE_REGISTRY)?;
+    Ok(())
+}
+
+#[test]
+// Define a function "g" inside "f", which defines a function which adds to
+// the graph. Tests that the resources of the graph of "g" aren't being
+// constrained to match those of the FuncDefn node.
+fn funcdefn() -> Result<(), Box<dyn Error>> {
+    use crate::builder::{Container, Dataflow};
+    use crate::values::Value;
+
+    let mut builder = ModuleBuilder::new();
+    let mut func_builder = builder.define_function(
+        "F",
+        FunctionType::new(vec![NAT], vec![NAT])
+            .with_extension_delta(&ExtensionSet::singleton(&A))
+            .pure(),
+    )?;
+
+    let mut nested_func_builder = func_builder.define_function(
+        "G",
+        FunctionType::new(vec![NAT], vec![NAT])
+            .with_extension_delta(&ExtensionSet::singleton(&A))
+            .pure(),
+    )?;
+
+    let [w] = nested_func_builder.input_wires_arr();
+    let lift = nested_func_builder.add_dataflow_op(
+        ops::LeafOp::Lift {
+            type_row: type_row![NAT],
+            new_extension: A,
+        },
+        [w],
+    )?;
+    let [w] = lift.outputs_arr();
+    let g_id = nested_func_builder.finish_with_outputs([w])?;
+
+    let int_value: Value = ConstUsize::new(42).into();
+    let k_node = func_builder.add_constant(ops::Const::new(int_value, USIZE_T).unwrap(), None)?;
+    let k = func_builder.load_const(&k_node)?;
+    let call = func_builder.call(g_id.handle(), [k])?;
+    let [w] = call.outputs_arr();
+    func_builder.finish_with_outputs([w])?;
+
+    let hugr = builder.finish_hugr(&PRELUDE_REGISTRY)?;
+    infer_extensions(&hugr)?;
+    Ok(())
+}