fix: FuncDefns don't require that their extensions match their children

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));
+                    }
                 }
             }
 

src/extension/infer/test.rs

@@ -2,10 +2,15 @@ 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 +967,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(())
+}

src/hugr/validate.rs

@@ -200,11 +200,15 @@ impl<'a, 'b> ValidationContext<'a, 'b> {
         // Secondly that the node has correct children
         self.validate_children(node, node_type)?;
 
-        // If this is a container with I/O nodes, check that the extension they
-        // define match the extensions of the container.
-        if let Some([input, output]) = self.hugr.get_io(node) {
-            self.extension_validator
-                .validate_io_extensions(node, input, output)?;
+        // FuncDefns have no resources since they're static nodes, but the
+        // functions they define can have any extension delta.
+        if node_type.tag() != OpTag::FuncDefn {
+            // If this is a container with I/O nodes, check that the extension they
+            // define match the extensions of the container.
+            if let Some([input, output]) = self.hugr.get_io(node) {
+                self.extension_validator
+                    .validate_io_extensions(node, input, output)?;
+            }
         }
 
         Ok(())

src/hugr/validate/test.rs

@@ -563,7 +563,7 @@ fn extensions_mismatch() -> Result<(), BuildError> {
     assert_matches!(
         handle,
         Err(ValidationError::ExtensionError(
-            ExtensionError::ParentIOExtensionMismatch { .. }
+            ExtensionError::TgtExceedsSrcExtensionsAtPort { .. }
         ))
     );
     Ok(())
@@ -750,3 +750,51 @@ fn invalid_types() {
         SignatureError::TypeArgMismatch(TypeArgError::WrongNumberArgs(2, 1))
     );
 }
+
+#[test]
+fn parent_io_mismatch() {
+    let mut hugr = Hugr::new(NodeType::new_pure(ops::DFG {
+        signature: FunctionType::new(type_row![USIZE_T], type_row![USIZE_T]),
+    }));
+
+    let input = hugr
+        .add_node_with_parent(
+            hugr.root(),
+            NodeType::new_pure(ops::Input {
+                types: type_row![USIZE_T],
+            }),
+        )
+        .unwrap();
+    let output = hugr
+        .add_node_with_parent(
+            hugr.root(),
+            NodeType::new(
+                ops::Output {
+                    types: type_row![USIZE_T],
+                },
+                ExtensionSet::singleton(&XB),
+            ),
+        )
+        .unwrap();
+
+    let lift = hugr
+        .add_node_with_parent(
+            hugr.root(),
+            NodeType::new_pure(ops::LeafOp::Lift {
+                type_row: type_row![USIZE_T],
+                new_extension: XB,
+            }),
+        )
+        .unwrap();
+
+    hugr.connect(input, 0, lift, 0).unwrap();
+    hugr.connect(lift, 0, output, 0).unwrap();
+
+    let result = hugr.validate(&PRELUDE_REGISTRY);
+    assert_matches!(
+        result,
+        Err(ValidationError::ExtensionError(
+            ExtensionError::ParentIOExtensionMismatch { .. }
+        ))
+    );
+}