fix: reuse generated operation functions (#1995)

## Summary

- make cached `OpFunctionMap` definitions public while HUGR linking runs
- reuse `NodeTemplate::call_to_function` so `OnMultiDefn::UseTarget`
coalesces repeated definitions
- rely on the existing qsystem hiding phase to restore generated helpers
to private visibility
- add generic and wrapped-barrier regressions asserting repeated
operations share one definition

## Root cause

`OpFunctionMap` cached one generated HUGR per operation shape, but its
replacement template embedded a private function HUGR at every call
site. Name linking only coalesces public definitions, so
`OnMultiDefn::UseTarget` had no effect and repeated barriers produced
thousands of equivalent wrapper functions. Wide wrappers then amplified
LLVM IPSCCP memory use.

The fix belongs in `func_as_node_template`: cached helpers are now
linker-visible when passed through the existing `call_to_function` path.
`QSystemRebasePass` already records pre-existing public functions and
hides newly introduced helpers after lowering, so the final HUGR retains
private linkage without adding a parallel call-template mechanism.

## Validation

- `cargo check -p tket -p tket-qsystem`
- `cargo test -p tket --lib` (604 passed, 2 ignored)
- `cargo test -p tket-qsystem --lib` (167 passed, 1 ignored)
- `just check`: formatting, Ruff, mypy, Cargo check/docs, and pg-libs
tests passed; the full local gate remains blocked by unrelated existing
nightly Clippy findings and macOS hugrenv `@rpath` resolution in
aggregate nextest/Python discovery
- `git diff --check origin/main...HEAD`

Author: qartik

tket-qsystem/src/extension/qsystem/barrier.rs

@@ -139,4 +139,29 @@ mod test {
             }
         }
     }
+
+    #[test]
+    fn repeated_barriers_share_wrapped_function() {
+        let mut b =
+            DFGBuilder::new(hugr::types::Signature::new_endo(vec![qb_t(), qb_t()])).unwrap();
+        let first = b.add_barrier(b.input_wires()).unwrap();
+        let second = b.add_barrier(first.outputs()).unwrap();
+        let mut h = b.finish_hugr_with_outputs(second.outputs()).unwrap();
+
+        lower_tk2_ops(&mut h, Preserve::Public, QSystemPlatform::Helios).unwrap();
+        h.validate().unwrap();
+
+        let wrapped_function = h
+            .children(h.module_root())
+            .filter(|&node| {
+                h.get_optype(node).as_func_defn().is_some_and(|op| {
+                    op.func_name()
+                        .contains(wrapped_barrier::WRAPPED_BARRIER_NAME.as_str())
+                })
+            })
+            .exactly_one()
+            .ok()
+            .expect("identical barriers should share one wrapped function");
+        assert_eq!(h.output_neighbours(wrapped_function).count(), 2);
+    }
 }

tket/src/passes/utils/unpack_container.rs

@@ -541,13 +541,15 @@ mod tests {
     use super::*;
     use hugr::{
         HugrView,
-        builder::{DFGBuilder, DataflowHugr as _},
+        builder::{DFGBuilder, DataflowHugr as _, FunctionBuilder},
         extension::prelude::{bool_t, option_type, qb_t, usize_t},
         std_extensions::collections::array::array_type,
         types::Signature,
     };
     use rstest::rstest;
 
+    use crate::passes::{ComposablePass, ReplaceTypes};
+
     #[test]
     fn test_container_factory_creation() {
         let analyzer = TypeUnpacker::for_qubits();
@@ -571,6 +573,52 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn repeated_ops_share_lowering_functions() -> Result<(), BuildError> {
+        let factory = UnpackContainerBuilder::new(TypeUnpacker::for_qubits());
+        let option_qb_type = Type::from(option_type([qb_t()]));
+        let option_usize_type = Type::from(option_type([usize_t()]));
+        let mut builder = FunctionBuilder::new(
+            "main",
+            Signature::new_endo([option_qb_type, option_usize_type]),
+        )?;
+
+        let qb_input = builder.input().out_wire(0);
+        let unwrapped = factory.unpack_option(&mut builder, qb_input, &qb_t())?;
+        let wrapped = factory.repack_option(&mut builder, unwrapped, &qb_t())?;
+        let unwrapped = factory.unpack_option(&mut builder, wrapped, &qb_t())?;
+        let qb_output = factory.repack_option(&mut builder, unwrapped, &qb_t())?;
+
+        let usize_input = builder.input().out_wire(1);
+        let unwrapped = factory.unpack_option(&mut builder, usize_input, &usize_t())?;
+        let usize_output = factory.repack_option(&mut builder, unwrapped, &usize_t())?;
+        let mut hugr = builder.finish_hugr_with_outputs([qb_output, usize_output])?;
+
+        let mut lowerer = ReplaceTypes::new_empty();
+        factory
+            .into_function_map()
+            .register_operation_replacements(&mut hugr, &mut lowerer);
+        lowerer.run(&mut hugr).unwrap();
+        hugr.validate().unwrap();
+
+        let lowering_functions = hugr
+            .children(hugr.module_root())
+            .filter(|&node| {
+                hugr.get_optype(node)
+                    .as_func_defn()
+                    .is_some_and(|op| op.func_name() != "main")
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(lowering_functions.len(), 4);
+        let mut use_counts = lowering_functions
+            .iter()
+            .map(|&function| hugr.output_neighbours(function).count())
+            .collect::<Vec<_>>();
+        use_counts.sort_unstable();
+        assert_eq!(use_counts, [1, 1, 2, 2]);
+        Ok(())
+    }
+
     #[rstest]
     #[case::array(Array, ARRAY_UNPACK, ARRAY_REPACK)]
     #[case::borrow_array(BorrowArray, BARRAY_UNPACK, BARRAY_REPACK)]

tket/src/passes/utils/unpack_container/op_function_map.rs

@@ -3,18 +3,15 @@
 use crate::passes::monomorphize::mangle_name;
 use crate::passes::{ReplaceTypes, replace_types::NodeTemplate};
 use hugr::HugrView;
-use hugr::builder::{Container, Dataflow, HugrBuilder};
-use hugr::hugr::linking::OnMultiDefn;
-use hugr::ops::handle::{FuncID, NodeHandle};
 use hugr::{
     Hugr, Node, Wire,
     builder::{BuildError, DataflowHugr, FunctionBuilder},
     hugr::hugrmut::HugrMut,
-    ops::{DataflowOpTrait, ExtensionOp},
+    ops::{DataflowOpTrait, ExtensionOp, OpType},
     types::TypeArg,
 };
 use hugr_core::Visibility;
-use hugr_core::hugr::linking::NameLinkingPolicy;
+use hugr_core::hugr::internal::HugrMutInternals;
 use indexmap::IndexMap;
 use std::{cell::RefCell, ops::Deref};
 
@@ -79,7 +76,16 @@ impl OpFunctionMap {
         if self.map.borrow().contains_key(&key) {
             return Ok(());
         }
-        let name = mangle_name(op.def().name(), mangle_args);
+        // Definitions are made public while name linking deduplicates them, so
+        // the symbol must distinguish every cached operation instance. Keep
+        // caller-selected arguments for readability, and append the complete
+        // operation arguments to guarantee uniqueness.
+        let name_args = mangle_args
+            .iter()
+            .chain(op.args())
+            .cloned()
+            .collect::<Vec<_>>();
+        let name = mangle_name(op.def().name(), name_args);
         let sig = op.signature().deref().clone();
         let mut func_b = FunctionBuilder::new(name, sig)?;
         // insert None as a placeholder to avoid cyclic recursion in func_builder call
@@ -122,42 +128,27 @@ impl OpFunctionMap {
         _hugr: &mut impl HugrMut<Node = Node>,
         lowerer: &mut ReplaceTypes,
     ) {
-        // Use the centralized cache for all operation replacements
         for (op, func_def) in self.into_function_iter() {
             lowerer.set_replace_op(&op, func_as_node_template(func_def));
         }
     }
 }
 
+/// Given a HUGR with a function definition as entrypoint, constructs a
+/// [`NodeTemplate::LinkedHugr`] that produces a call to the function.
+fn func_as_node_template(mut func_def: Hugr) -> NodeTemplate {
+    let entrypoint = func_def.entrypoint();
+    let OpType::FuncDefn(func) = func_def.optype_mut(entrypoint) else {
+        panic!("OpFunctionMap entries must be function definitions");
+    };
+    *func.visibility_mut() = Visibility::Public;
+
+    NodeTemplate::call_to_function(func_def, &[])
+        .expect("OpFunctionMap entries must be monomorphic function definitions")
+}
+
 impl Default for OpFunctionMap {
     fn default() -> Self {
         Self::new()
     }
 }
-
-/// Given a hugr with a function definition as entrypoint, constructs a
-/// [`NodeTemplate::LinkedHugr`] that produces a call to the function.
-//
-// TODO: Use [`NodeTemplate::call_to_function`] once it gets released in `hugr 0.25.6`.
-fn func_as_node_template(func_def: Hugr) -> NodeTemplate {
-    // Create a replacement hugr for the op nodes: Add a `call` node in the `func_def` hugr and set it as entrypoint.
-    let func_signature = func_def.inner_function_type().unwrap().into_owned();
-
-    // Build a new hugr and insert the function definition into it
-    let mut b = FunctionBuilder::new_vis("", func_signature, Visibility::Private).unwrap();
-    let func_id = FuncID::<true>::from(
-        b.module_root_builder()
-            .add_hugr(func_def)
-            .inserted_entrypoint,
-    );
-
-    // Build a call to the function in the new separate function.
-    let call = b.call(&func_id, &[], b.input_wires()).unwrap();
-    let mut call_hugr = b.finish_hugr_with_outputs(call.outputs()).unwrap();
-    call_hugr.set_entrypoint(call.node());
-
-    NodeTemplate::LinkedHugr(
-        Box::new(call_hugr),
-        NameLinkingPolicy::default().on_multiple_defn(OnMultiDefn::UseTarget),
-    )
-}