fix(tket2-py): Replace with_hugr helpers with with_circ, keeping …

…the circuit parent. (#438)

Avoids unwrapping the hugr inside `Tk2Circuit`.
This caused problems with circuits with non-root parents. Circuits
defined in guppy ended up pointing to the root module instead of the
chosen function definition.

Adds the first rust-side test for `tket2-py`.
I'm pretty sure I'll have to install the python interpreter in the CI
rust test job, but let's see the error first.

tket2-py/Cargo.toml

@@ -19,7 +19,9 @@ test = false
 bench = false
 
 [dependencies]
-tket2 = { path = "../tket2", version = "0.1.0-alpha.1", features = ["portmatching"] }
+tket2 = { path = "../tket2", version = "0.1.0-alpha.1", features = [
+    "portmatching",
+] }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 tket-json-rs = { workspace = true, features = ["pyo3"] }
@@ -31,3 +33,7 @@ derive_more = { workspace = true }
 itertools = { workspace = true }
 portmatching = { workspace = true }
 strum = { workspace = true }
+
+[dev-dependencies]
+rstest = { workspace = true }
+cool_asserts = { workspace = true }

tket2-py/src/circuit.rs

@@ -23,7 +23,7 @@ use tket_json_rs::circuit_json::SerialCircuit;
 use crate::utils::create_py_exception;
 use crate::utils::ConvertPyErr;
 
-pub use self::convert::{try_update_hugr, try_with_hugr, update_hugr, with_hugr, CircuitType};
+pub use self::convert::{try_update_circ, try_with_circ, update_circ, with_circ, CircuitType};
 pub use self::cost::PyCircuitCost;
 use self::tk2circuit::Dfg;
 pub use self::tk2circuit::Tk2Circuit;
@@ -88,19 +88,19 @@ create_py_exception!(
 /// Run the validation checks on a circuit.
 #[pyfunction]
 pub fn validate_circuit(c: &Bound<PyAny>) -> PyResult<()> {
-    try_with_hugr(c, |hugr, _| hugr.validate(&REGISTRY))
+    try_with_circ(c, |circ, _| circ.hugr().validate(&REGISTRY))
 }
 
 /// Return a Graphviz DOT string representation of the circuit.
 #[pyfunction]
 pub fn render_circuit_dot(c: &Bound<PyAny>) -> PyResult<String> {
-    with_hugr(c, |hugr, _| hugr.dot_string())
+    with_circ(c, |hugr, _| hugr.dot_string())
 }
 
 /// Return a Mermaid diagram representation of the circuit.
 #[pyfunction]
 pub fn render_circuit_mermaid(c: &Bound<PyAny>) -> PyResult<String> {
-    with_hugr(c, |hugr, _| hugr.mermaid_string())
+    with_circ(c, |hugr, _| hugr.mermaid_string())
 }
 
 /// A [`hugr::Node`] wrapper for Python.

tket2-py/src/circuit/convert.rs

@@ -1,4 +1,4 @@
-//! Utilities for calling Hugr functions on generic python objects.
+//! Utilities for calling tket2 Circuit functions on generic python objects.
 
 use std::borrow::Borrow;
 
@@ -41,74 +41,98 @@ pub enum CircuitType {
 }
 
 impl CircuitType {
-    /// Converts a `Hugr` into the format indicated by the flag.
-    pub fn convert(self, py: Python, hugr: Hugr) -> PyResult<Bound<PyAny>> {
+    /// Converts a circuit into the format indicated by the flag.
+    pub fn convert(self, py: Python, circ: Circuit) -> PyResult<Bound<PyAny>> {
         match self {
-            CircuitType::Tket1 => SerialCircuit::encode(&hugr.into())
-                .convert_pyerrs()?
-                .to_tket1(py),
-            CircuitType::Tket2 => Ok(Bound::new(py, Tk2Circuit { circ: hugr.into() })?.into_any()),
+            CircuitType::Tket1 => SerialCircuit::encode(&circ).convert_pyerrs()?.to_tket1(py),
+            CircuitType::Tket2 => Ok(Bound::new(py, Tk2Circuit { circ })?.into_any()),
         }
     }
 }
 
-/// Apply a fallible function expecting a hugr on a python circuit.
+/// Apply a fallible function expecting a tket2 circuit on a python object.
 ///
 /// This method supports both `pytket.Circuit` and `Tk2Circuit` python objects.
-pub fn try_with_hugr<T, E, F>(circ: &Bound<PyAny>, f: F) -> PyResult<T>
+pub fn try_with_circ<T, E, F>(circ: &Bound<PyAny>, f: F) -> PyResult<T>
 where
     E: ConvertPyErr<Output = PyErr>,
-    F: FnOnce(Hugr, CircuitType) -> Result<T, E>,
+    F: FnOnce(Circuit, CircuitType) -> Result<T, E>,
 {
     let (circ, typ) = match Tk2Circuit::extract_bound(circ) {
-        // hugr circuit
+        // tket2 circuit
         Ok(t2circ) => (t2circ.circ, CircuitType::Tket2),
         // tket1 circuit
         Err(_) => (
             SerialCircuit::from_tket1(circ)?.decode().convert_pyerrs()?,
             CircuitType::Tket1,
         ),
     };
-    (f)(circ.into_hugr(), typ).map_err(|e| e.convert_pyerrs())
+    (f)(circ, typ).map_err(|e| e.convert_pyerrs())
 }
 
-/// Apply a function expecting a hugr on a python circuit.
+/// Apply a function expecting a tket2 circuit on a python object.
 ///
 /// This method supports both `pytket.Circuit` and `Tk2Circuit` python objects.
-pub fn with_hugr<T, F>(circ: &Bound<PyAny>, f: F) -> PyResult<T>
+pub fn with_circ<T, F>(circ: &Bound<PyAny>, f: F) -> PyResult<T>
 where
-    F: FnOnce(Hugr, CircuitType) -> T,
+    F: FnOnce(Circuit, CircuitType) -> T,
 {
-    try_with_hugr(circ, |hugr, typ| Ok::<T, PyErr>((f)(hugr, typ)))
+    try_with_circ(circ, |circ, typ| Ok::<T, PyErr>((f)(circ, typ)))
 }
 
-/// Apply a fallible hugr-to-hugr function on a python circuit, and return the modified circuit.
+/// Apply a fallible circuit-to-circuit function on a python object, and return the modified result.
 ///
 /// This method supports both `pytket.Circuit` and `Tk2Circuit` python objects.
-/// The returned Hugr is converted to the matching python object.
-pub fn try_update_hugr<'py, E, F>(circ: &Bound<'py, PyAny>, f: F) -> PyResult<Bound<'py, PyAny>>
+/// The returned [`Circuit`] is converted back into the circuit representation used in the input.
+pub fn try_update_circ<'py, E, F>(circ: &Bound<'py, PyAny>, f: F) -> PyResult<Bound<'py, PyAny>>
 where
     E: ConvertPyErr<Output = PyErr>,
-    F: FnOnce(Hugr, CircuitType) -> Result<Hugr, E>,
+    F: FnOnce(Circuit, CircuitType) -> Result<Circuit, E>,
 {
     let py = circ.py();
-    try_with_hugr(circ, |hugr, typ| {
-        let hugr = f(hugr, typ).map_err(|e| e.convert_pyerrs())?;
-        typ.convert(py, hugr)
+    try_with_circ(circ, |circ, typ| {
+        let circ = f(circ, typ).map_err(|e| e.convert_pyerrs())?;
+        typ.convert(py, circ)
     })
 }
 
-/// Apply a hugr-to-hugr function on a python circuit, and return the modified circuit.
+/// Apply a circuit-to-circuit function on a python object, and return the modified result.
 ///
 /// This method supports both `pytket.Circuit` and `Tk2Circuit` python objects.
-/// The returned Hugr is converted to the matching python object.
-pub fn update_hugr<'py, F>(circ: &Bound<'py, PyAny>, f: F) -> PyResult<Bound<'py, PyAny>>
+/// The returned [`Circuit`] is converted back into the circuit representation used in the input.
+pub fn update_circ<'py, F>(circ: &Bound<'py, PyAny>, f: F) -> PyResult<Bound<'py, PyAny>>
 where
-    F: FnOnce(Hugr, CircuitType) -> Hugr,
+    F: FnOnce(Circuit, CircuitType) -> Circuit,
 {
     let py = circ.py();
-    try_with_hugr(circ, |hugr, typ| {
-        let hugr = f(hugr, typ);
-        typ.convert(py, hugr)
+    try_with_circ(circ, |circ, typ| {
+        let circ = f(circ, typ);
+        typ.convert(py, circ)
     })
 }
+
+#[cfg(test)]
+mod test {
+    use crate::utils::test::make_module_tk2_circuit;
+
+    use super::*;
+    use cool_asserts::assert_matches;
+    use pyo3::prelude::*;
+    use rstest::rstest;
+
+    #[rstest]
+    fn test_with_circ() -> PyResult<()> {
+        pyo3::prepare_freethreaded_python();
+        Python::with_gil(|py| {
+            let circ = make_module_tk2_circuit(py)?;
+
+            with_circ(&circ, |circ, typ| {
+                assert_eq!(typ, CircuitType::Tket2);
+
+                let parent_optype = circ.hugr().get_optype(circ.parent());
+                assert_matches!(parent_optype, OpType::FuncDefn(_));
+            })
+        })?;
+        Ok(())
+    }
+}

tket2-py/src/circuit/tk2circuit.rs

@@ -31,7 +31,7 @@ use crate::rewrite::PyCircuitRewrite;
 use crate::types::PyHugrType;
 use crate::utils::{into_vec, ConvertPyErr};
 
-use super::{cost, with_hugr, PyCircuitCost, PyNode, PyWire};
+use super::{cost, with_circ, PyCircuitCost, PyNode, PyWire};
 
 /// A circuit in tket2 format.
 ///
@@ -69,7 +69,7 @@ impl Tk2Circuit {
     #[new]
     pub fn new(circ: &Bound<PyAny>) -> PyResult<Self> {
         Ok(Self {
-            circ: with_hugr(circ, |hugr, _| hugr)?.into(),
+            circ: with_circ(circ, |hugr, _| hugr)?,
         })
     }
 

tket2-py/src/optimiser.rs

@@ -8,7 +8,7 @@ use tket2::optimiser::badger::BadgerOptions;
 use tket2::optimiser::{BadgerLogger, DefaultBadgerOptimiser};
 use tket2::Circuit;
 
-use crate::circuit::update_hugr;
+use crate::circuit::update_circ;
 
 /// The module definition
 pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
@@ -96,10 +96,7 @@ impl PyBadgerOptimiser {
             split_circuit: split_circ.unwrap_or(false),
             queue_size: queue_size.unwrap_or(100),
         };
-        update_hugr(circ, |circ, _| {
-            self.optimise(circ.into(), log_progress, options)
-                .into_hugr()
-        })
+        update_circ(circ, |circ, _| self.optimise(circ, log_progress, options))
     }
 }
 

tket2-py/src/passes.rs

@@ -6,11 +6,11 @@ use std::{cmp::min, convert::TryInto, fs, num::NonZeroUsize, path::PathBuf};
 
 use pyo3::{prelude::*, types::IntoPyDict};
 use tket2::optimiser::badger::BadgerOptions;
-use tket2::{op_matches, passes::apply_greedy_commutation, Circuit, Tk2Op};
+use tket2::{op_matches, passes::apply_greedy_commutation, Tk2Op};
 
 use crate::utils::{create_py_exception, ConvertPyErr};
 use crate::{
-    circuit::{try_update_hugr, try_with_hugr},
+    circuit::{try_update_circ, try_with_circ},
     optimiser::PyBadgerOptimiser,
 };
 
@@ -45,10 +45,9 @@ create_py_exception!(
 #[pyfunction]
 fn greedy_depth_reduce<'py>(circ: &Bound<'py, PyAny>) -> PyResult<(Bound<'py, PyAny>, u32)> {
     let py = circ.py();
-    try_with_hugr(circ, |h, typ| {
-        let mut circ: Circuit = h.into();
+    try_with_circ(circ, |mut circ, typ| {
         let n_moves = apply_greedy_commutation(&mut circ).convert_pyerrs()?;
-        let circ = typ.convert(py, circ.into_hugr())?;
+        let circ = typ.convert(py, circ)?;
         PyResult::Ok((circ, n_moves))
     })
 }
@@ -124,8 +123,7 @@ fn badger_optimise<'py>(
         _ => unreachable!(),
     };
     // Optimise
-    try_update_hugr(circ, |hugr, _| {
-        let mut circ: Circuit = hugr.into();
+    try_update_circ(circ, |mut circ, _| {
         let n_cx = circ
             .commands()
             .filter(|c| op_matches(c.optype(), Tk2Op::CX))
@@ -150,6 +148,6 @@ fn badger_optimise<'py>(
             };
             circ = optimiser.optimise(circ, log_file, options);
         }
-        PyResult::Ok(circ.into_hugr())
+        PyResult::Ok(circ)
     })
 }

tket2-py/src/passes/chunks.rs

@@ -4,18 +4,17 @@ use derive_more::From;
 use pyo3::exceptions::PyAttributeError;
 use pyo3::prelude::*;
 use tket2::passes::CircuitChunks;
-use tket2::Circuit;
 
 use crate::circuit::CircuitType;
-use crate::circuit::{try_with_hugr, with_hugr};
+use crate::circuit::{try_with_circ, with_circ};
 use crate::utils::ConvertPyErr;
 
 /// Split a circuit into chunks of a given size.
 #[pyfunction]
 pub fn chunks(c: &Bound<PyAny>, max_chunk_size: usize) -> PyResult<PyCircuitChunks> {
-    with_hugr(c, |hugr, typ| {
+    with_circ(c, |hugr, typ| {
         // TODO: Detect if the circuit is in tket1 format or Tk2Circuit.
-        let chunks = CircuitChunks::split(&hugr.into(), max_chunk_size);
+        let chunks = CircuitChunks::split(&hugr, max_chunk_size);
         (chunks, typ).into()
     })
 }
@@ -40,21 +39,20 @@ impl PyCircuitChunks {
     /// Reassemble the chunks into a circuit.
     fn reassemble<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let circ = self.clone().chunks.reassemble().convert_pyerrs()?;
-        self.original_type.convert(py, circ.into_hugr())
+        self.original_type.convert(py, circ)
     }
 
     /// Returns clones of the split circuits.
     fn circuits<'py>(&self, py: Python<'py>) -> PyResult<Vec<Bound<'py, PyAny>>> {
         self.chunks
             .iter()
-            .map(|circ| self.original_type.convert(py, circ.clone().into_hugr()))
+            .map(|circ| self.original_type.convert(py, circ.clone()))
             .collect()
     }
 
     /// Replaces a chunk's circuit with an updated version.
     fn update_circuit(&mut self, index: usize, new_circ: &Bound<PyAny>) -> PyResult<()> {
-        try_with_hugr(new_circ, |hugr, _| {
-            let circ: Circuit = hugr.into();
+        try_with_circ(new_circ, |circ, _| {
             let circuit_sig = circ.circuit_signature();
             let chunk_sig = self.chunks[index].circuit_signature();
             if circuit_sig.input() != chunk_sig.input()

tket2-py/src/pattern/portmatching.rs

@@ -9,7 +9,7 @@ use pyo3::{prelude::*, types::PyIterator};
 
 use tket2::portmatching::{CircuitPattern, PatternMatch, PatternMatcher};
 
-use crate::circuit::{try_with_hugr, with_hugr, PyNode};
+use crate::circuit::{try_with_circ, with_circ, PyNode};
 
 /// A pattern that match a circuit exactly
 ///
@@ -30,9 +30,7 @@ impl PyCircuitPattern {
     /// Construct a pattern from a TKET1 circuit
     #[new]
     pub fn from_circuit(circ: &Bound<PyAny>) -> PyResult<Self> {
-        let pattern = try_with_hugr(circ, |circ, _| {
-            CircuitPattern::try_from_circuit(&circ.into())
-        })?;
+        let pattern = try_with_circ(circ, |circ, _| CircuitPattern::try_from_circuit(&circ))?;
         Ok(pattern.into())
     }
 
@@ -82,19 +80,16 @@ impl PyPatternMatcher {
 
     /// Find one convex match in a circuit.
     pub fn find_match(&self, circ: &Bound<PyAny>) -> PyResult<Option<PyPatternMatch>> {
-        with_hugr(circ, |circ, _| {
-            self.matcher
-                .find_matches_iter(&circ.into())
-                .next()
-                .map(Into::into)
+        with_circ(circ, |circ, _| {
+            self.matcher.find_matches_iter(&circ).next().map(Into::into)
         })
     }
 
     /// Find all convex matches in a circuit.
     pub fn find_matches(&self, circ: &Bound<PyAny>) -> PyResult<Vec<PyPatternMatch>> {
-        with_hugr(circ, |circ, _| {
+        with_circ(circ, |circ, _| {
             self.matcher
-                .find_matches(&circ.into())
+                .find_matches(&circ)
                 .into_iter()
                 .map_into()
                 .collect()