feat: pauli propagation use case example

tket2-py/src/circuit.rs

@@ -5,17 +5,24 @@ pub mod convert;
 pub mod cost;
 
 use derive_more::{From, Into};
+use hugr::extension::prelude::{BOOL_T, QB_T};
+use hugr::hugr::IdentList;
+use hugr::ops::custom::{ExtensionOp, OpaqueOp};
+use hugr::ops::{CustomOp, OpName, OpType};
+use hugr::types::{CustomType, FunctionType, Type, TypeBound};
 use pyo3::prelude::*;
 use std::fmt;
 
-use hugr::{Hugr, HugrView};
-use tket2::extension::REGISTRY;
+use hugr::{type_row, Hugr, HugrView, PortIndex};
+use tket2::extension::{LINEAR_BIT, REGISTRY};
 use tket2::json::TKETDecode;
 use tket2::rewrite::CircuitRewrite;
 use tket_json_rs::circuit_json::SerialCircuit;
 
 use crate::utils::create_py_exception;
+use crate::utils::ConvertPyErr;
 
+use self::convert::{into_vec, Dfg};
 pub use self::convert::{try_update_hugr, try_with_hugr, update_hugr, with_hugr, Tk2Circuit};
 pub use self::cost::PyCircuitCost;
 pub use tket2::{Pauli, Tk2Op};
@@ -24,13 +31,18 @@ pub use tket2::{Pauli, Tk2Op};
 pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
     let m = PyModule::new_bound(py, "_circuit")?;
     m.add_class::<Tk2Circuit>()?;
+    m.add_class::<Dfg>()?;
     m.add_class::<PyNode>()?;
+    m.add_class::<PyWire>()?;
     m.add_class::<PyCircuitCost>()?;
     m.add_class::<Tk2Op>()?;
+    m.add_class::<PyCustom>()?;
+    m.add_class::<PyHugrType>()?;
     m.add_class::<Pauli>()?;
 
     m.add_function(wrap_pyfunction!(validate_hugr, &m)?)?;
     m.add_function(wrap_pyfunction!(to_hugr_dot, &m)?)?;
+    m.add_function(wrap_pyfunction!(to_hugr_mermaid, &m)?)?;
 
     m.add("HugrError", py.get_type_bound::<PyHugrError>())?;
     m.add("BuildError", py.get_type_bound::<PyBuildError>())?;
@@ -86,6 +98,12 @@ pub fn to_hugr_dot(c: &Bound<PyAny>) -> PyResult<String> {
     with_hugr(c, |hugr, _| hugr.dot_string())
 }
 
+/// Return a Mermaid diagram representation of the circuit.
+#[pyfunction]
+pub fn to_hugr_mermaid(c: &Bound<PyAny>) -> PyResult<String> {
+    with_hugr(c, |hugr, _| hugr.mermaid_string())
+}
+
 /// A [`hugr::Node`] wrapper for Python.
 #[pyclass]
 #[pyo3(name = "Node")]
@@ -108,10 +126,203 @@ impl fmt::Debug for PyNode {
     }
 }
 
+#[pyclass]
+struct WireIter {
+    node: PyNode,
+    current: usize,
+}
+
+#[pymethods]
+impl WireIter {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<PyWire> {
+        slf.current += 1;
+        Some(slf.node.__getitem__(slf.current - 1).unwrap())
+    }
+}
+
 #[pymethods]
 impl PyNode {
     /// A string representation of the pattern.
     pub fn __repr__(&self) -> String {
         format!("{:?}", self)
     }
+
+    fn __getitem__(&self, idx: usize) -> PyResult<PyWire> {
+        Ok(hugr::Wire::new(self.node, idx).into())
+    }
+
+    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<WireIter>> {
+        let iter = WireIter {
+            current: 0,
+            node: *slf,
+        };
+        Py::new(slf.py(), iter)
+    }
+
+    fn outs(&self, n: usize) -> Vec<PyWire> {
+        (0..n)
+            .map(|i| hugr::Wire::new(self.node, i).into())
+            .collect()
+    }
+}
+
+/// A [`hugr::Node`] wrapper for Python.
+#[pyclass]
+#[pyo3(name = "Wire")]
+#[repr(transparent)]
+#[derive(From, Into, PartialEq, Eq, Hash, Clone, Copy)]
+pub struct PyWire {
+    /// Rust representation of the node
+    pub wire: hugr::Wire,
+}
+
+impl fmt::Display for PyWire {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.wire.fmt(f)
+    }
+}
+
+impl fmt::Debug for PyWire {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.wire.fmt(f)
+    }
+}
+
+#[pymethods]
+impl PyWire {
+    /// A string representation of the pattern.
+    pub fn __repr__(&self) -> String {
+        format!("{:?}", self)
+    }
+
+    fn node(&self) -> PyNode {
+        self.wire.node().into()
+    }
+
+    fn port(&self) -> usize {
+        self.wire.source().index()
+    }
+}
+
+#[pyclass]
+#[pyo3(name = "CustomOp")]
+#[repr(transparent)]
+#[derive(From, Into, PartialEq, Clone)]
+struct PyCustom(CustomOp);
+
+impl fmt::Debug for PyCustom {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl From<PyCustom> for OpType {
+    fn from(op: PyCustom) -> Self {
+        op.0.into()
+    }
+}
+
+#[pymethods]
+impl PyCustom {
+    #[new]
+    fn new(
+        extension: &str,
+        op_name: &str,
+        input_types: Vec<PyHugrType>,
+        output_types: Vec<PyHugrType>,
+    ) -> PyResult<Self> {
+        Ok(CustomOp::new_opaque(OpaqueOp::new(
+            IdentList::new(extension).unwrap(),
+            op_name,
+            Default::default(),
+            [],
+            FunctionType::new(into_vec(input_types), into_vec(output_types)),
+        ))
+        .into())
+    }
+
+    fn to_custom(&self) -> Self {
+        self.clone()
+    }
+    pub fn __repr__(&self) -> String {
+        format!("{:?}", self)
+    }
+
+    fn name(&self) -> String {
+        self.0.name().to_string()
+    }
+}
+#[pyclass]
+#[pyo3(name = "TypeBound")]
+#[derive(PartialEq, Clone, Debug)]
+enum PyTypeBound {
+    Any,
+    Copyable,
+    Eq,
+}
+
+impl From<PyTypeBound> for TypeBound {
+    fn from(bound: PyTypeBound) -> Self {
+        match bound {
+            PyTypeBound::Any => TypeBound::Any,
+            PyTypeBound::Copyable => TypeBound::Copyable,
+            PyTypeBound::Eq => TypeBound::Eq,
+        }
+    }
+}
+
+impl From<TypeBound> for PyTypeBound {
+    fn from(bound: TypeBound) -> Self {
+        match bound {
+            TypeBound::Any => PyTypeBound::Any,
+            TypeBound::Copyable => PyTypeBound::Copyable,
+            TypeBound::Eq => PyTypeBound::Eq,
+        }
+    }
+}
+
+#[pyclass]
+#[pyo3(name = "HugrType")]
+#[repr(transparent)]
+#[derive(From, Into, PartialEq, Clone)]
+struct PyHugrType(Type);
+
+impl fmt::Debug for PyHugrType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+#[pymethods]
+impl PyHugrType {
+    #[new]
+    fn new(extension: &str, type_name: &str, bound: PyTypeBound) -> Self {
+        Self(Type::new_extension(CustomType::new_simple(
+            type_name.into(),
+            IdentList::new(extension).unwrap(),
+            bound.into(),
+        )))
+    }
+    #[staticmethod]
+    fn qubit() -> Self {
+        Self(QB_T)
+    }
+
+    #[staticmethod]
+    fn linear_bit() -> Self {
+        Self(LINEAR_BIT.to_owned())
+    }
+
+    #[staticmethod]
+    fn bool() -> Self {
+        Self(BOOL_T)
+    }
+
+    pub fn __repr__(&self) -> String {
+        format!("{:?}", self)
+    }
 }

tket2-py/src/circuit/convert.rs

@@ -1,14 +1,20 @@
 //! Utilities for calling Hugr functions on generic python objects.
 
-use hugr::ops::OpType;
+use hugr::builder::{CircuitBuilder, DFGBuilder, Dataflow, DataflowHugr};
+use hugr::extension::prelude::QB_T;
+use hugr::ops::handle::NodeHandle;
+use hugr::ops::{CustomOp, OpType};
+use hugr::types::{FunctionType, Type};
+use itertools::Itertools;
 use pyo3::exceptions::{PyAttributeError, PyValueError};
 use pyo3::types::PyAnyMethods;
 use pyo3::{
-    pyclass, pymethods, Bound, FromPyObject, PyAny, PyErr, PyResult, PyTypeInfo, Python, ToPyObject,
+    pyclass, pymethods, Bound, FromPyObject, PyAny, PyErr, PyObject, PyRefMut, PyResult,
+    PyTypeInfo, Python, ToPyObject,
 };
 
 use derive_more::From;
-use hugr::{Hugr, HugrView};
+use hugr::{Hugr, HugrView, Wire};
 use serde::Serialize;
 use tket2::circuit::CircuitHash;
 use tket2::extension::REGISTRY;
@@ -20,7 +26,7 @@ use tket_json_rs::circuit_json::SerialCircuit;
 use crate::rewrite::PyCircuitRewrite;
 use crate::utils::ConvertPyErr;
 
-use super::{cost, PyCircuitCost};
+use super::{cost, PyCircuitCost, PyCustom, PyHugrType, PyNode, PyWire};
 
 /// A circuit in tket2 format.
 ///
@@ -147,6 +153,43 @@ impl Tk2Circuit {
     pub fn __deepcopy__(&self, _memo: Bound<PyAny>) -> PyResult<Self> {
         Ok(self.clone())
     }
+
+    fn node_op(&self, node: PyNode) -> PyResult<PyCustom> {
+        let custom: CustomOp = self
+            .hugr
+            .get_optype(node.node)
+            .clone()
+            .try_into()
+            .map_err(|e| {
+                PyErr::new::<PyValueError, _>(format!(
+                    "Could not convert circuit operation to a `CustomOp`: {e}"
+                ))
+            })?;
+
+        Ok(custom.into())
+    }
+
+    fn node_inputs(&self, node: PyNode) -> Vec<PyWire> {
+        self.hugr
+            .all_linked_outputs(node.node)
+            .map(|(n, p)| Wire::new(n, p).into())
+            .collect()
+    }
+
+    fn node_outputs(&self, node: PyNode) -> Vec<PyWire> {
+        self.hugr
+            .node_outputs(node.node)
+            .map(|p| Wire::new(node.node, p).into())
+            .collect()
+    }
+
+    fn input_node(&self) -> PyNode {
+        self.hugr.input().into()
+    }
+
+    fn output_node(&self) -> PyNode {
+        self.hugr.output().into()
+    }
 }
 impl Tk2Circuit {
     /// Tries to extract a Tk2Circuit from a python object.
@@ -157,6 +200,47 @@ impl Tk2Circuit {
     }
 }
 
+#[pyclass]
+#[derive(Clone, Debug, PartialEq, From)]
+pub(super) struct Dfg {
+    /// Rust representation of the circuit.
+    builder: DFGBuilder<Hugr>,
+}
+#[pymethods]
+impl Dfg {
+    #[new]
+    fn new(input_types: Vec<PyHugrType>, output_types: Vec<PyHugrType>) -> PyResult<Self> {
+        let builder = DFGBuilder::new(FunctionType::new(
+            into_vec(input_types),
+            into_vec(output_types),
+        ))
+        .convert_pyerrs()?;
+        Ok(Self { builder })
+    }
+
+    fn inputs(&self) -> Vec<PyWire> {
+        self.builder.input_wires().map_into().collect()
+    }
+
+    fn add_op(&mut self, op: PyCustom, inputs: Vec<PyWire>) -> PyResult<PyNode> {
+        let custom: CustomOp = op.into();
+        self.builder
+            .add_dataflow_op(custom, inputs.into_iter().map_into())
+            .convert_pyerrs()
+            .map(|d| d.node().into())
+    }
+
+    fn finish(&mut self, outputs: Vec<PyWire>) -> PyResult<Tk2Circuit> {
+        Ok(Tk2Circuit {
+            hugr: self
+                .builder
+                .clone()
+                .finish_hugr_with_outputs(outputs.into_iter().map_into(), &REGISTRY)
+                .convert_pyerrs()?,
+        })
+    }
+}
+
 /// A flag to indicate the encoding of a circuit.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub enum CircuitType {
@@ -236,3 +320,7 @@ where
         typ.convert(py, hugr)
     })
 }
+
+pub(super) fn into_vec<T, S: From<T>>(v: impl IntoIterator<Item = T>) -> Vec<S> {
+    v.into_iter().map_into().collect()
+}