From 62da0911095b3e86256b195c137c083d5b1ab17c Mon Sep 17 00:00:00 2001 From: Douglas Wilson <141026920+doug-q@users.noreply.github.com> Date: Fri, 12 Jul 2024 09:10:48 +0100 Subject: [PATCH] feat(tket2-hseries): Add `tket2.futures` Hugr extension (#471) --- Cargo.lock | 5 + tket2-hseries/Cargo.toml | 11 +- tket2-hseries/src/extension.rs | 3 + tket2-hseries/src/extension/futures.rs | 303 +++++++++++++++++++++++++ tket2-hseries/src/lib.rs | 2 + 5 files changed, 322 insertions(+), 2 deletions(-) create mode 100644 tket2-hseries/src/extension.rs create mode 100644 tket2-hseries/src/extension/futures.rs diff --git a/Cargo.lock b/Cargo.lock index e22426d5..b7b0066f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1840,10 +1840,15 @@ name = "tket2-hseries" version = "0.1.0-alpha.1" dependencies = [ "clap", + "cool_asserts", "hugr", "hugr-cli", "lazy_static", + "serde", "serde_json", + "smol_str", + "strum", + "strum_macros", "tket2", ] diff --git a/tket2-hseries/Cargo.toml b/tket2-hseries/Cargo.toml index 3cc08261..1291f3ab 100644 --- a/tket2-hseries/Cargo.toml +++ b/tket2-hseries/Cargo.toml @@ -20,10 +20,17 @@ cli = ["dep:hugr-cli", "dep:clap"] [dependencies] hugr.workspace = true hugr-cli = { workspace = true, optional = true } -clap = { workspace = true, optional = true, features = ["derive"] } tket2 = { path = "../tket2", version = "0.1.0-alpha.1" } -serde_json.workspace = true +clap = { workspace = true, optional = true, features = ["derive"] } lazy_static.workspace = true +serde = { workspace = true, features = ["derive"] } +serde_json.workspace = true +smol_str.workspace = true +strum.workspace = true +strum_macros.workspace = true + +[dev-dependencies] +cool_asserts.workspace = true [lints] workspace = true diff --git a/tket2-hseries/src/extension.rs b/tket2-hseries/src/extension.rs new file mode 100644 index 00000000..7ec3fbdc --- /dev/null +++ b/tket2-hseries/src/extension.rs @@ -0,0 +1,3 @@ +//! This module defines the Hugr extensions used by tket2-hseries. + +pub mod futures; diff --git a/tket2-hseries/src/extension/futures.rs b/tket2-hseries/src/extension/futures.rs new file mode 100644 index 00000000..0265e95d --- /dev/null +++ b/tket2-hseries/src/extension/futures.rs @@ -0,0 +1,303 @@ +//! This module defines the Hugr extension used to represent futures. +//! +//! `Future` is a linear type representing a value that will be available in +//! the future. It can be consumed by `Read`, returning a `t`. It can be +//! duplicated by `Dup`, and discarded with `Free`. +use hugr::{ + builder::{BuildError, Dataflow}, + extension::{ + simple_op::{try_from_name, MakeExtensionOp, MakeOpDef, MakeRegisteredOp, OpLoadError}, + ExtensionBuildError, ExtensionId, ExtensionRegistry, OpDef, SignatureFunc, TypeDef, + }, + ops::{custom::ExtensionOp, CustomOp, OpType}, + types::{ + type_param::TypeParam, CustomType, FunctionType, PolyFuncType, Type, TypeArg, TypeBound, + }, + Extension, Wire, +}; +use lazy_static::lazy_static; +use serde::{Deserialize, Serialize}; +use smol_str::SmolStr; +use strum_macros::{EnumIter, EnumString, IntoStaticStr}; + +/// The ID of the `tket2.futures` extension. +pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("tket2.futures"); + +lazy_static! { + /// The "tket2.futures" extension. + pub static ref EXTENSION: Extension = { + let mut ext = Extension::new(EXTENSION_ID); + let _ = add_future_type_def(&mut ext).unwrap(); + + FutureOp::load_all_ops(&mut ext).unwrap(); + ext + }; + + /// Extension registry including the "tket2.futures" extension. + pub static ref REGISTRY: ExtensionRegistry = ExtensionRegistry::try_new([ + EXTENSION.to_owned() + ]).unwrap(); + + /// The name of the `Future` type. + pub static ref FUTURE_TYPE_NAME: SmolStr = SmolStr::new_inline("Future"); +} + +fn add_future_type_def(ext: &mut Extension) -> Result<&TypeDef, ExtensionBuildError> { + ext.add_type( + FUTURE_TYPE_NAME.to_owned(), + vec![TypeBound::Any.into()], + "A value that is computed asynchronously".into(), + TypeBound::Any.into(), + ) +} + +/// Returns a `Future` [CustomType]. +pub fn future_custom_type(t: Type) -> CustomType { + CustomType::new( + FUTURE_TYPE_NAME.to_owned(), + vec![t.into()], + EXTENSION_ID, + TypeBound::Any, + ) +} + +/// Returns a `Future` [Type]. +pub fn future_type(t: Type) -> Type { + future_custom_type(t).into() +} + +#[derive( + Clone, + Copy, + Debug, + Serialize, + Deserialize, + Hash, + PartialEq, + Eq, + PartialOrd, + Ord, + EnumIter, + IntoStaticStr, + EnumString, +)] +#[allow(missing_docs)] +#[non_exhaustive] +/// Simple enum of "tket2.futures" operations. +pub enum FutureOp { + Read, + Dup, + Free, +} + +impl MakeOpDef for FutureOp { + fn signature(&self) -> SignatureFunc { + let t_param = TypeParam::from(TypeBound::Any); + let t_type = Type::new_var_use(0, TypeBound::Any); + let future_type = future_type(t_type.clone()); + match self { + FutureOp::Read => { + PolyFuncType::new([t_param], FunctionType::new(future_type, t_type)).into() + } + FutureOp::Dup => PolyFuncType::new( + [t_param], + FunctionType::new(future_type.clone(), vec![future_type.clone(), future_type]), + ) + .into(), + FutureOp::Free => { + PolyFuncType::new([t_param], FunctionType::new(future_type.clone(), vec![])).into() + } + } + } + + fn extension(&self) -> ExtensionId { + EXTENSION_ID + } + + fn from_def(op_def: &OpDef) -> Result { + try_from_name(op_def.name(), &EXTENSION_ID) + } + + fn description(&self) -> String { + match self { + FutureOp::Read => "Read a value from a Future, consuming it".into(), + FutureOp::Dup => { + "Duplicate a Future. The original Future is consumed and two Futures are returned" + .into() + } + FutureOp::Free => "Consume a future without reading it.".into(), + } + } +} + +impl<'a> From<&'a ConcreteFutureOp> for &'static str { + fn from(value: &ConcreteFutureOp) -> Self { + value.op.into() + } +} + +/// Concrete "tket2.futures" operations with type set. +struct ConcreteFutureOp { + op: FutureOp, + typ: Type, +} + +fn concrete_future_op_type_args( + args: &[TypeArg], +) -> Result { + match args { + [TypeArg::Type { ty }] => Ok(ty.clone()), + _ => Err(OpLoadError::InvalidArgs( + hugr::extension::SignatureError::InvalidTypeArgs, + )), + } +} + +impl MakeExtensionOp for ConcreteFutureOp { + fn from_extension_op( + ext_op: &ExtensionOp, + ) -> Result + where + Self: Sized, + { + let op = FutureOp::from_def(ext_op.def())?; + let typ = concrete_future_op_type_args(ext_op.args())?; + Ok(Self { op, typ }) + } + + fn type_args(&self) -> Vec { + vec![self.typ.clone().into()] + } +} + +impl MakeRegisteredOp for ConcreteFutureOp { + fn extension_id(&self) -> ExtensionId { + EXTENSION_ID + } + + fn registry<'s, 'r: 's>(&'s self) -> &'r ExtensionRegistry { + ®ISTRY + } +} + +impl TryFrom<&OpType> for FutureOp { + type Error = (); + + fn try_from(value: &OpType) -> Result { + let Some(custom_op) = value.as_custom_op() else { + Err(())? + }; + match custom_op { + CustomOp::Extension(ext) => Self::from_extension_op(ext).ok(), + CustomOp::Opaque(opaque) => try_from_name(opaque.name(), &EXTENSION_ID).ok(), + } + .ok_or(()) + } +} + +impl TryFrom<&OpType> for ConcreteFutureOp { + type Error = (); + + fn try_from(value: &OpType) -> Result { + (|| { + let op = value.try_into().ok()?; + let typ = concrete_future_op_type_args(value.as_custom_op()?.args()).ok()?; + Some(Self { op, typ }) + })() + .ok_or(()) + } +} + +/// An extension trait for [Dataflow] providing methods to add "tket2.futures" +/// operations. +pub trait FutureOpBuilder: Dataflow { + /// Add a "tket2.futures.Read" op. + fn add_read(&mut self, lifted: Wire, typ: Type) -> Result<[Wire; 1], BuildError> { + Ok(self + .add_dataflow_op( + ConcreteFutureOp { + op: FutureOp::Read, + typ, + }, + [lifted], + )? + .outputs_arr()) + } + + /// Add a "tket2.futures.Dup" op. + fn add_dup(&mut self, lifted: Wire, typ: Type) -> Result<[Wire; 2], BuildError> { + Ok(self + .add_dataflow_op( + ConcreteFutureOp { + op: FutureOp::Dup, + typ, + }, + [lifted], + )? + .outputs_arr()) + } + + /// Add a "tket2.futures.Free" op. + fn add_free(&mut self, lifted: Wire, typ: Type) -> Result<(), BuildError> { + let op = self.add_dataflow_op( + ConcreteFutureOp { + op: FutureOp::Free, + typ, + }, + [lifted], + )?; + assert!(op.outputs().len() == 0); + Ok(()) + } +} + +impl FutureOpBuilder for D {} + +#[cfg(test)] +pub(crate) mod test { + use cool_asserts::assert_matches; + use hugr::{ + builder::{Dataflow, DataflowHugr, FunctionBuilder}, + ops::NamedOp, + }; + use std::sync::Arc; + use strum::IntoEnumIterator; + + use super::*; + + fn get_opdef(op: impl NamedOp) -> Option<&'static Arc> { + EXTENSION.get_op(&op.name()) + } + + #[test] + fn create_extension() { + assert_eq!(EXTENSION.name(), &EXTENSION_ID); + + for o in FutureOp::iter() { + assert_eq!(FutureOp::from_def(get_opdef(o).unwrap()), Ok(o)); + } + } + + #[test] + fn circuit() { + let t_param = TypeParam::from(TypeBound::Any); + let t = Type::new_var_use(0, TypeBound::Any); + let future_type = future_type(t.clone()); + + let hugr = { + let mut func_builder = FunctionBuilder::new( + "circuit", + PolyFuncType::new(vec![t_param], FunctionType::new(future_type, t.clone())), + ) + .unwrap(); + let [future_w] = func_builder.input_wires_arr(); + let [future_w, lazy_dup_w] = func_builder.add_dup(future_w, t.clone()).unwrap(); + func_builder.add_free(future_w, t.clone()).unwrap(); + let [t_w] = func_builder.add_read(lazy_dup_w, t).unwrap(); + func_builder + .finish_hugr_with_outputs([t_w], ®ISTRY) + .unwrap() + }; + assert_matches!(hugr.validate(®ISTRY), Ok(_)); + } +} diff --git a/tket2-hseries/src/lib.rs b/tket2-hseries/src/lib.rs index 36aa948e..d640c673 100644 --- a/tket2-hseries/src/lib.rs +++ b/tket2-hseries/src/lib.rs @@ -6,6 +6,8 @@ use hugr::Hugr; #[cfg(feature = "cli")] pub mod cli; +pub mod extension; + /// Modify a [Hugr] into a form that is acceptable for ingress into an H-series. /// /// Returns an error if this cannot be done.