perf: Return Cow<Signature> where possible (#1743)

``` group cows serialization ----- ---- ------------- circuit_roundtrip/capnp/0 1.02 5.3±0.52µs ? ?/sec 1.00 5.2±0.09µs ? ?/sec circuit_roundtrip/capnp/1 1.00 12.5±0.74µs ? ?/sec 1.14 14.2±0.16µs ? ?/sec circuit_roundtrip/capnp/10 1.00 25.9±0.39µs ? ?/sec 1.31 34.0±0.30µs ? ?/sec circuit_roundtrip/capnp/100 1.00 155.8±2.01µs ? ?/sec 1.47 229.6±2.46µs ? ?/sec circuit_roundtrip/capnp/1000 1.00 1483.7±64.06µs ? ?/sec 1.43 2.1±0.04ms ? ?/sec circuit_roundtrip/capnp/10000 1.00 15.9±0.26ms ? ?/sec 1.49 23.7±1.67ms ? ?/sec circuit_roundtrip/capnp/100000 1.00 192.8±2.26ms ? ?/sec 1.37 264.3±5.44ms ? ?/sec circuit_roundtrip/capnp/200000 1.00 411.5±4.77ms ? ?/sec 1.36 560.4±15.52ms ? ?/sec circuit_roundtrip/capnp/300000 1.00 648.9±5.40ms ? ?/sec 1.34 871.5±10.63ms ? ?/sec circuit_roundtrip/json/0 1.00 8.0±0.18µs ? ?/sec 1.00 8.0±0.14µs ? ?/sec circuit_roundtrip/json/1 1.00 19.0±0.27µs ? ?/sec 1.16 22.1±0.33µs ? ?/sec circuit_roundtrip/json/10 1.00 50.0±0.68µs ? ?/sec 1.26 63.0±0.94µs ? ?/sec circuit_roundtrip/json/100 1.00 342.5±4.81µs ? ?/sec 1.30 446.4±4.61µs ? ?/sec circuit_roundtrip/json/1000 1.00 3.4±0.03ms ? ?/sec 1.31 4.5±0.35ms ? ?/sec circuit_roundtrip/json/10000 1.00 37.2±0.73ms ? ?/sec 1.28 47.5±0.45ms ? ?/sec circuit_roundtrip/json/100000 1.00 399.3±8.46ms ? ?/sec 1.27 507.7±16.13ms ? ?/sec circuit_roundtrip/json/1000000 1.00 4.5±0.26s ? ?/sec 1.20 5.4±0.06s ? ?/sec ``` Closes #1741
CQCL · Dec 10, 2024 · 7a49f05 · 7a49f05
1 parent f2215e9
commit 7a49f05
Show file tree

Hide file tree

Showing 27 changed files with 152 additions and 92 deletions.
diff --git a/hugr-core/src/builder/cfg.rs b/hugr-core/src/builder/cfg.rs
@@ -391,7 +391,7 @@ impl<B: AsMut<Hugr> + AsRef<Hugr>> BlockBuilder<B> {
             .get_optype(block_n)
             .as_dataflow_block()
             .unwrap();
-        let signature = block_op.inner_signature();
+        let signature = block_op.inner_signature().into_owned();
         let db = DFGBuilder::create_with_io(base, block_n, signature)?;
         Ok(BlockBuilder::from_dfg_builder(db))
     }

diff --git a/hugr-core/src/builder/conditional.rs b/hugr-core/src/builder/conditional.rs
@@ -107,7 +107,7 @@ impl<B: AsMut<Hugr> + AsRef<Hugr>> ConditionalBuilder<B> {
             .clone()
             .try_into()
             .expect("Parent node does not have Conditional optype.");
-        let extension_delta = cond.signature().extension_reqs;
+        let extension_delta = cond.signature().extension_reqs.clone();
         let inputs = cond
             .case_input_row(case)
             .ok_or(ConditionalBuildError::NotCase { conditional, case })?;

diff --git a/hugr-core/src/builder/dataflow.rs b/hugr-core/src/builder/dataflow.rs
@@ -248,7 +248,7 @@ impl FunctionBuilder<Hugr> {
             .get_optype(parent)
             .as_func_defn()
             .expect("FunctionBuilder node must be a FuncDefn");
-        let signature = old_optype.inner_signature();
+        let signature = old_optype.inner_signature().into_owned();
         let name = old_optype.name.clone();
         self.hugr_mut()
             .replace_op(

diff --git a/hugr-core/src/extension/prelude.rs b/hugr-core/src/extension/prelude.rs
@@ -1022,8 +1022,8 @@ mod test {
         let op = Lift::new(type_row![Type::UNIT], ExtensionSet::singleton(XA));
         let optype: OpType = op.clone().into();
         assert_eq!(
-            optype.dataflow_signature().unwrap(),
-            Signature::new_endo(type_row![Type::UNIT])
+            optype.dataflow_signature().unwrap().as_ref(),
+            &Signature::new_endo(type_row![Type::UNIT])
                 .with_extension_delta(XA)
                 .with_prelude()
         );

diff --git a/hugr-core/src/extension/resolution/ops.rs b/hugr-core/src/extension/resolution/ops.rs
@@ -98,8 +98,8 @@ pub(crate) fn resolve_op_extensions<'e>(
             node,
             extension: opaque.extension().clone(),
             op: def.name().clone(),
-            computed: ext_op.signature().clone(),
-            stored: opaque.signature().clone(),
+            computed: ext_op.signature().into_owned(),
+            stored: opaque.signature().into_owned(),
         }
         .into());
     };

diff --git a/hugr-core/src/hugr/rewrite/outline_cfg.rs b/hugr-core/src/hugr/rewrite/outline_cfg.rs
@@ -71,7 +71,7 @@ impl OutlineCfg {
                     }
                 }
             }
-            extension_delta = extension_delta.union(o.signature().extension_reqs);
+            extension_delta = extension_delta.union(o.signature().extension_reqs.clone());
             let external_succs = h.output_neighbours(n).filter(|s| !self.blocks.contains(s));
             match external_succs.at_most_one() {
                 Ok(None) => (), // No external successors

diff --git a/hugr-core/src/hugr/rewrite/replace.rs b/hugr-core/src/hugr/rewrite/replace.rs
@@ -615,9 +615,9 @@ mod test {
                 },
                 op_sig.input()
             );
-            h.simple_entry_builder_exts(op_sig.output, 1, op_sig.extension_reqs.clone())?
+            h.simple_entry_builder_exts(op_sig.output.clone(), 1, op_sig.extension_reqs.clone())?
         } else {
-            h.simple_block_builder(op_sig, 1)?
+            h.simple_block_builder(op_sig.into_owned(), 1)?
         };
         let op: OpType = op.into();
         let op = bb.add_dataflow_op(op, bb.input_wires())?;

diff --git a/hugr-core/src/hugr/validate.rs b/hugr-core/src/hugr/validate.rs
@@ -16,7 +16,7 @@ use crate::ops::custom::{ExtensionOp, OpaqueOpError};
 use crate::ops::validate::{ChildrenEdgeData, ChildrenValidationError, EdgeValidationError};
 use crate::ops::{FuncDefn, NamedOp, OpName, OpParent, OpTag, OpTrait, OpType, ValidateOp};
 use crate::types::type_param::TypeParam;
-use crate::types::{EdgeKind, Signature};
+use crate::types::EdgeKind;
 use crate::{Direction, Hugr, Node, Port};
 
 use super::views::{HierarchyView, HugrView, SiblingGraph};
@@ -61,7 +61,7 @@ impl Hugr {
                 return Err(ValidationError::ExtensionsNotInferred { node: parent });
             }
             let parent_extensions = match parent_op.inner_function_type() {
-                Some(Signature { extension_reqs, .. }) => extension_reqs,
+                Some(s) => s.extension_reqs.clone(),
                 None => match parent_op.tag() {
                     OpTag::Cfg | OpTag::Conditional => parent_op.extension_delta(),
                     // ModuleRoot holds but does not execute its children, so allow any extensions

diff --git a/hugr-core/src/hugr/validate/test.rs b/hugr-core/src/hugr/validate/test.rs
@@ -784,7 +784,7 @@ fn test_polymorphic_call() -> Result<(), Box<dyn std::error::Error>> {
     let exp_fun_ty = Signature::new(vec![utou(PRELUDE_ID), int_pair.clone()], int_pair)
         .with_extension_delta(EXT_ID)
         .with_prelude();
-    assert_eq!(call_ty, exp_fun_ty);
+    assert_eq!(call_ty.as_ref(), &exp_fun_ty);
     Ok(())
 }
 

diff --git a/hugr-core/src/hugr/views.rs b/hugr-core/src/hugr/views.rs
@@ -11,6 +11,8 @@ pub mod sibling_subgraph;
 #[cfg(test)]
 mod tests;
 
+use std::borrow::Cow;
+
 pub use self::petgraph::PetgraphWrapper;
 use self::render::RenderConfig;
 pub use descendants::DescendantsGraph;
@@ -311,7 +313,7 @@ pub trait HugrView: HugrInternals {
     ///
     /// In contrast to [`poly_func_type`][HugrView::poly_func_type], this
     /// method always return a concrete [`Signature`].
-    fn inner_function_type(&self) -> Option<Signature> {
+    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
         self.root_type().inner_function_type()
     }
 
@@ -408,7 +410,7 @@ pub trait HugrView: HugrInternals {
 
     /// Get the "signature" (incoming and outgoing types) of a node, non-Value
     /// kind ports will be missing.
-    fn signature(&self, node: Node) -> Option<Signature> {
+    fn signature(&self, node: Node) -> Option<Cow<'_, Signature>> {
         self.get_optype(node).dataflow_signature()
     }
 

diff --git a/hugr-core/src/hugr/views/descendants.rs b/hugr-core/src/hugr/views/descendants.rs
@@ -173,6 +173,8 @@ where
 
 #[cfg(test)]
 pub(super) mod test {
+    use std::borrow::Cow;
+
     use rstest::rstest;
 
     use crate::extension::prelude::{qb_t, usize_t};
@@ -241,7 +243,7 @@ pub(super) mod test {
 
         let inner_region: DescendantsGraph = DescendantsGraph::try_new(&hugr, inner)?;
         assert_eq!(
-            inner_region.inner_function_type(),
+            inner_region.inner_function_type().map(Cow::into_owned),
             Some(Signature::new(vec![usize_t()], vec![usize_t()]))
         );
         assert_eq!(inner_region.node_count(), 3);

diff --git a/hugr-core/src/hugr/views/sibling.rs b/hugr-core/src/hugr/views/sibling.rs
@@ -332,6 +332,8 @@ impl<Root: NodeHandle> HugrMut for SiblingMut<'_, Root> {}
 
 #[cfg(test)]
 mod test {
+    use std::borrow::Cow;
+
     use rstest::rstest;
 
     use crate::builder::test::simple_dfg_hugr;
@@ -376,7 +378,7 @@ mod test {
         );
 
         assert_eq!(
-            inner_region.inner_function_type(),
+            inner_region.inner_function_type().map(Cow::into_owned),
             Some(Signature::new(vec![usize_t()], vec![usize_t()]))
         );
         assert_eq!(inner_region.node_count(), 3);
@@ -487,7 +489,7 @@ mod test {
     fn flat_mut(mut simple_dfg_hugr: Hugr) {
         simple_dfg_hugr.validate().unwrap();
         let root = simple_dfg_hugr.root();
-        let signature = simple_dfg_hugr.inner_function_type().unwrap().clone();
+        let signature = simple_dfg_hugr.inner_function_type().unwrap().into_owned();
 
         let sib_mut = SiblingMut::<CfgID>::try_new(&mut simple_dfg_hugr, root);
         assert_eq!(
@@ -517,7 +519,11 @@ mod test {
     fn sibling_mut_covariance(mut simple_dfg_hugr: Hugr) {
         let root = simple_dfg_hugr.root();
         let case_nodetype = crate::ops::Case {
-            signature: simple_dfg_hugr.root_type().dataflow_signature().unwrap(),
+            signature: simple_dfg_hugr
+                .root_type()
+                .dataflow_signature()
+                .unwrap()
+                .into_owned(),
         };
         let mut sib_mut = SiblingMut::<DfgID>::try_new(&mut simple_dfg_hugr, root).unwrap();
         // As expected, we cannot replace the root with a Case

diff --git a/hugr-core/src/hugr/views/sibling_subgraph.rs b/hugr-core/src/hugr/views/sibling_subgraph.rs
@@ -388,7 +388,7 @@ impl SiblingSubgraph {
         {
             return Err(InvalidReplacement::InvalidSignature {
                 expected: self.signature(hugr),
-                actual: dfg_optype.dataflow_signature(),
+                actual: dfg_optype.dataflow_signature().map(|s| s.into_owned()),
             });
         }
 

diff --git a/hugr-core/src/ops.rs b/hugr-core/src/ops.rs
@@ -12,6 +12,8 @@ pub mod validate;
 use crate::extension::resolution::{
     collect_op_extension, collect_op_types_extensions, ExtensionCollectionError,
 };
+use std::borrow::Cow;
+
 use crate::extension::simple_op::MakeExtensionOp;
 use crate::extension::{ExtensionId, ExtensionRegistry, ExtensionSet};
 use crate::types::{EdgeKind, Signature};
@@ -377,7 +379,7 @@ pub trait OpTrait {
     /// The signature of the operation.
     ///
     /// Only dataflow operations have a signature, otherwise returns None.
-    fn dataflow_signature(&self) -> Option<Signature> {
+    fn dataflow_signature(&self) -> Option<Cow<'_, Signature>> {
         None
     }
 
@@ -440,13 +442,13 @@ pub trait OpParent {
     /// sibling graph.
     ///
     /// Non-container ops like `FuncDecl` return `None` even though they represent a function.
-    fn inner_function_type(&self) -> Option<Signature> {
+    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
         None
     }
 }
 
 impl<T: DataflowParent> OpParent for T {
-    fn inner_function_type(&self) -> Option<Signature> {
+    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
         Some(DataflowParent::inner_signature(self))
     }
 }

diff --git a/hugr-core/src/ops/constant.rs b/hugr-core/src/ops/constant.rs
@@ -2,6 +2,7 @@
 
 mod custom;
 
+use std::borrow::Cow;
 use std::collections::hash_map::DefaultHasher; // Moves into std::hash in Rust 1.76.
 use std::hash::{Hash, Hasher};
 
@@ -348,12 +349,15 @@ pub enum ConstTypeError {
 }
 
 /// Hugrs (even functions) inside Consts must be monomorphic
-fn mono_fn_type(h: &Hugr) -> Result<Signature, ConstTypeError> {
+fn mono_fn_type(h: &Hugr) -> Result<Cow<'_, Signature>, ConstTypeError> {
     let err = || ConstTypeError::NotMonomorphicFunction {
         hugr_root_type: h.root_type().clone(),
     };
     if let Some(pf) = h.poly_func_type() {
-        return pf.try_into().map_err(|_| err());
+        match pf.try_into() {
+            Ok(sig) => return Ok(Cow::Owned(sig)),
+            Err(_) => return Err(err()),
+        };
     }
 
     h.inner_function_type().ok_or_else(err)
@@ -367,7 +371,7 @@ impl Value {
             Self::Sum(Sum { sum_type, .. }) => sum_type.clone().into(),
             Self::Function { hugr } => {
                 let func_type = mono_fn_type(hugr).unwrap_or_else(|e| panic!("{}", e));
-                Type::new_function(func_type)
+                Type::new_function(func_type.into_owned())
             }
         }
     }

diff --git a/hugr-core/src/ops/controlflow.rs b/hugr-core/src/ops/controlflow.rs
@@ -1,5 +1,7 @@
 //! Control flow operations.
 
+use std::borrow::Cow;
+
 use crate::extension::ExtensionSet;
 use crate::types::{EdgeKind, Signature, Type, TypeRow};
 use crate::Direction;
@@ -31,10 +33,13 @@ impl DataflowOpTrait for TailLoop {
         "A tail-controlled loop"
     }
 
-    fn signature(&self) -> Signature {
+    fn signature(&self) -> Cow<'_, Signature> {
+        // TODO: Store a cached signature
         let [inputs, outputs] =
             [&self.just_inputs, &self.just_outputs].map(|row| row.extend(self.rest.iter()));
-        Signature::new(inputs, outputs).with_extension_delta(self.extension_delta.clone())
+        Cow::Owned(
+            Signature::new(inputs, outputs).with_extension_delta(self.extension_delta.clone()),
+        )
     }
 }
 
@@ -64,9 +69,12 @@ impl TailLoop {
 }
 
 impl DataflowParent for TailLoop {
-    fn inner_signature(&self) -> Signature {
-        Signature::new(self.body_input_row(), self.body_output_row())
-            .with_extension_delta(self.extension_delta.clone())
+    fn inner_signature(&self) -> Cow<'_, Signature> {
+        // TODO: Store a cached signature
+        Cow::Owned(
+            Signature::new(self.body_input_row(), self.body_output_row())
+                .with_extension_delta(self.extension_delta.clone()),
+        )
     }
 }
 
@@ -92,13 +100,16 @@ impl DataflowOpTrait for Conditional {
         "HUGR conditional operation"
     }
 
-    fn signature(&self) -> Signature {
+    fn signature(&self) -> Cow<'_, Signature> {
+        // TODO: Store a cached signature
         let mut inputs = self.other_inputs.clone();
         inputs
             .to_mut()
             .insert(0, Type::new_sum(self.sum_rows.clone()));
-        Signature::new(inputs, self.outputs.clone())
-            .with_extension_delta(self.extension_delta.clone())
+        Cow::Owned(
+            Signature::new(inputs, self.outputs.clone())
+                .with_extension_delta(self.extension_delta.clone()),
+        )
     }
 }
 
@@ -126,8 +137,8 @@ impl DataflowOpTrait for CFG {
         "A dataflow node defined by a child CFG"
     }
 
-    fn signature(&self) -> Signature {
-        self.signature.clone()
+    fn signature(&self) -> Cow<'_, Signature> {
+        Cow::Borrowed(&self.signature)
     }
 }
 
@@ -172,13 +183,16 @@ impl StaticTag for ExitBlock {
 }
 
 impl DataflowParent for DataflowBlock {
-    fn inner_signature(&self) -> Signature {
+    fn inner_signature(&self) -> Cow<'_, Signature> {
+        // TODO: Store a cached signature
         // The node outputs a Sum before the data outputs of the block node
         let sum_type = Type::new_sum(self.sum_rows.clone());
         let mut node_outputs = vec![sum_type];
         node_outputs.extend_from_slice(&self.other_outputs);
-        Signature::new(self.inputs.clone(), TypeRow::from(node_outputs))
-            .with_extension_delta(self.extension_delta.clone())
+        Cow::Owned(
+            Signature::new(self.inputs.clone(), TypeRow::from(node_outputs))
+                .with_extension_delta(self.extension_delta.clone()),
+        )
     }
 }
 
@@ -280,8 +294,8 @@ impl StaticTag for Case {
 }
 
 impl DataflowParent for Case {
-    fn inner_signature(&self) -> Signature {
-        self.signature.clone()
+    fn inner_signature(&self) -> Cow<'_, Signature> {
+        Cow::Borrowed(&self.signature)
     }
 }