feat: Move python bindings of structs to tket-py

Cargo.toml

@@ -17,9 +17,10 @@ license-file = "LICENCE"
 [workspace.dependencies]
 
 tket2 = { path = "./tket2" }
-quantinuum-hugr = { git = "https://github.com/CQCL-DEV/hugr", rev = "b71cae6" }
+quantinuum-hugr = { git = "https://github.com/CQCL/hugr", rev = "b71cae6" }
 portgraph = { version = "0.10" }
 pyo3 = { version = "0.20" }
 itertools = { version = "0.11.0" }
-tket-json-rs = "0.2.0"
+tket-json-rs = { version = "0.3.0" }
 tracing = "0.1.37"
+portmatching = { git = "https://github.com/lmondada/portmatching", rev = "738c91c" }

tket2-py/Cargo.toml

@@ -22,3 +22,4 @@ pyo3 = { workspace = true, features = ["extension-module"] }
 num_cpus = "1.16.0"
 derive_more = "0.99.17"
 itertools = { workspace = true }
+portmatching = { workspace = true }

tket2-py/src/circuit.rs

@@ -1,27 +1,31 @@
 //! Circuit-related functionality and utilities.
 #![allow(unused)]
 
+pub mod convert;
+
+use derive_more::{From, Into};
 use pyo3::prelude::*;
+use std::fmt;
 
 use hugr::{Hugr, HugrView};
 use tket2::extension::REGISTRY;
 use tket2::json::TKETDecode;
-use tket2::passes::CircuitChunks;
 use tket2::rewrite::CircuitRewrite;
 use tket_json_rs::circuit_json::SerialCircuit;
 
+pub use self::convert::{try_update_hugr, try_with_hugr, update_hugr, with_hugr, T2Circuit};
+
 /// The module definition
 pub fn module(py: Python) -> PyResult<&PyModule> {
     let m = PyModule::new(py, "_circuit")?;
     m.add_class::<T2Circuit>()?;
+    m.add_class::<PyNode>()?;
     m.add_class::<tket2::T2Op>()?;
     m.add_class::<tket2::Pauli>()?;
-    m.add_class::<tket2::passes::CircuitChunks>()?;
 
     m.add_function(wrap_pyfunction!(validate_hugr, m)?)?;
     m.add_function(wrap_pyfunction!(to_hugr_dot, m)?)?;
     m.add_function(wrap_pyfunction!(to_hugr, m)?)?;
-    m.add_function(wrap_pyfunction!(chunks, m)?)?;
 
     m.add("HugrError", py.get_type::<hugr::hugr::PyHugrError>())?;
     m.add("BuildError", py.get_type::<hugr::builder::PyBuildError>())?;
@@ -41,72 +45,50 @@ pub fn module(py: Python) -> PyResult<&PyModule> {
     Ok(m)
 }
 
-/// Apply a fallible function expecting a hugr on a pytket circuit.
-pub fn try_with_hugr<T, E, F>(circ: Py<PyAny>, f: F) -> PyResult<T>
-where
-    E: Into<PyErr>,
-    F: FnOnce(Hugr) -> Result<T, E>,
-{
-    let hugr = SerialCircuit::_from_tket1(circ).decode()?;
-    (f)(hugr).map_err(|e| e.into())
-}
-
-/// Apply a function expecting a hugr on a pytket circuit.
-pub fn with_hugr<T, F: FnOnce(Hugr) -> T>(circ: Py<PyAny>, f: F) -> PyResult<T> {
-    try_with_hugr(circ, |hugr| Ok::<T, PyErr>((f)(hugr)))
-}
-
-/// Apply a hugr-to-hugr function on a pytket circuit, and return the modified circuit.
-pub fn try_update_hugr<E: Into<PyErr>, F: FnOnce(Hugr) -> Result<Hugr, E>>(
-    circ: Py<PyAny>,
-    f: F,
-) -> PyResult<Py<PyAny>> {
-    let hugr = try_with_hugr(circ, f)?;
-    SerialCircuit::encode(&hugr)?.to_tket1()
-}
-
-/// Apply a hugr-to-hugr function on a pytket circuit, and return the modified circuit.
-pub fn update_hugr<F: FnOnce(Hugr) -> Hugr>(circ: Py<PyAny>, f: F) -> PyResult<Py<PyAny>> {
-    let hugr = with_hugr(circ, f)?;
-    SerialCircuit::encode(&hugr)?.to_tket1()
-}
-
+/// Run the validation checks on a circuit.
 #[pyfunction]
 pub fn validate_hugr(c: Py<PyAny>) -> PyResult<()> {
     try_with_hugr(c, |hugr| hugr.validate(&REGISTRY))
 }
 
+/// Return a Graphviz DOT string representation of the circuit.
 #[pyfunction]
 pub fn to_hugr_dot(c: Py<PyAny>) -> PyResult<String> {
     with_hugr(c, |hugr| hugr.dot_string())
 }
 
+/// Downcast a python object to a [`Hugr`].
 #[pyfunction]
 pub fn to_hugr(c: Py<PyAny>) -> PyResult<Hugr> {
     with_hugr(c, |hugr| hugr)
 }
 
-#[pyfunction]
-pub fn chunks(c: Py<PyAny>, max_chunk_size: usize) -> PyResult<CircuitChunks> {
-    with_hugr(c, |hugr| CircuitChunks::split(&hugr, max_chunk_size))
-}
-
+/// A [`hugr::Node`] wrapper for Python.
 #[pyclass]
-/// A manager for tket 2 operations on a tket 1 Circuit.
-pub struct T2Circuit(pub Hugr);
+#[pyo3(name = "Node")]
+#[repr(transparent)]
+#[derive(From, Into, PartialEq, Eq, Hash, Clone, Copy)]
+pub struct PyNode {
+    /// Rust representation of the node
+    pub node: hugr::Node,
+}
 
-#[pymethods]
-impl T2Circuit {
-    #[new]
-    fn from_circuit(circ: PyObject) -> PyResult<Self> {
-        Ok(Self(to_hugr(circ)?))
+impl fmt::Display for PyNode {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.node.fmt(f)
     }
+}
 
-    fn finish(&self) -> PyResult<PyObject> {
-        SerialCircuit::encode(&self.0)?.to_tket1()
+impl fmt::Debug for PyNode {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.node.fmt(f)
     }
+}
 
-    fn apply_match(&mut self, rw: CircuitRewrite) {
-        rw.apply(&mut self.0).expect("Apply error.");
+#[pymethods]
+impl PyNode {
+    /// A string representation of the pattern.
+    pub fn __repr__(&self) -> String {
+        format!("{:?}", self)
     }
 }

tket2-py/src/circuit/convert.rs

@@ -0,0 +1,89 @@
+//! Utilities for calling Hugr functions on generic python objects.
+
+use pyo3::{prelude::*, PyTypeInfo};
+
+use hugr::{Hugr, HugrView};
+use tket2::extension::REGISTRY;
+use tket2::json::TKETDecode;
+use tket2::passes::CircuitChunks;
+use tket2::rewrite::CircuitRewrite;
+use tket_json_rs::circuit_json::SerialCircuit;
+
+/// A manager for tket 2 operations on a tket 1 Circuit.
+#[pyclass]
+#[derive(Clone, Debug, PartialEq)]
+pub struct T2Circuit {
+    /// Rust representation of the circuit.
+    pub hugr: Hugr,
+}
+
+#[pymethods]
+impl T2Circuit {
+    #[new]
+    fn from_circuit(circ: PyObject) -> PyResult<Self> {
+        Ok(Self {
+            hugr: super::to_hugr(circ)?,
+        })
+    }
+
+    fn finish(&self) -> PyResult<PyObject> {
+        SerialCircuit::encode(&self.hugr)?.to_tket1_with_gil()
+    }
+
+    fn apply_match(&mut self, rw: CircuitRewrite) {
+        rw.apply(&mut self.hugr).expect("Apply error.");
+    }
+}
+impl T2Circuit {
+    /// Tries to extract a T2Circuit from a python object.
+    ///
+    /// Returns an error if the py object is not a T2Circuit.
+    pub fn try_extract(circ: Py<PyAny>) -> PyResult<Self> {
+        Python::with_gil(|py| circ.as_ref(py).extract::<T2Circuit>())
+    }
+}
+
+/// Apply a fallible function expecting a hugr on a pytket circuit.
+pub fn try_with_hugr<T, E, F>(circ: Py<PyAny>, f: F) -> PyResult<T>
+where
+    E: Into<PyErr>,
+    F: FnOnce(Hugr) -> Result<T, E>,
+{
+    let hugr = Python::with_gil(|py| -> PyResult<Hugr> {
+        let circ = circ.as_ref(py);
+        match T2Circuit::extract(circ) {
+            // hugr circuit
+            Ok(t2circ) => Ok(t2circ.hugr),
+            // tket1 circuit
+            Err(_) => Ok(SerialCircuit::from_tket1(circ)?.decode()?),
+        }
+    })?;
+    (f)(hugr).map_err(|e| e.into())
+}
+
+/// Apply a function expecting a hugr on a pytket circuit.
+pub fn with_hugr<T, F>(circ: Py<PyAny>, f: F) -> PyResult<T>
+where
+    F: FnOnce(Hugr) -> T,
+{
+    try_with_hugr(circ, |hugr| Ok::<T, PyErr>((f)(hugr)))
+}
+
+/// Apply a hugr-to-hugr function on a pytket circuit, and return the modified circuit.
+pub fn try_update_hugr<E, F>(circ: Py<PyAny>, f: F) -> PyResult<Py<PyAny>>
+where
+    E: Into<PyErr>,
+    F: FnOnce(Hugr) -> Result<Hugr, E>,
+{
+    let hugr = try_with_hugr(circ, f)?;
+    SerialCircuit::encode(&hugr)?.to_tket1_with_gil()
+}
+
+/// Apply a hugr-to-hugr function on a pytket circuit, and return the modified circuit.
+pub fn update_hugr<F>(circ: Py<PyAny>, f: F) -> PyResult<Py<PyAny>>
+where
+    F: FnOnce(Hugr) -> Hugr,
+{
+    let hugr = with_hugr(circ, f)?;
+    SerialCircuit::encode(&hugr)?.to_tket1_with_gil()
+}

tket2-py/src/lib.rs

@@ -1,10 +1,10 @@
 //! Python bindings for TKET2.
 #![warn(missing_docs)]
 
-mod circuit;
-mod optimiser;
-mod passes;
-mod pattern;
+pub mod circuit;
+pub mod optimiser;
+pub mod passes;
+pub mod pattern;
 
 use pyo3::prelude::*;
 

tket2-py/src/passes.rs

@@ -1,3 +1,7 @@
+//! Passes for optimising circuits.
+
+pub mod chunks;
+
 use std::{cmp::min, convert::TryInto, fs, num::NonZeroUsize, path::PathBuf};
 
 use pyo3::{prelude::*, types::IntoPyDict};
@@ -16,7 +20,8 @@ pub fn module(py: Python) -> PyResult<&PyModule> {
     let m = PyModule::new(py, "_passes")?;
     m.add_function(wrap_pyfunction!(greedy_depth_reduce, m)?)?;
     m.add_function(wrap_pyfunction!(badger_optimise, m)?)?;
-    m.add_class::<tket2::T2Op>()?;
+    m.add_class::<self::chunks::PyCircuitChunks>()?;
+    m.add_function(wrap_pyfunction!(self::chunks::chunks, m)?)?;
     m.add(
         "PullForwardError",
         py.get_type::<tket2::passes::PyPullForwardError>(),
@@ -28,7 +33,7 @@ pub fn module(py: Python) -> PyResult<&PyModule> {
 fn greedy_depth_reduce(py_c: PyObject) -> PyResult<(PyObject, u32)> {
     try_with_hugr(py_c, |mut h| {
         let n_moves = apply_greedy_commutation(&mut h)?;
-        let py_c = SerialCircuit::encode(&h)?.to_tket1()?;
+        let py_c = SerialCircuit::encode(&h)?.to_tket1_with_gil()?;
         PyResult::Ok((py_c, n_moves))
     })
 }

tket2-py/src/passes/chunks.rs

@@ -0,0 +1,69 @@
+//! Circuit chunking utilities.
+
+use derive_more::From;
+use pyo3::exceptions::PyAttributeError;
+use pyo3::prelude::*;
+use tket2::json::TKETDecode;
+use tket2::passes::CircuitChunks;
+use tket2::Circuit;
+use tket_json_rs::circuit_json::SerialCircuit;
+
+use crate::circuit::{with_hugr, T2Circuit};
+
+/// Split a circuit into chunks of a given size.
+#[pyfunction]
+pub fn chunks(c: Py<PyAny>, max_chunk_size: usize) -> PyResult<PyCircuitChunks> {
+    with_hugr(c, |hugr| {
+        // TODO: Detect if the circuit is in tket1 format or T2Circuit.
+        let is_tket1 = true;
+        let chunks = CircuitChunks::split(&hugr, max_chunk_size);
+        (chunks, is_tket1).into()
+    })
+}
+
+/// A pattern that match a circuit exactly
+///
+/// Python equivalent of [`CircuitChunks`].
+///
+/// [`CircuitChunks`]: tket2::passes::chunks::CircuitChunks
+#[pyclass]
+#[pyo3(name = "CircuitChunks")]
+#[derive(Debug, Clone, From)]
+pub struct PyCircuitChunks {
+    /// Rust representation of the circuit chunks.
+    pub chunks: CircuitChunks,
+    /// Whether to reassemble the circuit in the tket1 format.
+    pub in_tket1: bool,
+}
+
+#[pymethods]
+impl PyCircuitChunks {
+    /// Reassemble the chunks into a circuit.
+    fn reassemble(&self) -> PyResult<Py<PyAny>> {
+        let hugr = self.clone().chunks.reassemble()?;
+        Python::with_gil(|py| match self.in_tket1 {
+            true => Ok(SerialCircuit::encode(&hugr)?.to_tket1(py)?.into_py(py)),
+            false => Ok(T2Circuit { hugr }.into_py(py)),
+        })
+    }
+
+    /// Returns clones of the split circuits.
+    fn circuits(&self) -> PyResult<Vec<Py<PyAny>>> {
+        self.chunks
+            .iter()
+            .map(|hugr| SerialCircuit::encode(hugr)?.to_tket1_with_gil())
+            .collect()
+    }
+
+    /// Replaces a chunk's circuit with an updated version.
+    fn update_circuit(&mut self, index: usize, new_circ: Py<PyAny>) -> PyResult<()> {
+        let hugr = SerialCircuit::from_tket1_with_gil(new_circ)?.decode()?;
+        if hugr.circuit_signature() != self.chunks[index].circuit_signature() {
+            return Err(PyAttributeError::new_err(
+                "The new circuit has a different signature.",
+            ));
+        }
+        self.chunks[index] = hugr;
+        Ok(())
+    }
+}