From 3893f5a22dc1011f6ceb9f3f9a95fede28fa97c8 Mon Sep 17 00:00:00 2001
From: Luca Mondada <luca@mondada.net>
Date: Wed, 24 Jul 2024 17:21:18 +0200
Subject: [PATCH] feat: Expose advanced timeout options to tket2-py

---
 tket2-py/src/passes.rs   | 16 +++++++++++++---
 tket2-py/tket2/passes.py |  7 +++++--
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/tket2-py/src/passes.rs b/tket2-py/src/passes.rs
index a747738ab..ce521bc05 100644
--- a/tket2-py/src/passes.rs
+++ b/tket2-py/src/passes.rs
@@ -102,18 +102,27 @@ fn lower_to_pytket<'py>(circ: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>>
 /// optimising. This can be deactivated by setting `rebase` to `false`, in which
 /// case the circuit is expected to be in the Nam gate set.
 ///
-/// Will use at most `max_threads` threads (plus a constant) and take at most
-/// `timeout` seconds (plus a constant). Default to the number of cpus and
-/// 15min respectively.
+/// Will use at most `max_threads` threads (plus a constant). Defaults to the
+/// number of CPUs available.
+///
+/// The optimisation will terminate at the first of the following timeout
+/// criteria, if set:
+/// - `timeout` seconds (default: 15min) have elapsed since the start of the
+///    optimisation
+/// - `progress_timeout` (default: None) seconds have elapsed since progress
+///    in the cost function was last made
+/// - `max_circuit_cnt` (default: None) circuits have been explored.
 ///
 /// Log files will be written to the directory `log_dir` if specified.
 #[pyfunction]
+#[allow(clippy::too_many_arguments)]
 fn badger_optimise<'py>(
     circ: &Bound<'py, PyAny>,
     optimiser: &PyBadgerOptimiser,
     max_threads: Option<NonZeroUsize>,
     timeout: Option<u64>,
     progress_timeout: Option<u64>,
+    max_circuit_cnt: Option<usize>,
     log_dir: Option<PathBuf>,
     rebase: Option<bool>,
 ) -> PyResult<Bound<'py, PyAny>> {
@@ -165,6 +174,7 @@ fn badger_optimise<'py>(
                 progress_timeout,
                 n_threads: n_threads.try_into().unwrap(),
                 split_circuit: true,
+                max_circuit_cnt,
                 ..Default::default()
             };
             circ = optimiser.optimise(circ, log_file, options);
diff --git a/tket2-py/tket2/passes.py b/tket2-py/tket2/passes.py
index e63a50876..098ce15df 100644
--- a/tket2-py/tket2/passes.py
+++ b/tket2-py/tket2/passes.py
@@ -35,6 +35,7 @@ def badger_pass(
     max_threads: Optional[int] = None,
     timeout: Optional[int] = None,
     progress_timeout: Optional[int] = None,
+    max_circuit_cnt: Optional[int] = None,
     log_dir: Optional[Path] = None,
     rebase: bool = False,
 ) -> BasePass:
@@ -44,8 +45,9 @@ def badger_pass(
     `compile-rewriter <https://github.com/CQCL/tket2/tree/main/badger-optimiser>`_
     utility. If `rewriter` is not specified, a default one will be used.
 
-    The arguments `max_threads`, `timeout`, `log_dir` and `rebase` are optional
-    and will be passed on to the Badger optimiser if provided."""
+    The arguments `max_threads`, `timeout`, `progress_timeout`, `max_circuit_cnt`,
+    `log_dir` and `rebase` are optional and will be passed on to the Badger
+    optimiser if provided."""
     if rewriter is None:
         with resources.as_file(
             resources.files("tket2").joinpath("data/nam_6_3.rwr")
@@ -61,6 +63,7 @@ def apply(circuit: Circuit) -> Circuit:
             max_threads=max_threads,
             timeout=timeout,
             progress_timeout=progress_timeout,
+            max_circuit_cnt=max_circuit_cnt,
             log_dir=log_dir,
             rebase=rebase,
         )