feat: Add CircuitReplacer

tket-py/src/rewrite.rs

@@ -83,20 +83,8 @@ pub enum PyRewriter {
     Vec(Vec<PyRewriter>),
 }
 
-// impl<H: HugrView<Node = Node>> Rewriter<H> for PyRewriter {
-//     fn get_rewrites(&self, circ: &H) -> Vec<CircuitRewrite> {
-//         match self {
-//             Self::ECC(ecc) => ecc.0.get_rewrites(circ),
-//             Self::Vec(rewriters) => rewriters
-//                 .iter()
-//                 .flat_map(|r| r.get_rewrites(circ))
-//                 .collect(),
-//         }
-//     }
-// }
-
 impl<H: HugrView<Node = Node>> Rewriter<ResourceScope<H>> for PyRewriter {
-    fn get_rewrites(&self, circ: &ResourceScope<H>) -> Vec<CircuitRewrite<<H>::Node>> {
+    fn get_rewrites(&self, circ: &ResourceScope<H>) -> Vec<CircuitRewrite<H::Node>> {
         match self {
             Self::ECC(ecc) => ecc.0.get_rewrites(circ),
             Self::Vec(rewriters) => rewriters
@@ -157,7 +145,8 @@ impl PyECCRewriter {
         )?))
     }
 
-    /// Returns a list of circuit rewrites that can be applied to the given Tk2Circuit.
+    /// Returns a list of circuit rewrites that can be applied to the given
+    /// Tk2Circuit.
     pub fn get_rewrites(&self, circ: &Tk2Circuit) -> Vec<PyCircuitRewrite> {
         self.0
             .get_rewrites(&circ.circ)

tket/examples/badger-hadamard-opt.rs

@@ -0,0 +1,140 @@
+//! Using Badger to perform Hadamard cancellation.
+
+use hugr::{
+    builder::{endo_sig, DFGBuilder, Dataflow, DataflowHugr},
+    extension::prelude::qb_t,
+    HugrView,
+};
+use tket::{
+    op_matches,
+    optimiser::BadgerOptimiser,
+    resource::{CircuitUnit, ResourceScope},
+    rewrite::{
+        matcher::{CircuitMatcher, MatchContext, MatchOutcome},
+        replacer::CircuitReplacer,
+        strategy::LexicographicCostFunction,
+        MatchReplaceRewriter,
+    },
+    Circuit, Subcircuit, TketOp,
+};
+
+/// A matcher that matches two Hadamard gates in a row.
+#[derive(Clone, Copy, Debug)]
+struct TwoHMatcher;
+
+/// A replacement that replaces two Hadamard gates in a row with the identity.
+#[derive(Clone, Copy, Debug)]
+struct HadamardCancellation;
+
+/// State to keep track of how much has been matched so far.
+#[derive(Clone, Copy, Default, PartialEq, Eq, Hash)]
+enum PartialMatchState {
+    /// No hadamard matched so far.
+    #[default]
+    NoMatch,
+    /// One hadamard matched so far.
+    MatchedOne,
+}
+
+impl CircuitMatcher for TwoHMatcher {
+    type PartialMatchInfo = PartialMatchState;
+    type MatchInfo = ();
+
+    fn match_tket_op<H: HugrView>(
+        &self,
+        op: tket::TketOp,
+        _op_args: &[CircuitUnit<H::Node>],
+        match_context: MatchContext<Self::PartialMatchInfo, H>,
+    ) -> MatchOutcome<Self::PartialMatchInfo, Self::MatchInfo> {
+        if op != TketOp::H {
+            // We are not intersted in matching this op
+            return MatchOutcome::stop();
+        }
+
+        match match_context.match_info {
+            PartialMatchState::NoMatch => {
+                // Making progress! Proceed to matching the second Hadamard
+                MatchOutcome::default().proceed(PartialMatchState::MatchedOne)
+            }
+            PartialMatchState::MatchedOne => {
+                // We have matched two Hadamards, so we can report the match
+                MatchOutcome::default()
+                    .complete(()) // report a full match
+                    .skip(PartialMatchState::MatchedOne) // consider skipping
+                                                         // this op (and match
+                                                         // another one)
+            }
+        }
+    }
+}
+
+impl CircuitReplacer<()> for HadamardCancellation {
+    fn replace_match<H: hugr::HugrView>(
+        &self,
+        subcircuit: &Subcircuit<H::Node>,
+        circuit: &ResourceScope<H>,
+        _match_info: (),
+    ) -> Vec<tket::Circuit> {
+        let hugr = circuit.hugr();
+        // subgraph should be a pair of Hadamards
+        assert_eq!(subcircuit.node_count(circuit), 2);
+        assert!(subcircuit
+            .nodes(circuit)
+            .all(|n| op_matches(hugr.get_optype(n), TketOp::H)));
+
+        // The right hand side of the rewrite is just an empty one-qubit circuit
+        let h = DFGBuilder::new(endo_sig(qb_t())).unwrap();
+        let inps = h.input_wires();
+        let empty_circ = h.finish_hugr_with_outputs(inps).unwrap();
+
+        vec![Circuit::new(empty_circ)]
+    }
+}
+
+/// A 4-qubit circuit made of three layers
+///  - layer of 4x Hadamards on each qubit,
+///  - layer of CX gates between pairs of qubits,
+///  - layer of 4x Hadamards on each qubit,
+fn h_cx_h() -> Circuit {
+    let mut h = DFGBuilder::new(endo_sig(vec![qb_t(); 4])).unwrap();
+    let qbs = h.input_wires();
+    let mut circ = h.as_circuit(qbs);
+
+    for _ in 0..4 {
+        for i in 0..4 {
+            circ.append(TketOp::H, [i]).unwrap();
+        }
+    }
+
+    for i in (0..4).step_by(2) {
+        circ.append(TketOp::CX, [i, i + 1]).unwrap();
+    }
+
+    for _ in 0..4 {
+        for i in 0..4 {
+            circ.append(TketOp::H, [i]).unwrap();
+        }
+    }
+
+    let qbs = circ.finish();
+    Circuit::new(h.finish_hugr_with_outputs(qbs).unwrap())
+}
+
+fn main() {
+    let rewriter = MatchReplaceRewriter::new(TwoHMatcher, HadamardCancellation);
+
+    let optimiser =
+        BadgerOptimiser::new(rewriter, LexicographicCostFunction::default_cx_strategy());
+
+    let circuit = h_cx_h();
+
+    let optimised = optimiser.optimise(&circuit, Default::default());
+
+    // Only CX gates are left
+    assert_eq!(optimised.num_operations(), 2);
+    assert!(optimised
+        .operations()
+        .all(|cmd| op_matches(cmd.optype(), TketOp::CX)));
+
+    println!("Success!");
+}

tket/examples/circuit-matcher.rs

@@ -103,7 +103,7 @@ impl CircuitMatcher for CliffordMatcher {
 fn main() {
     const CIRCUIT: &str = r#"{"bits": [], "commands": [{"args": [["q", [0]], ["q", [1]]], "op": {"type": "CX"}}, {"args": [["q", [0]]], "op": {"params": ["0.5"], "type": "Rz"}}, {"args": [["q", [1]]], "op": {"type": "V"}}, {"args": [["q", [0]], ["q", [2]]], "op": {"type": "CX"}}, {"args": [["q", [0]], ["q", [1]]], "op": {"type": "CX"}}, {"args": [["q", [2]]], "op": {"type": "S"}}, {"args": [["q", [0]]], "op": {"params": ["0.111"], "type": "Rz"}}, {"args": [["q", [2]]], "op": {"type": "T"}}, {"args": [["q", [1]], ["q", [2]]], "op": {"type": "CX"}}, {"args": [["q", [1]]], "op": {"type": "T"}}, {"args": [["q", [2]]], "op": {"type": "S"}}, {"args": [["q", [0]], ["q", [1]]], "op": {"type": "CX"}}], "created_qubits": [], "discarded_qubits": [], "implicit_permutation": [[["q", [0]], ["q", [0]]], [["q", [1]], ["q", [1]]], [["q", [2]], ["q", [2]]]], "phase": "0.0", "qubits": [["q", [0]], ["q", [1]], ["q", [2]]]}"#;
     let ser_circ: SerialCircuit = serde_json::from_str(CIRCUIT).unwrap();
-    let circuit = ResourceScope::from_circuit(ser_circ.decode().unwrap());
+    let circuit = ResourceScope::from_circuit(ser_circ.decode(Default::default()).unwrap());
 
     let matcher = CliffordMatcher {
         allowed_num_cx: 2..4,

tket/src/optimiser/badger.rs

@@ -1,15 +1,15 @@
 //! Badger circuit optimiser.
 //!
 //! This module implements the Badger circuit optimiser. It relies on a rewriter
-//! and a RewriteStrategy instance to repeatedly rewrite a circuit and optimising
-//! it according to some cost metric (typically gate count).
+//! and a RewriteStrategy instance to repeatedly rewrite a circuit and
+//! optimising it according to some cost metric (typically gate count).
 //!
-//! The optimiser is implemented as a priority queue of circuits to be processed.
-//! On top of the queue are the circuits with the lowest cost. They are popped
-//! from the queue and replaced by the new circuits obtained from the rewriter
-//! and the rewrite strategy. A hash of every circuit computed is stored to
-//! detect and ignore duplicates. The priority queue is truncated whenever
-//! it gets too large.
+//! The optimiser is implemented as a priority queue of circuits to be
+//! processed. On top of the queue are the circuits with the lowest cost. They
+//! are popped from the queue and replaced by the new circuits obtained from the
+//! rewriter and the rewrite strategy. A hash of every circuit computed is
+//! stored to detect and ignore duplicates. The priority queue is truncated
+//! whenever it gets too large.
 
 mod eq_circ_class;
 pub mod log;
@@ -49,7 +49,8 @@ pub struct BadgerOptions {
     ///
     /// Defaults to `None`, which means no timeout.
     pub progress_timeout: Option<u64>,
-    /// The maximum number of circuits to process before stopping the optimisation.
+    /// The maximum number of circuits to process before stopping the
+    /// optimisation.
     ///
     /// For data parallel multi-threading, (split_circuit=true), applies on a
     /// per-thread basis, otherwise applies globally.
@@ -60,13 +61,14 @@ pub struct BadgerOptions {
     ///
     /// Defaults to `1`.
     pub n_threads: NonZeroUsize,
-    /// Whether to split the circuit into chunks and process each in a separate thread.
+    /// Whether to split the circuit into chunks and process each in a separate
+    /// thread.
     ///
-    /// If this option is set to `true`, the optimiser will split the circuit into `n_threads`
-    /// chunks.
+    /// If this option is set to `true`, the optimiser will split the circuit
+    /// into `n_threads` chunks.
     ///
-    /// If this option is set to `false`, the optimiser will run parallel searches on the whole
-    /// circuit.
+    /// If this option is set to `false`, the optimiser will run parallel
+    /// searches on the whole circuit.
     ///
     /// Defaults to `false`.
     pub split_circuit: bool,
@@ -99,10 +101,11 @@ impl Default for BadgerOptions {
 ///
 /// Optimisation is done by maintaining a priority queue of circuits and
 /// always processing the circuit with the lowest cost first. Rewrites are
-/// computed for that circuit and all new circuit obtained are added to the queue.
+/// computed for that circuit and all new circuit obtained are added to the
+/// queue.
 ///
-/// There are a single-threaded and two multi-threaded versions of the optimiser,
-/// controlled by setting the [`BadgerOptions::n_threads`] and
+/// There are a single-threaded and two multi-threaded versions of the
+/// optimiser, controlled by setting the [`BadgerOptions::n_threads`] and
 /// [`BadgerOptions::split_circuit`] fields.
 ///
 /// [Quartz]: https://arxiv.org/abs/2204.09033
@@ -366,7 +369,8 @@ where
         best_circ
     }
 
-    /// Run the Badger optimiser on a circuit, with data parallel multithreading.
+    /// Run the Badger optimiser on a circuit, with data parallel
+    /// multithreading.
     ///
     /// Split the circuit into chunks and process each in a separate thread.
     #[tracing::instrument(target = "badger::metrics", skip(self, circ, logger))]
@@ -509,7 +513,8 @@ mod badger_default {
             Ok(BadgerOptimiser::new(rewriter, strategy))
         }
 
-        /// An optimiser minimising Rz gate count using a precompiled binary rewriter.
+        /// An optimiser minimising Rz gate count using a precompiled binary
+        /// rewriter.
         #[cfg(feature = "binary-eccs")]
         pub fn rz_opt_with_rewriter_binary(
             rewriter_path: impl AsRef<Path>,
@@ -578,7 +583,8 @@ mod tests {
     /// ```
     const NON_COMPOSABLE: &str = r#"{"phase":"0.0","commands":[{"op":{"type":"CX","n_qb":2,"signature":["Q","Q"]},"args":[["q",[4]],["q",[1]]]},{"op":{"type":"CX","n_qb":2,"signature":["Q","Q"]},"args":[["q",[1]],["q",[2]]]},{"op":{"type":"U3","params":["0.5","0","0.5"],"signature":["Q"]},"args":[["q",[1]]]},{"op":{"type":"CX","n_qb":2,"signature":["Q","Q"]},"args":[["q",[3]],["q",[4]]]},{"op":{"type":"CX","n_qb":2,"signature":["Q","Q"]},"args":[["q",[4]],["q",[0]]]},{"op":{"type":"CX","n_qb":2,"signature":["Q","Q"]},"args":[["q",[0]],["q",[2]]]},{"op":{"type":"CX","n_qb":2,"signature":["Q","Q"]},"args":[["q",[0]],["q",[2]]]},{"op":{"type":"CX","n_qb":2,"signature":["Q","Q"]},"args":[["q",[3]],["q",[1]]]}],"qubits":[["q",[0]],["q",[1]],["q",[2]],["q",[3]],["q",[4]]],"bits":[],"implicit_permutation":[[["q",[0]],["q",[0]]],[["q",[1]],["q",[1]]],[["q",[2]],["q",[2]]],[["q",[3]],["q",[3]]],[["q",[4]],["q",[4]]]]}"#;
 
-    /// A circuit that would trigger non-composable rewrites, if we applied them blindly from nam_6_3 matches.
+    /// A circuit that would trigger non-composable rewrites, if we applied them
+    /// blindly from nam_6_3 matches.
     #[fixture]
     fn non_composable_rw_hugr() -> Circuit {
         load_tk1_json_str(NON_COMPOSABLE, DecodeOptions::new()).unwrap()

tket/src/optimiser/badger/worker.rs

@@ -2,11 +2,10 @@
 
 use std::thread::{self, JoinHandle};
 
-use crate::circuit::cost::CircuitCost;
 use crate::circuit::CircuitHash;
-use crate::resource::ResourceScope;
 use crate::rewrite::strategy::RewriteStrategy;
 use crate::rewrite::Rewriter;
+use crate::{circuit::cost::CircuitCost, resource::ResourceScope};
 
 use super::pqueue_worker::{StatePQueueChannels, Work};
 

tket/src/resource.rs

@@ -54,7 +54,7 @@ use hugr::{
 };
 pub use interval::{Interval, InvalidInterval};
 use itertools::Itertools;
-pub use scope::{ResourceScope, ResourceScopeConfig};
+pub use scope::{CircuitRewriteError, ResourceScope, ResourceScopeConfig};
 pub use types::{CircuitUnit, Position, ResourceAllocator, ResourceId};
 
 use crate::{
@@ -82,7 +82,8 @@ impl<H: HugrMut> ResourceScope<H> {
 
     /// Register a rewrite applied to the circuit.
     ///
-    /// Returns `true` if the rewrite was successfully registered, or `false` if it was ignored.
+    /// Returns `true` if the rewrite was successfully registered, or `false` if
+    /// it was ignored.
     #[inline]
     pub fn add_rewrite_trace(&mut self, rewrite: impl Into<RewriteTrace>) -> bool {
         self.as_circuit_mut().add_rewrite_trace(rewrite)

tket/src/resource/scope.rs

@@ -4,8 +4,10 @@
 //! tracking within a specific region of a HUGR, computing resource paths and
 //! providing efficient lookup of circuit units associated with ports.
 
-use std::collections::BTreeSet;
-use std::{cmp, iter};
+mod patch;
+pub use patch::CircuitRewriteError;
+
+use std::{cmp, collections::BTreeSet, iter};
 
 use crate::resource::flow::{DefaultResourceFlow, ResourceFlow};
 use crate::resource::types::{CircuitUnit, PortMap};
@@ -185,17 +187,15 @@ impl<H: HugrView> ResourceScope<H> {
         self.hugr
     }
 
-    pub(crate) fn hugr_mut(&mut self) -> &mut H {
-        &mut self.hugr
-    }
-
     /// Wrap the underlying HUGR in a Circuit as reference.
     pub fn as_circuit(&self) -> Circuit<&H> {
         Circuit::new(self.hugr())
     }
 
-    pub(crate) fn as_circuit_mut(&mut self) -> Circuit<&mut H> {
-        Circuit::new(self.hugr_mut())
+    /// Careful: this will not update the circuit units, so do not modify
+    /// the HUGR using this.
+    pub(super) fn as_circuit_mut(&mut self) -> Circuit<&mut H> {
+        Circuit::new(&mut self.hugr)
     }
 
     /// Get the underlying subgraph, or `None` if the circuit is empty.