feat: Generalise array op builers to arbitrary array kinds

Author: mark-koch

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

@@ -15,14 +15,15 @@ use std::sync::Arc;
 use delegate::delegate;
 use lazy_static::lazy_static;
 
+use crate::builder::{BuildError, Dataflow};
 use crate::extension::resolution::{ExtensionResolutionError, WeakExtensionRegistry};
-use crate::extension::simple_op::{MakeOpDef, MakeRegisteredOp};
+use crate::extension::simple_op::{HasConcrete, MakeOpDef, MakeRegisteredOp};
 use crate::extension::{ExtensionId, ExtensionSet, SignatureError, TypeDef, TypeDefBound};
 use crate::ops::constant::{CustomConst, ValueName};
 use crate::ops::{ExtensionOp, OpName};
 use crate::types::type_param::{TypeArg, TypeParam};
 use crate::types::{CustomCheckFailure, Type, TypeBound, TypeName};
-use crate::Extension;
+use crate::{Extension, Wire};
 
 pub use array_clone::{GenericArrayClone, GenericArrayCloneDef, ARRAY_CLONE_OP_ID};
 pub use array_conversion::{Direction, GenericArrayConvert, GenericArrayConvertDef, FROM, INTO};
@@ -32,7 +33,8 @@ pub use array_op::{GenericArrayOp, GenericArrayOpDef};
 pub use array_repeat::{GenericArrayRepeat, GenericArrayRepeatDef, ARRAY_REPEAT_OP_ID};
 pub use array_scan::{GenericArrayScan, GenericArrayScanDef, ARRAY_SCAN_OP_ID};
 pub use array_value::GenericArrayValue;
-pub use op_builder::ArrayOpBuilder;
+
+use op_builder::GenericArrayOpBuilder;
 
 /// Reported unique name of the array type.
 pub const ARRAY_TYPENAME: TypeName = TypeName::new_inline("array");
@@ -170,6 +172,236 @@ pub fn new_array_op(element_ty: Type, size: u64) -> ExtensionOp {
     op.to_extension_op().unwrap()
 }
 
+/// Trait for building array operations in a dataflow graph.
+pub trait ArrayOpBuilder: GenericArrayOpBuilder {
+    /// Adds a new array operation to the dataflow graph and return the wire
+    /// representing the new array.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `values` - An iterator over the values to initialize the array with.
+    ///
+    /// # Errors
+    ///
+    /// If building the operation fails.
+    ///
+    /// # Returns
+    ///
+    /// The wire representing the new array.
+    fn add_new_array(
+        &mut self,
+        elem_ty: Type,
+        values: impl IntoIterator<Item = Wire>,
+    ) -> Result<Wire, BuildError> {
+        self.add_new_generic_array::<Array>(elem_ty, values)
+    }
+
+    /// Adds an array clone operation to the dataflow graph and return the wires
+    /// representing the originala and cloned array.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `size` - The size of the array.
+    /// * `input` - The wire representing the array.
+    ///
+    /// # Errors
+    ///
+    /// If building the operation fails.
+    ///
+    /// # Returns
+    ///
+    /// The wires representing the original and cloned array.
+    fn add_array_clone(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+    ) -> Result<(Wire, Wire), BuildError> {
+        self.add_generic_array_clone::<Array>(elem_ty, size, input)
+    }
+
+    /// Adds an array discard operation to the dataflow graph.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `size` - The size of the array.
+    /// * `input` - The wire representing the array.
+    ///
+    /// # Errors
+    ///
+    /// If building the operation fails.
+    fn add_array_discard(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+    ) -> Result<(), BuildError> {
+        self.add_generic_array_discard::<Array>(elem_ty, size, input)
+    }
+
+    /// Adds an array get operation to the dataflow graph.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `size` - The size of the array.
+    /// * `input` - The wire representing the array.
+    /// * `index` - The wire representing the index to get.
+    ///
+    /// # Errors
+    ///
+    /// If building the operation fails.
+    ///
+    /// # Returns
+    ///
+    /// * The wire representing the value at the specified index in the array
+    /// * The wire representing the array
+    fn add_array_get(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+        index: Wire,
+    ) -> Result<(Wire, Wire), BuildError> {
+        self.add_generic_array_get::<Array>(elem_ty, size, input, index)
+    }
+
+    /// Adds an array set operation to the dataflow graph.
+    ///
+    /// This operation sets the value at a specified index in the array.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `size` - The size of the array.
+    /// * `input` - The wire representing the array.
+    /// * `index` - The wire representing the index to set.
+    /// * `value` - The wire representing the value to set at the specified index.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if building the operation fails.
+    ///
+    /// # Returns
+    ///
+    /// The wire representing the updated array after the set operation.
+    fn add_array_set(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+        index: Wire,
+        value: Wire,
+    ) -> Result<Wire, BuildError> {
+        self.add_generic_array_set::<Array>(elem_ty, size, input, index, value)
+    }
+
+    /// Adds an array swap operation to the dataflow graph.
+    ///
+    /// This operation swaps the values at two specified indices in the array.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `size` - The size of the array.
+    /// * `input` - The wire representing the array.
+    /// * `index1` - The wire representing the first index to swap.
+    /// * `index2` - The wire representing the second index to swap.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if building the operation fails.
+    ///
+    /// # Returns
+    ///
+    /// The wire representing the updated array after the swap operation.
+    fn add_array_swap(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+        index1: Wire,
+        index2: Wire,
+    ) -> Result<Wire, BuildError> {
+        let op = GenericArrayOpDef::<Array>::swap.instantiate(&[size.into(), elem_ty.into()])?;
+        let [out] = self
+            .add_dataflow_op(op, vec![input, index1, index2])?
+            .outputs_arr();
+        Ok(out)
+    }
+
+    /// Adds an array pop-left operation to the dataflow graph.
+    ///
+    /// This operation removes the leftmost element from the array.
+    ///
+    /// # 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
+    ///
+    /// The wire representing the Option<elemty, array<SIZE-1, elemty>>
+    fn add_array_pop_left(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+    ) -> Result<Wire, BuildError> {
+        self.add_generic_array_pop_left::<Array>(elem_ty, size, input)
+    }
+
+    /// Adds an array pop-right operation to the dataflow graph.
+    ///
+    /// This operation removes the rightmost element from the array.
+    ///
+    /// # 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
+    ///
+    /// The wire representing the Option<elemty, array<SIZE-1, elemty>>
+    fn add_array_pop_right(
+        &mut self,
+        elem_ty: Type,
+        size: u64,
+        input: Wire,
+    ) -> Result<Wire, BuildError> {
+        self.add_generic_array_pop_right::<Array>(elem_ty, size, input)
+    }
+
+    /// Adds an operation to discard an empty array from the dataflow graph.
+    ///
+    /// # Arguments
+    ///
+    /// * `elem_ty` - The type of the elements in the array.
+    /// * `input` - The wire representing the array.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if building the operation fails.
+    fn add_array_discard_empty(&mut self, elem_ty: Type, input: Wire) -> Result<(), BuildError> {
+        self.add_generic_array_discard_empty::<Array>(elem_ty, input)
+    }
+}
+
+impl<D: Dataflow> ArrayOpBuilder for D {}
+
 #[cfg(test)]
 mod test {
     use crate::builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr};