From 8a3fdb9ba55720d24e6f666f295933ed2ae17acd Mon Sep 17 00:00:00 2001
From: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>
Date: Tue, 16 Jan 2024 11:08:41 +0000
Subject: [PATCH] feat: add qalloc, qfree, reset ops

Closes #283
---
 tket2/src/json/op.rs |  7 ++++++-
 tket2/src/ops.rs     | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/tket2/src/json/op.rs b/tket2/src/json/op.rs
index 7ae81832..07f09082 100644
--- a/tket2/src/json/op.rs
+++ b/tket2/src/json/op.rs
@@ -191,6 +191,7 @@ impl From<&JsonOp> for OpType {
             JsonOpType::ZZMax => Tk2Op::ZZMax.into(),
             JsonOpType::ZZPhase => Tk2Op::ZZPhase.into(),
             JsonOpType::CZ => Tk2Op::CZ.into(),
+            JsonOpType::Reset => Tk2Op::Reset.into(),
             JsonOpType::noop => LeafOp::Noop { ty: QB_T }.into(),
             _ => LeafOp::CustomOp(Box::new(json_op.as_opaque_op())).into(),
         }
@@ -228,12 +229,16 @@ impl TryFrom<&OpType> for JsonOp {
                 Tk2Op::RxF64 => JsonOpType::Rx,
                 // TODO: Use a TK2 opaque op once we update the tket-json-rs dependency.
                 Tk2Op::AngleAdd => {
-                    unimplemented!("Serialising AngleAdd not supported. Are all constant folded?")
+                    unimplemented!("Serialising AngleAdd not supported. Are all constants folded?")
                 }
                 Tk2Op::TK1 => JsonOpType::TK1,
                 Tk2Op::PhasedX => JsonOpType::PhasedX,
                 Tk2Op::ZZPhase => JsonOpType::ZZPhase,
                 Tk2Op::CZ => JsonOpType::CZ,
+                Tk2Op::Reset => JsonOpType::Reset,
+                Tk2Op::QAlloc | Tk2Op::QFree => {
+                    unimplemented!("TKET1 does not support dynamic qubit allocation/discarding.")
+                }
             }
         } else if let LeafOp::CustomOp(b) = leaf {
             let ext = (*b).as_ref();
diff --git a/tket2/src/ops.rs b/tket2/src/ops.rs
index 87fc90d4..84f33540 100644
--- a/tket2/src/ops.rs
+++ b/tket2/src/ops.rs
@@ -65,6 +65,9 @@ pub enum Tk2Op {
     AngleAdd,
     CZ,
     TK1,
+    QAlloc,
+    QFree,
+    Reset,
 }
 
 /// Whether an op is a given Tk2Op.
@@ -99,7 +102,9 @@ impl MakeOpDef for Tk2Op {
         let one_qb_row = type_row![QB_T];
         let two_qb_row = type_row![QB_T, QB_T];
         match self {
-            H | T | S | X | Y | Z | Tdg | Sdg => FunctionType::new(one_qb_row.clone(), one_qb_row),
+            H | T | S | X | Y | Z | Tdg | Sdg | Reset => {
+                FunctionType::new(one_qb_row.clone(), one_qb_row)
+            }
             CX | ZZMax | CZ => FunctionType::new(two_qb_row.clone(), two_qb_row),
             ZZPhase => FunctionType::new(type_row![QB_T, QB_T, FLOAT64_TYPE], two_qb_row),
             Measure => FunctionType::new(one_qb_row, type_row![QB_T, BOOL_T]),
@@ -113,6 +118,8 @@ impl MakeOpDef for Tk2Op {
                 type_row![QB_T, FLOAT64_TYPE, FLOAT64_TYPE, FLOAT64_TYPE],
                 one_qb_row,
             ),
+            QAlloc => FunctionType::new(type_row![], one_qb_row),
+            QFree => FunctionType::new(one_qb_row, type_row![]),
         }
         .into()
     }
@@ -159,7 +166,7 @@ impl Tk2Op {
         match self {
             H | CX | T | S | X | Y | Z | Tdg | Sdg | ZZMax | RzF64 | RxF64 | PhasedX | ZZPhase
             | CZ | TK1 => true,
-            AngleAdd | Measure => false,
+            AngleAdd | Measure | QAlloc | QFree | Reset => false,
         }
     }
 }
@@ -269,6 +276,7 @@ pub(crate) mod test {
 
     use hugr::extension::simple_op::MakeOpDef;
     use hugr::ops::OpName;
+    use hugr::CircuitUnit;
     use hugr::{extension::OpDef, Hugr};
     use rstest::{fixture, rstest};
     use strum::IntoEnumIterator;
@@ -303,4 +311,26 @@ pub(crate) mod test {
     fn check_t2_bell(t2_bell_circuit: Hugr) {
         assert_eq!(t2_bell_circuit.commands().count(), 2);
     }
+
+    #[test]
+    fn ancilla_circ() {
+        let h = build_simple_circuit(1, |circ| {
+            let empty: [CircuitUnit; 0] = []; // requires type annotation
+            let ancilla = circ.append_with_outputs(Tk2Op::QAlloc, empty)?[0];
+            let ancilla = circ.append_with_outputs(Tk2Op::Reset, [ancilla])?[0];
+
+            let ancilla = circ.append_with_outputs(
+                Tk2Op::CX,
+                [CircuitUnit::Linear(0), CircuitUnit::Wire(ancilla)],
+            )?[0];
+            let ancilla = circ.append_with_outputs(Tk2Op::Measure, [ancilla])?[0];
+            circ.append_and_consume(Tk2Op::QFree, [ancilla])?;
+
+            Ok(())
+        })
+        .unwrap();
+
+        // 5 commands: alloc, reset, cx, measure, free
+        assert_eq!(h.commands().count(), 5);
+    }
 }