diff --git a/src/extension/prelude.rs b/src/extension/prelude.rs index 210b83fac..d01ebc149 100644 --- a/src/extension/prelude.rs +++ b/src/extension/prelude.rs @@ -5,15 +5,16 @@ use smol_str::SmolStr; use crate::{ extension::{ExtensionId, TypeDefBound}, + ops::LeafOp, types::{ type_param::{TypeArg, TypeParam}, - CustomCheckFailure, CustomType, Type, TypeBound, + CustomCheckFailure, CustomType, FunctionType, Type, TypeBound, }, values::{CustomConst, KnownTypeConst}, Extension, }; -use super::ExtensionRegistry; +use super::{ExtensionRegistry, EMPTY_REG}; /// Name of prelude extension. pub const PRELUDE_ID: ExtensionId = ExtensionId::new_unchecked("prelude"); @@ -39,6 +40,23 @@ lazy_static! { ) .unwrap(); + prelude + .add_op_custom_sig_simple( + SmolStr::new_inline(NEW_ARRAY_OP_ID), + "Create a new array from elements".to_string(), + vec![TypeParam::Type(TypeBound::Any), TypeParam::max_nat()], + |args: &[TypeArg]| { + let [TypeArg::Type { ty }, TypeArg::BoundedNat { n }] = args else { + panic!("should have been checked already.") + }; + Ok(FunctionType::new( + vec![ty.clone(); *n as usize], + vec![array_type(ty.clone(), *n)], + )) + }, + ) + .unwrap(); + prelude .add_type( SmolStr::new_inline("qubit"), @@ -70,18 +88,36 @@ pub const USIZE_T: Type = Type::new_extension(USIZE_CUSTOM_T); /// Boolean type - Sum of two units. pub const BOOL_T: Type = Type::new_simple_predicate(2); -/// Initialize a new array of type `typ` of length `size` -pub fn new_array(typ: Type, size: u64) -> Type { +/// Initialize a new array of element type `element_ty` of length `size` +pub fn array_type(element_ty: Type, size: u64) -> Type { let array_def = PRELUDE.get_type("array").unwrap(); let custom_t = array_def .instantiate_concrete(vec![ - TypeArg::Type { ty: typ }, + TypeArg::Type { ty: element_ty }, TypeArg::BoundedNat { n: size }, ]) .unwrap(); Type::new_extension(custom_t) } +/// Name of the operation in the prelude for creating new arrays. +pub const NEW_ARRAY_OP_ID: &str = "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) -> LeafOp { + PRELUDE + .instantiate_extension_op( + NEW_ARRAY_OP_ID, + vec![ + TypeArg::Type { ty: element_ty }, + TypeArg::BoundedNat { n: size }, + ], + &EMPTY_REG, + ) + .unwrap() + .into() +} + pub(crate) const ERROR_TYPE: Type = Type::new_extension(CustomType::new_simple( smol_str::SmolStr::new_inline("error"), PRELUDE_ID, @@ -117,3 +153,28 @@ impl CustomConst for ConstUsize { impl KnownTypeConst for ConstUsize { const TYPE: CustomType = USIZE_CUSTOM_T; } + +#[cfg(test)] +mod test { + use crate::builder::{DFGBuilder, Dataflow, DataflowHugr}; + + use super::*; + + #[test] + /// Test building a HUGR involving a new_array operation. + fn test_new_array() { + let mut b = DFGBuilder::new(FunctionType::new( + vec![QB_T, QB_T], + vec![array_type(QB_T, 2)], + )) + .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(); + } +} diff --git a/src/types/serialize.rs b/src/types/serialize.rs index a3cfe0867..f70cc57b0 100644 --- a/src/types/serialize.rs +++ b/src/types/serialize.rs @@ -4,7 +4,7 @@ use super::custom::CustomType; use super::FunctionType; -use crate::extension::prelude::{new_array, QB_T, USIZE_T}; +use crate::extension::prelude::{array_type, QB_T, USIZE_T}; use crate::ops::AliasDecl; use crate::types::primitive::PrimType; @@ -51,7 +51,7 @@ impl From for Type { SerSimpleType::G(sig) => Type::new_function(*sig), SerSimpleType::Tuple { inner } => Type::new_tuple(inner), SerSimpleType::Sum(sum) => sum.into(), - SerSimpleType::Array { inner, len } => new_array((*inner).into(), len), + SerSimpleType::Array { inner, len } => array_type((*inner).into(), len), SerSimpleType::Opaque(custom) => Type::new_extension(custom), SerSimpleType::Alias(a) => Type::new_alias(a), }