feat: add qsystem op for measure leaked (#924)

Closes: #905 

Corresponding eldarion PR:
quantinuum-dev/eldarion#309

Author: narmlu

justfile

@@ -52,6 +52,10 @@ recompile-eccs:
 gen-extensions:
     cargo run -p tket2-hseries gen-extensions -o tket2-exts/src/tket2_exts/data
 
+# Interactively update snapshot tests (requires `cargo-insta`)
+update-snapshots:
+    cargo insta review
+
 # Runs a rust and a python command, depending on the `language` variable.
 #
 # If `language` is set to `rust` or `python`, only run the command for that language.

tket2-exts/src/tket2_exts/data/tket2/qsystem.json

@@ -1,5 +1,5 @@
 {
-  "version": "0.4.0",
+  "version": "0.4.1",
   "name": "tket2.qsystem",
   "types": {},
   "operations": {
@@ -37,6 +37,47 @@
       },
       "binary": false
     },
+    "LazyMeasureLeaked": {
+      "extension": "tket2.qsystem",
+      "name": "LazyMeasureLeaked",
+      "description": "Measure a qubit (return 0 or 1) or detect leakage (return 2).",
+      "signature": {
+        "params": [],
+        "body": {
+          "input": [
+            {
+              "t": "Q"
+            }
+          ],
+          "output": [
+            {
+              "t": "Opaque",
+              "extension": "tket2.futures",
+              "id": "Future",
+              "args": [
+                {
+                  "tya": "Type",
+                  "ty": {
+                    "t": "Opaque",
+                    "extension": "arithmetic.int.types",
+                    "id": "int",
+                    "args": [
+                      {
+                        "tya": "BoundedNat",
+                        "n": 6
+                      }
+                    ],
+                    "bound": "C"
+                  }
+                }
+              ],
+              "bound": "A"
+            }
+          ]
+        }
+      },
+      "binary": false
+    },
     "LazyMeasureReset": {
       "extension": "tket2.qsystem",
       "name": "LazyMeasureReset",

tket2-hseries/src/extension/qsystem.rs

@@ -21,6 +21,7 @@ use hugr::{
         arithmetic::{
             float_ops::FloatOps,
             float_types::{float64_type, ConstF64, EXTENSION as FLOAT_TYPES},
+            int_types::int_type,
         },
         collections::array::{array_type_parametric, ArrayOpBuilder},
     },
@@ -45,7 +46,7 @@ pub use lower::{check_lowered, lower_tk2_op, LowerTk2Error, LowerTket2ToQSystemP
 /// The "tket2.qsystem" extension id.
 pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("tket2.qsystem");
 /// The "tket2.qsystem" extension version.
-pub const EXTENSION_VERSION: Version = Version::new(0, 4, 0);
+pub const EXTENSION_VERSION: Version = Version::new(0, 4, 1);
 
 lazy_static! {
     /// The "tket2.qsystem" extension.
@@ -95,6 +96,7 @@ pub enum QSystemOp {
     QFree,
     Reset,
     MeasureReset,
+    LazyMeasureLeaked,
 }
 
 impl MakeOpDef for QSystemOp {
@@ -108,6 +110,7 @@ impl MakeOpDef for QSystemOp {
         let two_qb_row = TypeRow::from(vec![qb_t(), qb_t()]);
         match self {
             LazyMeasure => Signature::new(qb_t(), future_type(bool_t())),
+            LazyMeasureLeaked => Signature::new(qb_t(), future_type(int_type(6))),
             LazyMeasureReset => Signature::new(qb_t(), vec![qb_t(), future_type(bool_t())]),
             Reset => Signature::new(one_qb_row.clone(), one_qb_row),
             ZZPhase => Signature::new(vec![qb_t(), qb_t(), float64_type()], two_qb_row),
@@ -144,6 +147,9 @@ impl MakeOpDef for QSystemOp {
             QSystemOp::QFree => "Free a qubit (lose track of it).",
             QSystemOp::Reset => "Reset a qubit to the Z |0> eigenstate.",
             QSystemOp::MeasureReset => "Measure a qubit and reset it to the Z |0> eigenstate.",
+            QSystemOp::LazyMeasureLeaked => {
+                "Measure a qubit (return 0 or 1) or detect leakage (return 2)."
+            }
             QSystemOp::LazyMeasureReset => {
                 "Lazily measure a qubit and reset it to the Z |0> eigenstate."
             }
@@ -227,6 +233,13 @@ pub trait QSystemOpBuilder: Dataflow + UnwrapBuilder + ArrayOpBuilder {
             .out_wire(0))
     }
 
+    /// Add a "tket2.qsystem.LazyMeasureLeaked" op.
+    fn add_lazy_measure_leaked(&mut self, qb: Wire) -> Result<Wire, BuildError> {
+        Ok(self
+            .add_dataflow_op(QSystemOp::LazyMeasureLeaked, [qb])?
+            .out_wire(0))
+    }
+
     /// Add a "tket2.qsystem.LazyMeasureReset" op.
     fn add_lazy_measure_reset(&mut self, qb: Wire) -> Result<[Wire; 2], BuildError> {
         Ok(self
@@ -568,6 +581,19 @@ mod test {
         assert_matches!(hugr.validate(), Ok(_));
     }
 
+    #[test]
+    fn leaked() {
+        let hugr = {
+            let mut func_builder =
+                FunctionBuilder::new("leaked", Signature::new(qb_t(), vec![int_type(6)])).unwrap();
+            let [qb] = func_builder.input_wires_arr();
+            let lazy_i = func_builder.add_lazy_measure_leaked(qb).unwrap();
+            let [i] = func_builder.add_read(lazy_i, int_type(6)).unwrap();
+            func_builder.finish_hugr_with_outputs([i]).unwrap()
+        };
+        assert_matches!(hugr.validate(), Ok(_));
+    }
+
     #[test]
     fn all_ops() {
         let hugr = {

tket2-hseries/src/llvm/futures.rs

@@ -5,6 +5,7 @@ use anyhow::{anyhow, Result};
 use hugr::extension::prelude::bool_t;
 use hugr::extension::simple_op::MakeExtensionOp;
 use hugr::ops::{ExtensionOp, Value};
+use hugr::std_extensions::arithmetic::int_types::INT_TYPES;
 use hugr::types::TypeArg;
 use hugr::{HugrView, Node};
 use hugr_llvm::custom::CodegenExtension;
@@ -36,10 +37,13 @@ impl CodegenExtension for FuturesCodegenExtension {
                 (futures::EXTENSION_ID, FUTURE_TYPE_NAME.to_owned()),
                 |session, hugr_type| {
                     match (hugr_type.name().as_str(), hugr_type.args()) {
-                        // For now, we only support future bools
+                        // For now, we only support future bools and ints
                         ("Future", [TypeArg::Type { ty }]) if *ty == bool_t() => {
                             Ok(future_type(session.iw_context()))
                         }
+                        ("Future", [TypeArg::Type { ty }]) if *ty == INT_TYPES[6] => {
+                            Ok(future_type(session.iw_context()))
+                        }
                         _ => Err(anyhow!(
                             "FuturesCodegenExtension: Unsupported type: {}",
                             hugr_type
@@ -74,6 +78,10 @@ impl<'c, H: HugrView<Node = Node>> FuturesEmitter<'c, '_, '_, H> {
         self.iw_context().bool_type()
     }
 
+    fn ll_uint_type(&self) -> IntType<'c> {
+        self.iw_context().i64_type()
+    }
+
     fn builder(&self) -> &Builder<'c> {
         self.0.builder()
     }
@@ -101,10 +109,19 @@ impl<'c, H: HugrView<Node = Node>> FuturesEmitter<'c, '_, '_, H> {
         self.0.get_extern_func("___read_future_bool", func_type)
     }
 
+    fn get_func_read_uint(&self) -> Result<FunctionValue<'c>> {
+        let func_type = self
+            .ll_uint_type()
+            .fn_type(&[self.ll_future_type().into()], false);
+        self.0.get_extern_func("___read_future_uint", func_type)
+    }
+
     fn emit(&mut self, args: EmitOpArgs<'c, '_, ExtensionOp, H>) -> Result<()> {
-        let op = FutureOp::from_optype(&args.node().generalise()).unwrap().op;
-        match op {
-            FutureOpDef::Dup => {
+        let future_op = FutureOp::from_optype(&args.node().generalise()).unwrap();
+        let op = &future_op.op;
+        let typ = &future_op.typ;
+        match (op, typ) {
+            (FutureOpDef::Dup, _) => {
                 let [arg] = args
                     .inputs
                     .try_into()
@@ -114,7 +131,7 @@ impl<'c, H: HugrView<Node = Node>> FuturesEmitter<'c, '_, '_, H> {
                     .build_call(func, &[arg.into()], "inc_refcount")?;
                 args.outputs.finish(self.builder(), [arg, arg])
             }
-            FutureOpDef::Free => {
+            (FutureOpDef::Free, _) => {
                 let [arg] = args
                     .inputs
                     .try_into()
@@ -124,7 +141,7 @@ impl<'c, H: HugrView<Node = Node>> FuturesEmitter<'c, '_, '_, H> {
                     .build_call(func, &[arg.into()], "dec_refcount")?;
                 args.outputs.finish(self.builder(), [])
             }
-            FutureOpDef::Read => {
+            (FutureOpDef::Read, ty) if *ty == bool_t() => {
                 let [arg] = args
                     .inputs
                     .try_into()
@@ -146,29 +163,57 @@ impl<'c, H: HugrView<Node = Node>> FuturesEmitter<'c, '_, '_, H> {
                     .build_select(result_i1, true_val, false_val, "measure")?;
                 args.outputs.finish(self.builder(), [result])
             }
+            (FutureOpDef::Read, ty) if *ty == INT_TYPES[6] => {
+                let [arg] = args
+                    .inputs
+                    .try_into()
+                    .map_err(|_| anyhow!("Read expects a single input"))?;
+                let read_func = self.get_func_read_uint()?;
+                let result = self
+                    .builder()
+                    .build_call(read_func, &[arg.into()], "read_uint")?
+                    .try_as_basic_value()
+                    .unwrap_left();
+                let dec_refcount_func = self.get_func_dec_refcount()?;
+                self.builder()
+                    .build_call(dec_refcount_func, &[arg.into()], "dec_refcount")?;
+                args.outputs.finish(self.builder(), [result])
+            }
+            _ => Err(anyhow!(
+                "Unsupported future operation: {:?} with type: {}",
+                op,
+                typ
+            )),
         }
     }
 }
 
 #[cfg(test)]
 mod test {
     use hugr::extension::simple_op::MakeRegisteredOp;
+    use hugr::std_extensions::arithmetic::int_types::int_type;
     use hugr_llvm::check_emission;
     use hugr_llvm::test::llvm_ctx;
     use hugr_llvm::test::single_op_hugr;
     use hugr_llvm::test::TestContext;
 
     use super::*;
     #[rstest::rstest]
-    #[case::read(1,FutureOp { op: FutureOpDef::Read, typ: bool_t() })]
-    #[case::dup(2,FutureOp { op: FutureOpDef::Dup, typ: bool_t() })]
-    #[case::free(3,FutureOp { op: FutureOpDef::Free, typ: bool_t() })]
+    #[case::read_bool(1,FutureOp { op: FutureOpDef::Read, typ: bool_t() })]
+    #[case::dup_bool(2,FutureOp { op: FutureOpDef::Dup, typ: bool_t() })]
+    #[case::free_bool(3,FutureOp { op: FutureOpDef::Free, typ: bool_t() })]
+    #[case::read_int(4,FutureOp { op: FutureOpDef::Read, typ: int_type(6) })]
+    #[case::dup_int(5,FutureOp { op: FutureOpDef::Dup, typ: int_type(6) })]
+    #[case::free_int(6,FutureOp { op: FutureOpDef::Free, typ: int_type(6) })]
     fn emit_futures_codegen(
         #[case] _i: i32,
         #[with(_i)] mut llvm_ctx: TestContext,
         #[case] op: FutureOp,
     ) {
-        llvm_ctx.add_extensions(|ceb| ceb.add_extension(FuturesCodegenExtension));
+        llvm_ctx.add_extensions(|ceb| {
+            ceb.add_extension(FuturesCodegenExtension)
+                .add_default_int_extensions()
+        });
         let hugr = single_op_hugr(op.to_extension_op().unwrap().into());
         check_emission!(hugr, llvm_ctx);
     }

tket2-hseries/src/llvm/qsystem.rs

@@ -52,6 +52,7 @@ enum RuntimeFunction {
     QAlloc,
     QFree,
     Measure,
+    LazyMeasureLeaked,
     LazyMeasure,
     Reset,
 }
@@ -65,6 +66,7 @@ impl RuntimeFunction {
             RuntimeFunction::QAlloc => "___qalloc",
             RuntimeFunction::QFree => "___qfree",
             RuntimeFunction::Measure => "___measure",
+            RuntimeFunction::LazyMeasureLeaked => "___lazy_measure_leaked",
             RuntimeFunction::LazyMeasure => "___lazy_measure",
             RuntimeFunction::Reset => "___reset",
         }
@@ -94,6 +96,9 @@ impl RuntimeFunction {
             RuntimeFunction::QAlloc => qb_type.fn_type(&[], false),
             RuntimeFunction::QFree => iwc.void_type().fn_type(&[qb_type.into()], false),
             RuntimeFunction::Measure => iwc.bool_type().fn_type(&[qb_type.into()], false),
+            RuntimeFunction::LazyMeasureLeaked => {
+                future_type(iwc).fn_type(&[qb_type.into()], false)
+            }
             RuntimeFunction::LazyMeasure => future_type(iwc).fn_type(&[qb_type.into()], false),
             RuntimeFunction::Reset => iwc.void_type().fn_type(&[qb_type.into()], false),
         }
@@ -220,6 +225,28 @@ impl<PCG: PreludeCodegen> QSystemCodegenExtension<PCG> {
                 )?;
                 args.outputs.finish(builder, [result])
             }
+            // Measure qubit in Z basis or detect leakage, not forcing to a boolean
+            QSystemOp::LazyMeasureLeaked => {
+                let builder = context.builder();
+                let [qb] = args
+                    .inputs
+                    .try_into()
+                    .map_err(|_| anyhow!("LazyMeasureLeaked expects one input"))?;
+                let result = builder
+                    .build_call(
+                        self.runtime_func(context, RuntimeFunction::LazyMeasureLeaked)?,
+                        &[qb.into()],
+                        "lazy_measure_leaked",
+                    )?
+                    .try_as_basic_value()
+                    .unwrap_left();
+                builder.build_call(
+                    self.runtime_func(context, RuntimeFunction::QFree)?,
+                    &[qb.into()],
+                    "qfree",
+                )?;
+                args.outputs.finish(builder, [result])
+            }
             QSystemOp::LazyMeasureReset => {
                 let builder = context.builder();
                 let [qb] = args
@@ -336,6 +363,7 @@ mod test {
     #[case::qfree(7, QSystemOp::QFree)]
     #[case::reset(8, QSystemOp::Reset)]
     #[case::measure_reset(9, QSystemOp::MeasureReset)]
+    #[case::lazy_measure_leaked(10, QSystemOp::LazyMeasureLeaked)]
     fn emit_qsystem_codegen(
         #[case] _i: i32,
         #[with(_i)] mut llvm_ctx: TestContext,

tket2-hseries/src/llvm/snapshots/tket2_hseries__llvm__futures__test__emit_futures_codegen@llvm14_4.snap

@@ -0,0 +1,20 @@
+---
+source: tket2-hseries/src/llvm/futures.rs
+expression: mod_str
+---
+; ModuleID = 'test_context'
+source_filename = "test_context"
+
+define i64 @_hl.main.1(i64 %0) {
+alloca_block:
+  br label %entry_block
+
+entry_block:                                      ; preds = %alloca_block
+  %read_uint = call i64 @___read_future_uint(i64 %0)
+  call void @___dec_future_refcount(i64 %0)
+  ret i64 %read_uint
+}
+
+declare i64 @___read_future_uint(i64)
+
+declare void @___dec_future_refcount(i64)

tket2-hseries/src/llvm/snapshots/tket2_hseries__llvm__futures__test__emit_futures_codegen@llvm14_5.snap

@@ -0,0 +1,19 @@
+---
+source: tket2-hseries/src/llvm/futures.rs
+expression: mod_str
+---
+; ModuleID = 'test_context'
+source_filename = "test_context"
+
+define { i64, i64 } @_hl.main.1(i64 %0) {
+alloca_block:
+  br label %entry_block
+
+entry_block:                                      ; preds = %alloca_block
+  call void @___inc_future_refcount(i64 %0)
+  %mrv = insertvalue { i64, i64 } undef, i64 %0, 0
+  %mrv6 = insertvalue { i64, i64 } %mrv, i64 %0, 1
+  ret { i64, i64 } %mrv6
+}
+
+declare void @___inc_future_refcount(i64)