1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

#![deny(unsafe_op_in_unsafe_fn)]
#![cfg_attr(not(any(doc, test, feature = "std")), no_std)]

use core::mem::{needs_drop, MaybeUninit};
use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};

#[cfg(test)]
mod tests;

mod iter;
use iter::{uninit_array, Slice};

pub struct Array<T, const N: usize>(pub [T; N]);

fn binop_impl<T, U, O, const N: usize>(
    lhs: [T; N],
    rhs: [U; N],
    op: impl Fn(T, U) -> O + Copy,
) -> [O; N] {
    if !needs_drop::<T>() && !needs_drop::<U>() && !needs_drop::<O>() {
        // SAFETY:
        // we've just checked that T, U and O are non-drop types
        unsafe { binop_impl_copy(lhs, rhs, op) }
    } else {
        binop_impl_drop(lhs, rhs, op)
    }
}

fn binop_impl_drop<T, U, O, const N: usize>(
    lhs: [T; N],
    rhs: [U; N],
    op: impl Fn(T, U) -> O + Copy,
) -> [O; N] {
    let mut lhs = Slice::full(lhs);
    let mut rhs = Slice::full(rhs);
    let mut output = Slice::new();

    for _ in 0..N {
        unsafe {
            let lhs = lhs.pop_front_unchecked();
            let rhs = rhs.pop_front_unchecked();
            output.push_unchecked(op(lhs, rhs));
        }
    }

    unsafe { output.output() }
}

/// # Safety
/// must only be called if T, U and O are Copy types (no drop needed)
unsafe fn binop_impl_copy<T, U, O, const N: usize>(
    lhs: [T; N],
    rhs: [U; N],
    op: impl Fn(T, U) -> O + Copy,
) -> [O; N] {
    // SAFETY:
    // we will not read from output, and caller ensures that O is non-drop
    let mut output: [MaybeUninit<O>; N] = uninit_array();

    for i in 0..N {
        unsafe {
            let lhs = core::ptr::read(&lhs[i]);
            let rhs = core::ptr::read(&rhs[i]);
            output[i].write(op(lhs, rhs));
        }
    }

    unsafe { core::ptr::read(&output as *const [MaybeUninit<O>; N] as *const [O; N]) }
}

fn binop_assign_impl<T, U, const N: usize>(
    lhs: &mut [T; N],
    rhs: [U; N],
    op: impl Fn(&mut T, U) + Copy,
) {
    let mut rhs = Slice::full(rhs);

    for i in 0..N {
        // SAFETY:
        // Will only be called a maximum of N times
        unsafe { op(lhs.get_unchecked_mut(i), rhs.pop_front_unchecked()) }
    }
}

macro_rules! binop {
    ($trait:ident, $method:ident) => {
        impl<T, U, const N: usize> $trait<[U; N]> for Array<T, N>
        where
            T: $trait<U>,
        {
            type Output = [T::Output; N];

            fn $method(self, rhs: [U; N]) -> Self::Output {
                binop_impl(self.0, rhs, T::$method)
            }
        }

        impl<T, U, const N: usize> $trait<Array<U, N>> for Array<T, N>
        where
            T: $trait<U>,
        {
            type Output = Array<T::Output, N>;

            fn $method(self, rhs: Array<U, N>) -> Self::Output {
                Array(binop_impl(self.0, rhs.0, T::$method))
            }
        }
    };
}

macro_rules! binop_assign {
    ($trait:ident, $method:ident) => {
        impl<T, U, const N: usize> $trait<[U; N]> for Array<T, N>
        where
            T: $trait<U>,
        {
            fn $method(&mut self, rhs: [U; N]) {
                binop_assign_impl(&mut self.0, rhs, T::$method)
            }
        }

        impl<T, U, const N: usize> $trait<Array<U, N>> for Array<T, N>
        where
            T: $trait<U>,
        {
            fn $method(&mut self, rhs: Array<U, N>) {
                binop_assign_impl(&mut self.0, rhs.0, T::$method)
            }
        }
    };
}

binop!(Add, add);
binop!(Mul, mul);
binop!(Div, div);
binop!(Sub, sub);

binop_assign!(AddAssign, add_assign);
binop_assign!(MulAssign, mul_assign);
binop_assign!(DivAssign, div_assign);
binop_assign!(SubAssign, sub_assign);