feat: add a tket2.qsystem.utils extension with GetCurrentShot (#772)

Closes: #767

DEVELOPMENT.md

@@ -28,8 +28,8 @@ shell by setting up [direnv](https://devenv.sh/automatic-shell-activation/).
 
 To setup the environment manually you will need:
 
-- Just: https://just.systems/
-- Rust `>=1.75`: https://www.rust-lang.org/tools/install
+- Just: <https://just.systems/>
+- Rust `>=1.75`: <https://www.rust-lang.org/tools/install>
 - uv `>=0.3`: docs.astral.sh/uv/getting-started/installation
 
 Once you have these installed, install the required python dependencies and setup pre-commit hooks with:
@@ -101,6 +101,7 @@ line-by-line coverage report on
 [codecov](https://app.codecov.io/gh/CQCL/tket2/commits?branch=All%20branches).
 
 To run the coverage checks locally, first install `cargo-llvm-cov`.
+
 ```bash
 cargo install cargo-llvm-cov
 ```

tket2-exts/README.md

@@ -12,8 +12,7 @@ This is an auxiliary Python package containing HUGR extension definitions for `t
 operations and types.
 
 This package is intended to be used as an internal dependency for `tket2`.
-See https://pypi.org/project/tket2/ for the main package.
-
+See <https://pypi.org/project/tket2/> for the main package.
 
 ## Install
 
@@ -29,9 +28,8 @@ See [DEVELOPMENT.md] information on how to develop and contribute to this packag
 
   [DEVELOPMENT.md]: https://github.com/CQCL/tket2/blob/main/DEVELOPMENT.md
 
-
 ## License
 
-This project is licensed under Apache License, Version 2.0 ([LICENCE][] or http://www.apache.org/licenses/LICENSE-2.0).
+This project is licensed under Apache License, Version 2.0 ([LICENCE][] or <http://www.apache.org/licenses/LICENSE-2.0>).
 
   [LICENCE]: https://github.com/CQCL/tket2/blob/main/LICENCE

tket2-exts/src/tket2_exts/__init__.py

@@ -26,6 +26,11 @@ def qsystem() -> Extension:
     return load_extension("tket2.qsystem")
 
 
+@functools.cache
+def qsystem_utils() -> Extension:
+    return load_extension("tket2.qsystem.utils")
+
+
 @functools.cache
 def quantum() -> Extension:
     return load_extension("tket2.quantum")

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

@@ -0,0 +1,38 @@
+{
+  "version": "0.1.0",
+  "name": "tket2.qsystem.utils",
+  "runtime_reqs": [
+    "prelude"
+  ],
+  "types": {},
+  "values": {},
+  "operations": {
+    "GetCurrentShot": {
+      "extension": "tket2.qsystem.utils",
+      "name": "GetCurrentShot",
+      "description": "Get current shot number.",
+      "signature": {
+        "params": [],
+        "body": {
+          "input": [],
+          "output": [
+            {
+              "t": "Opaque",
+              "extension": "arithmetic.int.types",
+              "id": "int",
+              "args": [
+                {
+                  "tya": "BoundedNat",
+                  "n": 6
+                }
+              ],
+              "bound": "C"
+            }
+          ],
+          "runtime_reqs": []
+        }
+      },
+      "binary": false
+    }
+  }
+}

tket2-hseries/README.md

@@ -26,7 +26,7 @@ See [DEVELOPMENT.md][] for instructions on setting up the development environmen
 
 ## License
 
-This project is licensed under Apache License, Version 2.0 ([LICENSE][] or http://www.apache.org/licenses/LICENSE-2.0).
+This project is licensed under Apache License, Version 2.0 ([LICENSE][] or <http://www.apache.org/licenses/LICENSE-2.0>).
 
   [msrv]: https://img.shields.io/badge/rust-1.75.0%2B-blue.svg
   [LICENSE]: https://github.com/CQCL/tket2/blob/main/LICENCE

tket2-hseries/src/bin/tket2-hseries.rs

@@ -14,6 +14,7 @@ fn main() {
                 tket2_hseries::extension::futures::EXTENSION.to_owned(),
                 tket2_hseries::extension::result::EXTENSION.to_owned(),
                 tket2_hseries::extension::wasm::EXTENSION.to_owned(),
+                tket2_hseries::extension::utils::EXTENSION.to_owned(),
             ]);
 
             args.run_dump(&reg);

tket2-hseries/src/extension.rs

@@ -2,4 +2,5 @@
 pub mod futures;
 pub mod qsystem;
 pub mod result;
+pub mod utils;
 pub mod wasm;

tket2-hseries/src/extension/utils.rs

@@ -0,0 +1,154 @@
+//! This module contains utility functions for Quantinuum systems.
+
+use std::sync::{Arc, Weak};
+
+use derive_more::derive::Display;
+use hugr::{
+    builder::{BuildError, Dataflow},
+    extension::{
+        prelude::UnwrapBuilder,
+        simple_op::{try_from_name, MakeOpDef, MakeRegisteredOp},
+        ExtensionId, ExtensionRegistry, ExtensionSet, OpDef, SignatureFunc, Version, PRELUDE,
+    },
+    std_extensions::arithmetic::int_types::int_type,
+    type_row,
+    types::Signature,
+    Extension, Wire,
+};
+use lazy_static::lazy_static;
+use strum::{EnumIter, EnumString, IntoStaticStr};
+
+/// The extension ID for the utils extension.
+pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("tket2.qsystem.utils");
+/// The version of the "tket2.qsystem.utils" extension.
+pub const EXTENSION_VERSION: Version = Version::new(0, 1, 0);
+
+lazy_static! {
+    /// The "tket2.qsystem.utils" extension.
+    pub static ref EXTENSION: Arc<Extension> = {
+         Extension::new_arc(EXTENSION_ID, EXTENSION_VERSION, |ext, ext_ref| {
+            ext.add_requirements(ExtensionSet::from_iter([
+                PRELUDE.name(),
+            ].into_iter().cloned()));
+            UtilsOp::load_all_ops( ext, ext_ref).unwrap();
+        })
+    };
+
+    /// Extension registry including the "tket2.qsystem.utils" extension and
+    /// dependencies.
+    pub static ref REGISTRY: ExtensionRegistry = ExtensionRegistry::new([
+        EXTENSION.to_owned(),
+        PRELUDE.to_owned(),
+    ]);
+}
+
+#[derive(
+    Clone,
+    Copy,
+    Debug,
+    serde::Serialize,
+    serde::Deserialize,
+    Hash,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    EnumIter,
+    IntoStaticStr,
+    EnumString,
+    Display,
+)]
+/// The operations provided by the utils extension.
+pub enum UtilsOp {
+    /// fn get_current_shot() -> usize
+    GetCurrentShot,
+}
+
+impl MakeOpDef for UtilsOp {
+    fn init_signature(&self, _extension_ref: &std::sync::Weak<Extension>) -> SignatureFunc {
+        match self {
+            UtilsOp::GetCurrentShot => Signature::new(type_row![], int_type(6)),
+        }
+        .into()
+    }
+
+    fn from_def(op_def: &OpDef) -> Result<Self, hugr::extension::simple_op::OpLoadError> {
+        try_from_name(op_def.name(), op_def.extension_id())
+    }
+
+    fn extension(&self) -> ExtensionId {
+        EXTENSION_ID
+    }
+
+    fn extension_ref(&self) -> std::sync::Weak<Extension> {
+        Arc::downgrade(&EXTENSION)
+    }
+
+    fn description(&self) -> String {
+        match self {
+            UtilsOp::GetCurrentShot => "Get current shot number.",
+        }
+        .to_string()
+    }
+}
+
+impl MakeRegisteredOp for UtilsOp {
+    fn extension_id(&self) -> ExtensionId {
+        EXTENSION_ID
+    }
+
+    fn extension_ref(&self) -> Weak<Extension> {
+        Arc::downgrade(&EXTENSION)
+    }
+}
+
+/// An extension trait for [Dataflow] providing methods to add
+/// "tket2.qsystem.utils" operations.
+pub trait UtilsOpBuilder: Dataflow + UnwrapBuilder {
+    /// Add a "tket2.qsystem.utils.GetCurrentShot" op.
+    fn add_get_current_shot(&mut self) -> Result<Wire, BuildError> {
+        Ok(self
+            .add_dataflow_op(UtilsOp::GetCurrentShot, [])?
+            .out_wire(0))
+    }
+}
+
+impl<D: Dataflow> UtilsOpBuilder for D {}
+
+#[cfg(test)]
+mod test {
+    use std::sync::Arc;
+
+    use hugr::builder::{DataflowHugr, FunctionBuilder};
+    use hugr::ops::NamedOp;
+    use strum::IntoEnumIterator;
+
+    use super::*;
+
+    fn get_opdef(op: impl NamedOp) -> Option<&'static Arc<OpDef>> {
+        EXTENSION.get_op(&op.name())
+    }
+
+    #[test]
+    fn create_extension() {
+        assert_eq!(EXTENSION.name(), &EXTENSION_ID);
+
+        for o in UtilsOp::iter() {
+            assert_eq!(UtilsOp::from_def(get_opdef(o).unwrap()), Ok(o));
+        }
+    }
+
+    #[test]
+    fn get_current_shot() {
+        let hugr = {
+            let mut func_builder = FunctionBuilder::new(
+                "get_current_shot",
+                Signature::new(vec![], vec![int_type(6)]),
+            )
+            .unwrap();
+            let shot = func_builder.add_get_current_shot().unwrap();
+            func_builder.finish_hugr_with_outputs([shot]).unwrap()
+        };
+        hugr.validate().unwrap()
+    }
+}

tket2-py/tket2/extensions/__init__.py

@@ -1,4 +1,4 @@
-from tket2_exts import rotation, futures, qsystem, quantum, result
+from tket2_exts import rotation, futures, qsystem, qsystem_utils, quantum, result
 
 
-__all__ = ["rotation", "futures", "qsystem", "quantum", "result"]
+__all__ = ["rotation", "futures", "qsystem", "qsystem_utils", "quantum", "result"]