feat!: Allow static inputs to extension operations

hugr-core/src/builder/build_traits.rs

@@ -191,7 +191,28 @@ pub trait Dataflow: Container {
 
         Ok(outs.into())
     }
+    /// Add a dataflow [`OpType`] to the sibling graph, wiring up the `input_wires` to the
+    /// incoming ports of the resulting node, and the `static_wires` to the static input ports.
+    ///
+    /// # Errors
+    ///
+    /// Returns a [`BuildError::OperationWiring`] error if the `input_wires` or `static_wires` cannot be connected.
+    fn add_dataflow_op_with_static(
+        &mut self,
+        nodetype: impl Into<OpType>,
+        input_wires: impl IntoIterator<Item = Wire>,
+        static_wires: impl IntoIterator<Item = Wire>,
+    ) -> Result<BuildHandle<DataflowOpID>, BuildError> {
+        let op: OpType = nodetype.into();
+        let static_in_ports = op.static_input_ports();
+        let handle = self.add_dataflow_op(op, input_wires)?;
 
+        for (src, in_port) in static_wires.into_iter().zip(static_in_ports) {
+            self.hugr_mut()
+                .connect(src.node(), src.source(), handle.node(), in_port);
+        }
+        Ok(handle)
+    }
     /// Insert a hugr-defined op to the sibling graph, wiring up the
     /// `input_wires` to the incoming ports of the resulting root node.
     ///
@@ -672,7 +693,7 @@ pub trait Dataflow: Container {
             }
         };
         let op: OpType = ops::Call::try_new(type_scheme, type_args, exts)?.into();
-        let const_in_port = op.static_input_port().unwrap();
+        let const_in_port = op.static_input_ports()[0];
         let op_id = self.add_dataflow_op(op, input_wires)?;
         let src_port = self.hugr_mut().num_outputs(function.node()) - 1;
 

hugr-core/src/export.rs

@@ -151,7 +151,7 @@ impl<'a> Context<'a> {
     /// Get the node that declares or defines the function associated with the given
     /// node via the static input. Returns `None` if the node is not connected to a function.
     fn connected_function(&self, node: Node) -> Option<Node> {
-        let func_node = self.hugr.static_source(node)?;
+        let func_node = *self.hugr.static_sources(node).first()?;
 
         match self.hugr.get_optype(func_node) {
             OpType::FuncDecl(_) => Some(func_node),
@@ -417,7 +417,7 @@ impl<'a> Context<'a> {
         use std::collections::hash_map::Entry;
 
         let poly_func_type = match opdef.signature_func() {
-            SignatureFunc::PolyFuncType(poly_func_type) => poly_func_type,
+            SignatureFunc::PolyFuncType(op_def_sig) => op_def_sig.poly_func_type(),
             _ => return self.make_named_global_ref(opdef.extension(), opdef.name()),
         };
 

hugr-core/src/extension.rs

@@ -22,8 +22,8 @@ use crate::types::{Signature, TypeNameRef};
 
 mod op_def;
 pub use op_def::{
-    CustomSignatureFunc, CustomValidator, LowerFunc, OpDef, SignatureFromArgs, SignatureFunc,
-    ValidateJustArgs, ValidateTypeArgs,
+    CustomSignatureFunc, CustomValidator, ExtOpSignature, LowerFunc, OpDef, SignatureFromArgs,
+    SignatureFunc, ValidateJustArgs, ValidateTypeArgs,
 };
 mod type_def;
 pub use type_def::{TypeDef, TypeDefBound};

hugr-core/src/extension/declarative/signature.rs

@@ -14,7 +14,7 @@ use smol_str::SmolStr;
 use crate::extension::prelude::PRELUDE_ID;
 use crate::extension::{ExtensionSet, SignatureFunc, TypeDef};
 use crate::types::type_param::TypeParam;
-use crate::types::{CustomType, FuncValueType, PolyFuncTypeRV, Type, TypeRowRV};
+use crate::types::{CustomType, FuncValueType, OpDefSignature, Type, TypeRowRV};
 use crate::Extension;
 
 use super::{DeclarationContext, ExtensionDeclarationError};
@@ -56,7 +56,7 @@ impl SignatureDeclaration {
             extension_reqs: self.extensions.clone(),
         };
 
-        let poly_func = PolyFuncTypeRV::new(op_params, body);
+        let poly_func = OpDefSignature::new(op_params, body);
         Ok(poly_func.into())
     }
 }