fix(qsystem): handle barrier lowering for all array kinds (#1024)

ss2165 · web-flow · commit 0e553218c0f7 · 2025-08-08T14:14:12.000Z
Closes #1023
diff --git a/tket-qsystem/src/extension/qsystem/barrier.rs b/tket-qsystem/src/extension/qsystem/barrier.rs
@@ -8,6 +8,8 @@ mod test {
     use crate::extension::qsystem::{self, lower_tk2_op};
     use hugr::builder::{Dataflow, DataflowHugr};
     use hugr::extension::prelude::Barrier;
+    use hugr::std_extensions::collections::borrow_array::borrow_array_type;
+    use hugr::std_extensions::collections::value_array::value_array_type;
     use hugr::{
         builder::DFGBuilder,
         extension::prelude::{bool_t, option_type, qb_t},
@@ -32,6 +34,8 @@ mod test {
     // bare option of qubit is ignored
     #[case(vec![qb_t(), option_type(qb_t()).into()], 1, false)]
     #[case(vec![array_type(2, bool_t())], 0, false)]
+    #[case(vec![value_array_type(2, option_type(qb_t()).into())], 2, false)]
+    #[case(vec![borrow_array_type(2, qb_t())], 2, false)]
     // special case, single array of qubits is passed directly to op without unpacking
     #[case(vec![array_type(3, qb_t())], 1, true)]
     #[case(vec![qb_t(), array_type(2, qb_t()), array_type(2, array_type(2, qb_t()))], 7, false)]
@@ -93,6 +97,7 @@ mod test {
                 let mut analyzer = QTypeAnalyzer::new();
                 let tuple_type = hugr::types::Type::new_tuple(type_row);
                 assert!(!analyzer.is_qubit_container(&tuple_type));
+                assert_eq!(num_qb, 0);
                 return;
             }
             h.single_linked_input(run_barr_func_n.unwrap(), 0)
diff --git a/tket-qsystem/src/extension/qsystem/barrier/barrier_inserter.rs b/tket-qsystem/src/extension/qsystem/barrier/barrier_inserter.rs
@@ -69,7 +69,9 @@ impl BarrierInserter {
         target: Target,
     ) -> Option<Result<(), LowerTk2Error>> {
         // Check if this is an array of qubits
-        let size = is_qubit_array(typ)?;
+        let size = is_qubit_array::<hugr::std_extensions::collections::array::Array>(typ)?;
+
+        // TODO if other array type, convert
 
         // Build and insert the barrier
         Some(match build_runtime_barrier_op(size) {
diff --git a/tket-qsystem/src/extension/qsystem/barrier/barrier_ops.rs b/tket-qsystem/src/extension/qsystem/barrier/barrier_ops.rs
@@ -9,7 +9,11 @@ use hugr::{
         Extension,
     },
     ops::{DataflowOpTrait, ExtensionOp, OpName},
-    std_extensions::collections::array::{array_type, array_type_parametric, ArrayOpBuilder},
+    std_extensions::collections::{
+        array::{array_type, op_builder::GenericArrayOpBuilder, Array, ArrayKind},
+        borrow_array::BorrowArray,
+        value_array::ValueArray,
+    },
     types::{
         type_param::TypeParam, FuncValueType, PolyFuncTypeRV, Signature, SumType, Type, TypeArg,
         TypeBound, TypeRV,
@@ -62,6 +66,43 @@ pub struct BarrierOperationFactory {
     type_analyzer: QTypeAnalyzer,
 }
 
+fn generic_array_unpack_sig<AK: ArrayKind>() -> PolyFuncTypeRV {
+    PolyFuncTypeRV::new(
+        vec![
+            TypeParam::max_nat_type(),
+            TypeParam::RuntimeType(TypeBound::Linear),
+            TypeParam::new_list_type(TypeBound::Linear),
+        ],
+        FuncValueType::new(
+            AK::ty_parametric(
+                TypeArg::new_var_use(0, TypeParam::max_nat_type()),
+                Type::new_var_use(1, TypeBound::Linear),
+            )
+            .unwrap(),
+            TypeRV::new_row_var_use(2, TypeBound::Linear),
+        ),
+    )
+}
+
+/// Helper function to add array operations for any ArrayKind to the extension
+fn add_array_ops<AK: ArrayKind>(
+    ext: &mut Extension,
+    ext_ref: &std::sync::Weak<Extension>,
+    unpack_name: OpName,
+    repack_name: OpName,
+) -> Result<(), hugr::extension::ExtensionBuildError> {
+    let array_unpack_sig = generic_array_unpack_sig::<AK>();
+    // pack some wires into an array
+    ext.add_op(
+        repack_name,
+        Default::default(),
+        invert_sig(&array_unpack_sig),
+        ext_ref,
+    )?;
+    // unpack an array into some wires
+    ext.add_op(unpack_name, Default::default(), array_unpack_sig, ext_ref)?;
+    Ok(())
+}
 impl BarrierOperationFactory {
     /// Temporary extension name.
     pub(super) const TEMP_EXT_NAME: hugr::hugr::IdentList =
@@ -72,6 +113,10 @@ impl BarrierOperationFactory {
     pub(super) const WRAPPED_BARRIER: OpName = OpName::new_static("wrapped_barrier");
     pub(super) const ARRAY_UNPACK: OpName = OpName::new_static("array_unpack");
     pub(super) const ARRAY_REPACK: OpName = OpName::new_static("array_repack");
+    pub(super) const VARRAY_UNPACK: OpName = OpName::new_static("varray_unpack");
+    pub(super) const VARRAY_REPACK: OpName = OpName::new_static("varray_repack");
+    pub(super) const BARRAY_UNPACK: OpName = OpName::new_static("barray_unpack");
+    pub(super) const BARRAY_REPACK: OpName = OpName::new_static("barray_repack");
     pub(super) const TUPLE_UNPACK: OpName = OpName::new_static("tuple_unpack");
     pub(super) const TUPLE_REPACK: OpName = OpName::new_static("tuple_repack");
 
@@ -122,35 +167,17 @@ impl BarrierOperationFactory {
                     ext_ref,
                 )
                 .unwrap();
-                let array_unpack_sig = PolyFuncTypeRV::new(
-                    vec![
-                        TypeParam::max_nat_type(),
-                        TypeParam::RuntimeType(TypeBound::Linear),
-                        TypeParam::new_list_type(TypeBound::Linear),
-                    ],
-                    FuncValueType::new(
-                        array_type_parametric(
-                            TypeArg::new_var_use(0, TypeParam::max_nat_type()),
-                            Type::new_var_use(1, TypeBound::Linear),
-                        )
-                        .unwrap(),
-                        TypeRV::new_row_var_use(2, TypeBound::Linear),
-                    ),
-                );
-                // pack some wires into an array
-                ext.add_op(
-                    Self::ARRAY_REPACK,
-                    Default::default(),
-                    invert_sig(&array_unpack_sig),
-                    ext_ref,
-                )
-                .unwrap();
-                // unpack an array into some wires
-                ext.add_op(
-                    Self::ARRAY_UNPACK,
-                    Default::default(),
-                    array_unpack_sig,
+
+                // Add array operations for all ArrayKind types
+                add_array_ops::<Array>(ext, ext_ref, Self::ARRAY_UNPACK, Self::ARRAY_REPACK)
+                    .unwrap();
+                add_array_ops::<ValueArray>(ext, ext_ref, Self::VARRAY_UNPACK, Self::VARRAY_REPACK)
+                    .unwrap();
+                add_array_ops::<BorrowArray>(
+                    ext,
                     ext_ref,
+                    Self::BARRAY_UNPACK,
+                    Self::BARRAY_REPACK,
                 )
                 .unwrap();
 
@@ -303,46 +330,47 @@ impl BarrierOperationFactory {
         )
     }
 
-    /// Unpack an array into individual wires
-    pub fn unpack_array(
+    /// Generic array unpacking using the ArrayKind trait
+    fn unpack_array<AK: ArrayKind>(
         &mut self,
         builder: &mut impl Dataflow,
         array_wire: Wire,
         size: u64,
         elem_ty: &Type,
+        op_name: &OpName,
     ) -> Result<Vec<Wire>, BuildError> {
-        // Calculate arguments for the array unpacking
-        let args = match self.array_args(size, elem_ty) {
+        let args = match self.array_args::<AK>(size, elem_ty) {
             Some(args) => args,
             None => return Ok(vec![array_wire]), // Not a qubit-containing array
         };
 
         let outputs = self.apply_cached_operation(
             builder,
-            &Self::ARRAY_UNPACK,
+            op_name,
             args.clone(),
             &args[..2],
             [array_wire],
             |slf, func_b| {
                 let w = func_b.input().out_wire(0);
-                let elems = func_b.add_array_unpack(elem_ty.clone(), size, w)?;
-                let unpacked: Vec<_> = elems
+                let elems = func_b.add_generic_array_unpack::<AK>(elem_ty.clone(), size, w)?;
+
+                let result: Vec<_> = elems
                     .into_iter()
                     .map(|wire| slf.unpack_container(func_b, elem_ty, wire))
                     .collect::<Result<Vec<_>, _>>()?
                     .concat();
-                Ok(unpacked)
+                Ok(result)
             },
         )?;
 
         Ok(outputs.collect())
     }
 
     /// Helper function for array arguments
-    fn array_args(&mut self, size: u64, elem_ty: &Type) -> Option<[TypeArg; 3]> {
+    fn array_args<AK: ArrayKind>(&mut self, size: u64, elem_ty: &Type) -> Option<[TypeArg; 3]> {
         let row = self
             .type_analyzer
-            .unpack_type(&array_type(size, elem_ty.clone()))?;
+            .unpack_type(&AK::ty(size, elem_ty.clone()))?;
         let args = [
             size.into(),
             elem_ty.clone().into(),
@@ -351,15 +379,16 @@ impl BarrierOperationFactory {
         Some(args)
     }
 
-    /// Repack wires into an array
-    pub fn repack_array(
+    /// Generic array repacking using the ArrayKind trait
+    fn repack_array<AK: ArrayKind>(
         &mut self,
         builder: &mut impl Dataflow,
         elem_wires: impl IntoIterator<Item = Wire>,
         size: u64,
         elem_ty: &Type,
+        op_name: &OpName,
     ) -> Result<Wire, BuildError> {
-        let args = match self.array_args(size, elem_ty) {
+        let args = match self.array_args::<AK>(size, elem_ty) {
             Some(args) => args,
             None => {
                 return Ok(elem_wires
@@ -373,7 +402,7 @@ impl BarrierOperationFactory {
 
         let mut outputs = self.apply_cached_operation(
             builder,
-            &Self::ARRAY_REPACK,
+            op_name,
             args.clone(),
             &args[..2],
             elem_wires,
@@ -385,8 +414,8 @@ impl BarrierOperationFactory {
                     .chunks(inner_row_len)
                     .map(|chunk| slf.repack_container(func_b, elem_ty, chunk.to_vec()))
                     .collect();
-                let array_wire = func_b.add_new_array(elem_ty.clone(), elems?)?;
 
+                let array_wire = func_b.add_new_generic_array::<AK>(elem_ty.clone(), elems?)?;
                 Ok(vec![array_wire])
             },
         )?;
@@ -528,9 +557,23 @@ impl BarrierOperationFactory {
         if is_opt_qb(typ) {
             return Ok(vec![self.unpack_option(builder, container_wire)?]);
         }
-        if let Some((n, elem_ty)) = typ.as_extension().and_then(array_args) {
-            return self.unpack_array(builder, container_wire, n, elem_ty);
+        macro_rules! handle_array_type {
+            ($array_kind:ty, $unpack_op:expr) => {
+                if let Some((n, elem_ty)) = typ.as_extension().and_then(array_args::<$array_kind>) {
+                    return self.unpack_array::<$array_kind>(
+                        builder,
+                        container_wire,
+                        n,
+                        elem_ty,
+                        &$unpack_op,
+                    );
+                }
+            };
         }
+
+        handle_array_type!(Array, Self::ARRAY_UNPACK);
+        handle_array_type!(ValueArray, Self::VARRAY_UNPACK);
+        handle_array_type!(BorrowArray, Self::BARRAY_UNPACK);
         if let Some(row) = typ.as_sum().and_then(SumType::as_tuple) {
             let row: hugr::types::TypeRow =
                 row.clone().try_into().expect("unexpected row variable.");
@@ -556,10 +599,24 @@ impl BarrierOperationFactory {
             debug_assert!(unpacked_wires.len() == 1);
             return self.repack_option(builder, unpacked_wires[0]);
         }
-        if let Some((n, elem_ty)) = typ.as_extension().and_then(array_args) {
-            return self.repack_array(builder, unpacked_wires, n, elem_ty);
+        macro_rules! handle_array_type {
+            ($array_kind:ty, $repack_op:expr) => {
+                if let Some((n, elem_ty)) = typ.as_extension().and_then(array_args::<$array_kind>) {
+                    return self.repack_array::<$array_kind>(
+                        builder,
+                        unpacked_wires,
+                        n,
+                        elem_ty,
+                        &$repack_op,
+                    );
+                }
+            };
         }
 
+        handle_array_type!(Array, Self::ARRAY_REPACK);
+        handle_array_type!(ValueArray, Self::VARRAY_REPACK);
+        handle_array_type!(BorrowArray, Self::BARRAY_REPACK);
+
         if let Some(row) = typ.as_sum().and_then(SumType::as_tuple) {
             let row: hugr::types::TypeRow =
                 row.clone().try_into().expect("unexpected row variable.");
@@ -586,9 +643,9 @@ pub fn build_runtime_barrier_op(array_size: u64) -> Result<Hugr, BuildError> {
 
 #[cfg(test)]
 mod tests {
-    use hugr::{extension::prelude::bool_t, HugrView};
-
     use super::*;
+    use hugr::{extension::prelude::bool_t, HugrView};
+    use rstest::rstest;
 
     #[test]
     fn test_barrier_op_factory_creation() {
@@ -610,21 +667,48 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn test_array_unpack_repack() -> Result<(), BuildError> {
+    #[rstest]
+    #[case::array(
+        Array,
+        BarrierOperationFactory::ARRAY_UNPACK,
+        BarrierOperationFactory::ARRAY_REPACK
+    )]
+    #[case::value_array(
+        ValueArray,
+        BarrierOperationFactory::VARRAY_UNPACK,
+        BarrierOperationFactory::VARRAY_REPACK
+    )]
+    #[case::borrow_array(
+        BorrowArray,
+        BarrierOperationFactory::BARRAY_UNPACK,
+        BarrierOperationFactory::BARRAY_REPACK
+    )]
+    fn test_array_unpack_repack<AK: ArrayKind>(
+        #[case] _kind: AK,
+        #[case] unpack_op: OpName,
+        #[case] repack_op: OpName,
+    ) -> Result<(), BuildError> {
         let mut factory = BarrierOperationFactory::new();
-        let array_size = 3;
-        let array_type = array_type(array_size, qb_t());
+        let array_size = 2;
 
-        let mut builder = DFGBuilder::new(Signature::new_endo(array_type))?;
+        // Create the specific array type
+        let array_type = AK::ty(array_size, qb_t());
 
+        // Build a dataflow graph that unpacks and repacks the array
+        let mut builder = DFGBuilder::new(Signature::new_endo(array_type))?;
         let input = builder.input().out_wire(0);
-        let unpacked = factory.unpack_array(&mut builder, input, array_size, &qb_t())?;
-        assert_eq!(unpacked.len(), array_size as usize);
 
-        let repacked = factory.repack_array(&mut builder, unpacked, array_size, &qb_t())?;
+        // Unpack the array
+        let unpacked =
+            factory.unpack_array::<AK>(&mut builder, input, array_size, &qb_t(), &unpack_op)?;
+
+        // Repack the array
+        let repacked =
+            factory.repack_array::<AK>(&mut builder, unpacked, array_size, &qb_t(), &repack_op)?;
+
         let hugr = builder.finish_hugr_with_outputs([repacked])?;
         assert!(hugr.validate().is_ok());
+
         Ok(())
     }
 
diff --git a/tket-qsystem/src/extension/qsystem/barrier/qtype_analyzer.rs b/tket-qsystem/src/extension/qsystem/barrier/qtype_analyzer.rs
