From 5842c19ccced47ade6bb359f7f964e38476fdc99 Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Thu, 29 Aug 2024 14:51:51 +0100 Subject: [PATCH 1/7] refactor: put array in submodule of prelude --- hugr-core/src/extension/prelude.rs | 63 +++------------- hugr-core/src/extension/prelude/array.rs | 95 ++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 53 deletions(-) create mode 100644 hugr-core/src/extension/prelude/array.rs diff --git a/hugr-core/src/extension/prelude.rs b/hugr-core/src/extension/prelude.rs index 40b4caddb..c9277fde6 100644 --- a/hugr-core/src/extension/prelude.rs +++ b/hugr-core/src/extension/prelude.rs @@ -11,7 +11,8 @@ use crate::extension::{ ConstFold, ExtensionId, ExtensionSet, OpDef, SignatureError, SignatureFunc, TypeDefBound, }; use crate::ops::constant::{CustomCheckFailure, CustomConst, ValueName}; -use crate::ops::{ExtensionOp, NamedOp, OpName, Value}; +use crate::ops::OpName; +use crate::ops::{NamedOp, Value}; use crate::types::type_param::{TypeArg, TypeParam}; use crate::types::{ CustomType, FuncValueType, PolyFuncType, PolyFuncTypeRV, Signature, SumType, Type, TypeBound, @@ -22,30 +23,11 @@ use crate::{type_row, Extension}; use strum_macros::{EnumIter, EnumString, IntoStaticStr}; -use super::{ExtensionRegistry, SignatureFromArgs}; -struct ArrayOpCustom; +use super::ExtensionRegistry; -const MAX: &[TypeParam; 1] = &[TypeParam::max_nat()]; -impl SignatureFromArgs for ArrayOpCustom { - fn compute_signature(&self, arg_values: &[TypeArg]) -> Result { - let [TypeArg::BoundedNat { n }] = *arg_values else { - return Err(SignatureError::InvalidTypeArgs); - }; - let elem_ty_var = Type::new_var_use(0, TypeBound::Any); - - let var_arg_row = vec![elem_ty_var.clone(); n as usize]; - let other_row = vec![array_type(TypeArg::BoundedNat { n }, elem_ty_var.clone())]; - - Ok(PolyFuncTypeRV::new( - vec![TypeBound::Any.into()], - FuncValueType::new(var_arg_row, other_row), - )) - } - - fn static_params(&self) -> &[TypeParam] { - MAX - } -} +/// Array type and operations. +pub mod array; +pub use array::{array_type, new_array_op, NEW_ARRAY_OP_ID}; /// Name of prelude extension. pub const PRELUDE_ID: ExtensionId = ExtensionId::new_unchecked("prelude"); @@ -84,9 +66,9 @@ lazy_static! { .unwrap(); prelude .add_op( - NEW_ARRAY_OP_ID, + array::NEW_ARRAY_OP_ID, "Create a new array from elements".to_string(), - ArrayOpCustom, + array::ArrayOpCustom, ) .unwrap(); @@ -152,18 +134,6 @@ pub const USIZE_T: Type = Type::new_extension(USIZE_CUSTOM_T); /// Boolean type - Sum of two units. pub const BOOL_T: Type = Type::new_unit_sum(2); -/// Initialize a new array of element type `element_ty` of length `size` -pub fn array_type(size: impl Into, element_ty: Type) -> Type { - let array_def = PRELUDE.get_type("array").unwrap(); - let custom_t = array_def - .instantiate(vec![size.into(), element_ty.into()]) - .unwrap(); - Type::new_extension(custom_t) -} - -/// Name of the operation in the prelude for creating new arrays. -pub const NEW_ARRAY_OP_ID: OpName = OpName::new_inline("new_array"); - /// Name of the prelude panic operation. /// /// This operation can have any input and any output wires; it is instantiated @@ -175,20 +145,6 @@ pub const NEW_ARRAY_OP_ID: OpName = OpName::new_inline("new_array"); /// satisfied. pub const PANIC_OP_ID: OpName = OpName::new_inline("panic"); -/// Initialize a new array op of element type `element_ty` of length `size` -pub fn new_array_op(element_ty: Type, size: u64) -> ExtensionOp { - PRELUDE - .instantiate_extension_op( - &NEW_ARRAY_OP_ID, - vec![ - TypeArg::BoundedNat { n: size }, - TypeArg::Type { ty: element_ty }, - ], - &PRELUDE_REGISTRY, - ) - .unwrap() -} - /// Name of the string type. pub const STRING_TYPE_NAME: TypeName = TypeName::new_inline("string"); @@ -934,10 +890,11 @@ impl MakeRegisteredOp for Lift { #[cfg(test)] mod test { + use crate::builder::inout_sig; use crate::std_extensions::arithmetic::float_ops::FLOAT_OPS_REGISTRY; use crate::std_extensions::arithmetic::float_types::{ConstF64, FLOAT64_TYPE}; use crate::{ - builder::{endo_sig, inout_sig, DFGBuilder, Dataflow, DataflowHugr}, + builder::{endo_sig, DFGBuilder, Dataflow, DataflowHugr}, utils::test_quantum_extension::cx_gate, Hugr, Wire, }; diff --git a/hugr-core/src/extension/prelude/array.rs b/hugr-core/src/extension/prelude/array.rs new file mode 100644 index 000000000..842ac7653 --- /dev/null +++ b/hugr-core/src/extension/prelude/array.rs @@ -0,0 +1,95 @@ +use crate::extension::SignatureFromArgs; +use crate::ops::ExtensionOp; +use crate::ops::OpName; +use crate::types::FuncValueType; + +use crate::types::TypeBound; + +use crate::types::Type; + +use crate::extension::SignatureError; + +use crate::types::PolyFuncTypeRV; + +use crate::types::type_param::TypeArg; + +use super::{PRELUDE, PRELUDE_REGISTRY}; +use crate::types::type_param::TypeParam; +pub(super) struct ArrayOpCustom; + +pub(crate) const MAX: &[TypeParam; 1] = &[TypeParam::max_nat()]; + +impl SignatureFromArgs for ArrayOpCustom { + fn compute_signature(&self, arg_values: &[TypeArg]) -> Result { + let [TypeArg::BoundedNat { n }] = *arg_values else { + return Err(SignatureError::InvalidTypeArgs); + }; + let elem_ty_var = Type::new_var_use(0, TypeBound::Any); + + let var_arg_row = vec![elem_ty_var.clone(); n as usize]; + let other_row = vec![array_type(TypeArg::BoundedNat { n }, elem_ty_var.clone())]; + + Ok(PolyFuncTypeRV::new( + vec![TypeBound::Any.into()], + FuncValueType::new(var_arg_row, other_row), + )) + } + + fn static_params(&self) -> &[TypeParam] { + MAX + } +} + +/// Initialize a new array of element type `element_ty` of length `size` +pub fn array_type(size: impl Into, element_ty: Type) -> Type { + let array_def = PRELUDE.get_type("array").unwrap(); + let custom_t = array_def + .instantiate(vec![size.into(), element_ty.into()]) + .unwrap(); + Type::new_extension(custom_t) +} + +/// Name of the operation in the prelude for creating new arrays. +pub const NEW_ARRAY_OP_ID: OpName = OpName::new_inline("new_array"); + +/// Initialize a new array op of element type `element_ty` of length `size` +pub fn new_array_op(element_ty: Type, size: u64) -> ExtensionOp { + PRELUDE + .instantiate_extension_op( + &NEW_ARRAY_OP_ID, + vec![ + TypeArg::BoundedNat { n: size }, + TypeArg::Type { ty: element_ty }, + ], + &PRELUDE_REGISTRY, + ) + .unwrap() +} + +#[cfg(test)] +mod tests { + use crate::{ + builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr}, + extension::prelude::QB_T, + }; + + use super::*; + + #[test] + /// Test building a HUGR involving a new_array operation. + fn test_new_array() { + let mut b = DFGBuilder::new(inout_sig( + vec![QB_T, QB_T], + array_type(TypeArg::BoundedNat { n: 2 }, QB_T), + )) + .unwrap(); + + let [q1, q2] = b.input_wires_arr(); + + let op = new_array_op(QB_T, 2); + + let out = b.add_dataflow_op(op, [q1, q2]).unwrap(); + + b.finish_prelude_hugr_with_outputs(out.outputs()).unwrap(); + } +} From 7a3f0ac7d42de708e5d5b21acecdb95caee4e6da Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Thu, 29 Aug 2024 17:13:30 +0100 Subject: [PATCH 2/7] feat: introduce enum/struct for array ops --- hugr-core/src/extension/prelude.rs | 12 +- hugr-core/src/extension/prelude/array.rs | 167 +++++++++++++++++++---- 2 files changed, 145 insertions(+), 34 deletions(-) diff --git a/hugr-core/src/extension/prelude.rs b/hugr-core/src/extension/prelude.rs index c9277fde6..952c8a7e2 100644 --- a/hugr-core/src/extension/prelude.rs +++ b/hugr-core/src/extension/prelude.rs @@ -27,7 +27,7 @@ use super::ExtensionRegistry; /// Array type and operations. pub mod array; -pub use array::{array_type, new_array_op, NEW_ARRAY_OP_ID}; +pub use array::{array_type, new_array_op, ArrayOp, ArrayOpDef, ARRAY_TYPE_NAME, NEW_ARRAY_OP_ID}; /// Name of prelude extension. pub const PRELUDE_ID: ExtensionId = ExtensionId::new_unchecked("prelude"); @@ -58,19 +58,12 @@ lazy_static! { ) .unwrap(); prelude.add_type( - TypeName::new_inline("array"), + TypeName::new_inline(ARRAY_TYPE_NAME), vec![ TypeParam::max_nat(), TypeBound::Any.into()], "array".into(), TypeDefBound::from_params(vec![1] ), ) .unwrap(); - prelude - .add_op( - array::NEW_ARRAY_OP_ID, - "Create a new array from elements".to_string(), - array::ArrayOpCustom, - ) - .unwrap(); prelude .add_type( @@ -107,6 +100,7 @@ lazy_static! { TupleOpDef::load_all_ops(&mut prelude).unwrap(); NoopDef.add_to_extension(&mut prelude).unwrap(); LiftDef.add_to_extension(&mut prelude).unwrap(); + array::ArrayOpDef::load_all_ops(&mut prelude).unwrap(); prelude }; /// An extension registry containing only the prelude diff --git a/hugr-core/src/extension/prelude/array.rs b/hugr-core/src/extension/prelude/array.rs index 842ac7653..dc0669aee 100644 --- a/hugr-core/src/extension/prelude/array.rs +++ b/hugr-core/src/extension/prelude/array.rs @@ -1,5 +1,18 @@ +use strum_macros::EnumIter; +use strum_macros::EnumString; +use strum_macros::IntoStaticStr; + +use crate::extension::simple_op::MakeExtensionOp; +use crate::extension::simple_op::MakeOpDef; +use crate::extension::simple_op::MakeRegisteredOp; +use crate::extension::simple_op::OpLoadError; +use crate::extension::ExtensionId; +use crate::extension::OpDef; use crate::extension::SignatureFromArgs; +use crate::extension::SignatureFunc; +use crate::extension::TypeDef; use crate::ops::ExtensionOp; +use crate::ops::NamedOp; use crate::ops::OpName; use crate::types::FuncValueType; @@ -13,26 +26,33 @@ use crate::types::PolyFuncTypeRV; use crate::types::type_param::TypeArg; +use super::PRELUDE_ID; use super::{PRELUDE, PRELUDE_REGISTRY}; use crate::types::type_param::TypeParam; -pub(super) struct ArrayOpCustom; -pub(crate) const MAX: &[TypeParam; 1] = &[TypeParam::max_nat()]; +/// Array operation definitions. +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, EnumIter, IntoStaticStr, EnumString)] +#[allow(non_camel_case_types, missing_docs)] +#[non_exhaustive] +pub enum ArrayOpDef { + new_array, +} +const MAX: &[TypeParam; 1] = &[TypeParam::max_nat()]; -impl SignatureFromArgs for ArrayOpCustom { +impl SignatureFromArgs for ArrayOpDef { fn compute_signature(&self, arg_values: &[TypeArg]) -> Result { let [TypeArg::BoundedNat { n }] = *arg_values else { return Err(SignatureError::InvalidTypeArgs); }; let elem_ty_var = Type::new_var_use(0, TypeBound::Any); - - let var_arg_row = vec![elem_ty_var.clone(); n as usize]; - let other_row = vec![array_type(TypeArg::BoundedNat { n }, elem_ty_var.clone())]; - - Ok(PolyFuncTypeRV::new( - vec![TypeBound::Any.into()], - FuncValueType::new(var_arg_row, other_row), - )) + let array_ty = array_type(TypeArg::BoundedNat { n }, elem_ty_var.clone()); + let poly_func_ty = match self { + ArrayOpDef::new_array => PolyFuncTypeRV::new( + vec![TypeBound::Any.into()], + FuncValueType::new(vec![elem_ty_var.clone(); n as usize], array_ty), + ), + }; + Ok(poly_func_ty) } fn static_params(&self) -> &[TypeParam] { @@ -40,13 +60,104 @@ impl SignatureFromArgs for ArrayOpCustom { } } +impl ArrayOpDef { + /// Instantiate a new array operation with the given element type and array size. + pub fn to_concrete(self, elem_ty: Type, size: u64) -> ArrayOp { + ArrayOp { + def: self, + elem_ty, + size, + } + } +} + +impl MakeOpDef for ArrayOpDef { + fn from_def(op_def: &OpDef) -> Result + where + Self: Sized, + { + crate::extension::simple_op::try_from_name(op_def.name(), op_def.extension()) + } + + fn signature(&self) -> SignatureFunc { + (*self).into() // implements SignatureFromArgs + } + + fn extension(&self) -> ExtensionId { + PRELUDE_ID + } + + fn description(&self) -> String { + match self { + ArrayOpDef::new_array => "Create a new array from elements", + } + .into() + } +} + +#[derive(Clone, Debug, PartialEq)] +/// Concrete array operation. +pub struct ArrayOp { + /// The operation definition. + pub def: ArrayOpDef, + /// The element type of the array. + pub elem_ty: Type, + /// The size of the array. + pub size: u64, +} + +impl NamedOp for ArrayOp { + fn name(&self) -> OpName { + self.def.name() + } +} + +impl MakeExtensionOp for ArrayOp { + fn from_extension_op(ext_op: &ExtensionOp) -> Result + where + Self: Sized, + { + let def = ArrayOpDef::from_def(ext_op.def())?; + let [TypeArg::BoundedNat { n }, TypeArg::Type { ty }] = ext_op.args() else { + return Err(SignatureError::InvalidTypeArgs.into()); + }; + + Ok(def.to_concrete(ty.clone(), *n)) + } + + fn type_args(&self) -> Vec { + vec![ + TypeArg::BoundedNat { n: self.size }, + TypeArg::Type { + ty: self.elem_ty.clone(), + }, + ] + } +} + +impl MakeRegisteredOp for ArrayOp { + fn extension_id(&self) -> ExtensionId { + PRELUDE_ID + } + + fn registry<'s, 'r: 's>(&'s self) -> &'r crate::extension::ExtensionRegistry { + &PRELUDE_REGISTRY + } +} + +/// Name of the array type in the prelude. +pub const ARRAY_TYPE_NAME: &str = "array"; + +fn array_type_def() -> &'static TypeDef { + PRELUDE.get_type(ARRAY_TYPE_NAME).unwrap() +} /// Initialize a new array of element type `element_ty` of length `size` pub fn array_type(size: impl Into, element_ty: Type) -> Type { - let array_def = PRELUDE.get_type("array").unwrap(); - let custom_t = array_def + let array_def = array_type_def(); + array_def .instantiate(vec![size.into(), element_ty.into()]) - .unwrap(); - Type::new_extension(custom_t) + .unwrap() + .into() } /// Name of the operation in the prelude for creating new arrays. @@ -54,27 +165,33 @@ pub const NEW_ARRAY_OP_ID: OpName = OpName::new_inline("new_array"); /// Initialize a new array op of element type `element_ty` of length `size` pub fn new_array_op(element_ty: Type, size: u64) -> ExtensionOp { - PRELUDE - .instantiate_extension_op( - &NEW_ARRAY_OP_ID, - vec![ - TypeArg::BoundedNat { n: size }, - TypeArg::Type { ty: element_ty }, - ], - &PRELUDE_REGISTRY, - ) - .unwrap() + let op = ArrayOpDef::new_array.to_concrete(element_ty, size); + op.to_extension_op().unwrap() } #[cfg(test)] mod tests { + use strum::IntoEnumIterator; + use crate::{ builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr}, extension::prelude::QB_T, + ops::OpType, }; use super::*; + #[test] + fn test_array_ops() { + for def in ArrayOpDef::iter() { + let op = def.to_concrete(QB_T, 2); + let ext_op = op.clone().to_extension_op().unwrap(); + let optype: OpType = ext_op.into(); + let new_op: ArrayOp = optype.cast().unwrap(); + assert_eq!(new_op, op); + } + } + #[test] /// Test building a HUGR involving a new_array operation. fn test_new_array() { From 2e15594c2f49cdbb793e0eaec721f535daee0a47 Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Fri, 30 Aug 2024 14:35:14 +0100 Subject: [PATCH 3/7] implement remaining array ops --- hugr-core/src/extension/prelude/array.rs | 160 +++++++++++++++++++++-- 1 file changed, 146 insertions(+), 14 deletions(-) diff --git a/hugr-core/src/extension/prelude/array.rs b/hugr-core/src/extension/prelude/array.rs index dc0669aee..c51fa352f 100644 --- a/hugr-core/src/extension/prelude/array.rs +++ b/hugr-core/src/extension/prelude/array.rs @@ -2,6 +2,9 @@ use strum_macros::EnumIter; use strum_macros::EnumString; use strum_macros::IntoStaticStr; +use crate::extension::prelude::either_type; +use crate::extension::prelude::option_type; +use crate::extension::prelude::USIZE_T; use crate::extension::simple_op::MakeExtensionOp; use crate::extension::simple_op::MakeOpDef; use crate::extension::simple_op::MakeRegisteredOp; @@ -14,6 +17,7 @@ use crate::extension::TypeDef; use crate::ops::ExtensionOp; use crate::ops::NamedOp; use crate::ops::OpName; +use crate::type_row; use crate::types::FuncValueType; use crate::types::TypeBound; @@ -25,6 +29,7 @@ use crate::extension::SignatureError; use crate::types::PolyFuncTypeRV; use crate::types::type_param::TypeArg; +use crate::Extension; use super::PRELUDE_ID; use super::{PRELUDE, PRELUDE_REGISTRY}; @@ -36,6 +41,12 @@ use crate::types::type_param::TypeParam; #[non_exhaustive] pub enum ArrayOpDef { new_array, + get, + set, + swap, + pop_left, + pop_right, + discard_empty, } const MAX: &[TypeParam; 1] = &[TypeParam::max_nat()]; @@ -46,11 +57,21 @@ impl SignatureFromArgs for ArrayOpDef { }; let elem_ty_var = Type::new_var_use(0, TypeBound::Any); let array_ty = array_type(TypeArg::BoundedNat { n }, elem_ty_var.clone()); + let params = vec![TypeBound::Any.into()]; let poly_func_ty = match self { ArrayOpDef::new_array => PolyFuncTypeRV::new( - vec![TypeBound::Any.into()], + params, FuncValueType::new(vec![elem_ty_var.clone(); n as usize], array_ty), ), + ArrayOpDef::pop_left | ArrayOpDef::pop_right => { + let popped_array_ty = + array_type(TypeArg::BoundedNat { n: n - 1 }, elem_ty_var.clone()); + PolyFuncTypeRV::new( + params, + FuncValueType::new(array_ty, vec![elem_ty_var, popped_array_ty]), + ) + } + _ => unreachable!("Other operations should not need custom computation."), }; Ok(poly_func_ty) } @@ -63,12 +84,73 @@ impl SignatureFromArgs for ArrayOpDef { impl ArrayOpDef { /// Instantiate a new array operation with the given element type and array size. pub fn to_concrete(self, elem_ty: Type, size: u64) -> ArrayOp { + if self == ArrayOpDef::discard_empty { + debug_assert_eq!( + size, 0, + "discard_empty should only be called on empty arrays" + ); + } ArrayOp { def: self, elem_ty, size, } } + + /// To avoid recursion when defining the extension, take the type definition as an argument. + fn signature_from_def(&self, array_def: &TypeDef) -> SignatureFunc { + use ArrayOpDef::*; + if let new_array | pop_left | pop_right = self { + // implements SignatureFromArgs + // signature computed dynamically, so can rely on type definition in extension. + (*self).into() + } else { + let size_var = TypeArg::new_var_use(0, TypeParam::max_nat()); + let elem_ty_var = Type::new_var_use(1, TypeBound::Any); + let array_ty = instantiate(array_def, size_var.clone(), elem_ty_var.clone()); + let standard_params = vec![TypeParam::max_nat(), TypeBound::Any.into()]; + + match self { + get => { + let params = vec![TypeParam::max_nat(), TypeBound::Copyable.into()]; + let copy_elem_ty = Type::new_var_use(1, TypeBound::Copyable); + let copy_array_ty = instantiate(array_def, size_var, copy_elem_ty.clone()); + let option_type: Type = option_type(copy_elem_ty).into(); + PolyFuncTypeRV::new( + params, + FuncValueType::new(vec![copy_array_ty, USIZE_T], option_type), + ) + } + set => { + let result_row = vec![elem_ty_var.clone(), array_ty.clone()]; + let result_type: Type = either_type(result_row.clone(), result_row).into(); + PolyFuncTypeRV::new( + standard_params, + FuncValueType::new( + vec![array_ty.clone(), USIZE_T, elem_ty_var], + result_type, + ), + ) + } + swap => { + let result_type: Type = either_type(array_ty.clone(), array_ty.clone()).into(); + PolyFuncTypeRV::new( + standard_params, + FuncValueType::new(vec![array_ty, USIZE_T, USIZE_T], result_type), + ) + } + discard_empty => PolyFuncTypeRV::new( + vec![TypeBound::Any.into()], + FuncValueType::new( + instantiate(array_def, 0, Type::new_var_use(0, TypeBound::Any)), + type_row![], + ), + ), + new_array | pop_left | pop_right => unreachable!(), + } + .into() + } + } } impl MakeOpDef for ArrayOpDef { @@ -80,7 +162,7 @@ impl MakeOpDef for ArrayOpDef { } fn signature(&self) -> SignatureFunc { - (*self).into() // implements SignatureFromArgs + self.signature_from_def(array_type_def()) } fn extension(&self) -> ExtensionId { @@ -90,9 +172,32 @@ impl MakeOpDef for ArrayOpDef { fn description(&self) -> String { match self { ArrayOpDef::new_array => "Create a new array from elements", + ArrayOpDef::get => "Get an element from an array", + ArrayOpDef::set => "Set an element in an array", + ArrayOpDef::swap => "Swap two elements in an array", + ArrayOpDef::pop_left => "Pop an element from the left of an array", + ArrayOpDef::pop_right => "Pop an element from the right of an array", + ArrayOpDef::discard_empty => "Discard an empty array", } .into() } + + /// Add an operation implemented as an [MakeOpDef], which can provide the data + /// required to define an [OpDef], to an extension. + // + // This method is re-defined here since we need to pass the list type def while computing the signature, + // to avoid recursive loops initializing the extension. + fn add_to_extension( + &self, + extension: &mut Extension, + ) -> Result<(), crate::extension::ExtensionBuildError> { + let sig = self.signature_from_def(extension.get_type(ARRAY_TYPE_NAME).unwrap()); + let def = extension.add_op(self.name(), self.description(), sig)?; + + self.post_opdef(def); + + Ok(()) + } } #[derive(Clone, Debug, PartialEq)] @@ -118,20 +223,33 @@ impl MakeExtensionOp for ArrayOp { Self: Sized, { let def = ArrayOpDef::from_def(ext_op.def())?; - let [TypeArg::BoundedNat { n }, TypeArg::Type { ty }] = ext_op.args() else { - return Err(SignatureError::InvalidTypeArgs.into()); + + let (ty, size) = match (def, ext_op.args()) { + (ArrayOpDef::discard_empty, [TypeArg::Type { ty }]) => (ty.clone(), 0), + (_, [TypeArg::BoundedNat { n }, TypeArg::Type { ty }]) => (ty.clone(), *n), + _ => return Err(SignatureError::InvalidTypeArgs.into()), }; - Ok(def.to_concrete(ty.clone(), *n)) + Ok(def.to_concrete(ty.clone(), size)) } fn type_args(&self) -> Vec { - vec![ - TypeArg::BoundedNat { n: self.size }, - TypeArg::Type { - ty: self.elem_ty.clone(), - }, - ] + use ArrayOpDef::*; + let ty_arg = TypeArg::Type { + ty: self.elem_ty.clone(), + }; + match self.def { + discard_empty => { + debug_assert_eq!( + self.size, 0, + "discard_empty should only be called on empty arrays" + ); + vec![ty_arg] + } + new_array | pop_left | pop_right | get | set | swap => { + vec![TypeArg::BoundedNat { n: self.size }, ty_arg] + } + } } } @@ -153,7 +271,14 @@ fn array_type_def() -> &'static TypeDef { } /// Initialize a new array of element type `element_ty` of length `size` pub fn array_type(size: impl Into, element_ty: Type) -> Type { - let array_def = array_type_def(); + instantiate(array_type_def(), size, element_ty) +} + +fn instantiate( + array_def: &TypeDef, + size: impl Into, + element_ty: crate::types::TypeBase, +) -> crate::types::TypeBase { array_def .instantiate(vec![size.into(), element_ty.into()]) .unwrap() @@ -175,7 +300,7 @@ mod tests { use crate::{ builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr}, - extension::prelude::QB_T, + extension::prelude::{BOOL_T, QB_T}, ops::OpType, }; @@ -184,7 +309,14 @@ mod tests { #[test] fn test_array_ops() { for def in ArrayOpDef::iter() { - let op = def.to_concrete(QB_T, 2); + dbg!(def); + let ty = if def == ArrayOpDef::get { BOOL_T } else { QB_T }; + let size = if def == ArrayOpDef::discard_empty { + 0 + } else { + 2 + }; + let op = def.to_concrete(ty, size); let ext_op = op.clone().to_extension_op().unwrap(); let optype: OpType = ext_op.into(); let new_op: ArrayOp = optype.cast().unwrap(); From b2d3f2ec8c5114eb74fa4c107a74e4145847c12b Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Fri, 30 Aug 2024 15:29:50 +0100 Subject: [PATCH 4/7] test all signatures --- hugr-core/src/extension/prelude/array.rs | 147 ++++++++++++++++++++--- 1 file changed, 130 insertions(+), 17 deletions(-) diff --git a/hugr-core/src/extension/prelude/array.rs b/hugr-core/src/extension/prelude/array.rs index c51fa352f..7aca0be0d 100644 --- a/hugr-core/src/extension/prelude/array.rs +++ b/hugr-core/src/extension/prelude/array.rs @@ -5,10 +5,9 @@ use strum_macros::IntoStaticStr; use crate::extension::prelude::either_type; use crate::extension::prelude::option_type; use crate::extension::prelude::USIZE_T; -use crate::extension::simple_op::MakeExtensionOp; -use crate::extension::simple_op::MakeOpDef; -use crate::extension::simple_op::MakeRegisteredOp; -use crate::extension::simple_op::OpLoadError; +use crate::extension::simple_op::{ + HasConcrete, HasDef, MakeExtensionOp, MakeOpDef, MakeRegisteredOp, OpLoadError, +}; use crate::extension::ExtensionId; use crate::extension::OpDef; use crate::extension::SignatureFromArgs; @@ -68,7 +67,10 @@ impl SignatureFromArgs for ArrayOpDef { array_type(TypeArg::BoundedNat { n: n - 1 }, elem_ty_var.clone()); PolyFuncTypeRV::new( params, - FuncValueType::new(array_ty, vec![elem_ty_var, popped_array_ty]), + FuncValueType::new( + array_ty, + Type::from(option_type(vec![elem_ty_var, popped_array_ty])), + ), ) } _ => unreachable!("Other operations should not need custom computation."), @@ -223,14 +225,7 @@ impl MakeExtensionOp for ArrayOp { Self: Sized, { let def = ArrayOpDef::from_def(ext_op.def())?; - - let (ty, size) = match (def, ext_op.args()) { - (ArrayOpDef::discard_empty, [TypeArg::Type { ty }]) => (ty.clone(), 0), - (_, [TypeArg::BoundedNat { n }, TypeArg::Type { ty }]) => (ty.clone(), *n), - _ => return Err(SignatureError::InvalidTypeArgs.into()), - }; - - Ok(def.to_concrete(ty.clone(), size)) + def.instantiate(ext_op.args()) } fn type_args(&self) -> Vec { @@ -263,6 +258,24 @@ impl MakeRegisteredOp for ArrayOp { } } +impl HasDef for ArrayOp { + type Def = ArrayOpDef; +} + +impl HasConcrete for ArrayOpDef { + type Concrete = ArrayOp; + + fn instantiate(&self, type_args: &[TypeArg]) -> Result { + let (ty, size) = match (self, type_args) { + (ArrayOpDef::discard_empty, [TypeArg::Type { ty }]) => (ty.clone(), 0), + (_, [TypeArg::BoundedNat { n }, TypeArg::Type { ty }]) => (ty.clone(), *n), + _ => return Err(SignatureError::InvalidTypeArgs.into()), + }; + + Ok(self.to_concrete(ty.clone(), size)) + } +} + /// Name of the array type in the prelude. pub const ARRAY_TYPE_NAME: &str = "array"; @@ -301,7 +314,7 @@ mod tests { use crate::{ builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr}, extension::prelude::{BOOL_T, QB_T}, - ops::OpType, + ops::{OpTrait, OpType}, }; use super::*; @@ -309,7 +322,6 @@ mod tests { #[test] fn test_array_ops() { for def in ArrayOpDef::iter() { - dbg!(def); let ty = if def == ArrayOpDef::get { BOOL_T } else { QB_T }; let size = if def == ArrayOpDef::discard_empty { 0 @@ -317,8 +329,7 @@ mod tests { 2 }; let op = def.to_concrete(ty, size); - let ext_op = op.clone().to_extension_op().unwrap(); - let optype: OpType = ext_op.into(); + let optype: OpType = op.clone().into(); let new_op: ArrayOp = optype.cast().unwrap(); assert_eq!(new_op, op); } @@ -341,4 +352,106 @@ mod tests { b.finish_prelude_hugr_with_outputs(out.outputs()).unwrap(); } + + #[test] + fn test_get() { + let size = 2; + let element_ty = BOOL_T; + let op = ArrayOpDef::get.to_concrete(element_ty.clone(), size); + + let optype: OpType = op.into(); + + let sig = optype.dataflow_signature().unwrap(); + + assert_eq!( + sig.io(), + ( + &vec![array_type(size, element_ty.clone()), USIZE_T].into(), + &vec![option_type(element_ty.clone()).into()].into() + ) + ); + } + + #[test] + fn test_set() { + let size = 2; + let element_ty = BOOL_T; + let op = ArrayOpDef::set.to_concrete(element_ty.clone(), size); + + let optype: OpType = op.into(); + + let sig = optype.dataflow_signature().unwrap(); + let array_ty = array_type(size, element_ty.clone()); + let result_row = vec![element_ty.clone(), array_ty.clone()]; + assert_eq!( + sig.io(), + ( + &vec![array_ty.clone(), USIZE_T, element_ty.clone()].into(), + &vec![either_type(result_row.clone(), result_row).into()].into() + ) + ); + } + + #[test] + fn test_swap() { + let size = 2; + let element_ty = BOOL_T; + let op = ArrayOpDef::swap.to_concrete(element_ty.clone(), size); + + let optype: OpType = op.into(); + + let sig = optype.dataflow_signature().unwrap(); + let array_ty = array_type(size, element_ty.clone()); + assert_eq!( + sig.io(), + ( + &vec![array_ty.clone(), USIZE_T, USIZE_T].into(), + &vec![either_type(array_ty.clone(), array_ty).into()].into() + ) + ); + } + + #[test] + fn test_pops() { + let size = 2; + let element_ty = BOOL_T; + for op in [ArrayOpDef::pop_left, ArrayOpDef::pop_right].iter() { + let op = op.to_concrete(element_ty.clone(), size); + + let optype: OpType = op.into(); + + let sig = optype.dataflow_signature().unwrap(); + assert_eq!( + sig.io(), + ( + &vec![array_type(size, element_ty.clone())].into(), + &vec![option_type(vec![ + element_ty.clone(), + array_type(size - 1, element_ty.clone()) + ]) + .into()] + .into() + ) + ); + } + } + + #[test] + fn test_discard_empty() { + let size = 0; + let element_ty = BOOL_T; + let op = ArrayOpDef::discard_empty.to_concrete(element_ty.clone(), size); + + let optype: OpType = op.into(); + + let sig = optype.dataflow_signature().unwrap(); + + assert_eq!( + sig.io(), + ( + &vec![array_type(size, element_ty.clone())].into(), + &type_row![] + ) + ); + } } From e9e0f6194388427684b641d84d248752e84e9716 Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Fri, 30 Aug 2024 15:38:11 +0100 Subject: [PATCH 5/7] params docstring --- hugr-core/src/extension/prelude/array.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/hugr-core/src/extension/prelude/array.rs b/hugr-core/src/extension/prelude/array.rs index 7aca0be0d..540f8d086 100644 --- a/hugr-core/src/extension/prelude/array.rs +++ b/hugr-core/src/extension/prelude/array.rs @@ -47,7 +47,9 @@ pub enum ArrayOpDef { pop_right, discard_empty, } -const MAX: &[TypeParam; 1] = &[TypeParam::max_nat()]; + +/// Static parameters for array operations. Includes array size. Type is part of the type scheme. +const STATIC_SIZE_PARAM: &[TypeParam; 1] = &[TypeParam::max_nat()]; impl SignatureFromArgs for ArrayOpDef { fn compute_signature(&self, arg_values: &[TypeArg]) -> Result { @@ -79,7 +81,7 @@ impl SignatureFromArgs for ArrayOpDef { } fn static_params(&self) -> &[TypeParam] { - MAX + STATIC_SIZE_PARAM } } From feb706ee93b14a587edb91db5fe5fd7774e9d2b8 Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Fri, 30 Aug 2024 15:39:58 +0100 Subject: [PATCH 6/7] update serialized prelude --- hugr-py/src/hugr/std/_json_defs/prelude.json | 358 +++++++++++++++++++ specification/std_extensions/prelude.json | 358 +++++++++++++++++++ 2 files changed, 716 insertions(+) diff --git a/hugr-py/src/hugr/std/_json_defs/prelude.json b/hugr-py/src/hugr/std/_json_defs/prelude.json index 3ec1d8699..014ba3ede 100644 --- a/hugr-py/src/hugr/std/_json_defs/prelude.json +++ b/hugr-py/src/hugr/std/_json_defs/prelude.json @@ -221,6 +221,112 @@ }, "binary": false }, + "discard_empty": { + "extension": "prelude", + "name": "discard_empty", + "description": "Discard an empty array", + "signature": { + "params": [ + { + "tp": "Type", + "b": "A" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "BoundedNat", + "n": 0 + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 0, + "b": "A" + } + } + ], + "bound": "A" + } + ], + "output": [], + "extension_reqs": [] + } + }, + "binary": false + }, + "get": { + "extension": "prelude", + "name": "get", + "description": "Get an element from an array", + "signature": { + "params": [ + { + "tp": "BoundedNat", + "bound": null + }, + { + "tp": "Type", + "b": "C" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "C" + } + } + ], + "bound": "C" + }, + { + "t": "I" + } + ], + "output": [ + { + "t": "Sum", + "s": "General", + "rows": [ + [], + [ + { + "t": "V", + "i": 1, + "b": "C" + } + ] + ] + } + ], + "extension_reqs": [] + } + }, + "binary": false + }, "new_array": { "extension": "prelude", "name": "new_array", @@ -276,6 +382,20 @@ }, "binary": false }, + "pop_left": { + "extension": "prelude", + "name": "pop_left", + "description": "Pop an element from the left of an array", + "signature": null, + "binary": true + }, + "pop_right": { + "extension": "prelude", + "name": "pop_right", + "description": "Pop an element from the right of an array", + "signature": null, + "binary": true + }, "print": { "extension": "prelude", "name": "print", @@ -297,6 +417,244 @@ } }, "binary": false + }, + "set": { + "extension": "prelude", + "name": "set", + "description": "Set an element in an array", + "signature": { + "params": [ + { + "tp": "BoundedNat", + "bound": null + }, + { + "tp": "Type", + "b": "A" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + }, + { + "t": "I" + }, + { + "t": "V", + "i": 1, + "b": "A" + } + ], + "output": [ + { + "t": "Sum", + "s": "General", + "rows": [ + [ + { + "t": "V", + "i": 1, + "b": "A" + }, + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ], + [ + { + "t": "V", + "i": 1, + "b": "A" + }, + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ] + ] + } + ], + "extension_reqs": [] + } + }, + "binary": false + }, + "swap": { + "extension": "prelude", + "name": "swap", + "description": "Swap two elements in an array", + "signature": { + "params": [ + { + "tp": "BoundedNat", + "bound": null + }, + { + "tp": "Type", + "b": "A" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + }, + { + "t": "I" + }, + { + "t": "I" + } + ], + "output": [ + { + "t": "Sum", + "s": "General", + "rows": [ + [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ], + [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ] + ] + } + ], + "extension_reqs": [] + } + }, + "binary": false } } } diff --git a/specification/std_extensions/prelude.json b/specification/std_extensions/prelude.json index 3ec1d8699..014ba3ede 100644 --- a/specification/std_extensions/prelude.json +++ b/specification/std_extensions/prelude.json @@ -221,6 +221,112 @@ }, "binary": false }, + "discard_empty": { + "extension": "prelude", + "name": "discard_empty", + "description": "Discard an empty array", + "signature": { + "params": [ + { + "tp": "Type", + "b": "A" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "BoundedNat", + "n": 0 + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 0, + "b": "A" + } + } + ], + "bound": "A" + } + ], + "output": [], + "extension_reqs": [] + } + }, + "binary": false + }, + "get": { + "extension": "prelude", + "name": "get", + "description": "Get an element from an array", + "signature": { + "params": [ + { + "tp": "BoundedNat", + "bound": null + }, + { + "tp": "Type", + "b": "C" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "C" + } + } + ], + "bound": "C" + }, + { + "t": "I" + } + ], + "output": [ + { + "t": "Sum", + "s": "General", + "rows": [ + [], + [ + { + "t": "V", + "i": 1, + "b": "C" + } + ] + ] + } + ], + "extension_reqs": [] + } + }, + "binary": false + }, "new_array": { "extension": "prelude", "name": "new_array", @@ -276,6 +382,20 @@ }, "binary": false }, + "pop_left": { + "extension": "prelude", + "name": "pop_left", + "description": "Pop an element from the left of an array", + "signature": null, + "binary": true + }, + "pop_right": { + "extension": "prelude", + "name": "pop_right", + "description": "Pop an element from the right of an array", + "signature": null, + "binary": true + }, "print": { "extension": "prelude", "name": "print", @@ -297,6 +417,244 @@ } }, "binary": false + }, + "set": { + "extension": "prelude", + "name": "set", + "description": "Set an element in an array", + "signature": { + "params": [ + { + "tp": "BoundedNat", + "bound": null + }, + { + "tp": "Type", + "b": "A" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + }, + { + "t": "I" + }, + { + "t": "V", + "i": 1, + "b": "A" + } + ], + "output": [ + { + "t": "Sum", + "s": "General", + "rows": [ + [ + { + "t": "V", + "i": 1, + "b": "A" + }, + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ], + [ + { + "t": "V", + "i": 1, + "b": "A" + }, + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ] + ] + } + ], + "extension_reqs": [] + } + }, + "binary": false + }, + "swap": { + "extension": "prelude", + "name": "swap", + "description": "Swap two elements in an array", + "signature": { + "params": [ + { + "tp": "BoundedNat", + "bound": null + }, + { + "tp": "Type", + "b": "A" + } + ], + "body": { + "input": [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + }, + { + "t": "I" + }, + { + "t": "I" + } + ], + "output": [ + { + "t": "Sum", + "s": "General", + "rows": [ + [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ], + [ + { + "t": "Opaque", + "extension": "prelude", + "id": "array", + "args": [ + { + "tya": "Variable", + "idx": 0, + "cached_decl": { + "tp": "BoundedNat", + "bound": null + } + }, + { + "tya": "Type", + "ty": { + "t": "V", + "i": 1, + "b": "A" + } + } + ], + "bound": "A" + } + ] + ] + } + ], + "extension_reqs": [] + } + }, + "binary": false } } } From 43e515d98bfab98424ebfec0db96ea2db5a873c1 Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Fri, 30 Aug 2024 15:54:58 +0100 Subject: [PATCH 7/7] better unreachable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Agustín Borgna <121866228+aborgna-q@users.noreply.github.com> --- hugr-core/src/extension/prelude/array.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/hugr-core/src/extension/prelude/array.rs b/hugr-core/src/extension/prelude/array.rs index 540f8d086..a15bf23cc 100644 --- a/hugr-core/src/extension/prelude/array.rs +++ b/hugr-core/src/extension/prelude/array.rs @@ -75,7 +75,10 @@ impl SignatureFromArgs for ArrayOpDef { ), ) } - _ => unreachable!("Other operations should not need custom computation."), + _ => unreachable!( + "Operation {} should not need custom computation.", + self.name() + ), }; Ok(poly_func_ty) }