perf: Avoid cloning signatures for simple port checks

hugr-core/src/hugr/views.rs

@@ -483,26 +483,22 @@ pub trait HugrView: HugrInternals {
         self.get_optype(node).dataflow_signature()
     }
 
-    /// Iterator over all outgoing ports that have Value type, along
-    /// with corresponding types.
+    /// Iterator over all ports in a direction that have Value type, along with
+    /// their corresponding types.
     fn value_types(&self, node: Self::Node, dir: Direction) -> impl Iterator<Item = (Port, Type)> {
-        let sig = self.signature(node).unwrap_or_default();
-        self.node_ports(node, dir)
-            .filter_map(move |port| sig.port_type(port).map(|typ| (port, typ.clone())))
+        self.get_optype(node).value_types(dir)
     }
 
     /// Iterator over all incoming ports that have Value type, along
     /// with corresponding types.
     fn in_value_types(&self, node: Self::Node) -> impl Iterator<Item = (IncomingPort, Type)> {
-        self.value_types(node, Direction::Incoming)
-            .map(|(p, t)| (p.as_incoming().unwrap(), t))
+        self.get_optype(node).value_input_types()
     }
 
     /// Iterator over all outgoing ports that have Value type, along
     /// with corresponding types.
     fn out_value_types(&self, node: Self::Node) -> impl Iterator<Item = (OutgoingPort, Type)> {
-        self.value_types(node, Direction::Outgoing)
-            .map(|(p, t)| (p.as_outgoing().unwrap(), t))
+        self.get_optype(node).value_output_types()
     }
 
     /// Returns the set of extensions used by the HUGR.

hugr-core/src/hugr/views/sibling_subgraph.rs

@@ -648,16 +648,16 @@ impl<N: HugrNode> SiblingSubgraph<N> {
             .iter()
             .map(|part| {
                 let &(n, p) = part.iter().next().expect("is non-empty");
-                let sig = hugr.signature(n).expect("must have dataflow signature");
-                sig.port_type(p).cloned().expect("must be dataflow edge")
+                let op = hugr.get_optype(n);
+                op.value_input_type(p).expect("must be dataflow edge")
             })
             .collect_vec();
         let output = self
             .outputs
             .iter()
             .map(|&(n, p)| {
-                let sig = hugr.signature(n).expect("must have dataflow signature");
-                sig.port_type(p).cloned().expect("must be dataflow edge")
+                let op = hugr.get_optype(n);
+                op.value_output_type(p).expect("must be dataflow edge")
             })
             .collect_vec();
 
@@ -1212,12 +1212,14 @@ fn get_edge_type<H: HugrView, P: Into<Port> + Copy>(
     ports: &[(H::Node, P)],
 ) -> Option<Type> {
     let &(n, p) = ports.first()?;
-    let edge_t = hugr.signature(n)?.port_type(p)?.clone();
+    let op = hugr.get_optype(n);
+    let edge_t = op.value_port_type(p.into())?.clone();
     ports
         .iter()
         .all(|&(n, p)| {
-            hugr.signature(n)
-                .is_some_and(|s| s.port_type(p) == Some(&edge_t))
+            hugr.get_optype(n)
+                .value_port_type(p.into())
+                .is_some_and(|t| t == edge_t)
         })
         .then_some(edge_t)
 }

hugr-core/src/ops.rs

@@ -112,7 +112,7 @@ use std::cmp::Ordering;
 
 use crate::extension::simple_op::MakeExtensionOp;
 use crate::extension::{ExtensionId, ExtensionRegistry};
-use crate::types::{EdgeKind, Signature, Substitution};
+use crate::types::{EdgeKind, Signature, Substitution, Type};
 use crate::{Direction, Node, OutgoingPort, Port};
 use crate::{IncomingPort, PortIndex};
 use handle::NodeHandle;
@@ -316,14 +316,13 @@ impl OpType {
     /// See [`OpType::dataflow_signature`], [`OpType::static_port_kind`], and
     /// [`OpType::other_port_kind`].
     pub fn port_kind(&self, port: impl Into<Port>) -> Option<EdgeKind> {
-        let signature = self.dataflow_signature().unwrap_or_default();
         let port: Port = port.into();
         let dir = port.direction();
-        let port_count = signature.port_count(dir);
+        let port_count = self.value_port_count(dir);
 
         // Dataflow ports
         if port.index() < port_count {
-            return signature.port_type(port).cloned().map(EdgeKind::Value);
+            return OpTrait::value_port_type(self, port).map(EdgeKind::Value);
         }
 
         // Constant port
@@ -410,6 +409,27 @@ impl OpType {
         (0..self.value_port_count(dir)).map(move |i| Port::new(dir, i))
     }
 
+    /// Return the dataflow value ports and their types for the given direction.
+    #[inline]
+    pub fn value_types(&self, dir: Direction) -> impl Iterator<Item = (Port, Type)> {
+        self.value_ports(dir)
+            .map(|port| (port, self.value_port_type(port).unwrap()))
+    }
+
+    /// Return the dataflow value input ports and their types.
+    #[inline]
+    pub fn value_input_types(&self) -> impl Iterator<Item = (IncomingPort, Type)> {
+        self.value_types(Direction::Incoming)
+            .map(|(port, typ)| (port.as_incoming().unwrap(), typ))
+    }
+
+    /// Return the dataflow value output ports and their types.
+    #[inline]
+    pub fn value_output_types(&self) -> impl Iterator<Item = (OutgoingPort, Type)> {
+        self.value_types(Direction::Outgoing)
+            .map(|(port, typ)| (port.as_outgoing().unwrap(), typ))
+    }
+
     /// Return the dataflow value input ports for the given direction.
     #[inline]
     #[must_use]
@@ -430,8 +450,7 @@ impl OpType {
     #[inline]
     #[must_use]
     pub fn value_port_count(&self, dir: portgraph::Direction) -> usize {
-        self.dataflow_signature()
-            .map_or(0, |sig| sig.port_count(dir))
+        OpTrait::value_port_count(self, dir)
     }
 
     /// The number of Value input ports.
@@ -577,6 +596,37 @@ pub trait OpTrait: Sized + Clone {
         None
     }
 
+    /// Returns the type of a value port.
+    ///
+    /// Implementations may override this to avoid constructing a complete
+    /// [`Signature`] when only one port type is needed.
+    fn value_port_type(&self, port: Port) -> Option<Type> {
+        self.dataflow_signature()?.port_type(port).cloned()
+    }
+
+    /// Returns the type of an input value port.
+    ///
+    /// Shorthand for `value_port_type(port.into())`.
+    fn value_input_type(&self, port: IncomingPort) -> Option<Type> {
+        self.value_port_type(port.into())
+    }
+
+    /// Returns the type of an output value port.
+    ///
+    /// Shorthand for `value_port_type(port.into())`.
+    fn value_output_type(&self, port: OutgoingPort) -> Option<Type> {
+        self.value_port_type(port.into())
+    }
+
+    /// Returns the number of value ports in one direction.
+    ///
+    /// Implementations may override this to avoid constructing a complete
+    /// [`Signature`] when only its size is needed.
+    fn value_port_count(&self, dir: Direction) -> usize {
+        self.dataflow_signature()
+            .map_or(0, |signature| signature.port_count(dir))
+    }
+
     /// The edge kind for the non-dataflow inputs of the operation,
     /// not described by the signature.
     ///
@@ -694,3 +744,99 @@ macro_rules! impl_validate_op {
 }
 
 use impl_validate_op;
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::types::{PolyFuncType, Type};
+    use rstest::rstest;
+
+    fn signature() -> Signature {
+        Signature::new([Type::UNIT, Type::new_unit_sum(2)], [Type::new_unit_sum(3)])
+    }
+
+    fn input() -> OpType {
+        Input {
+            types: signature().input,
+        }
+        .into()
+    }
+
+    fn output() -> OpType {
+        Output {
+            types: signature().output,
+        }
+        .into()
+    }
+
+    fn call_indirect() -> OpType {
+        CallIndirect {
+            signature: signature(),
+        }
+        .into()
+    }
+
+    fn load_constant() -> OpType {
+        LoadConstant {
+            datatype: Type::new_unit_sum(2),
+        }
+        .into()
+    }
+
+    fn load_function() -> OpType {
+        let instantiation = signature();
+        LoadFunction {
+            func_sig: PolyFuncType::new(Vec::new(), instantiation.clone()),
+            type_args: Vec::new(),
+            instantiation,
+        }
+        .into()
+    }
+
+    fn tag() -> OpType {
+        Tag::new(1, vec![vec![Type::UNIT].into(), signature().input]).into()
+    }
+
+    fn tail_loop() -> OpType {
+        TailLoop {
+            just_inputs: vec![Type::UNIT].into(),
+            just_outputs: vec![Type::new_unit_sum(2)].into(),
+            rest: vec![Type::new_unit_sum(3)].into(),
+        }
+        .into()
+    }
+
+    fn conditional() -> OpType {
+        Conditional {
+            sum_rows: vec![vec![Type::UNIT].into(), vec![Type::new_unit_sum(2)].into()],
+            other_inputs: signature().input,
+            outputs: signature().output,
+        }
+        .into()
+    }
+
+    /// Borrow-first queries must remain equivalent to the public signature API.
+    #[rstest]
+    #[case::input(input())]
+    #[case::output(output())]
+    #[case::call_indirect(call_indirect())]
+    #[case::load_constant(load_constant())]
+    #[case::load_function(load_function())]
+    #[case::tag(tag())]
+    #[case::tail_loop(tail_loop())]
+    #[case::conditional(conditional())]
+    fn value_ports_match_signature(#[case] op: OpType) {
+        let signature = op.dataflow_signature().expect("dataflow operation");
+
+        for dir in [Direction::Incoming, Direction::Outgoing] {
+            assert_eq!(op.value_port_count(dir), signature.port_count(dir));
+            for index in 0..signature.port_count(dir) {
+                let port = Port::new(dir, index);
+                assert_eq!(
+                    op.port_kind(port),
+                    signature.port_type(port).cloned().map(EdgeKind::Value)
+                );
+            }
+        }
+    }
+}

hugr-core/src/ops/controlflow.rs

@@ -2,8 +2,8 @@
 
 use std::borrow::Cow;
 
-use crate::Direction;
 use crate::types::{EdgeKind, Signature, Type, TypeRow, TypeRowLike};
+use crate::{Direction, Port, PortIndex};
 
 use super::OpTag;
 use super::dataflow::{DataflowOpTrait, DataflowParent};
@@ -37,6 +37,35 @@ impl DataflowOpTrait for TailLoop {
         Cow::Owned(Signature::new(inputs, outputs))
     }
 
+    fn value_port_type(&self, port: Port) -> Option<Type> {
+        let (head, tail) = match port.direction() {
+            Direction::Incoming => (&self.just_inputs, &self.rest),
+            Direction::Outgoing => (&self.just_outputs, &self.rest),
+        };
+        let mut index = port.index();
+
+        // The op value ports are defined as the concatenation of `head` and `tail`.
+        // See if `index` falls within any of those segments.
+        if index < head.len() {
+            return Some(head[index].clone());
+        }
+        index -= head.len();
+
+        if index < tail.len() {
+            return Some(tail[index].clone());
+        }
+
+        None
+    }
+
+    fn value_port_count(&self, dir: Direction) -> usize {
+        let directional = match dir {
+            Direction::Incoming => &self.just_inputs,
+            Direction::Outgoing => &self.just_outputs,
+        };
+        directional.len() + self.rest.len()
+    }
+
     fn substitute(&self, subst: &crate::types::Substitution) -> Self {
         Self {
             just_inputs: self.just_inputs.substitute(subst),
@@ -114,6 +143,21 @@ impl DataflowOpTrait for Conditional {
         Cow::Owned(Signature::new(inputs, self.outputs.clone()))
     }
 
+    fn value_port_type(&self, port: Port) -> Option<Type> {
+        match port.direction() {
+            Direction::Incoming if port.index() == 0 => Some(Type::new_sum(self.sum_rows.clone())),
+            Direction::Incoming => self.other_inputs.get(port.index() - 1).cloned(),
+            Direction::Outgoing => self.outputs.get(port.index()).cloned(),
+        }
+    }
+
+    fn value_port_count(&self, dir: Direction) -> usize {
+        match dir {
+            Direction::Incoming => self.other_inputs.len() + 1,
+            Direction::Outgoing => self.outputs.len(),
+        }
+    }
+
     fn substitute(&self, subst: &crate::types::Substitution) -> Self {
         Self {
             sum_rows: self.sum_rows.iter().map(|r| r.substitute(subst)).collect(),