feat!: runtime entrypoint arguments via generic tket.argreader op

qis-compiler/pyproject.toml

@@ -77,6 +77,8 @@ cache-keys = [
     { file = "pyproject.toml" },
     { file = "Cargo.toml" },
     { file = "rust/**/*.rs" },
+    { file = "../tket/**/*.rs" },
+    { file = "../tket-qsystem/**/*.rs" },
 ]
 
 [tool.cibuildwheel]

qis-compiler/rust/lib.rs

@@ -1,7 +1,7 @@
 //! The compiler for HUGR to QIS
 pub mod array;
 
-use anyhow::{Result, anyhow};
+use anyhow::{Result, anyhow, bail};
 use hugr::envelope::EnvelopeConfig;
 use hugr::llvm::CodegenExtsBuilder;
 use hugr::llvm::custom::CodegenExtsMap;
@@ -36,8 +36,9 @@ use tket_qsystem::llvm::array_utils::ArrayLowering;
 pub use tket_qsystem::llvm::futures::FuturesCodegenExtension;
 use tket_qsystem::llvm::globals::GlobalsCodegenExtension;
 use tket_qsystem::llvm::{
-    debug::DebugCodegenExtension, prelude::QISPreludeCodegen, qsystem::QSystemCodegenExtension,
-    random::RandomCodegenExtension, result::ResultsCodegenExtension, utils::UtilsCodegenExtension,
+    argreader::ArgReaderCodegenExtension, debug::DebugCodegenExtension, prelude::QISPreludeCodegen,
+    qsystem::QSystemCodegenExtension, random::RandomCodegenExtension,
+    result::ResultsCodegenExtension, utils::UtilsCodegenExtension,
 };
 use tracing::{Level, event, instrument};
 use utils::read_hugr_envelope;
@@ -135,11 +136,15 @@ fn codegen_extensions(platform: qsystem::QSystemPlatform) -> CodegenExtsMap<'sta
         .add_extension(ResultsCodegenExtension::new(
             SeleneHeapArrayCodegen::LOWERING,
         ))
-        .add_extension(RotationCodegenExtension::new(pcg))
+        .add_extension(RotationCodegenExtension::new(pcg.clone()))
         .add_extension(UtilsCodegenExtension)
         // State results use standard arrays.
         .add_extension(DebugCodegenExtension::new(SeleneHeapArrayCodegen::LOWERING))
         .add_extension(gpu::GpuCodegen)
+        // Argument reading uses standard arrays.
+        .add_extension(ArgReaderCodegenExtension::new(
+            SeleneHeapArrayCodegen::LOWERING,
+        ))
         .finish()
 }
 
@@ -217,10 +222,10 @@ fn get_entry_point_name(namer: &Namer, hugr: &impl HugrView<Node = Node>) -> Res
             .as_func_defn()
             .ok_or_else(|| anyhow!("Entry point node is not a function definition"))?;
         if func_defn.inner_signature().input_count() != 0 {
-            return Err(anyhow!(
+            bail!(
                 "Entry point function must have no input parameters (found {})",
                 func_defn.inner_signature().input_count()
-            ));
+            );
         }
         (func_defn.func_name().as_ref(), hugr.entrypoint())
     };

tket-exts/src/tket_exts/__init__.py

@@ -19,6 +19,7 @@
 from tket_exts.tket.result import ResultExtension
 from tket_exts.tket.wasm import WasmExtension
 from tket_exts.tket.measurement import MeasurementExtension
+from tket_exts.tket.argreader import ArgReaderExtension
 
 from hugr.ext import ExtensionRegistry
 from tket_exts import tket
@@ -45,6 +46,7 @@
     "global_phase",
     "globals",
     "measurement",
+    "argreader",
 ]
 
 debug: DebugExtension = tket.debug.DebugExtension()
@@ -64,6 +66,7 @@
 global_phase: GlobalPhaseExtension = tket.global_phase.GlobalPhaseExtension()
 globals: GlobalsExtension = tket.globals.GlobalsExtension()
 measurement: MeasurementExtension = tket.measurement.MeasurementExtension()
+argreader: ArgReaderExtension = tket.argreader.ArgReaderExtension()
 
 
 def tket_registry() -> ExtensionRegistry:
@@ -92,6 +95,7 @@ def tket_registry() -> ExtensionRegistry:
         tket.modifier.ModifierExtension(),
         tket.global_phase.GlobalPhaseExtension(),
         tket.measurement.MeasurementExtension(),
+        tket.argreader.ArgReaderExtension(),
     ]
 
     registry = ExtensionRegistry()

tket-exts/src/tket_exts/data/tket/argreader.json

@@ -0,0 +1,34 @@
+{
+  "version": "0.1.0",
+  "name": "tket.argreader",
+  "types": {},
+  "operations": {
+    "read_arg": {
+      "extension": "tket.argreader",
+      "name": "read_arg",
+      "description": "Read a runtime argument of the given type identified by a string tag",
+      "signature": {
+        "params": [
+          {
+            "tp": "String"
+          },
+          {
+            "tp": "Type",
+            "b": "A"
+          }
+        ],
+        "body": {
+          "input": [],
+          "output": [
+            {
+              "t": "V",
+              "i": 1,
+              "b": "A"
+            }
+          ]
+        }
+      },
+      "binary": false
+    }
+  }
+}

tket-exts/src/tket_exts/tket/__init__.py

@@ -13,6 +13,7 @@
     modifier,
     global_phase,
     globals,
+    argreader,
 )
 
 __all__ = [
@@ -30,4 +31,5 @@
     "modifier",
     "global_phase",
     "globals",
+    "argreader",
 ]

tket-exts/src/tket_exts/tket/argreader.py

@@ -0,0 +1,50 @@
+"""Argument reader extension operations."""
+
+import functools
+from typing import List
+
+from hugr.ext import Extension, OpDef, TypeDef
+from hugr.ops import ExtOp
+from hugr.tys import FunctionType, StringArg, Type, TypeTypeArg
+
+from ._util import TketExtension, load_extension
+
+
+class ArgReaderExtension(TketExtension):
+    """Operations for reading runtime entrypoint arguments."""
+
+    @functools.cache
+    def __call__(self) -> Extension:
+        """Returns the argreader extension"""
+        return load_extension("tket.argreader")
+
+    def TYPES(self) -> List[TypeDef]:
+        """Return the types defined by this extension"""
+        return []
+
+    def OPS(self) -> List[OpDef]:
+        """Return the operations defined by this extension"""
+        return [
+            self.read_arg_def,
+        ]
+
+    @functools.cached_property
+    def read_arg_def(self) -> OpDef:
+        """Read a runtime argument of the given type identified by a string tag.
+
+        This is the generic operation definition. For the instantiated operation, see
+        `read_arg`.
+        """
+        return self().get_op("read_arg")
+
+    def read_arg(self, tag: str, ty: Type) -> ExtOp:
+        """Read a runtime argument of type ``ty`` identified by ``tag``.
+
+        Args:
+            tag: String tag identifying the argument (matches the provider's key).
+            ty: The HUGR type of the argument to read.
+        """
+        return self.read_arg_def.instantiate(
+            [StringArg(tag), TypeTypeArg(ty)],
+            FunctionType(input=[], output=[ty]),
+        )

tket-exts/tests/test_validate_exts.py

@@ -218,6 +218,15 @@ def ext_wasm() -> Tuple[TketExtension, List[ExtType], List[ExtOp]]:
     )
 
 
+def ext_argreader() -> Tuple[TketExtension, List[ExtType], List[ExtOp]]:
+    ext = tket_exts.argreader
+    return (
+        ext,
+        [],
+        [ext.read_arg("test", Bool)],
+    )
+
+
 @pytest.mark.parametrize(
     "ext_vals",
     [
@@ -235,6 +244,7 @@ def ext_wasm() -> Tuple[TketExtension, List[ExtType], List[ExtOp]]:
         ext_result,
         ext_rotation,
         ext_wasm,
+        ext_argreader,
     ],
 )
 def test_exported_extension(

tket-qsystem/src/extension.rs

@@ -7,6 +7,7 @@ use hugr::extension::ExtensionRegistry;
 pub mod classical_compute;
 pub use classical_compute::gpu;
 pub use classical_compute::wasm;
+pub mod argreader;
 pub mod futures;
 pub mod globals;
 pub mod qsystem;
@@ -29,7 +30,7 @@ pub static REGISTRY: LazyLock<ExtensionRegistry> = LazyLock::new(|| {
 /// The list intentionally excludes std and base tket extensions so callers can
 /// combine it with [`tket::extension::tket_extensions`] or another base
 /// registry without duplicating either list.
-pub fn qsystem_extensions() -> [Arc<Extension>; 10] {
+pub fn qsystem_extensions() -> [Arc<Extension>; 11] {
     [
         gpu::EXTENSION.to_owned(),
         qsystem::EXTENSION.to_owned(),
@@ -41,5 +42,6 @@ pub fn qsystem_extensions() -> [Arc<Extension>; 10] {
         utils::EXTENSION.to_owned(),
         wasm::EXTENSION.to_owned(),
         globals::EXTENSION.to_owned(),
+        argreader::EXTENSION.to_owned(),
     ]
 }