array_bin_ops/
simple.rs

1use core::mem::{needs_drop, MaybeUninit};
2use core::ops::{Add, Div, Mul, Sub};
3
4use crate::iter::uninit_array;
5use crate::{iter::Slice, Array};
6
7impl<T, const N: usize> Array<T, N> {
8    pub fn zip_map<U, O>(self, rhs: [U; N], mut op: impl FnMut(T, U) -> O) -> [O; N] {
9        if needs_drop::<T>() || needs_drop::<U>() || needs_drop::<O>() {
10            let mut lhs = Slice::full(self.0);
11            let mut rhs = Slice::full(rhs);
12            let mut output = Slice::new();
13
14            for _ in 0..N {
15                unsafe {
16                    let lhs = lhs.pop_front_unchecked();
17                    let rhs = rhs.pop_front_unchecked();
18                    output.push_unchecked(op(lhs, rhs));
19                }
20            }
21
22            unsafe { output.output() }
23        } else {
24            // SAFETY:
25            // we will not read from output, and caller ensures that O is non-drop
26            let mut output: [MaybeUninit<O>; N] = uninit_array();
27
28            for i in 0..N {
29                unsafe {
30                    let lhs = core::ptr::read(&self.0[i]);
31                    let rhs = core::ptr::read(&rhs[i]);
32                    output[i].write(op(lhs, rhs));
33                }
34            }
35
36            unsafe { core::ptr::read(&output as *const [MaybeUninit<O>; N] as *const [O; N]) }
37        }
38    }
39}
40
41macro_rules! binop {
42    ($trait:ident, $method:ident) => {
43        impl<T, U, const N: usize> $trait<[U; N]> for Array<T, N>
44        where
45            T: $trait<U>,
46        {
47            type Output = [T::Output; N];
48
49            fn $method(self, rhs: [U; N]) -> Self::Output {
50                self.zip_map(rhs, T::$method)
51            }
52        }
53
54        impl<T, U, const N: usize> $trait<Array<U, N>> for Array<T, N>
55        where
56            T: $trait<U>,
57        {
58            type Output = Array<T::Output, N>;
59
60            fn $method(self, rhs: Array<U, N>) -> Self::Output {
61                Array(self.zip_map(rhs.0, T::$method))
62            }
63        }
64    };
65}
66
67binop!(Add, add);
68binop!(Mul, mul);
69binop!(Div, div);
70binop!(Sub, sub);
array_bin_ops/simple.rs

array_bin_ops/
simple.rs
