feat: add inlining annotation + pass

tket-py/src/passes.rs

@@ -1,6 +1,7 @@
 //! Passes for optimising circuits.
 
 pub mod chunks;
+mod inline_always;
 mod inline_funcs;
 mod scope;
 pub mod tket1;
@@ -27,12 +28,14 @@ pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
     m.add_function(wrap_pyfunction!(greedy_depth_reduce, &m)?)?;
     m.add_function(wrap_pyfunction!(badger_optimise, &m)?)?;
     m.add_function(wrap_pyfunction!(normalize_guppy, &m)?)?;
+    m.add_function(wrap_pyfunction!(self::inline_always::inline_always, &m)?)?;
     m.add_function(wrap_pyfunction!(self::inline_funcs::inline_functions, &m)?)?;
     m.add_class::<self::chunks::PyCircuitChunks>()?;
     m.add_function(wrap_pyfunction!(self::chunks::chunks, &m)?)?;
     m.add_function(wrap_pyfunction!(self::tket1::tket1_pass, &m)?)?;
     m.add_function(wrap_pyfunction!(resolve_modifiers, &m)?)?;
     m.add("PullForwardError", py.get_type::<PyPullForwardError>())?;
+    m.add("InlineAlwaysError", py.get_type::<PyInlineAlwaysError>())?;
     m.add(
         "InlineFunctionsError",
         py.get_type::<PyInlineFunctionsError>(),
@@ -59,6 +62,12 @@ create_py_exception!(
     "Errors from the modifer resolver pass."
 );
 
+create_py_exception!(
+    tket::passes::inline_always::InlineAlwaysError,
+    PyInlineAlwaysError,
+    "Error from `InlineAlwaysPass`"
+);
+
 create_py_exception!(
     tket::passes::inline_funcs::InlineFuncsError,
     PyInlineFunctionsError,

tket-py/src/passes/inline_always.rs

@@ -0,0 +1,28 @@
+//! Python bindings for the `InlineFunctions` pass.
+
+use pyo3::prelude::*;
+
+use tket::passes::{ComposablePass, WithScope};
+
+use super::PyPassScope;
+use crate::state::CompilationState;
+use crate::utils::ConvertPyErr;
+
+/// Inline acyclic function calls below the selected scope.
+///
+/// Parameters:
+/// - heuristic: Heuristic used to choose which non-recursive functions to
+///   inline. Defaults to `tket.passes.MaxSize(64)`.
+/// - follow_inline_hints: Whether to follow compiler hints for inlining
+///   functions.
+#[pyfunction]
+#[pyo3(signature = (circ, *, scope = None))]
+pub(super) fn inline_always(
+    circ: &mut CompilationState,
+    scope: Option<PyPassScope>,
+) -> PyResult<()> {
+    let py_scope = scope.unwrap_or_default();
+    let pass = tket::passes::InlineAlwaysPass::default_with_scope(py_scope.scope);
+    pass.run(&mut circ.hugr).convert_pyerrs()?;
+    Ok(())
+}

tket-py/test/test_pass.py

@@ -7,6 +7,8 @@
 from tket.passes import (
     _badger_optimise,
     _greedy_depth_reduce,
+    InlineAlwaysError,
+    InlineAlwaysPass,
     InlineFunctions,
     inline_funcs,
     NormalizeGuppy,
@@ -23,6 +25,7 @@
 from tket.passes import PytketHugrPass
 from pytket.passes import CliffordSimp, SquashRzPhasedX, SequencePass
 from hugr.build.base import Hugr
+import hugr.tys as tys
 
 import numpy as np
 import pytest
@@ -285,6 +288,52 @@ def test_modifier_execution() -> None:
             np.testing.assert_allclose(computed_statevector, expected_statevector)
 
 
+@pytest.mark.parametrize("annotate", [True, False])
+def test_inline_always(annotate: bool) -> None:
+    import hugr.ops as ops
+    from hugr.build.dfg import Dfg
+
+    d = Dfg(tys.Tuple(tys.Qubit, tys.Qubit))
+
+    f_id = d.module_root_builder().define_function(
+        "id",
+        [tys.Qubit],
+    )
+    f_id.set_outputs(f_id.input_node[0])
+
+    if annotate:
+        f_id.metadata["tket.inline"] = "always"
+
+    (tup,) = d.inputs()
+    (q1, q2) = d.add(ops.UnpackTuple()(tup))
+    call1 = d.call(f_id, q1)
+    call2 = d.call(f_id, q2)
+    (tup,) = d.add(ops.MakeTuple()(call1, call2))
+
+    d.set_outputs(tup)
+
+    # validate(d.hugr)
+
+    InlineAlwaysPass()(d.hugr)
+    # validate(d.hugr)
+    assert _count_ops(d.hugr, "Call") == 0 if annotate else 2
+
+
+def test_inline_always_cycle() -> None:
+    from hugr.build.function import Module
+
+    mod = Module()
+
+    f_recursive = mod.define_function("recurse", [tys.Qubit])
+    f_recursive.declare_outputs([tys.Qubit])
+    call = f_recursive.call(f_recursive, f_recursive.input_node[0])
+    f_recursive.set_outputs(call)
+
+    f_recursive.metadata["tket.inline"] = "always"
+    with pytest.raises(InlineAlwaysError):
+        InlineAlwaysPass()(mod.hugr)
+
+
 def test_inline_functions() -> None:
     hugr = _hugr_from_path("test_files/guppy_examples/fn_calls.hugr")
 

tket-py/tket/_tket/passes.pyi

@@ -43,6 +43,13 @@ def normalize_guppy(
     - remove_redundant_order_edges: Whether to remove redundant order edges.
     """
 
+def inline_always(
+    circ: CompilationState,
+    *,
+    scope: PassScope = GlobalScope.PRESERVE_PUBLIC,
+) -> None:
+    """Inline functions marked with the `inline="always"` decorator below the selected scope."""
+
 def inline_functions(
     circ: CompilationState,
     *,

tket-py/tket/passes/__init__.py

@@ -12,7 +12,7 @@
 from tket import _state
 from . import inline_funcs
 from .._tket import passes as _passes, optimiser as _optimiser
-
+from .._tket.passes import InlineAlwaysError  # type: ignore
 from hugr.passes.composable import (
     ComposablePass,
     ComposedPass,
@@ -25,6 +25,8 @@
 __all__ = [
     "PytketHugrPass",
     "PassResult",
+    "InlineAlwaysError",
+    "InlineAlwaysPass",
     "InlineFuncsHeuristic",
     "InlineFunctions",
     "NormalizeGuppy",
@@ -147,6 +149,39 @@ def _run_tk(self, program: _state.CompilationState) -> _state.CompilationState:
         return program
 
 
+@dataclass
+class InlineAlwaysPass(ComposablePass):
+    """Inline functions marked with the `inline="always"` decorator below the selected scope."""
+
+    _scope: PassScope = GlobalScope.PRESERVE_PUBLIC
+
+    def run(self, hugr: Hugr, *, inplace: bool = True) -> PassResult:
+        return implement_pass_run(
+            self,
+            hugr=hugr,
+            inplace=inplace,
+            copy_call=lambda h: self._inline_always(h, inplace),
+        )
+
+    def with_scope(self, scope: PassScope) -> InlineAlwaysPass:
+        """Set the scope of this pass and return self."""
+        self._scope = scope
+        return self
+
+    def _inline_always(self, hugr: Hugr, inplace: bool) -> PassResult:
+        tk_program = _state.CompilationState.from_python(hugr)
+
+        _passes.inline_always(
+            tk_program._inner,
+            scope=self._scope,
+        )
+
+        package = tk_program.to_python()
+        return PassResult.for_pass(
+            self, hugr=package.modules[0], inplace=inplace, result=None
+        )
+
+
 @dataclass
 class InlineFunctions(ComposablePass):
     """Inline acyclic function calls below the selected scope.

tket/src/metadata.rs

@@ -12,6 +12,24 @@ impl Metadata for MaxQubits {
     type Type<'hugr> = u32;
 }
 
+/// Metadata that may be supplied for a function to indicate
+/// that/when calls to it should be inlined.
+#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "lowercase")]
+#[non_exhaustive]
+pub enum InlineAnnotation {
+    /// Always inline calls to this function.
+    ///
+    /// If this cannot be done, an error will be raised.
+    Always,
+}
+
+impl Metadata for InlineAnnotation {
+    type Type<'hugr> = Self;
+
+    const KEY: &'static str = "tket.inline";
+}
+
 /// Metadata key for traced rewrites that were applied during circuit transformation.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct CircuitRewriteTraces;

tket/src/passes.rs

@@ -42,6 +42,8 @@ pub mod inline_dfgs;
 pub use inline_dfgs::InlineDFGsPass;
 
 // Inline function calls.
+pub mod inline_always;
+pub use inline_always::InlineAlwaysPass;
 pub mod inline_funcs;
 pub use inline_funcs::InlineFunctionsPass;
 #[expect(deprecated)]