afarray/
ext.rs

1use std::fmt;
2use std::iter::FromIterator;
3use std::marker::PhantomData;
4use std::ops::*;
5
6use arrayfire as af;
7use async_trait::async_trait;
8use destream::{de, en};
9use futures::{future, stream, Stream};
10use number_general::{DType, NumberType};
11use serde::de::{Deserialize, Deserializer};
12use serde::ser::{Serialize, Serializer};
13
14use super::Complex;
15
16/// Convenience methods defining the base value of a reduce operation on an `ArrayExt`.
17pub trait HasArrayExt
18where
19    Self: af::HasAfEnum,
20{
21    /// The base value of a `product` operation.
22    fn product_one() -> <Self as af::HasAfEnum>::ProductOutType;
23
24    /// The base value of a `sum` operation.
25    fn zero_sum() -> <Self as af::HasAfEnum>::AggregateOutType;
26
27    /// The base value of a comparison (min or max).
28    fn zero_cmp() -> Self;
29}
30
31macro_rules! has_array_ext {
32    ($t:ty, $one:expr, $zero_sum:expr, $zero_cmp:expr) => {
33        impl HasArrayExt for $t {
34            fn product_one() -> <Self as af::HasAfEnum>::ProductOutType {
35                $one
36            }
37
38            fn zero_sum() -> <Self as af::HasAfEnum>::AggregateOutType {
39                $zero_sum
40            }
41
42            fn zero_cmp() -> Self {
43                $zero_cmp
44            }
45        }
46    };
47}
48
49has_array_ext!(bool, true, 0, false);
50has_array_ext!(u8, 1, 0, 0);
51has_array_ext!(u16, 1, 0, 0);
52has_array_ext!(u32, 1, 0, 0);
53has_array_ext!(u64, 1, 0, 0);
54has_array_ext!(i16, 1, 0, 0);
55has_array_ext!(i32, 1, 0, 0);
56has_array_ext!(i64, 1, 0, 0);
57has_array_ext!(f32, 1., 0., 0.);
58has_array_ext!(f64, 1., 0., 0.);
59has_array_ext!(
60    Complex<f32>,
61    Complex::new(1., 1.),
62    Complex::new(0., 0.),
63    Complex::new(0., 0.)
64);
65has_array_ext!(
66    Complex<f64>,
67    Complex::new(1., 1.),
68    Complex::new(0., 0.),
69    Complex::new(0., 0.)
70);
71
72/// Defines common access methods for instance of [`ArrayExt`].
73pub trait ArrayInstance: Deref<Target = af::Array<Self::DType>> + DerefMut {
74    type DType: af::HasAfEnum;
75
76    /// Cast this instance into an [`af::Array`] with type `T`.
77    fn af_cast<T: af::HasAfEnum>(&self) -> af::Array<T> {
78        self.cast()
79    }
80
81    /// How many elements are in this `ArrayInstance`.
82    fn len(&self) -> usize {
83        self.elements()
84    }
85
86    /// Get the values specified by the given [`af::Indexer`].
87    fn get(&self, index: af::Indexer) -> af::Array<Self::DType> {
88        af::index_gen(self, index)
89    }
90
91    /// Set the values specified by the given [`af::Indexer`] to the corresponding values in `T`.
92    fn set<T: ArrayInstance<DType = Self::DType>>(&mut self, index: &af::Indexer, other: &T) {
93        af::assign_gen(self, index, other);
94    }
95
96    /// Set the value at the specified index to `value`.
97    fn set_at(&mut self, offset: usize, value: Self::DType) {
98        let seq = af::seq!(offset as i32, offset as i32, 1);
99        af::assign_seq(self, &[seq], &af::Array::new(&[value], af::dim4!(1)));
100    }
101
102    /// Copy the data in this [`af::Array`] into a new `Vec`.
103    fn to_vec(&self) -> Vec<Self::DType>
104    where
105        Self::DType: Clone + Default,
106    {
107        let mut v = vec![Self::DType::default(); self.len()];
108        self.host(&mut v);
109        v
110    }
111}
112
113/// A wrapper around [`af::Array`] which defines common operations.
114#[derive(Clone)]
115pub struct ArrayExt<T: af::HasAfEnum>(af::Array<T>);
116
117impl<T: af::HasAfEnum + Default> ArrayExt<T> {
118    /// Construct a new `ArrayExt` with the given value and length.
119    pub fn constant(value: T, length: usize) -> Self
120    where
121        T: af::ConstGenerator<OutType = T>,
122    {
123        let dim = af::dim4!(length as u64);
124        Self(af::constant(value, dim))
125    }
126
127    /// Concatenate two instances of `ArrayExt<T>`.
128    pub fn concatenate(left: &Self, right: &Self) -> Self {
129        af::join(0, left, right).into()
130    }
131
132    /// Raise `e` to the power of `self`.
133    pub fn exp(&self) -> ArrayExt<T::UnaryOutType> {
134        af::exp(self).into()
135    }
136
137    /// Cast the values of this `ArrayExt` into a destination type `D`.
138    pub fn type_cast<D: af::HasAfEnum>(&self) -> ArrayExt<D> {
139        ArrayExt(self.af_cast())
140    }
141
142    /// Get the values specified by the given [`af::Indexer`].
143    pub fn get(&self, index: af::Indexer) -> Self {
144        Self(ArrayInstance::get(self, index))
145    }
146
147    /// Get the value at the given index.
148    pub fn get_value(&self, index: usize) -> T {
149        debug_assert!(index < self.elements());
150
151        let af_value = af::index(self, &[af::seq!(index as i32, index as i32, 1)]);
152        if af_value.elements() == 1 {
153            let mut value = vec![T::default()];
154            af_value.host(&mut value);
155            value.pop().unwrap()
156        } else {
157            panic!("no value at {}", index)
158        }
159    }
160
161    /// Deconstruct this `ArrayExt<T>` into its underlying [`af::Array`].
162    pub fn into_inner(self) -> af::Array<T> {
163        self.0
164    }
165
166    /// Return a slice of this `ArrayExt`.
167    ///
168    /// Panics: if `end` is out of bounds
169    pub fn slice(&self, start: usize, end: usize) -> Self {
170        af::index(self, &[af::seq!(start as i32, (end - 1) as i32, 1)]).into()
171    }
172
173    /// Split this `ArrayExt<T>` into two new instances at the given pivot.
174    pub fn split(&self, at: usize) -> (Self, Self) {
175        let left = af::seq!(0, at as i32, 1);
176        let right = af::seq!(at as i32, self.len() as i32, 1);
177        (
178            ArrayExt(af::index(self, &[left, af::seq!()])),
179            ArrayExt(af::index(self, &[right, af::seq!()])),
180        )
181    }
182
183    fn into_stream(self) -> impl Stream<Item = Vec<T>>
184    where
185        T: Clone,
186    {
187        stream::once(future::ready(self.to_vec()))
188    }
189
190    fn to_stream(&self) -> impl Stream<Item = Vec<T>>
191    where
192        T: Clone,
193    {
194        stream::once(future::ready(self.to_vec()))
195    }
196}
197
198impl<T: af::HasAfEnum + af::RealNumber + Clone + Default> ArrayExt<T> {
199    /// Return `true` if the elements of this `ArrayExt` are in sorted order.
200    pub fn is_sorted(&self) -> bool
201    where
202        T: af::RealNumber + Clone,
203    {
204        let sorted = self.sorted(true);
205        af::all_true_all(&af::eq(self.deref(), sorted.deref(), false)).0
206    }
207
208    /// Sort this `ArrayExt`.
209    pub fn sort(&mut self, ascending: bool) {
210        *self = self.sorted(ascending)
211    }
212
213    /// Compute the indices needed to sort this `ArrayExt`.
214    pub fn sort_index(&self, ascending: bool) -> (ArrayExt<T>, ArrayExt<u32>) {
215        let (sorted, indices) = af::sort_index(self, 0, ascending);
216        (sorted.into(), indices.into())
217    }
218
219    /// Return a sorted copy of this `ArrayExt`.
220    pub fn sorted(&self, ascending: bool) -> Self {
221        debug_assert_eq!(self.dims(), af::dim4!(self.len() as u64));
222        Self(af::sort(self, 0, ascending))
223    }
224
225    /// Return only the unique values from this `ArrayExt`.
226    ///
227    /// Pass `true` for `sorted` if this `ArrayExt` is known to be in sorted order.
228    pub fn unique(&self, sorted: bool) -> Self {
229        Self(af::set_unique(self, sorted))
230    }
231}
232
233impl<T: af::HasAfEnum + af::FloatingPoint + Default> ArrayExt<T> {
234    /// Construct a new `ArrayExt` with a random normal distribution.
235    pub fn random_normal(length: usize) -> Self {
236        let dim = af::dim4!(length as u64);
237        let engine = af::get_default_random_engine();
238        Self(af::random_normal(dim, &engine))
239    }
240
241    /// Construct a new `ArrayExt` with a uniform random distribution.
242    pub fn random_uniform(length: usize) -> Self {
243        let dim = af::dim4!(length as u64);
244        let engine = af::get_default_random_engine();
245        Self(af::random_uniform(dim, &engine))
246    }
247}
248
249impl<T: af::HasAfEnum> ArrayInstance for ArrayExt<T> {
250    type DType = T;
251}
252
253impl<T: af::HasAfEnum> Deref for ArrayExt<T> {
254    type Target = af::Array<T>;
255
256    fn deref(&self) -> &Self::Target {
257        &self.0
258    }
259}
260
261impl<T: af::HasAfEnum> DerefMut for ArrayExt<T> {
262    fn deref_mut(&mut self) -> &mut Self::Target {
263        &mut self.0
264    }
265}
266
267impl ArrayExt<bool> {
268    /// Logical not.
269    pub fn not(&self) -> Self {
270        ArrayExt(!self.deref())
271    }
272
273    /// Logical and.
274    pub fn and(&self, other: &Self) -> Self {
275        ArrayExt(af::and(self, other, batch(self, other)))
276    }
277
278    /// Logical or.
279    pub fn or(&self, other: &Self) -> Self {
280        ArrayExt(af::or(self, other, batch(self, other)))
281    }
282
283    /// Logical xor.
284    pub fn xor(&self, other: &Self) -> Self {
285        let batch = batch(self, other);
286        let one = af::or(self, other, batch);
287        let not_both = !(&af::and(self, other, batch));
288        let one_and_not_both = af::and(&one, &not_both, batch);
289        ArrayExt(one_and_not_both)
290    }
291}
292
293impl ArrayExt<Complex<f32>> {
294    /// Get the real component of this array.
295    pub fn re(&self) -> ArrayExt<f32> {
296        af::real(&self.0).into()
297    }
298
299    /// Get the imaginary component of this array.
300    pub fn im(&self) -> ArrayExt<f32> {
301        af::imag(&self.0).into()
302    }
303}
304
305impl ArrayExt<Complex<f64>> {
306    /// Get the real component of this array.
307    pub fn re(&self) -> ArrayExt<f64> {
308        af::real(&self.0).into()
309    }
310
311    /// Get the imaginary component of this array.
312    pub fn im(&self) -> ArrayExt<f64> {
313        af::imag(&self.0).into()
314    }
315}
316
317impl ArrayExt<u64> {
318    /// Construct a new `ArrayExt<u64>` with elements `start..end`.
319    pub fn range(start: u64, end: u64) -> Self {
320        let dims = af::dim4!((end - start));
321        let tile = af::dim4!(1);
322        let range: af::Array<u64> = af::iota(dims, tile).into();
323        if start == 0 {
324            range.into()
325        } else {
326            af::add(&range, &af::Array::new(&[start], af::dim4!(1)), true).into()
327        }
328    }
329}
330
331impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Add for ArrayExt<T>
332    where
333        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
334        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
335        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
336
337    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
338
339    fn add(self, other: Self) -> Self::Output {
340        ArrayExt(af::add(&self.0, &other.0, batch(&self, &other)))
341    }
342}
343
344impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Add for &ArrayExt<T>
345    where
346        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
347        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
348        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
349
350    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
351
352    fn add(self, other: Self) -> Self::Output {
353        ArrayExt(af::add(&self.0, &other.0, batch(self, other)))
354    }
355}
356
357impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Mul for ArrayExt<T>
358    where
359        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
360        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
361        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
362
363    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
364
365    fn mul(self, other: Self) -> Self::Output {
366        ArrayExt(af::mul(&self.0, &other.0, batch(&self, &other)))
367    }
368}
369
370impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Mul for &ArrayExt<T>
371    where
372        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
373        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
374        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
375
376    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
377
378    fn mul(self, other: Self) -> Self::Output {
379        ArrayExt(af::mul(&self.0, &other.0, batch(self, other)))
380    }
381}
382
383impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Div for ArrayExt<T>
384    where
385        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
386        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
387        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
388
389    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
390
391    fn div(self, other: Self) -> Self::Output {
392        ArrayExt(af::div(&self.0, &other.0, batch(&self, &other)))
393    }
394}
395
396impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Div for &ArrayExt<T>
397    where
398        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
399        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
400        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
401
402    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
403
404    fn div(self, other: Self) -> Self::Output {
405        ArrayExt(af::div(&self.0, &other.0, batch(self, other)))
406    }
407}
408
409impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> DivAssign for ArrayExt<T>
410where
411    <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
412    <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
413    <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum
414{
415
416    fn div_assign(&mut self, other: Self) {
417        let quotient = &*self / &other;
418        *self = quotient.type_cast();
419    }
420}
421
422impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Sub for ArrayExt<T>
423where
424    <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
425    <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
426    <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum,
427{
428
429    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
430
431    fn sub(self, other: Self) -> Self::Output {
432        ArrayExt(af::sub(&self.0, &other.0, batch(&self, &other)))
433    }
434}
435
436impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Sub for &ArrayExt<T>
437    where
438        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
439        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
440        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
441
442    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
443
444    fn sub(self, other: Self) -> Self::Output {
445        ArrayExt(af::sub(&self.0, &other.0, batch(self, other)))
446    }
447}
448
449impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> SubAssign for ArrayExt<T>
450    where
451        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
452        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
453        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
454
455    fn sub_assign(&mut self, other: Self) {
456        let diff = &*self - &other;
457        *self = diff.type_cast();
458    }
459}
460
461impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Rem for ArrayExt<T>
462    where
463        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
464        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
465        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
466
467    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
468
469    fn rem(self, other: Self) -> Self::Output {
470        ArrayExt(af::modulo(&self.0, &other.0, batch(&self, &other)))
471    }
472}
473
474impl<T: af::HasAfEnum + af::ImplicitPromote<T> + af::Convertable<OutType = T>> Rem for &ArrayExt<T>
475    where
476        <T as af::ImplicitPromote<T>>::Output: af::HasAfEnum + Default,
477        <T as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
478        <<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output: af::HasAfEnum, {
479
480    type Output = ArrayExt<<<T as af::Convertable>::OutType as af::ImplicitPromote<<T as af::Convertable>::OutType>>::Output>;
481
482    fn rem(self, other: Self) -> Self::Output {
483        ArrayExt(af::modulo(&self.0, &other.0, batch(self, other)))
484    }
485}
486
487impl From<(ArrayExt<f32>, ArrayExt<f32>)> for ArrayExt<Complex<f32>> {
488    fn from(elements: (ArrayExt<f32>, ArrayExt<f32>)) -> Self {
489        let (re, im) = elements;
490        Self(af::cplx2(&*re, &*im, false).cast())
491    }
492}
493
494impl From<(ArrayExt<f64>, ArrayExt<f64>)> for ArrayExt<Complex<f64>> {
495    fn from(elements: (ArrayExt<f64>, ArrayExt<f64>)) -> Self {
496        let (re, im) = elements;
497        Self(af::cplx2(&*re, &*im, false).cast())
498    }
499}
500
501impl<T: af::HasAfEnum> From<af::Array<T>> for ArrayExt<T> {
502    fn from(array: af::Array<T>) -> ArrayExt<T> {
503        ArrayExt(array)
504    }
505}
506
507impl<T: af::HasAfEnum> From<&[T]> for ArrayExt<T> {
508    fn from(values: &[T]) -> ArrayExt<T> {
509        let dim = af::dim4!(values.len() as u64);
510        ArrayExt(af::Array::new(values, dim))
511    }
512}
513
514impl<T: af::HasAfEnum> From<Vec<T>> for ArrayExt<T> {
515    fn from(values: Vec<T>) -> ArrayExt<T> {
516        Self::from(&values[..])
517    }
518}
519
520impl<T: af::HasAfEnum> FromIterator<T> for ArrayExt<T> {
521    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
522        let values = Vec::from_iter(iter);
523        let dim = af::dim4!(values.len() as u64);
524        ArrayExt(af::Array::new(&values, dim))
525    }
526}
527
528/// Defines an absolute value method `abs`.
529pub trait ArrayInstanceAbs: ArrayInstance {
530    type AbsValue: af::HasAfEnum;
531
532    /// Calculate the element-wise absolute value.
533    fn abs(&self) -> ArrayExt<Self::AbsValue>;
534}
535
536impl ArrayInstanceAbs for ArrayExt<Complex<f32>> {
537    type AbsValue = f32;
538
539    fn abs(&self) -> ArrayExt<f32> {
540        ArrayExt(af::abs(&*self))
541    }
542}
543
544impl ArrayInstanceAbs for ArrayExt<Complex<f64>> {
545    type AbsValue = f64;
546
547    fn abs(&self) -> ArrayExt<f64> {
548        ArrayExt(af::abs(&*self))
549    }
550}
551
552impl ArrayInstanceAbs for ArrayExt<f32> {
553    type AbsValue = f32;
554
555    fn abs(&self) -> ArrayExt<f32> {
556        ArrayExt(af::abs(&*self))
557    }
558}
559
560impl ArrayInstanceAbs for ArrayExt<f64> {
561    type AbsValue = f64;
562
563    fn abs(&self) -> ArrayExt<f64> {
564        ArrayExt(af::abs(&*self))
565    }
566}
567
568impl ArrayInstanceAbs for ArrayExt<i16> {
569    type AbsValue = i16;
570
571    fn abs(&self) -> ArrayExt<i16> {
572        ArrayExt(af::abs(&*self).cast())
573    }
574}
575
576impl ArrayInstanceAbs for ArrayExt<i32> {
577    type AbsValue = i32;
578
579    fn abs(&self) -> ArrayExt<i32> {
580        ArrayExt(af::abs(&*self).cast())
581    }
582}
583
584impl ArrayInstanceAbs for ArrayExt<i64> {
585    type AbsValue = i64;
586
587    fn abs(&self) -> ArrayExt<i64> {
588        ArrayExt(af::abs(&*self).cast())
589    }
590}
591
592/// Defines a rounding method `round`.
593pub trait ArrayInstanceRound: ArrayInstance {
594    type Round: af::HasAfEnum;
595
596    /// Round to the nearest integer, element-wise.
597    fn round(&self) -> ArrayExt<Self::Round>;
598}
599
600impl<T: af::HasAfEnum> ArrayInstanceRound for ArrayExt<T> {
601    type Round = T::AbsOutType;
602
603    fn round(&self) -> ArrayExt<Self::Round> {
604        ArrayExt(af::round(self))
605    }
606}
607
608/// Defines cumulative boolean operations `any` and `all`.
609pub trait ArrayInstanceAnyAll: ArrayInstance {
610    /// Returns `true` if all elements are nonzero.
611    fn all(&self) -> bool {
612        af::all_true_all(&*self).0
613    }
614
615    /// Returns `true` if any element is nonzero.
616    fn any(&self) -> bool {
617        af::any_true_all(&*self).0
618    }
619}
620
621impl ArrayInstanceAnyAll for ArrayExt<bool> {}
622impl ArrayInstanceAnyAll for ArrayExt<f32> {}
623impl ArrayInstanceAnyAll for ArrayExt<f64> {}
624impl ArrayInstanceAnyAll for ArrayExt<i16> {}
625impl ArrayInstanceAnyAll for ArrayExt<i32> {}
626impl ArrayInstanceAnyAll for ArrayExt<i64> {}
627impl ArrayInstanceAnyAll for ArrayExt<u8> {}
628impl ArrayInstanceAnyAll for ArrayExt<u16> {}
629impl ArrayInstanceAnyAll for ArrayExt<u32> {}
630impl ArrayInstanceAnyAll for ArrayExt<u64> {}
631
632impl ArrayInstanceAnyAll for ArrayExt<Complex<f32>> {
633    fn all(&self) -> bool {
634        af::all_true_all(&af::abs(&*self)).0
635    }
636
637    fn any(&self) -> bool {
638        let any = af::any_true_all(&*self);
639        any.0 || any.1
640    }
641}
642
643impl ArrayInstanceAnyAll for ArrayExt<Complex<f64>> {
644    fn all(&self) -> bool {
645        af::all_true_all(&af::abs(&*self)).0
646    }
647
648    fn any(&self) -> bool {
649        let any = af::any_true_all(&*self);
650        any.0 || any.1
651    }
652}
653
654/// Indexing operations for `ArrayExt`
655pub trait ArrayInstanceIndex: ArrayInstance {
656    fn argmax(&self) -> (usize, Self::DType);
657}
658
659macro_rules! array_index_real {
660    ($t:ty) => {
661        impl ArrayInstanceIndex for ArrayExt<$t> {
662            fn argmax(&self) -> (usize, $t) {
663                let (max, _, i) = af::imax_all(self);
664                (i as usize, max)
665            }
666        }
667    };
668}
669
670array_index_real!(bool);
671array_index_real!(f32);
672array_index_real!(f64);
673array_index_real!(i16);
674array_index_real!(i32);
675array_index_real!(i64);
676array_index_real!(u8);
677array_index_real!(u16);
678array_index_real!(u32);
679array_index_real!(u64);
680
681macro_rules! array_index_complex {
682    ($t:ty) => {
683        impl ArrayInstanceIndex for ArrayExt<$t> {
684            fn argmax(&self) -> (usize, $t) {
685                let (max_re, max_im, i) = af::imax_all(self);
686                (i as usize, <$t>::new(max_re, max_im))
687            }
688        }
689    };
690}
691
692array_index_complex!(Complex<f32>);
693array_index_complex!(Complex<f64>);
694
695/// Methods to check for infinite or non-numeric array elements.
696pub trait ArrayInstanceUnreal {
697    /// Element-wise check for infinite values.
698    fn is_infinite(&self) -> ArrayExt<bool>;
699
700    /// Element-wise check for non-numeric (NaN) values.
701    fn is_nan(&self) -> ArrayExt<bool>;
702}
703
704impl<T: af::HasAfEnum + af::ImplicitPromote<T>> ArrayInstanceUnreal for ArrayExt<T> {
705    fn is_infinite(&self) -> ArrayExt<bool> {
706        af::isinf(&*self).into()
707    }
708
709    fn is_nan(&self) -> ArrayExt<bool> {
710        af::isnan(&*self).into()
711    }
712}
713
714/// Defines element-wise comparison operations.
715pub trait ArrayInstanceCompare<U> {
716    /// Element-wise equality.
717    fn eq(&self, other: &U) -> ArrayExt<bool>;
718
719    /// Element-wise greater-than comparison.
720    fn gt(&self, other: &U) -> ArrayExt<bool>;
721
722    /// Element-wise greater-or-equal comparison
723    fn gte(&self, other: &U) -> ArrayExt<bool>;
724
725    /// Element-wise less-than comparison.
726    fn lt(&self, other: &U) -> ArrayExt<bool>;
727
728    /// Element-wise less-or-equal comparison.
729    fn lte(&self, other: &U) -> ArrayExt<bool>;
730
731    /// Element-wise inequality.
732    fn ne(&self, other: &U) -> ArrayExt<bool>;
733}
734
735impl<T, U> ArrayInstanceCompare<U> for ArrayExt<T>
736where
737    T: Clone + af::HasAfEnum + af::Convertable<OutType = T>,
738    U: af::Convertable<OutType = T>,
739    <T as af::Convertable>::OutType: af::ImplicitPromote<<U as af::Convertable>::OutType>,
740    <U as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
741{
742    fn eq(&self, other: &U) -> ArrayExt<bool> {
743        af::eq(self.deref(), other, true).into()
744    }
745
746    fn gt(&self, other: &U) -> ArrayExt<bool> {
747        af::gt(self.deref(), other, true).into()
748    }
749
750    fn gte(&self, other: &U) -> ArrayExt<bool> {
751        af::ge(self.deref(), other, true).into()
752    }
753
754    fn lt(&self, other: &U) -> ArrayExt<bool> {
755        af::lt(self.deref(), other, true).into()
756    }
757
758    fn lte(&self, other: &U) -> ArrayExt<bool> {
759        af::le(self.deref(), other, true).into()
760    }
761
762    fn ne(&self, other: &U) -> ArrayExt<bool> {
763        self.eq(other).not()
764    }
765}
766
767/// Defines the natural logarithm.
768pub trait ArrayInstanceNaturalLog<T>: ArrayInstance
769where
770    T: af::HasAfEnum,
771{
772    /// Calculate the element-wise natural logarithm.
773    fn ln(&self) -> ArrayExt<T::UnaryOutType>;
774}
775
776impl<T> ArrayInstanceNaturalLog<T> for ArrayExt<T>
777where
778    T: af::HasAfEnum,
779{
780    fn ln(&self) -> ArrayExt<T::UnaryOutType> {
781        af::log(self).into()
782    }
783}
784
785/// Defines a general logarithm.
786pub trait ArrayInstanceLog<T, U>: ArrayInstance
787where
788    T: af::HasAfEnum,
789    U: af::HasAfEnum,
790    ArrayExt<T>: ArrayInstanceNaturalLog<T>,
791    ArrayExt<U>: ArrayInstanceNaturalLog<U>,
792    ArrayExt<T::UnaryOutType>: Div<ArrayExt<U::UnaryOutType>>,
793{
794    /// Calculate the element-wise logarithm.
795    fn log(
796        &self,
797        base: &ArrayExt<U>,
798    ) -> <ArrayExt<T::UnaryOutType> as Div<ArrayExt<U::UnaryOutType>>>::Output;
799}
800
801impl<T, U> ArrayInstanceLog<T, U> for ArrayExt<T>
802where
803    T: af::HasAfEnum,
804    U: af::HasAfEnum,
805    Self: ArrayInstanceNaturalLog<T>,
806    ArrayExt<U>: ArrayInstanceNaturalLog<U>,
807    ArrayExt<T::UnaryOutType>: Div<ArrayExt<U::UnaryOutType>>,
808{
809    fn log(
810        &self,
811        base: &ArrayExt<U>,
812    ) -> <ArrayExt<T::UnaryOutType> as Div<ArrayExt<U::UnaryOutType>>>::Output {
813        self.ln() / base.ln()
814    }
815}
816
817/// Defines an exponentiation method `pow`.
818pub trait ArrayInstancePow<U>: ArrayInstance {
819    type Pow: af::HasAfEnum;
820
821    /// Calculate the element-wise exponentiation.
822    fn pow(&self, other: &U) -> ArrayExt<Self::Pow>;
823}
824
825impl<T, U> ArrayInstancePow<U> for ArrayExt<T>
826where
827    T: af::HasAfEnum + af::Convertable<OutType = T>,
828    U: af::Convertable<OutType = T>,
829    <T as af::Convertable>::OutType: af::ImplicitPromote<T> + af::ImplicitPromote<<U as af::Convertable>::OutType>,
830    <U as af::Convertable>::OutType: af::ImplicitPromote<<T as af::Convertable>::OutType>,
831    <<T as af::Convertable>::OutType as af::ImplicitPromote<<U as af::Convertable>::OutType>>::Output: af::HasAfEnum,
832{
833    type Pow = <<T as af::Convertable>::OutType as af::ImplicitPromote<<U as af::Convertable>::OutType>>::Output;
834
835    fn pow(&self, other: &U) -> ArrayExt<Self::Pow> {
836        ArrayExt(af::pow(self.deref(), other, true))
837    }
838}
839
840/// Defines common reduction operation `sum`.
841pub trait ArrayInstanceSum<T>: ArrayInstance
842where
843    T: af::HasAfEnum,
844    T::AggregateOutType: DType,
845{
846    type Sum: af::HasAfEnum;
847
848    /// Calculate the cumulative sum.
849    fn sum(&self) -> T::AggregateOutType;
850
851    /// The `NumberType` of the sum of this array.
852    fn sum_dtype() -> NumberType {
853        T::AggregateOutType::dtype()
854    }
855}
856
857/// Defines common reduction operation `product`.
858pub trait ArrayInstanceProduct<T>: ArrayInstance
859where
860    T: af::HasAfEnum,
861    T::ProductOutType: DType,
862{
863    type Product: af::HasAfEnum;
864
865    /// Calculate the cumulative product.
866    fn product(&self) -> T::ProductOutType;
867
868    /// The `NumberType` of the product of this array.
869    fn product_dtype() -> NumberType {
870        T::ProductOutType::dtype()
871    }
872}
873
874/// Defines common reduction operations `min` and `max`.
875pub trait ArrayInstanceMinMax<T>: ArrayInstance
876where
877    T: af::HasAfEnum,
878{
879    /// Find the maximum element.
880    fn max(&self) -> T;
881
882    /// Find the minimum element.
883    fn min(&self) -> T;
884}
885
886macro_rules! reduce_real {
887    ($t:ty) => {
888        impl ArrayInstanceSum<$t> for ArrayExt<$t> {
889            type Sum = <$t as af::HasAfEnum>::AggregateOutType;
890
891            fn sum(&self) -> Self::Sum {
892                af::sum_all(self).0
893            }
894        }
895
896        impl ArrayInstanceProduct<$t> for ArrayExt<$t> {
897            type Product = <$t as af::HasAfEnum>::ProductOutType;
898
899            fn product(&self) -> Self::Product {
900                af::product_all(self).0
901            }
902        }
903
904        impl ArrayInstanceMinMax<$t> for ArrayExt<$t> {
905            fn max(&self) -> $t {
906                af::max_all(self).0
907            }
908
909            fn min(&self) -> $t {
910                af::min_all(self).0
911            }
912        }
913    };
914}
915
916reduce_real!(bool);
917reduce_real!(u8);
918reduce_real!(u16);
919reduce_real!(u32);
920reduce_real!(u64);
921reduce_real!(i16);
922reduce_real!(i32);
923reduce_real!(i64);
924reduce_real!(f32);
925reduce_real!(f64);
926
927macro_rules! reduce_complex {
928    ($t:ty) => {
929        impl ArrayInstanceSum<Complex<$t>> for ArrayExt<Complex<$t>> {
930            type Sum = Complex<$t>;
931
932            fn sum(&self) -> Self::Sum {
933                let sum = af::sum_all(self);
934                Complex::new(sum.0, sum.1)
935            }
936        }
937
938        impl ArrayInstanceProduct<Complex<$t>> for ArrayExt<Complex<$t>> {
939            type Product = Complex<$t>;
940
941            fn product(&self) -> Self::Product {
942                let product = af::product_all(self);
943                Complex::new(product.0, product.1)
944            }
945        }
946
947        impl ArrayInstanceMinMax<Complex<$t>> for ArrayExt<Complex<$t>> {
948            fn max(&self) -> Complex<$t> {
949                let max = af::min_all(self);
950                Complex::new(max.0, max.1)
951            }
952
953            fn min(&self) -> Complex<$t> {
954                let min = af::min_all(self);
955                Complex::new(min.0, min.1)
956            }
957        }
958    };
959}
960
961reduce_complex!(f32);
962reduce_complex!(f64);
963
964macro_rules! unary {
965    ($t:ty, $f:ident) => {
966        fn $f(&self) -> ArrayExt<<$t>::UnaryOutType> {
967            af::$f(self).into()
968        }
969    };
970}
971
972pub trait ArrayInstanceTrig<T>: ArrayInstance<DType = T>
973where
974    T: af::HasAfEnum,
975{
976    unary!(T, sin);
977    unary!(T, asin);
978    unary!(T, sinh);
979    unary!(T, asinh);
980
981    unary!(T, cos);
982    unary!(T, acos);
983    unary!(T, cosh);
984    unary!(T, acosh);
985
986    unary!(T, tan);
987    unary!(T, atan);
988    unary!(T, tanh);
989    unary!(T, atanh);
990}
991
992impl<T: af::HasAfEnum + Default> ArrayInstanceTrig<T> for ArrayExt<T> {}
993
994impl<'de, T: af::HasAfEnum + Deserialize<'de> + 'de> Deserialize<'de> for ArrayExt<T>
995where
996    ArrayExt<T>: From<Vec<T>>,
997{
998    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
999        Vec::<T>::deserialize(deserializer).map(ArrayExt::from)
1000    }
1001}
1002
1003impl<T: af::HasAfEnum + Clone + Default + Serialize> Serialize for ArrayExt<T> {
1004    fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
1005        self.to_vec().serialize(s)
1006    }
1007}
1008
1009#[async_trait]
1010impl de::FromStream for ArrayExt<bool> {
1011    type Context = ();
1012
1013    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1014        decoder
1015            .decode_array_bool(ArrayExtVisitor::<bool>::default())
1016            .await
1017    }
1018}
1019
1020#[async_trait]
1021impl de::FromStream for ArrayExt<f32> {
1022    type Context = ();
1023
1024    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1025        decoder
1026            .decode_array_f32(ArrayExtVisitor::<f32>::default())
1027            .await
1028    }
1029}
1030
1031#[async_trait]
1032impl de::FromStream for ArrayExt<f64> {
1033    type Context = ();
1034
1035    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1036        decoder
1037            .decode_array_f64(ArrayExtVisitor::<f64>::default())
1038            .await
1039    }
1040}
1041
1042#[async_trait]
1043impl de::FromStream for ArrayExt<u8> {
1044    type Context = ();
1045
1046    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1047        decoder
1048            .decode_array_u8(ArrayExtVisitor::<u8>::default())
1049            .await
1050    }
1051}
1052
1053#[async_trait]
1054impl de::FromStream for ArrayExt<u16> {
1055    type Context = ();
1056
1057    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1058        decoder
1059            .decode_array_u16(ArrayExtVisitor::<u16>::default())
1060            .await
1061    }
1062}
1063
1064#[async_trait]
1065impl de::FromStream for ArrayExt<u32> {
1066    type Context = ();
1067
1068    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1069        decoder
1070            .decode_array_u32(ArrayExtVisitor::<u32>::default())
1071            .await
1072    }
1073}
1074
1075#[async_trait]
1076impl de::FromStream for ArrayExt<u64> {
1077    type Context = ();
1078
1079    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1080        decoder
1081            .decode_array_u64(ArrayExtVisitor::<u64>::default())
1082            .await
1083    }
1084}
1085
1086#[async_trait]
1087impl de::FromStream for ArrayExt<i16> {
1088    type Context = ();
1089
1090    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1091        decoder
1092            .decode_array_i16(ArrayExtVisitor::<i16>::default())
1093            .await
1094    }
1095}
1096
1097#[async_trait]
1098impl de::FromStream for ArrayExt<i32> {
1099    type Context = ();
1100
1101    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1102        decoder
1103            .decode_array_i32(ArrayExtVisitor::<i32>::default())
1104            .await
1105    }
1106}
1107
1108#[async_trait]
1109impl de::FromStream for ArrayExt<i64> {
1110    type Context = ();
1111
1112    async fn from_stream<D: de::Decoder>(_: (), decoder: &mut D) -> Result<Self, D::Error> {
1113        decoder
1114            .decode_array_i64(ArrayExtVisitor::<i64>::default())
1115            .await
1116    }
1117}
1118
1119impl<'en> en::IntoStream<'en> for ArrayExt<bool> {
1120    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1121        encoder.encode_array_bool(self.into_stream())
1122    }
1123}
1124
1125impl<'en> en::ToStream<'en> for ArrayExt<bool> {
1126    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1127        encoder.encode_array_bool(self.to_stream())
1128    }
1129}
1130
1131impl<'en> en::IntoStream<'en> for ArrayExt<f32> {
1132    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1133        encoder.encode_array_f32(self.into_stream())
1134    }
1135}
1136
1137impl<'en> en::ToStream<'en> for ArrayExt<f32> {
1138    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1139        encoder.encode_array_f32(self.to_stream())
1140    }
1141}
1142
1143impl<'en> en::IntoStream<'en> for ArrayExt<f64> {
1144    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1145        encoder.encode_array_f64(self.into_stream())
1146    }
1147}
1148
1149impl<'en> en::ToStream<'en> for ArrayExt<f64> {
1150    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1151        encoder.encode_array_f64(self.to_stream())
1152    }
1153}
1154
1155impl<'en> en::IntoStream<'en> for ArrayExt<u8> {
1156    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1157        encoder.encode_array_u8(self.into_stream())
1158    }
1159}
1160
1161impl<'en> en::ToStream<'en> for ArrayExt<u8> {
1162    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1163        encoder.encode_array_u8(self.to_stream())
1164    }
1165}
1166
1167impl<'en> en::IntoStream<'en> for ArrayExt<u16> {
1168    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1169        encoder.encode_array_u16(self.into_stream())
1170    }
1171}
1172
1173impl<'en> en::ToStream<'en> for ArrayExt<u16> {
1174    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1175        encoder.encode_array_u16(self.to_stream())
1176    }
1177}
1178
1179impl<'en> en::IntoStream<'en> for ArrayExt<u32> {
1180    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1181        encoder.encode_array_u32(self.into_stream())
1182    }
1183}
1184
1185impl<'en> en::ToStream<'en> for ArrayExt<u32> {
1186    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1187        encoder.encode_array_u32(self.to_stream())
1188    }
1189}
1190
1191impl<'en> en::IntoStream<'en> for ArrayExt<u64> {
1192    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1193        encoder.encode_array_u64(self.into_stream())
1194    }
1195}
1196
1197impl<'en> en::ToStream<'en> for ArrayExt<u64> {
1198    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1199        encoder.encode_array_u64(self.to_stream())
1200    }
1201}
1202
1203impl<'en> en::IntoStream<'en> for ArrayExt<i16> {
1204    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1205        encoder.encode_array_i16(self.into_stream())
1206    }
1207}
1208
1209impl<'en> en::ToStream<'en> for ArrayExt<i16> {
1210    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1211        encoder.encode_array_i16(self.to_stream())
1212    }
1213}
1214
1215impl<'en> en::IntoStream<'en> for ArrayExt<i32> {
1216    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1217        encoder.encode_array_i32(self.into_stream())
1218    }
1219}
1220
1221impl<'en> en::ToStream<'en> for ArrayExt<i32> {
1222    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1223        encoder.encode_array_i32(self.to_stream())
1224    }
1225}
1226
1227impl<'en> en::IntoStream<'en> for ArrayExt<i64> {
1228    fn into_stream<E: en::Encoder<'en>>(self, encoder: E) -> Result<E::Ok, E::Error> {
1229        encoder.encode_array_i64(self.into_stream())
1230    }
1231}
1232
1233impl<'en> en::ToStream<'en> for ArrayExt<i64> {
1234    fn to_stream<E: en::Encoder<'en>>(&'en self, encoder: E) -> Result<E::Ok, E::Error> {
1235        encoder.encode_array_i64(self.to_stream())
1236    }
1237}
1238
1239impl<T: af::HasAfEnum + fmt::Display + Default + Clone> fmt::Debug for ArrayExt<T> {
1240    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1241        write!(
1242            f,
1243            "ArrayExt<{}>({}): {}",
1244            std::any::type_name::<T>(),
1245            self.dims(),
1246            self
1247        )
1248    }
1249}
1250
1251impl<T: af::HasAfEnum + fmt::Display + Default + Clone> fmt::Display for ArrayExt<T> {
1252    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1253        let as_str: String = self
1254            .to_vec()
1255            .into_iter()
1256            .map(|n| n.to_string())
1257            .collect::<Vec<String>>()
1258            .join(",");
1259
1260        write!(f, "[{}]", as_str)
1261    }
1262}
1263
1264#[derive(Default)]
1265struct ArrayExtVisitor<T> {
1266    phantom: PhantomData<T>,
1267}
1268
1269impl<'a, T: af::HasAfEnum + Clone + Copy + Default + Send + 'a> ArrayExtVisitor<T>
1270where
1271    ArrayExt<T>: From<Vec<T>>,
1272{
1273    async fn visit_array<A: de::ArrayAccess<T>>(mut access: A) -> Result<ArrayExt<T>, A::Error> {
1274        let mut buf = [T::default(); 4096];
1275        let mut elements = Vec::new();
1276
1277        loop {
1278            let num = access.buffer(&mut buf).await?;
1279            if num == 0 {
1280                break;
1281            } else {
1282                elements.extend_from_slice(&buf[..num]);
1283            }
1284        }
1285
1286        Ok(ArrayExt::from(elements))
1287    }
1288}
1289
1290#[async_trait]
1291impl de::Visitor for ArrayExtVisitor<bool> {
1292    type Value = ArrayExt<bool>;
1293
1294    fn expecting() -> &'static str {
1295        "a boolean array"
1296    }
1297
1298    async fn visit_array_bool<A: de::ArrayAccess<bool>>(
1299        self,
1300        access: A,
1301    ) -> Result<Self::Value, A::Error> {
1302        Self::visit_array(access).await
1303    }
1304}
1305
1306#[async_trait]
1307impl de::Visitor for ArrayExtVisitor<f32> {
1308    type Value = ArrayExt<f32>;
1309
1310    fn expecting() -> &'static str {
1311        "a 32-bit float array"
1312    }
1313
1314    async fn visit_array_f32<A: de::ArrayAccess<f32>>(
1315        self,
1316        access: A,
1317    ) -> Result<Self::Value, A::Error> {
1318        Self::visit_array(access).await
1319    }
1320}
1321
1322#[async_trait]
1323impl de::Visitor for ArrayExtVisitor<f64> {
1324    type Value = ArrayExt<f64>;
1325
1326    fn expecting() -> &'static str {
1327        "a 64-bit float array"
1328    }
1329
1330    async fn visit_array_f64<A: de::ArrayAccess<f64>>(
1331        self,
1332        access: A,
1333    ) -> Result<Self::Value, A::Error> {
1334        Self::visit_array(access).await
1335    }
1336}
1337
1338#[async_trait]
1339impl de::Visitor for ArrayExtVisitor<u8> {
1340    type Value = ArrayExt<u8>;
1341
1342    fn expecting() -> &'static str {
1343        "an 8-bit unsigned integer array"
1344    }
1345
1346    async fn visit_array_u8<A: de::ArrayAccess<u8>>(
1347        self,
1348        access: A,
1349    ) -> Result<Self::Value, A::Error> {
1350        Self::visit_array(access).await
1351    }
1352}
1353
1354#[async_trait]
1355impl de::Visitor for ArrayExtVisitor<u16> {
1356    type Value = ArrayExt<u16>;
1357
1358    fn expecting() -> &'static str {
1359        "a 16-bit unsigned integer array"
1360    }
1361
1362    async fn visit_array_u16<A: de::ArrayAccess<u16>>(
1363        self,
1364        access: A,
1365    ) -> Result<Self::Value, A::Error> {
1366        Self::visit_array(access).await
1367    }
1368}
1369
1370#[async_trait]
1371impl de::Visitor for ArrayExtVisitor<u32> {
1372    type Value = ArrayExt<u32>;
1373
1374    fn expecting() -> &'static str {
1375        "a 32-bit unsigned integer array"
1376    }
1377
1378    async fn visit_array_u32<A: de::ArrayAccess<u32>>(
1379        self,
1380        access: A,
1381    ) -> Result<Self::Value, A::Error> {
1382        Self::visit_array(access).await
1383    }
1384}
1385
1386#[async_trait]
1387impl de::Visitor for ArrayExtVisitor<u64> {
1388    type Value = ArrayExt<u64>;
1389
1390    fn expecting() -> &'static str {
1391        "a 64-bit unsigned integer array"
1392    }
1393
1394    async fn visit_array_u64<A: de::ArrayAccess<u64>>(
1395        self,
1396        access: A,
1397    ) -> Result<Self::Value, A::Error> {
1398        Self::visit_array(access).await
1399    }
1400}
1401
1402#[async_trait]
1403impl de::Visitor for ArrayExtVisitor<i16> {
1404    type Value = ArrayExt<i16>;
1405
1406    fn expecting() -> &'static str {
1407        "a 16-bit integer array"
1408    }
1409
1410    async fn visit_array_i16<A: de::ArrayAccess<i16>>(
1411        self,
1412        access: A,
1413    ) -> Result<Self::Value, A::Error> {
1414        Self::visit_array(access).await
1415    }
1416}
1417
1418#[async_trait]
1419impl de::Visitor for ArrayExtVisitor<i32> {
1420    type Value = ArrayExt<i32>;
1421
1422    fn expecting() -> &'static str {
1423        "a 32-bit integer array"
1424    }
1425
1426    async fn visit_array_i32<A: de::ArrayAccess<i32>>(
1427        self,
1428        access: A,
1429    ) -> Result<Self::Value, A::Error> {
1430        Self::visit_array(access).await
1431    }
1432}
1433
1434#[async_trait]
1435impl de::Visitor for ArrayExtVisitor<i64> {
1436    type Value = ArrayExt<i64>;
1437
1438    fn expecting() -> &'static str {
1439        "a 64-bit integer array"
1440    }
1441
1442    async fn visit_array_i64<A: de::ArrayAccess<i64>>(
1443        self,
1444        access: A,
1445    ) -> Result<Self::Value, A::Error> {
1446        Self::visit_array(access).await
1447    }
1448}
1449
1450#[cfg(test)]
1451mod tests {
1452    use super::*;
1453
1454    #[test]
1455    fn test_range() {
1456        let range = ArrayExt::range(1, 10);
1457        assert_eq!(range.to_vec(), (1..10).collect::<Vec<u64>>())
1458    }
1459}
1460
1461#[inline]
1462fn batch<T: af::HasAfEnum>(this: &ArrayExt<T>, that: &ArrayExt<T>) -> bool {
1463    this.0.dims() != that.0.dims()
1464}
afarray/ext.rs

afarray/
ext.rs
