fix: Ops require their own extensions

src/builder/circuit.rs

@@ -137,22 +137,33 @@ mod test {
             test::{build_main, NAT, QB},
             Dataflow, DataflowSubContainer, Wire,
         },
-        extension::prelude::BOOL_T,
+        extension::{prelude::BOOL_T, ExtensionSet},
         ops::{custom::OpaqueOp, LeafOp},
         type_row,
         types::FunctionType,
-        utils::test_quantum_extension::{cx_gate, h_gate, measure},
+        utils::test_quantum_extension::{cx_gate, h_gate, measure, EXTENSION_ID},
     };
 
     #[test]
     fn simple_linear() {
         let build_res = build_main(
-            FunctionType::new(type_row![QB, QB], type_row![QB, QB]).into(),
+            FunctionType::new_endo(type_row![QB, QB])
+                .with_extension_delta(&ExtensionSet::singleton(&EXTENSION_ID))
+                .into(),
             |mut f_build| {
-                let wires = f_build.input_wires().collect();
+                let mut wires: [Wire; 2] = f_build.input_wires_arr();
+                [wires[1]] = f_build
+                    .add_dataflow_op(
+                        LeafOp::Lift {
+                            type_row: vec![QB].into(),
+                            new_extension: EXTENSION_ID,
+                        },
+                        [wires[1]],
+                    )?
+                    .outputs_arr();
 
                 let mut linear = CircuitBuilder {
-                    wires,
+                    wires: Vec::from(wires),
                     builder: &mut f_build,
                 };
 
@@ -184,10 +195,20 @@ mod test {
             .into(),
         );
         let build_res = build_main(
-            FunctionType::new(type_row![QB, QB, NAT], type_row![QB, QB, BOOL_T]).into(),
+            FunctionType::new(type_row![QB, QB, NAT], type_row![QB, QB, BOOL_T])
+                .with_extension_delta(&ExtensionSet::singleton(&EXTENSION_ID))
+                .into(),
             |mut f_build| {
                 let [q0, q1, angle]: [Wire; 3] = f_build.input_wires_arr();
-
+                let [angle] = f_build
+                    .add_dataflow_op(
+                        LeafOp::Lift {
+                            type_row: vec![NAT].into(),
+                            new_extension: EXTENSION_ID,
+                        },
+                        [angle],
+                    )?
+                    .outputs_arr();
                 let mut linear = f_build.as_circuit(vec![q0, q1]);
 
                 let measure_out = linear

src/builder/dataflow.rs

@@ -213,9 +213,9 @@ pub(crate) mod test {
     use crate::hugr::validate::InterGraphEdgeError;
     use crate::ops::{handle::NodeHandle, LeafOp, OpTag};
 
-    use crate::std_extensions::logic::test::and_op;
+    use crate::std_extensions::logic::{self, test::and_op};
     use crate::types::Type;
-    use crate::utils::test_quantum_extension::h_gate;
+    use crate::utils::test_quantum_extension::{self, h_gate};
     use crate::{
         builder::{
             test::{n_identity, BIT, NAT, QB},
@@ -235,13 +235,25 @@ pub(crate) mod test {
             let _f_id = {
                 let mut func_builder = module_builder.define_function(
                     "main",
-                    FunctionType::new(type_row![NAT, QB], type_row![NAT, QB]).into(),
+                    FunctionType::new(type_row![NAT, QB], type_row![NAT, QB])
+                        .with_extension_delta(&ExtensionSet::singleton(
+                            &test_quantum_extension::EXTENSION_ID,
+                        ))
+                        .into(),
                 )?;
 
                 let [int, qb] = func_builder.input_wires_arr();
 
                 let q_out = func_builder.add_dataflow_op(h_gate(), vec![qb])?;
-
+                let [int] = func_builder
+                    .add_dataflow_op(
+                        LeafOp::Lift {
+                            type_row: vec![NAT].into(),
+                            new_extension: test_quantum_extension::EXTENSION_ID,
+                        },
+                        [int],
+                    )?
+                    .outputs_arr();
                 let inner_builder = func_builder.dfg_builder(
                     FunctionType::new(type_row![NAT], type_row![NAT]),
                     None,
@@ -260,7 +272,7 @@ pub(crate) mod test {
     }
 
     // Scaffolding for copy insertion tests
-    fn copy_scaffold<F>(f: F, msg: &'static str) -> Result<(), BuildError>
+    fn copy_scaffold<F>(f: F, delta: &ExtensionSet, msg: &'static str) -> Result<(), BuildError>
     where
         F: FnOnce(FunctionBuilder<&mut Hugr>) -> Result<BuildHandle<FuncID<true>>, BuildError>,
     {
@@ -269,7 +281,9 @@ pub(crate) mod test {
 
             let f_build = module_builder.define_function(
                 "main",
-                FunctionType::new(type_row![BOOL_T], type_row![BOOL_T, BOOL_T]).into(),
+                FunctionType::new(type_row![BOOL_T], type_row![BOOL_T, BOOL_T])
+                    .with_extension_delta(delta)
+                    .into(),
             )?;
 
             f(f_build)?;
@@ -287,25 +301,47 @@ pub(crate) mod test {
                 let [b1] = f_build.input_wires_arr();
                 f_build.finish_with_outputs([b1, b1])
             },
+            &ExtensionSet::new(),
             "Copy input and output",
         )?;
 
+        let es = ExtensionSet::singleton(&logic::EXTENSION_ID);
         copy_scaffold(
             |mut f_build| {
                 let [b1] = f_build.input_wires_arr();
                 let xor = f_build.add_dataflow_op(and_op(), [b1, b1])?;
+                let [b1] = f_build
+                    .add_dataflow_op(
+                        LeafOp::Lift {
+                            type_row: vec![BOOL_T].into(),
+                            new_extension: logic::EXTENSION_ID,
+                        },
+                        [b1],
+                    )?
+                    .outputs_arr();
                 f_build.finish_with_outputs([xor.out_wire(0), b1])
             },
+            &es,
             "Copy input and use with binary function",
         )?;
 
         copy_scaffold(
             |mut f_build| {
                 let [b1] = f_build.input_wires_arr();
                 let xor1 = f_build.add_dataflow_op(and_op(), [b1, b1])?;
+                let [b1] = f_build
+                    .add_dataflow_op(
+                        LeafOp::Lift {
+                            type_row: vec![BOOL_T].into(),
+                            new_extension: logic::EXTENSION_ID,
+                        },
+                        [b1],
+                    )?
+                    .outputs_arr();
                 let xor2 = f_build.add_dataflow_op(and_op(), [b1, xor1.out_wire(0)])?;
                 f_build.finish_with_outputs([xor2.out_wire(0), b1])
             },
+            &es,
             "Copy multiple times",
         )?;
 

src/extension/infer.rs

@@ -648,11 +648,12 @@ impl UnificationContext {
     fn search_variable_deps(&self) -> HashSet<Meta> {
         let mut seen = HashSet::new();
         let mut new_variables: HashSet<Meta> = self.variables.clone();
+        let constraints_for_solved = HashSet::new();
         while !new_variables.is_empty() {
             new_variables = new_variables
                 .into_iter()
                 .filter(|m| seen.insert(*m))
-                .flat_map(|m| self.get_constraints(&m).unwrap())
+                .flat_map(|m| self.get_constraints(&m).unwrap_or(&constraints_for_solved))
                 .map(|c| match c {
                     Constraint::Plus(_, other) => self.resolve(*other),
                     Constraint::Equal(other) => self.resolve(*other),

src/extension/op_def.rs

@@ -223,7 +223,7 @@ impl SignatureFunc {
     ///
     /// This function will return an error if the type arguments are invalid or
     /// there is some error in type computation.
-    pub fn compute_signature(
+    fn compute_signature(
         &self,
         def: &OpDef,
         args: &[TypeArg],
@@ -245,10 +245,10 @@ impl SignatureFunc {
             }
         };
 
-        let res = pf.instantiate(args, exts)?;
-        // TODO bring this assert back once resource inference is done?
-        // https://github.com/CQCL/hugr/issues/388
-        // debug_assert!(res.extension_reqs.contains(def.extension()));
+        let mut res = pf.instantiate(args, exts)?;
+        res.extension_reqs = res
+            .extension_reqs
+            .union(&ExtensionSet::singleton(def.extension()));
         Ok(res)
     }
 }
@@ -541,10 +541,10 @@ mod test {
         let args = [TypeArg::BoundedNat { n: 3 }, USIZE_T.into()];
         assert_eq!(
             def.compute_signature(&args, &PRELUDE_REGISTRY),
-            Ok(FunctionType::new(
-                vec![USIZE_T; 3],
-                vec![Type::new_tuple(vec![USIZE_T; 3])]
-            ))
+            Ok(
+                FunctionType::new(vec![USIZE_T; 3], vec![Type::new_tuple(vec![USIZE_T; 3])])
+                    .with_extension_delta(&ExtensionSet::singleton(&EXT_ID))
+            )
         );
         assert_eq!(def.validate_args(&args, &PRELUDE_REGISTRY, &[]), Ok(()));
 
@@ -554,10 +554,10 @@ mod test {
         let args = [TypeArg::BoundedNat { n: 3 }, tyvar.clone().into()];
         assert_eq!(
             def.compute_signature(&args, &PRELUDE_REGISTRY),
-            Ok(FunctionType::new(
-                tyvars.clone(),
-                vec![Type::new_tuple(tyvars)]
-            ))
+            Ok(
+                FunctionType::new(tyvars.clone(), vec![Type::new_tuple(tyvars)])
+                    .with_extension_delta(&ExtensionSet::singleton(&EXT_ID))
+            )
         );
         def.validate_args(&args, &PRELUDE_REGISTRY, &[TypeBound::Eq.into()])
             .unwrap();
@@ -610,7 +610,8 @@ mod test {
         def.validate_args(&args, &EMPTY_REG, &decls).unwrap();
         assert_eq!(
             def.compute_signature(&args, &EMPTY_REG),
-            Ok(FunctionType::new_endo(vec![tv]))
+            Ok(FunctionType::new_endo(vec![tv])
+                .with_extension_delta(&ExtensionSet::singleton(&EXT_ID)))
         );
         Ok(())
     }

src/hugr/rewrite/replace.rs

@@ -477,12 +477,11 @@ mod test {
             .unwrap()
             .into();
         let just_list = TypeRow::from(vec![listy.clone()]);
+        let exset = ExtensionSet::singleton(&collections::EXTENSION_NAME);
         let intermed = TypeRow::from(vec![listy.clone(), USIZE_T]);
 
         let mut cfg = CFGBuilder::new(
-            // One might expect an extension_delta of "collections" here, but push/pop
-            // have an empty delta themselves, pending https://github.com/CQCL/hugr/issues/388
-            FunctionType::new_endo(just_list.clone()),
+            FunctionType::new_endo(just_list.clone()).with_extension_delta(&exset),
         )?;
 
         let pred_const = cfg.add_constant(ops::Const::unary_unit_sum())?;
@@ -628,13 +627,31 @@ mod test {
                 },
                 op_sig.input()
             );
-            h.simple_entry_builder(op_sig.output, 1, op_sig.extension_reqs.clone())?
+            h.simple_entry_builder(op_sig.output.clone(), 1, op_sig.extension_reqs.clone())?
         } else {
-            h.simple_block_builder(op_sig, 1)?
+            h.simple_block_builder(op_sig.clone(), 1)?
         };
         let op: OpType = op.into();
         let op = bb.add_dataflow_op(op, bb.input_wires())?;
-        let load_pred = bb.load_const(pred_const)?;
+        let mut load_pred = bb.load_const(pred_const)?;
+        let const_ty = bb
+            .hugr()
+            .get_optype(pred_const.node())
+            .as_const()
+            .unwrap()
+            .const_type()
+            .clone();
+        for e in op_sig.extension_reqs.iter() {
+            [load_pred] = bb
+                .add_dataflow_op(
+                    LeafOp::Lift {
+                        type_row: vec![const_ty.clone()].into(),
+                        new_extension: e.clone(),
+                    },
+                    [load_pred],
+                )?
+                .outputs_arr();
+        }
         bb.finish_with_outputs(load_pred, op.outputs())
     }
 

src/hugr/rewrite/simple_replace.rs

@@ -221,7 +221,7 @@ pub(in crate::hugr::rewrite) mod test {
         HugrBuilder, ModuleBuilder,
     };
     use crate::extension::prelude::BOOL_T;
-    use crate::extension::{EMPTY_REG, PRELUDE_REGISTRY};
+    use crate::extension::{ExtensionSet, EMPTY_REG, PRELUDE_REGISTRY};
     use crate::hugr::views::{HugrView, SiblingSubgraph};
     use crate::hugr::{Hugr, HugrMut, Rewrite};
     use crate::ops::dataflow::DataflowOpTrait;
@@ -230,7 +230,7 @@ pub(in crate::hugr::rewrite) mod test {
     use crate::std_extensions::logic::test::and_op;
     use crate::type_row;
     use crate::types::{FunctionType, Type};
-    use crate::utils::test_quantum_extension::{cx_gate, h_gate};
+    use crate::utils::test_quantum_extension::{cx_gate, h_gate, EXTENSION_ID};
     use crate::{IncomingPort, Node};
 
     use super::SimpleReplacement;
@@ -249,17 +249,20 @@ pub(in crate::hugr::rewrite) mod test {
     fn make_hugr() -> Result<Hugr, BuildError> {
         let mut module_builder = ModuleBuilder::new();
         let _f_id = {
+            let delta = ExtensionSet::singleton(&EXTENSION_ID);
             let mut func_builder = module_builder.define_function(
                 "main",
-                FunctionType::new(type_row![QB, QB, QB], type_row![QB, QB, QB]).into(),
+                FunctionType::new_endo(type_row![QB, QB, QB])
+                    .with_extension_delta(&delta)
+                    .into(),
             )?;
 
             let [qb0, qb1, qb2] = func_builder.input_wires_arr();
 
             let q_out = func_builder.add_dataflow_op(h_gate(), vec![qb2])?;
 
             let mut inner_builder = func_builder.dfg_builder(
-                FunctionType::new(type_row![QB, QB], type_row![QB, QB]),
+                FunctionType::new_endo(type_row![QB, QB]).with_extension_delta(&delta),
                 None,
                 [qb0, qb1],
             )?;

src/hugr/validate/test.rs

@@ -16,7 +16,7 @@ use crate::macros::const_extension_ids;
 use crate::ops::dataflow::IOTrait;
 use crate::ops::{self, Const, LeafOp, OpType};
 use crate::std_extensions::logic::test::{and_op, or_op};
-use crate::std_extensions::logic::{self, NotOp};
+use crate::std_extensions::logic::{self, NotOp, EXTENSION_ID};
 use crate::types::type_param::{TypeArg, TypeArgError, TypeParam};
 use crate::types::{CustomType, FunctionType, PolyFuncType, Type, TypeBound, TypeRow};
 use crate::values::Value;
@@ -360,17 +360,18 @@ fn cfg_children_restrictions() {
 
 #[test]
 fn test_ext_edge() -> Result<(), HugrError> {
-    let mut h = closed_dfg_root_hugr(FunctionType::new(
-        type_row![BOOL_T, BOOL_T],
-        type_row![BOOL_T],
-    ));
+    let delta = ExtensionSet::singleton(&EXTENSION_ID);
+    let mut h = closed_dfg_root_hugr(
+        FunctionType::new(type_row![BOOL_T, BOOL_T], type_row![BOOL_T])
+            .with_extension_delta(&delta),
+    );
     let [input, output] = h.get_io(h.root()).unwrap();
 
     // Nested DFG BOOL_T -> BOOL_T
     let sub_dfg = h.add_node_with_parent(
         h.root(),
         ops::DFG {
-            signature: FunctionType::new_endo(type_row![BOOL_T]),
+            signature: FunctionType::new_endo(type_row![BOOL_T]).with_extension_delta(&delta),
         },
     )?;
     // this Xor has its 2nd input unconnected
@@ -411,7 +412,10 @@ const_extension_ids! {
 
 #[test]
 fn test_local_const() -> Result<(), HugrError> {
-    let mut h = closed_dfg_root_hugr(FunctionType::new(type_row![BOOL_T], type_row![BOOL_T]));
+    let mut h = closed_dfg_root_hugr(
+        FunctionType::new_endo(type_row![BOOL_T])
+            .with_extension_delta(&ExtensionSet::singleton(&EXTENSION_ID)),
+    );
     let [input, output] = h.get_io(h.root()).unwrap();
     let and = h.add_node_with_parent(h.root(), and_op())?;
     h.connect(input, 0, and, 0)?;