array_bin_ops/
assign.rs

1use core::mem::needs_drop;
2use core::ops::{AddAssign, DivAssign, MulAssign, SubAssign};
3
4use crate::{iter::Slice, Array};
5
6impl<T, const N: usize> Array<T, N> {
7    pub fn zip_mut_map<U>(&mut self, rhs: [U; N], op: impl Fn(&mut T, U) + Copy) {
8        if needs_drop::<U>() {
9            let mut rhs = Slice::full(rhs);
10
11            for i in 0..N {
12                // SAFETY:
13                // Will only be called a maximum of N times
14                unsafe {
15                    let lhs = self.0.get_unchecked_mut(i);
16                    let rhs = rhs.pop_front_unchecked();
17                    op(lhs, rhs)
18                }
19            }
20        } else {
21            for i in 0..N {
22                // SAFETY:
23                // Will only be called a maximum of N times
24                unsafe {
25                    let lhs = self.0.get_unchecked_mut(i);
26                    let rhs = core::ptr::read(&rhs[i]);
27                    op(lhs, rhs)
28                }
29            }
30        }
31    }
32}
33
34macro_rules! binop_assign {
35    ($trait:ident, $method:ident) => {
36        impl<T, U, const N: usize> $trait<[U; N]> for Array<T, N>
37        where
38            T: $trait<U>,
39        {
40            fn $method(&mut self, rhs: [U; N]) {
41                self.zip_mut_map(rhs, T::$method)
42            }
43        }
44
45        impl<T, U, const N: usize> $trait<Array<U, N>> for Array<T, N>
46        where
47            T: $trait<U>,
48        {
49            fn $method(&mut self, rhs: Array<U, N>) {
50                self.zip_mut_map(rhs.0, T::$method)
51            }
52        }
53    };
54}
55
56binop_assign!(AddAssign, add_assign);
57binop_assign!(MulAssign, mul_assign);
58binop_assign!(DivAssign, div_assign);
59binop_assign!(SubAssign, sub_assign);