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}