feat: make FuncDecl/FuncDefn polymorphic

src/builder.rs

@@ -5,6 +5,7 @@ use thiserror::Error;
 #[cfg(feature = "pyo3")]
 use pyo3::{create_exception, exceptions::PyException, PyErr};
 
+use crate::extension::SignatureError;
 use crate::hugr::{HugrError, ValidationError};
 use crate::ops::handle::{BasicBlockID, CfgID, ConditionalID, DfgID, FuncID, TailLoopID};
 use crate::types::ConstTypeError;
@@ -43,6 +44,10 @@ pub enum BuildError {
     /// The constructed HUGR is invalid.
     #[error("The constructed HUGR is invalid: {0}.")]
     InvalidHUGR(#[from] ValidationError),
+    /// SignatureError in trying to construct a node (differs from
+    /// [ValidationError::SignatureError] in that we could not construct a node to report about)
+    #[error(transparent)]
+    SignatureError(#[from] SignatureError),
     /// Tried to add a malformed [Const]
     ///
     /// [Const]: crate::ops::constant::Const
@@ -100,7 +105,7 @@ pub(crate) mod test {
 
     use crate::hugr::{views::HugrView, HugrMut, NodeType};
     use crate::ops;
-    use crate::types::{FunctionType, Signature, Type};
+    use crate::types::{FunctionType, PolyFuncType, Type};
     use crate::{type_row, Hugr};
 
     use super::handle::BuildHandle;
@@ -123,7 +128,7 @@ pub(crate) mod test {
     }
 
     pub(super) fn build_main(
-        signature: Signature,
+        signature: PolyFuncType,
         f: impl FnOnce(FunctionBuilder<&mut Hugr>) -> Result<BuildHandle<FuncID<true>>, BuildError>,
     ) -> Result<Hugr, BuildError> {
         let mut module_builder = ModuleBuilder::new();

src/builder/build_traits.rs

@@ -20,7 +20,7 @@ use crate::{
 };
 
 use crate::extension::{ExtensionRegistry, ExtensionSet, PRELUDE_REGISTRY};
-use crate::types::{FunctionType, Signature, Type, TypeRow};
+use crate::types::{FunctionType, PolyFuncType, Type, TypeArg, TypeRow};
 
 use itertools::Itertools;
 
@@ -90,22 +90,19 @@ pub trait Container {
     fn define_function(
         &mut self,
         name: impl Into<String>,
-        signature: Signature,
+        signature: PolyFuncType,
     ) -> Result<FunctionBuilder<&mut Hugr>, BuildError> {
+        let body = signature.body.clone();
         let f_node = self.add_child_node(NodeType::new(
             ops::FuncDefn {
                 name: name.into(),
-                signature: signature.signature.clone(),
+                signature,
             },
-            signature.input_extensions.clone(),
+            ExtensionSet::new(),
         ))?;
 
-        let db = DFGBuilder::create_with_io(
-            self.hugr_mut(),
-            f_node,
-            signature.signature,
-            Some(signature.input_extensions),
-        )?;
+        let db =
+            DFGBuilder::create_with_io(self.hugr_mut(), f_node, body, Some(ExtensionSet::new()))?;
         Ok(FunctionBuilder::from_dfg_builder(db))
     }
 
@@ -598,11 +595,14 @@ pub trait Dataflow: Container {
     fn call<const DEFINED: bool>(
         &mut self,
         function: &FuncID<DEFINED>,
+        type_args: &[TypeArg],
         input_wires: impl IntoIterator<Item = Wire>,
+        // Sadly required as we substituting in type_args may result in recomputing bounds of types:
+        exts: &ExtensionRegistry,
     ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
         let hugr = self.hugr();
         let def_op = hugr.get_optype(function.node());
-        let signature = match def_op {
+        let type_scheme = match def_op {
             OpType::FuncDefn(ops::FuncDefn { signature, .. })
             | OpType::FuncDecl(ops::FuncDecl { signature, .. }) => signature.clone(),
             _ => {
@@ -612,6 +612,7 @@ pub trait Dataflow: Container {
                 })
             }
         };
+        let signature = type_scheme.instantiate(type_args, exts)?;
         let op: OpType = ops::Call { signature }.into();
         let const_in_port = op.static_input_port().unwrap();
         let op_id = self.add_dataflow_op(op, input_wires)?;

src/builder/cfg.rs

@@ -336,7 +336,7 @@ mod test {
         let build_result = {
             let mut module_builder = ModuleBuilder::new();
             let mut func_builder = module_builder
-                .define_function("main", FunctionType::new(vec![NAT], type_row![NAT]).pure())?;
+                .define_function("main", FunctionType::new(vec![NAT], type_row![NAT]).into())?;
             let _f_id = {
                 let [int] = func_builder.input_wires_arr();
 

src/builder/circuit.rs

@@ -147,7 +147,7 @@ mod test {
     #[test]
     fn simple_linear() {
         let build_res = build_main(
-            FunctionType::new(type_row![QB, QB], type_row![QB, QB]).pure(),
+            FunctionType::new(type_row![QB, QB], type_row![QB, QB]).into(),
             |mut f_build| {
                 let wires = f_build.input_wires().collect();
 
@@ -184,7 +184,7 @@ mod test {
             .into(),
         );
         let build_res = build_main(
-            FunctionType::new(type_row![QB, QB, NAT], type_row![QB, QB, BOOL_T]).pure(),
+            FunctionType::new(type_row![QB, QB, NAT], type_row![QB, QB, BOOL_T]).into(),
             |mut f_build| {
                 let [q0, q1, angle]: [Wire; 3] = f_build.input_wires_arr();
 

src/builder/conditional.rs

@@ -240,7 +240,7 @@ mod test {
             let mut module_builder = ModuleBuilder::new();
             let mut fbuild = module_builder.define_function(
                 "main",
-                FunctionType::new(type_row![NAT], type_row![NAT]).pure(),
+                FunctionType::new(type_row![NAT], type_row![NAT]).into(),
             )?;
             let tru_const = fbuild.add_constant(Const::true_val(), ExtensionSet::new())?;
             let _fdef = {

src/builder/dataflow.rs

@@ -7,7 +7,7 @@ use std::marker::PhantomData;
 use crate::hugr::{HugrView, NodeType, ValidationError};
 use crate::ops;
 
-use crate::types::{FunctionType, Signature};
+use crate::types::{FunctionType, PolyFuncType};
 
 use crate::extension::{ExtensionRegistry, ExtensionSet};
 use crate::Node;
@@ -146,21 +146,17 @@ impl FunctionBuilder<Hugr> {
     /// # Errors
     ///
     /// Error in adding DFG child nodes.
-    pub fn new(name: impl Into<String>, signature: Signature) -> Result<Self, BuildError> {
+    pub fn new(name: impl Into<String>, signature: PolyFuncType) -> Result<Self, BuildError> {
+        let body = signature.body.clone();
         let op = ops::FuncDefn {
-            signature: signature.clone().into(),
+            signature,
             name: name.into(),
         };
 
-        let base = Hugr::new(NodeType::new(op, signature.input_extensions.clone()));
+        let base = Hugr::new(NodeType::new_pure(op));
         let root = base.root();
 
-        let db = DFGBuilder::create_with_io(
-            base,
-            root,
-            signature.signature,
-            Some(signature.input_extensions),
-        )?;
+        let db = DFGBuilder::create_with_io(base, root, body, Some(ExtensionSet::new()))?;
         Ok(Self::from_dfg_builder(db))
     }
 }
@@ -239,7 +235,7 @@ 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]).pure(),
+                    FunctionType::new(type_row![NAT, QB], type_row![NAT, QB]).into(),
                 )?;
 
                 let [int, qb] = func_builder.input_wires_arr();
@@ -273,7 +269,7 @@ pub(crate) mod test {
 
             let f_build = module_builder.define_function(
                 "main",
-                FunctionType::new(type_row![BOOL_T], type_row![BOOL_T, BOOL_T]).pure(),
+                FunctionType::new(type_row![BOOL_T], type_row![BOOL_T, BOOL_T]).into(),
             )?;
 
             f(f_build)?;
@@ -323,7 +319,7 @@ pub(crate) mod test {
 
             let f_build = module_builder.define_function(
                 "main",
-                FunctionType::new(type_row![QB], type_row![QB, QB]).pure(),
+                FunctionType::new(type_row![QB], type_row![QB, QB]).into(),
             )?;
 
             let [q1] = f_build.input_wires_arr();
@@ -340,7 +336,7 @@ pub(crate) mod test {
         let builder = || -> Result<Hugr, BuildError> {
             let mut f_build = FunctionBuilder::new(
                 "main",
-                FunctionType::new(type_row![BIT], type_row![BIT]).pure(),
+                FunctionType::new(type_row![BIT], type_row![BIT]).into(),
             )?;
 
             let [i1] = f_build.input_wires_arr();
@@ -364,7 +360,7 @@ pub(crate) mod test {
     fn error_on_linear_inter_graph_edge() -> Result<(), BuildError> {
         let mut f_build = FunctionBuilder::new(
             "main",
-            FunctionType::new(type_row![QB], type_row![QB]).pure(),
+            FunctionType::new(type_row![QB], type_row![QB]).into(),
         )?;
 
         let [i1] = f_build.input_wires_arr();
@@ -408,7 +404,7 @@ pub(crate) mod test {
         let (dfg_node, f_node) = {
             let mut f_build = module_builder.define_function(
                 "main",
-                FunctionType::new(type_row![BIT], type_row![BIT]).pure(),
+                FunctionType::new(type_row![BIT], type_row![BIT]).into(),
             )?;
 
             let [i1] = f_build.input_wires_arr();

src/builder/module.rs

@@ -8,13 +8,11 @@ use crate::{
     extension::ExtensionRegistry,
     hugr::{hugrmut::sealed::HugrMutInternals, views::HugrView, ValidationError},
     ops,
-    types::{Type, TypeBound},
+    types::{PolyFuncType, Type, TypeBound},
 };
 
 use crate::ops::handle::{AliasID, FuncID, NodeHandle};
 
-use crate::types::Signature;
-
 use crate::Node;
 use smol_str::SmolStr;
 
@@ -86,16 +84,13 @@ impl<T: AsMut<Hugr> + AsRef<Hugr>> ModuleBuilder<T> {
                 op_desc: "crate::ops::OpType::FuncDecl",
             })?
             .clone();
-
+        let body = signature.body.clone();
         self.hugr_mut().replace_op(
             f_node,
-            NodeType::new_pure(ops::FuncDefn {
-                name,
-                signature: signature.clone(),
-            }),
+            NodeType::new_pure(ops::FuncDefn { name, signature }),
         )?;
 
-        let db = DFGBuilder::create_with_io(self.hugr_mut(), f_node, signature, None)?;
+        let db = DFGBuilder::create_with_io(self.hugr_mut(), f_node, body, None)?;
         Ok(FunctionBuilder::from_dfg_builder(db))
     }
 
@@ -108,17 +103,13 @@ impl<T: AsMut<Hugr> + AsRef<Hugr>> ModuleBuilder<T> {
     pub fn declare(
         &mut self,
         name: impl Into<String>,
-        signature: Signature,
+        signature: PolyFuncType,
     ) -> Result<FuncID<false>, BuildError> {
         // TODO add param names to metadata
-        let rs = signature.input_extensions.clone();
-        let declare_n = self.add_child_node(NodeType::new(
-            ops::FuncDecl {
-                signature: signature.into(),
-                name: name.into(),
-            },
-            rs,
-        ))?;
+        let declare_n = self.add_child_node(NodeType::new_pure(ops::FuncDecl {
+            signature,
+            name: name.into(),
+        }))?;
 
         Ok(declare_n.into())
     }
@@ -176,7 +167,7 @@ mod test {
             test::{n_identity, NAT},
             Dataflow, DataflowSubContainer,
         },
-        extension::EMPTY_REG,
+        extension::{EMPTY_REG, PRELUDE_REGISTRY},
         type_row,
         types::FunctionType,
     };
@@ -189,11 +180,11 @@ mod test {
 
             let f_id = module_builder.declare(
                 "main",
-                FunctionType::new(type_row![NAT], type_row![NAT]).pure(),
+                FunctionType::new(type_row![NAT], type_row![NAT]).into(),
             )?;
 
             let mut f_build = module_builder.define_declaration(&f_id)?;
-            let call = f_build.call(&f_id, f_build.input_wires())?;
+            let call = f_build.call(&f_id, &[], f_build.input_wires(), &PRELUDE_REGISTRY)?;
 
             f_build.finish_with_outputs(call.outputs())?;
             module_builder.finish_prelude_hugr()
@@ -216,7 +207,7 @@ mod test {
                     vec![qubit_state_type.get_alias_type()],
                     vec![qubit_state_type.get_alias_type()],
                 )
-                .pure(),
+                .into(),
             )?;
             n_identity(f_build)?;
             module_builder.finish_hugr(&EMPTY_REG)
@@ -232,16 +223,17 @@ mod test {
 
             let mut f_build = module_builder.define_function(
                 "main",
-                FunctionType::new(type_row![NAT], type_row![NAT, NAT]).pure(),
+                FunctionType::new(type_row![NAT], type_row![NAT, NAT]).into(),
             )?;
             let local_build = f_build.define_function(
                 "local",
-                FunctionType::new(type_row![NAT], type_row![NAT, NAT]).pure(),
+                FunctionType::new(type_row![NAT], type_row![NAT, NAT]).into(),
             )?;
             let [wire] = local_build.input_wires_arr();
             let f_id = local_build.finish_with_outputs([wire, wire])?;
 
-            let call = f_build.call(f_id.handle(), f_build.input_wires())?;
+            let call =
+                f_build.call(f_id.handle(), &[], f_build.input_wires(), &PRELUDE_REGISTRY)?;
 
             f_build.finish_with_outputs(call.outputs())?;
             module_builder.finish_prelude_hugr()

src/builder/tail_loop.rs

@@ -130,10 +130,19 @@ mod test {
             let mut fbuild = module_builder.define_function(
                 "main",
                 FunctionType::new(type_row![BIT], type_row![NAT])
-                    .with_input_extensions(ExtensionSet::singleton(&PRELUDE_ID)),
+                    .with_extension_delta(&ExtensionSet::singleton(&PRELUDE_ID))
+                    .into(),
             )?;
             let _fdef = {
-                let [b1] = fbuild.input_wires_arr();
+                let [b1] = fbuild
+                    .add_dataflow_op(
+                        ops::LeafOp::Lift {
+                            type_row: type_row![BIT],
+                            new_extension: PRELUDE_ID,
+                        },
+                        fbuild.input_wires(),
+                    )?
+                    .outputs_arr();
                 let loop_id = {
                     let mut loop_b =
                         fbuild.tail_loop_builder(vec![(BIT, b1)], vec![], type_row![NAT])?;

src/extension/infer/test.rs

@@ -955,7 +955,7 @@ fn simple_funcdefn() -> Result<(), Box<dyn Error>> {
         "F",
         FunctionType::new(vec![NAT], vec![NAT])
             .with_extension_delta(&ExtensionSet::singleton(&A))
-            .pure(),
+            .into(),
     )?;
 
     let [w] = func_builder.input_wires_arr();
@@ -979,7 +979,7 @@ fn funcdefn_signature_mismatch() -> Result<(), Box<dyn Error>> {
         "F",
         FunctionType::new(vec![NAT], vec![NAT])
             .with_extension_delta(&ExtensionSet::singleton(&A))
-            .pure(),
+            .into(),
     )?;
 
     let [w] = func_builder.input_wires_arr();
@@ -1013,7 +1013,7 @@ fn funcdefn_signature_mismatch2() -> Result<(), Box<dyn Error>> {
         "F",
         FunctionType::new(vec![NAT], vec![NAT])
             .with_extension_delta(&ExtensionSet::singleton(&A))
-            .pure(),
+            .into(),
     )?;
 
     let [w] = func_builder.input_wires_arr();