feat: add new_array operation to prelude (#544)

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();
+    }
+}

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<SerSimpleType> 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),
         }