fix: Ops require their own extensions

src/algorithm/nest_cfgs.rs

@@ -605,8 +605,8 @@ pub(crate) mod test {
         //               \-> right -/             \-<--<-/
         let mut cfg_builder = CFGBuilder::new(FunctionType::new(type_row![NAT], type_row![NAT]))?;
 
-        let pred_const = cfg_builder.add_constant(Const::unit_sum(0, 2), ExtensionSet::new())?; // Nothing here cares which
-        let const_unit = cfg_builder.add_constant(Const::unary_unit_sum(), ExtensionSet::new())?;
+        let pred_const = cfg_builder.add_constant(Const::unit_sum(0, 2))?; // Nothing here cares which
+        let const_unit = cfg_builder.add_constant(Const::unary_unit_sum())?;
 
         let entry = n_identity(
             cfg_builder.simple_entry_builder(type_row![NAT], 1, ExtensionSet::new())?,
@@ -887,8 +887,8 @@ pub(crate) mod test {
         separate: bool,
     ) -> Result<(Hugr, BasicBlockID, BasicBlockID), BuildError> {
         let mut cfg_builder = CFGBuilder::new(FunctionType::new(type_row![NAT], type_row![NAT]))?;
-        let pred_const = cfg_builder.add_constant(Const::unit_sum(0, 2), ExtensionSet::new())?; // Nothing here cares which
-        let const_unit = cfg_builder.add_constant(Const::unary_unit_sum(), ExtensionSet::new())?;
+        let pred_const = cfg_builder.add_constant(Const::unit_sum(0, 2))?; // Nothing here cares which
+        let const_unit = cfg_builder.add_constant(Const::unary_unit_sum())?;
 
         let entry = n_identity(
             cfg_builder.simple_entry_builder(type_row![NAT], 2, ExtensionSet::new())?,
@@ -929,8 +929,8 @@ pub(crate) mod test {
         cfg_builder: &mut CFGBuilder<T>,
         separate_headers: bool,
     ) -> Result<(BasicBlockID, BasicBlockID), BuildError> {
-        let pred_const = cfg_builder.add_constant(Const::unit_sum(0, 2), ExtensionSet::new())?; // Nothing here cares which
-        let const_unit = cfg_builder.add_constant(Const::unary_unit_sum(), ExtensionSet::new())?;
+        let pred_const = cfg_builder.add_constant(Const::unit_sum(0, 2))?; // Nothing here cares which
+        let const_unit = cfg_builder.add_constant(Const::unary_unit_sum())?;
 
         let entry = n_identity(
             cfg_builder.simple_entry_builder(type_row![NAT], 1, ExtensionSet::new())?,

src/builder/build_traits.rs

@@ -70,12 +70,8 @@ pub trait Container {
     ///
     /// This function will return an error if there is an error in adding the
     /// [`OpType::Const`] node.
-    fn add_constant(
-        &mut self,
-        constant: ops::Const,
-        extensions: impl Into<Option<ExtensionSet>>,
-    ) -> Result<ConstID, BuildError> {
-        let const_n = self.add_child_node(NodeType::new(constant, extensions.into()))?;
+    fn add_constant(&mut self, constant: ops::Const) -> Result<ConstID, BuildError> {
+        let const_n = self.add_child_node(NodeType::new(constant, ExtensionSet::new()))?;
 
         Ok(const_n.into())
     }
@@ -356,20 +352,16 @@ pub trait Dataflow: Container {
     fn load_const(&mut self, cid: &ConstID) -> Result<Wire, BuildError> {
         let const_node = cid.node();
         let nodetype = self.hugr().get_nodetype(const_node);
-        let input_extensions = nodetype.input_extensions().cloned();
         let op: ops::Const = nodetype
             .op()
             .clone()
             .try_into()
             .expect("ConstID does not refer to Const op.");
 
-        let load_n = self.add_dataflow_node(
-            NodeType::new(
-                ops::LoadConstant {
-                    datatype: op.const_type().clone(),
-                },
-                input_extensions,
-            ),
+        let load_n = self.add_dataflow_op(
+            ops::LoadConstant {
+                datatype: op.const_type().clone(),
+            },
             // Constant wire from the constant value node
             vec![Wire::new(const_node, OutgoingPort::from(0))],
         )?;
@@ -382,12 +374,8 @@ pub trait Dataflow: Container {
     /// # Errors
     ///
     /// This function will return an error if there is an error when adding the node.
-    fn add_load_const(
-        &mut self,
-        constant: ops::Const,
-        extensions: ExtensionSet,
-    ) -> Result<Wire, BuildError> {
-        let cid = self.add_constant(constant, extensions)?;
+    fn add_load_const(&mut self, constant: ops::Const) -> Result<Wire, BuildError> {
+        let cid = self.add_constant(constant)?;
         self.load_const(&cid)
     }
 

src/builder/cfg.rs

@@ -385,7 +385,7 @@ mod test {
         let mut middle_b = cfg_builder
             .simple_block_builder(FunctionType::new(type_row![NAT], type_row![NAT]), 1)?;
         let middle = {
-            let c = middle_b.add_load_const(ops::Const::unary_unit_sum(), ExtensionSet::new())?;
+            let c = middle_b.add_load_const(ops::Const::unary_unit_sum())?;
             let [inw] = middle_b.input_wires_arr();
             middle_b.finish_with_outputs(c, [inw])?
         };
@@ -398,8 +398,7 @@ mod test {
     #[test]
     fn test_dom_edge() -> Result<(), BuildError> {
         let mut cfg_builder = CFGBuilder::new(FunctionType::new(type_row![NAT], type_row![NAT]))?;
-        let sum_tuple_const =
-            cfg_builder.add_constant(ops::Const::unary_unit_sum(), ExtensionSet::new())?;
+        let sum_tuple_const = cfg_builder.add_constant(ops::Const::unary_unit_sum())?;
         let sum_variants = vec![type_row![]];
 
         let mut entry_b =
@@ -427,8 +426,7 @@ mod test {
     #[test]
     fn test_non_dom_edge() -> Result<(), BuildError> {
         let mut cfg_builder = CFGBuilder::new(FunctionType::new(type_row![NAT], type_row![NAT]))?;
-        let sum_tuple_const =
-            cfg_builder.add_constant(ops::Const::unary_unit_sum(), ExtensionSet::new())?;
+        let sum_tuple_const = cfg_builder.add_constant(ops::Const::unary_unit_sum())?;
         let sum_variants = vec![type_row![]];
         let mut middle_b = cfg_builder
             .simple_block_builder(FunctionType::new(type_row![NAT], type_row![NAT]), 1)?;

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/conditional.rs

@@ -242,7 +242,7 @@ mod test {
                 "main",
                 FunctionType::new(type_row![NAT], type_row![NAT]).into(),
             )?;
-            let tru_const = fbuild.add_constant(Const::true_val(), ExtensionSet::new())?;
+            let tru_const = fbuild.add_constant(Const::true_val())?;
             let _fdef = {
                 let const_wire = fbuild.load_const(&tru_const)?;
                 let [int] = fbuild.input_wires_arr();

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/builder/tail_loop.rs

@@ -109,10 +109,7 @@ mod test {
         let build_result: Result<Hugr, ValidationError> = {
             let mut loop_b = TailLoopBuilder::new(vec![], vec![BIT], vec![USIZE_T])?;
             let [i1] = loop_b.input_wires_arr();
-            let const_wire = loop_b.add_load_const(
-                ConstUsize::new(1).into(),
-                ExtensionSet::singleton(&PRELUDE_ID),
-            )?;
+            let const_wire = loop_b.add_load_const(ConstUsize::new(1).into())?;
 
             let break_wire = loop_b.make_break(loop_b.loop_signature()?.clone(), [const_wire])?;
             loop_b.set_outputs(break_wire, [i1])?;
@@ -148,8 +145,7 @@ mod test {
                         fbuild.tail_loop_builder(vec![(BIT, b1)], vec![], type_row![NAT])?;
                     let signature = loop_b.loop_signature()?.clone();
                     let const_val = Const::true_val();
-                    let const_wire =
-                        loop_b.add_load_const(Const::true_val(), ExtensionSet::new())?;
+                    let const_wire = loop_b.add_load_const(Const::true_val())?;
                     let lift_node = loop_b.add_dataflow_op(
                         ops::LeafOp::Lift {
                             type_row: vec![const_val.const_type().clone()].into(),
@@ -177,10 +173,7 @@ mod test {
                         let mut branch_1 = conditional_b.case_builder(1)?;
                         let [_b1] = branch_1.input_wires_arr();
 
-                        let wire = branch_1.add_load_const(
-                            ConstUsize::new(2).into(),
-                            ExtensionSet::singleton(&PRELUDE_ID),
-                        )?;
+                        let wire = branch_1.add_load_const(ConstUsize::new(2).into())?;
                         let break_wire = branch_1.make_break(signature, [wire])?;
                         branch_1.finish_with_outputs([break_wire])?;
 

src/extension/infer.rs

@@ -317,15 +317,11 @@ impl UnificationContext {
             match node_type.io_extensions() {
                 // Input extensions are open
                 None => {
-                    let c = if let Some(sig) = node_type.op_signature() {
-                        let delta = sig.extension_reqs;
-                        if delta.is_empty() {
-                            Constraint::Equal(m_input)
-                        } else {
-                            Constraint::Plus(delta, m_input)
-                        }
-                    } else {
+                    let delta = node_type.op().extension_delta();
+                    let c = if delta.is_empty() {
                         Constraint::Equal(m_input)
+                    } else {
+                        Constraint::Plus(delta, m_input)
                     };
                     self.add_constraint(m_output, c);
                 }
@@ -652,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),