hugr_core/std_extensions/collections/
array.rs

1//! Fixed-length array type and operations extension.
2
3mod array_clone;
4mod array_conversion;
5mod array_discard;
6mod array_kind;
7mod array_op;
8mod array_repeat;
9mod array_scan;
10mod array_value;
11pub mod op_builder;
12
13use std::sync::Arc;
14
15use delegate::delegate;
16use lazy_static::lazy_static;
17
18use crate::builder::{BuildError, Dataflow};
19use crate::extension::resolution::{ExtensionResolutionError, WeakExtensionRegistry};
20use crate::extension::simple_op::{HasConcrete, MakeOpDef, MakeRegisteredOp};
21use crate::extension::{ExtensionId, SignatureError, TypeDef, TypeDefBound};
22use crate::ops::constant::{CustomConst, ValueName};
23use crate::ops::{ExtensionOp, OpName};
24use crate::types::type_param::{TypeArg, TypeParam};
25use crate::types::{CustomCheckFailure, Type, TypeBound, TypeName};
26use crate::{Extension, Wire};
27
28pub use array_clone::{ARRAY_CLONE_OP_ID, GenericArrayClone, GenericArrayCloneDef};
29pub use array_conversion::{Direction, FROM, GenericArrayConvert, GenericArrayConvertDef, INTO};
30pub use array_discard::{ARRAY_DISCARD_OP_ID, GenericArrayDiscard, GenericArrayDiscardDef};
31pub use array_kind::ArrayKind;
32pub use array_op::{GenericArrayOp, GenericArrayOpDef};
33pub use array_repeat::{ARRAY_REPEAT_OP_ID, GenericArrayRepeat, GenericArrayRepeatDef};
34pub use array_scan::{ARRAY_SCAN_OP_ID, GenericArrayScan, GenericArrayScanDef};
35pub use array_value::GenericArrayValue;
36
37use op_builder::GenericArrayOpBuilder;
38
39/// Reported unique name of the array type.
40pub const ARRAY_TYPENAME: TypeName = TypeName::new_inline("array");
41/// Reported unique name of the array value.
42pub const ARRAY_VALUENAME: TypeName = TypeName::new_inline("array");
43/// Reported unique name of the extension
44pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("collections.array");
45/// Extension version.
46pub const VERSION: semver::Version = semver::Version::new(0, 1, 1);
47
48/// A linear, fixed-length collection of values.
49///
50/// Arrays are linear, even if their elements are copyable.
51#[derive(Clone, Copy, Debug, derive_more::Display, Eq, PartialEq, Default)]
52pub struct Array;
53
54impl ArrayKind for Array {
55    const EXTENSION_ID: ExtensionId = EXTENSION_ID;
56    const TYPE_NAME: TypeName = ARRAY_TYPENAME;
57    const VALUE_NAME: ValueName = ARRAY_VALUENAME;
58
59    fn extension() -> &'static Arc<Extension> {
60        &EXTENSION
61    }
62
63    fn type_def() -> &'static TypeDef {
64        EXTENSION.get_type(&ARRAY_TYPENAME).unwrap()
65    }
66}
67
68/// Array operation definitions.
69pub type ArrayOpDef = GenericArrayOpDef<Array>;
70/// Array clone operation definition.
71pub type ArrayCloneDef = GenericArrayCloneDef<Array>;
72/// Array discard operation definition.
73pub type ArrayDiscardDef = GenericArrayDiscardDef<Array>;
74/// Array repeat operation definition.
75pub type ArrayRepeatDef = GenericArrayRepeatDef<Array>;
76/// Array scan operation definition.
77pub type ArrayScanDef = GenericArrayScanDef<Array>;
78
79/// Array operations.
80pub type ArrayOp = GenericArrayOp<Array>;
81/// The array clone operation.
82pub type ArrayClone = GenericArrayClone<Array>;
83/// The array discard operation.
84pub type ArrayDiscard = GenericArrayDiscard<Array>;
85/// The array repeat operation.
86pub type ArrayRepeat = GenericArrayRepeat<Array>;
87/// The array scan operation.
88pub type ArrayScan = GenericArrayScan<Array>;
89
90/// An array extension value.
91pub type ArrayValue = GenericArrayValue<Array>;
92
93lazy_static! {
94    /// Extension for array operations.
95    pub static ref EXTENSION: Arc<Extension> = {
96        Extension::new_arc(EXTENSION_ID, VERSION, |extension, extension_ref| {
97            extension.add_type(
98                    ARRAY_TYPENAME,
99                    vec![ TypeParam::max_nat_type(), TypeBound::Linear.into()],
100                    "Fixed-length array".into(),
101                    // Default array is linear, even if the elements are copyable
102                    TypeDefBound::any(),
103                    extension_ref,
104                )
105                .unwrap();
106
107            ArrayOpDef::load_all_ops(extension, extension_ref).unwrap();
108            ArrayCloneDef::new().add_to_extension(extension, extension_ref).unwrap();
109            ArrayDiscardDef::new().add_to_extension(extension, extension_ref).unwrap();
110            ArrayRepeatDef::new().add_to_extension(extension, extension_ref).unwrap();
111            ArrayScanDef::new().add_to_extension(extension, extension_ref).unwrap();
112        })
113    };
114}
115
116impl ArrayValue {
117    /// Name of the constructor for creating constant arrays.
118    pub(crate) const CTR_NAME: &'static str = "collections.array.const";
119}
120
121#[typetag::serde(name = "ArrayValue")]
122impl CustomConst for ArrayValue {
123    delegate! {
124        to self {
125            fn name(&self) -> ValueName;
126            fn validate(&self) -> Result<(), CustomCheckFailure>;
127            fn update_extensions(
128                &mut self,
129                extensions: &WeakExtensionRegistry,
130            ) -> Result<(), ExtensionResolutionError>;
131            fn get_type(&self) -> Type;
132        }
133    }
134
135    fn equal_consts(&self, other: &dyn CustomConst) -> bool {
136        crate::ops::constant::downcast_equal_consts(self, other)
137    }
138}
139
140/// Gets the [`TypeDef`] for arrays. Note that instantiations are more easily
141/// created via [`array_type`] and [`array_type_parametric`]
142#[must_use]
143pub fn array_type_def() -> &'static TypeDef {
144    Array::type_def()
145}
146
147/// Instantiate a new array type given a size argument and element type.
148///
149/// This method is equivalent to [`array_type_parametric`], but uses concrete
150/// arguments types to ensure no errors are possible.
151#[must_use]
152pub fn array_type(size: u64, element_ty: Type) -> Type {
153    Array::ty(size, element_ty)
154}
155
156/// Instantiate a new array type given the size and element type parameters.
157///
158/// This is a generic version of [`array_type`].
159pub fn array_type_parametric(
160    size: impl Into<TypeArg>,
161    element_ty: impl Into<TypeArg>,
162) -> Result<Type, SignatureError> {
163    Array::ty_parametric(size, element_ty)
164}
165
166/// Name of the operation in the prelude for creating new arrays.
167pub const NEW_ARRAY_OP_ID: OpName = OpName::new_inline("new_array");
168
169/// Initialize a new array op of element type `element_ty` of length `size`
170#[must_use]
171pub fn new_array_op(element_ty: Type, size: u64) -> ExtensionOp {
172    let op = ArrayOpDef::new_array.to_concrete(element_ty, size);
173    op.to_extension_op().unwrap()
174}
175
176/// Trait for building array operations in a dataflow graph.
177pub trait ArrayOpBuilder: GenericArrayOpBuilder {
178    /// Adds a new array operation to the dataflow graph and return the wire
179    /// representing the new array.
180    ///
181    /// # Arguments
182    ///
183    /// * `elem_ty` - The type of the elements in the array.
184    /// * `values` - An iterator over the values to initialize the array with.
185    ///
186    /// # Errors
187    ///
188    /// If building the operation fails.
189    ///
190    /// # Returns
191    ///
192    /// The wire representing the new array.
193    fn add_new_array(
194        &mut self,
195        elem_ty: Type,
196        values: impl IntoIterator<Item = Wire>,
197    ) -> Result<Wire, BuildError> {
198        self.add_new_generic_array::<Array>(elem_ty, values)
199    }
200    /// Adds an array unpack operation to the dataflow graph.
201    ///
202    /// This operation unpacks an array into individual elements.
203    ///
204    /// # Arguments
205    ///
206    /// * `elem_ty` - The type of the elements in the array.
207    /// * `size` - The size of the array.
208    /// * `input` - The wire representing the array to unpack.
209    ///
210    /// # Errors
211    ///
212    /// If building the operation fails.
213    ///
214    /// # Returns
215    ///
216    /// A vector of wires representing the individual elements from the array.
217    fn add_array_unpack(
218        &mut self,
219        elem_ty: Type,
220        size: u64,
221        input: Wire,
222    ) -> Result<Vec<Wire>, BuildError> {
223        self.add_generic_array_unpack::<Array>(elem_ty, size, input)
224    }
225    /// Adds an array clone operation to the dataflow graph and return the wires
226    /// representing the original and cloned array.
227    ///
228    /// # Arguments
229    ///
230    /// * `elem_ty` - The type of the elements in the array.
231    /// * `size` - The size of the array.
232    /// * `input` - The wire representing the array.
233    ///
234    /// # Errors
235    ///
236    /// If building the operation fails.
237    ///
238    /// # Returns
239    ///
240    /// The wires representing the original and cloned array.
241    fn add_array_clone(
242        &mut self,
243        elem_ty: Type,
244        size: u64,
245        input: Wire,
246    ) -> Result<(Wire, Wire), BuildError> {
247        self.add_generic_array_clone::<Array>(elem_ty, size, input)
248    }
249
250    /// Adds an array discard operation to the dataflow graph.
251    ///
252    /// # Arguments
253    ///
254    /// * `elem_ty` - The type of the elements in the array.
255    /// * `size` - The size of the array.
256    /// * `input` - The wire representing the array.
257    ///
258    /// # Errors
259    ///
260    /// If building the operation fails.
261    fn add_array_discard(
262        &mut self,
263        elem_ty: Type,
264        size: u64,
265        input: Wire,
266    ) -> Result<(), BuildError> {
267        self.add_generic_array_discard::<Array>(elem_ty, size, input)
268    }
269
270    /// Adds an array get operation to the dataflow graph.
271    ///
272    /// # Arguments
273    ///
274    /// * `elem_ty` - The type of the elements in the array.
275    /// * `size` - The size of the array.
276    /// * `input` - The wire representing the array.
277    /// * `index` - The wire representing the index to get.
278    ///
279    /// # Errors
280    ///
281    /// If building the operation fails.
282    ///
283    /// # Returns
284    ///
285    /// * The wire representing the value at the specified index in the array
286    /// * The wire representing the array
287    fn add_array_get(
288        &mut self,
289        elem_ty: Type,
290        size: u64,
291        input: Wire,
292        index: Wire,
293    ) -> Result<(Wire, Wire), BuildError> {
294        self.add_generic_array_get::<Array>(elem_ty, size, input, index)
295    }
296
297    /// Adds an array set operation to the dataflow graph.
298    ///
299    /// This operation sets the value at a specified index in the array.
300    ///
301    /// # Arguments
302    ///
303    /// * `elem_ty` - The type of the elements in the array.
304    /// * `size` - The size of the array.
305    /// * `input` - The wire representing the array.
306    /// * `index` - The wire representing the index to set.
307    /// * `value` - The wire representing the value to set at the specified index.
308    ///
309    /// # Errors
310    ///
311    /// Returns an error if building the operation fails.
312    ///
313    /// # Returns
314    ///
315    /// The wire representing the updated array after the set operation.
316    fn add_array_set(
317        &mut self,
318        elem_ty: Type,
319        size: u64,
320        input: Wire,
321        index: Wire,
322        value: Wire,
323    ) -> Result<Wire, BuildError> {
324        self.add_generic_array_set::<Array>(elem_ty, size, input, index, value)
325    }
326
327    /// Adds an array swap operation to the dataflow graph.
328    ///
329    /// This operation swaps the values at two specified indices in the array.
330    ///
331    /// # Arguments
332    ///
333    /// * `elem_ty` - The type of the elements in the array.
334    /// * `size` - The size of the array.
335    /// * `input` - The wire representing the array.
336    /// * `index1` - The wire representing the first index to swap.
337    /// * `index2` - The wire representing the second index to swap.
338    ///
339    /// # Errors
340    ///
341    /// Returns an error if building the operation fails.
342    ///
343    /// # Returns
344    ///
345    /// The wire representing the updated array after the swap operation.
346    fn add_array_swap(
347        &mut self,
348        elem_ty: Type,
349        size: u64,
350        input: Wire,
351        index1: Wire,
352        index2: Wire,
353    ) -> Result<Wire, BuildError> {
354        let op = GenericArrayOpDef::<Array>::swap.instantiate(&[size.into(), elem_ty.into()])?;
355        let [out] = self
356            .add_dataflow_op(op, vec![input, index1, index2])?
357            .outputs_arr();
358        Ok(out)
359    }
360
361    /// Adds an array pop-left operation to the dataflow graph.
362    ///
363    /// This operation removes the leftmost element from the array.
364    ///
365    /// # Arguments
366    ///
367    /// * `elem_ty` - The type of the elements in the array.
368    /// * `size` - The size of the array.
369    /// * `input` - The wire representing the array.
370    ///
371    /// # Errors
372    ///
373    /// Returns an error if building the operation fails.
374    ///
375    /// # Returns
376    ///
377    /// The wire representing the Option<elemty, array<SIZE-1, elemty>>
378    fn add_array_pop_left(
379        &mut self,
380        elem_ty: Type,
381        size: u64,
382        input: Wire,
383    ) -> Result<Wire, BuildError> {
384        self.add_generic_array_pop_left::<Array>(elem_ty, size, input)
385    }
386
387    /// Adds an array pop-right operation to the dataflow graph.
388    ///
389    /// This operation removes the rightmost element from the array.
390    ///
391    /// # Arguments
392    ///
393    /// * `elem_ty` - The type of the elements in the array.
394    /// * `size` - The size of the array.
395    /// * `input` - The wire representing the array.
396    ///
397    /// # Errors
398    ///
399    /// Returns an error if building the operation fails.
400    ///
401    /// # Returns
402    ///
403    /// The wire representing the Option<elemty, array<SIZE-1, elemty>>
404    fn add_array_pop_right(
405        &mut self,
406        elem_ty: Type,
407        size: u64,
408        input: Wire,
409    ) -> Result<Wire, BuildError> {
410        self.add_generic_array_pop_right::<Array>(elem_ty, size, input)
411    }
412
413    /// Adds an operation to discard an empty array from the dataflow graph.
414    ///
415    /// # Arguments
416    ///
417    /// * `elem_ty` - The type of the elements in the array.
418    /// * `input` - The wire representing the array.
419    ///
420    /// # Errors
421    ///
422    /// Returns an error if building the operation fails.
423    fn add_array_discard_empty(&mut self, elem_ty: Type, input: Wire) -> Result<(), BuildError> {
424        self.add_generic_array_discard_empty::<Array>(elem_ty, input)
425    }
426}
427
428impl<D: Dataflow> ArrayOpBuilder for D {}
429
430#[cfg(test)]
431mod test {
432    use crate::builder::{DFGBuilder, Dataflow, DataflowHugr, inout_sig};
433    use crate::extension::prelude::qb_t;
434
435    use super::{array_type, new_array_op};
436
437    #[test]
438    /// Test building a HUGR involving a `new_array` operation.
439    fn test_new_array() {
440        let mut b =
441            DFGBuilder::new(inout_sig(vec![qb_t(), qb_t()], array_type(2, qb_t()))).unwrap();
442
443        let [q1, q2] = b.input_wires_arr();
444
445        let op = new_array_op(qb_t(), 2);
446
447        let out = b.add_dataflow_op(op, [q1, q2]).unwrap();
448
449        b.finish_hugr_with_outputs(out.outputs()).unwrap();
450    }
451}