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

#![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, Div, Mul, Sub};

#[cfg(test)]
mod tests;

pub mod assign;
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]) }
}

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))
            }
        }
    };
}

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