feat(core, llvm): add array unpack operations (#2339)

Dual of new_array

Closes #1947

Author: ss2165

hugr-core/src/std_extensions/collections/array.rs

@@ -43,7 +43,7 @@ pub const ARRAY_VALUENAME: TypeName = TypeName::new_inline("array");
 /// Reported unique name of the extension
 pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("collections.array");
 /// Extension version.
-pub const VERSION: semver::Version = semver::Version::new(0, 1, 0);
+pub const VERSION: semver::Version = semver::Version::new(0, 1, 1);
 
 /// A linear, fixed-length collection of values.
 ///
@@ -197,7 +197,31 @@ pub trait ArrayOpBuilder: GenericArrayOpBuilder {
     ) -> Result<Wire, BuildError> {
         self.add_new_generic_array::<Array>(elem_ty, values)
     }
-
+    /// Adds an array unpack operation to the dataflow graph.
+    ///
+    /// This operation unpacks an array into individual elements.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `size` - The size of the array.
+    /// * `input` - The wire representing the array to unpack.
+    ///
+    /// # Errors
+    ///
+    /// If building the operation fails.
+    ///
+    /// # Returns
+    ///
+    /// A vector of wires representing the individual elements from the array.
+    fn add_array_unpack(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+    ) -> Result<Vec<Wire>, BuildError> {
+        self.add_generic_array_unpack::<Array>(elem_ty, size, input)
+    }
     /// Adds an array clone operation to the dataflow graph and return the wires
     /// representing the originala and cloned array.
     ///

hugr-core/src/std_extensions/collections/array/array_op.rs

@@ -21,7 +21,7 @@ use crate::utils::Never;
 
 use super::array_kind::ArrayKind;
 
-/// Array operation definitions. Generic over the conrete array implementation.
+/// Array operation definitions. Generic over the concrete array implementation.
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, IntoStaticStr, EnumIter, EnumString)]
 #[allow(non_camel_case_types)]
 #[non_exhaustive]
@@ -58,6 +58,10 @@ pub enum GenericArrayOpDef<AK: ArrayKind> {
     /// references `AK` to ensure that the type parameter is used.
     #[strum(disabled)]
     _phantom(PhantomData<AK>, Never),
+    /// Unpacks an array into its individual elements:
+    /// `unpack<SIZE><elemty>: array<SIZE, elemty> -> (elemty)^SIZE`
+    /// where `SIZE` must be statically known (not a variable)
+    unpack,
 }
 
 /// Static parameters for array operations. Includes array size. Type is part of the type scheme.
@@ -76,6 +80,10 @@ impl<AK: ArrayKind> SignatureFromArgs for GenericArrayOpDef<AK> {
                 params,
                 FuncValueType::new(vec![elem_ty_var.clone(); n as usize], array_ty),
             ),
+            GenericArrayOpDef::unpack => PolyFuncTypeRV::new(
+                params,
+                FuncValueType::new(array_ty, vec![elem_ty_var.clone(); n as usize]),
+            ),
             GenericArrayOpDef::pop_left | GenericArrayOpDef::pop_right => {
                 let popped_array_ty = AK::ty(n - 1, elem_ty_var.clone());
                 PolyFuncTypeRV::new(
@@ -124,9 +132,9 @@ impl<AK: ArrayKind> GenericArrayOpDef<AK> {
         _extension_ref: &Weak<Extension>,
     ) -> SignatureFunc {
         use GenericArrayOpDef::{
-            _phantom, discard_empty, get, new_array, pop_left, pop_right, set, swap,
+            _phantom, discard_empty, get, new_array, pop_left, pop_right, set, swap, unpack,
         };
-        if let new_array | pop_left | pop_right = self {
+        if let new_array | unpack | pop_left | pop_right = self {
             // implements SignatureFromArgs
             // signature computed dynamically, so can rely on type definition in extension.
             (*self).into()
@@ -184,7 +192,7 @@ impl<AK: ArrayKind> GenericArrayOpDef<AK> {
                     ),
                 ),
                 _phantom(_, never) => match *never {},
-                new_array | pop_left | pop_right => unreachable!(),
+                new_array | unpack | pop_left | pop_right => unreachable!(),
             }
             .into()
         }
@@ -218,6 +226,7 @@ impl<AK: ArrayKind> MakeOpDef for GenericArrayOpDef<AK> {
     fn description(&self) -> String {
         match self {
             GenericArrayOpDef::new_array => "Create a new array from elements",
+            GenericArrayOpDef::unpack => "Unpack an array into its elements",
             GenericArrayOpDef::get => "Get an element from an array",
             GenericArrayOpDef::set => "Set an element in an array",
             GenericArrayOpDef::swap => "Swap two elements in an array",
@@ -250,7 +259,7 @@ impl<AK: ArrayKind> MakeOpDef for GenericArrayOpDef<AK> {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-/// Concrete array operation. Generic over the actual array implemenation.
+/// Concrete array operation. Generic over the actual array implementation.
 pub struct GenericArrayOp<AK: ArrayKind> {
     /// The operation definition.
     pub def: GenericArrayOpDef<AK>,
@@ -275,7 +284,7 @@ impl<AK: ArrayKind> MakeExtensionOp for GenericArrayOp<AK> {
 
     fn type_args(&self) -> Vec<TypeArg> {
         use GenericArrayOpDef::{
-            _phantom, discard_empty, get, new_array, pop_left, pop_right, set, swap,
+            _phantom, discard_empty, get, new_array, pop_left, pop_right, set, swap, unpack,
         };
         let ty_arg = TypeArg::Type {
             ty: self.elem_ty.clone(),
@@ -288,7 +297,7 @@ impl<AK: ArrayKind> MakeExtensionOp for GenericArrayOp<AK> {
                 );
                 vec![ty_arg]
             }
-            new_array | pop_left | pop_right | get | set | swap => {
+            new_array | unpack | pop_left | pop_right | get | set | swap => {
                 vec![TypeArg::BoundedNat { n: self.size }, ty_arg]
             }
             _phantom(_, never) => match never {},
@@ -379,6 +388,22 @@ mod tests {
         b.finish_hugr_with_outputs(out.outputs()).unwrap();
     }
 
+    #[rstest]
+    #[case(Array)]
+    #[case(ValueArray)]
+    /// Test building a HUGR involving an unpack operation.
+    fn test_unpack<AK: ArrayKind>(#[case] _kind: AK) {
+        let mut b = DFGBuilder::new(inout_sig(AK::ty(2, qb_t()), vec![qb_t(), qb_t()])).unwrap();
+
+        let [array] = b.input_wires_arr();
+
+        let op = GenericArrayOpDef::<AK>::unpack.to_concrete(qb_t(), 2);
+
+        let out = b.add_dataflow_op(op, [array]).unwrap();
+
+        b.finish_hugr_with_outputs(out.outputs()).unwrap();
+    }
+
     #[rstest]
     #[case(Array)]
     #[case(ValueArray)]

hugr-core/src/std_extensions/collections/array/op_builder.rs

@@ -49,6 +49,32 @@ pub trait GenericArrayOpBuilder: Dataflow {
         Ok(out)
     }
 
+    /// Adds an array unpack operation to the dataflow graph.
+    ///
+    /// This operation unpacks an array into individual elements.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `size` - The size of the array.
+    /// * `input` - The wire representing the array.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if building the operation fails.
+    ///
+    /// # Returns
+    ///
+    /// A vector of wires representing the individual elements of the array.
+    fn add_generic_array_unpack<AK: ArrayKind>(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+    ) -> Result<Vec<Wire>, BuildError> {
+        let op = GenericArrayOpDef::<AK>::unpack.instantiate(&[size.into(), elem_ty.into()])?;
+        Ok(self.add_dataflow_op(op, vec![input])?.outputs().collect())
+    }
     /// Adds an array clone operation to the dataflow graph and return the wires
     /// representing the originala and cloned array.
     ///
@@ -283,6 +309,17 @@ pub fn build_all_array_ops_generic<B: Dataflow, AK: ArrayKind>(mut builder: B) -
     let us0 = builder.add_load_value(ConstUsize::new(0));
     let us1 = builder.add_load_value(ConstUsize::new(1));
     let us2 = builder.add_load_value(ConstUsize::new(2));
+    let arr = builder
+        .add_new_generic_array::<AK>(usize_t(), [us1, us2])
+        .unwrap();
+
+    // Add array unpack operation
+    let [_us1, _us2] = builder
+        .add_generic_array_unpack::<AK>(usize_t(), 2, arr)
+        .unwrap()
+        .try_into()
+        .unwrap();
+
     let arr = builder
         .add_new_generic_array::<AK>(usize_t(), [us1, us2])
         .unwrap();

hugr-core/src/std_extensions/collections/value_array.rs

@@ -31,7 +31,7 @@ pub const VALUE_ARRAY_VALUENAME: TypeName = TypeName::new_inline("value_array");
 /// Reported unique name of the extension
 pub const EXTENSION_ID: ExtensionId = ExtensionId::new_static_unchecked("collections.value_array");
 /// Extension version.
-pub const VERSION: semver::Version = semver::Version::new(0, 1, 0);
+pub const VERSION: semver::Version = semver::Version::new(0, 1, 1);
 
 /// A fixed-length collection of values.
 ///

hugr-llvm/src/extension/collections/array.rs

@@ -455,6 +455,24 @@ pub fn emit_array_op<'c, H: HugrView<Node = Node>>(
             }
             outputs.finish(ctx.builder(), [array_v.into()])
         }
+        ArrayOpDef::unpack => {
+            let [array_v] = inputs
+                .try_into()
+                .map_err(|_| anyhow!("ArrayOpDef::unpack expects one argument"))?;
+            let (array_ptr, array_offset) = decompose_array_fat_pointer(builder, array_v)?;
+
+            let mut result = Vec::with_capacity(size as usize);
+            let usize_t = usize_ty(&ctx.typing_session());
+
+            for i in 0..size {
+                let idx = builder.build_int_add(array_offset, usize_t.const_int(i, false), "")?;
+                let elem_addr = unsafe { builder.build_in_bounds_gep(array_ptr, &[idx], "")? };
+                let elem_v = builder.build_load(elem_addr, "")?;
+                result.push(elem_v);
+            }
+
+            outputs.finish(ctx.builder(), result)
+        }
         ArrayOpDef::get => {
             let [array_v, index_v] = inputs
                 .try_into()
@@ -999,6 +1017,41 @@ mod test {
         check_emission!(hugr, llvm_ctx);
     }
 
+    // #[rstest]
+    // #[case(1, 2, 3)]
+    // #[case(0, 0, 0)]
+    // #[case(10, 20, 30)]
+    // fn exec_unpack_and_sum(mut exec_ctx: TestContext, #[case] a: u64, #[case] b: u64, #[case] expected: u64) {
+    //     let hugr = SimpleHugrConfig::new()
+    //         .with_extensions(exec_registry())
+    //         .with_outs(vec![usize_t()])
+    //         .finish(|mut builder| {
+    //             // Create an array with the test values
+    //             let values = vec![ConstUsize::new(a).into(), ConstUsize::new(b).into()];
+    //             let arr = builder.add_load_value(array::ArrayValue::new(usize_t(), values));
+
+    //             // Unpack the array
+    //             let [val_a, val_b] = builder.add_array_unpack(usize_t(), 2, arr).unwrap().try_into().unwrap();
+
+    //             // Add the values
+    //             let sum = {
+    //                 let int_ty = int_type(6);
+    //                 let a_int = builder.cast(val_a, int_ty.clone()).unwrap();
+    //                 let b_int = builder.cast(val_b, int_ty.clone()).unwrap();
+    //                 let sum_int = builder.add_iadd(6, a_int, b_int).unwrap();
+    //                 builder.cast(sum_int, usize_t()).unwrap()
+    //             };
+
+    //             builder.finish_hugr_with_outputs([sum]).unwrap()
+    //         });
+    //     exec_ctx.add_extensions(|cge| {
+    //         cge.add_default_prelude_extensions()
+    //             .add_default_array_extensions()
+    //             .add_default_int_extensions()
+    //     });
+    //     assert_eq!(expected, exec_ctx.exec_hugr_u64(hugr, "main"));
+    // }
+
     fn exec_registry() -> ExtensionRegistry {
         ExtensionRegistry::new([
             int_types::EXTENSION.to_owned(),
@@ -1398,6 +1451,51 @@ mod test {
         assert_eq!(expected, exec_ctx.exec_hugr_u64(hugr, "main"));
     }
 
+    #[rstest]
+    #[case(&[], 0)]
+    #[case(&[1, 2], 3)]
+    #[case(&[6, 6, 6], 18)]
+    fn exec_unpack(
+        mut exec_ctx: TestContext,
+        #[case] array_contents: &[u64],
+        #[case] expected: u64,
+    ) {
+        // We build a HUGR that:
+        // - Loads an array with the given contents
+        // - Unpacks all the elements
+        // - Returns the sum of the elements
+
+        let int_ty = int_type(6);
+        let hugr = SimpleHugrConfig::new()
+            .with_outs(int_ty.clone())
+            .with_extensions(exec_registry())
+            .finish(|mut builder| {
+                let array = array::ArrayValue::new(
+                    int_ty.clone(),
+                    array_contents
+                        .iter()
+                        .map(|&i| ConstInt::new_u(6, i).unwrap().into())
+                        .collect_vec(),
+                );
+                let array = builder.add_load_value(array);
+                let unpacked = builder
+                    .add_array_unpack(int_ty.clone(), array_contents.len() as u64, array)
+                    .unwrap();
+                let mut r = builder.add_load_value(ConstInt::new_u(6, 0).unwrap());
+                for elem in unpacked {
+                    r = builder.add_iadd(6, r, elem).unwrap();
+                }
+
+                builder.finish_hugr_with_outputs([r]).unwrap()
+            });
+        exec_ctx.add_extensions(|cge| {
+            cge.add_default_prelude_extensions()
+                .add_default_array_extensions()
+                .add_default_int_extensions()
+        });
+        assert_eq!(expected, exec_ctx.exec_hugr_u64(hugr, "main"));
+    }
+
     #[rstest]
     #[case(5, 42, 0)]
     #[case(5, 42, 1)]