impl_ops.rs

// Copyright 2014-2016 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use num_complex::Complex;

/// Elements that can be used as direct operands in arithmetic with arrays.
///
/// For example, `f64` is a `ScalarOperand` which means that for an array `a`,
/// arithmetic like `a + 1.0`, and, `a * 2.`, and `a += 3.` are allowed.
///
/// In the description below, let `A` be an array or array view,
/// let `B` be an array with owned data,
/// and let `C` be an array with mutable data.
///
/// `ScalarOperand` determines for which scalars `K` operations `&A @ K`, and `B @ K`,
/// and `C @= K` are defined, as ***right hand side operands***, for applicable
/// arithmetic operators (denoted `@`).
///
/// ***Left hand side*** scalar operands are not related to this trait
/// (they need one `impl` per concrete scalar type); but they are still
/// implemented for the same types, allowing operations
/// `K @ &A`, and `K @ B` for primitive numeric types `K`.
///
/// This trait ***does not*** limit which elements can be stored in an array in general.
/// Non-`ScalarOperand` types can still participate in arithmetic as array elements in
/// in array-array operations.
pub trait ScalarOperand: 'static + Clone {}
impl ScalarOperand for bool {}
impl ScalarOperand for i8 {}
impl ScalarOperand for u8 {}
impl ScalarOperand for i16 {}
impl ScalarOperand for u16 {}
impl ScalarOperand for i32 {}
impl ScalarOperand for u32 {}
impl ScalarOperand for i64 {}
impl ScalarOperand for u64 {}
impl ScalarOperand for i128 {}
impl ScalarOperand for u128 {}
impl ScalarOperand for isize {}
impl ScalarOperand for usize {}
impl ScalarOperand for f32 {}
impl ScalarOperand for f64 {}
impl ScalarOperand for Complex<f32> {}
impl ScalarOperand for Complex<f64> {}

macro_rules! impl_binary_op(
    ($trt:ident, $operator:tt, $mth:ident, $iop:tt, $doc:expr) => (
/// Perform elementwise
#[doc=$doc]
/// between `self` and `rhs`,
/// and return the result (based on `self`).
///
/// `self` must be an `Array` or `ArcArray`.
///
/// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
///
/// **Panics** if broadcasting isn’t possible.
impl<A, B, S, S2, D, E> $trt<ArrayBase<S2, E>> for ArrayBase<S, D>
where
    A: Clone + $trt<B, Output=A>,
    B: Clone,
    S: DataOwned<Elem=A> + DataMut,
    S2: Data<Elem=B>,
    D: Dimension,
    E: Dimension,
{
    type Output = ArrayBase<S, D>;
    fn $mth(self, rhs: ArrayBase<S2, E>) -> ArrayBase<S, D>
    {
        self.$mth(&rhs)
    }
}

/// Perform elementwise
#[doc=$doc]
/// between `self` and reference `rhs`,
/// and return the result (based on `self`).
///
/// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
///
/// **Panics** if broadcasting isn’t possible.
impl<'a, A, B, S, S2, D, E> $trt<&'a ArrayBase<S2, E>> for ArrayBase<S, D>
where
    A: Clone + $trt<B, Output=A>,
    B: Clone,
    S: DataOwned<Elem=A> + DataMut,
    S2: Data<Elem=B>,
    D: Dimension,
    E: Dimension,
{
    type Output = ArrayBase<S, D>;
    fn $mth(mut self, rhs: &ArrayBase<S2, E>) -> ArrayBase<S, D>
    {
        self.zip_mut_with(rhs, |x, y| {
            *x = x.clone() $operator y.clone();
        });
        self
    }
}

/// Perform elementwise
#[doc=$doc]
/// between references `self` and `rhs`,
/// and return the result as a new `Array`.
///
/// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
///
/// **Panics** if broadcasting isn’t possible.
impl<'a, A, B, S, S2, D, E> $trt<&'a ArrayBase<S2, E>> for &'a ArrayBase<S, D>
where
    A: Clone + $trt<B, Output=A>,
    B: Clone,
    S: Data<Elem=A>,
    S2: Data<Elem=B>,
    D: Dimension,
    E: Dimension,
{
    type Output = Array<A, D>;
    fn $mth(self, rhs: &'a ArrayBase<S2, E>) -> Array<A, D> {
        // FIXME: Can we co-broadcast arrays here? And how?
        self.to_owned().$mth(rhs)
    }
}

/// Perform elementwise
#[doc=$doc]
/// between `self` and the scalar `x`,
/// and return the result (based on `self`).
///
/// `self` must be an `Array` or `ArcArray`.
impl<A, S, D, B> $trt<B> for ArrayBase<S, D>
    where A: Clone + $trt<B, Output=A>,
          S: DataOwned<Elem=A> + DataMut,
          D: Dimension,
          B: ScalarOperand,
{
    type Output = ArrayBase<S, D>;
    fn $mth(mut self, x: B) -> ArrayBase<S, D> {
        self.unordered_foreach_mut(move |elt| {
            *elt = elt.clone() $operator x.clone();
        });
        self
    }
}

/// Perform elementwise
#[doc=$doc]
/// between the reference `self` and the scalar `x`,
/// and return the result as a new `Array`.
impl<'a, A, S, D, B, C> $trt<B> for &'a ArrayBase<S, D>
    where A: Clone + $trt<B, Output=C>,
          S: Data<Elem=A>,
          D: Dimension,
          B: ScalarOperand,
{
    type Output = Array<C, D>;
    fn $mth(self, x: B) -> Self::Output {
        self.map(move |elt| elt.clone() $operator x.clone())
    }
}
    );
);

macro_rules! impl_scalar_lhs_op {
    ($scalar:ty, $operator:tt, $trt:ident, $mth:ident, $doc:expr) => (
/// Perform elementwise
#[doc=$doc]
/// between the scalar `self` and array `rhs`,
/// and return the result (based on `self`).
impl<A, S, D> $trt<ArrayBase<S, D>> for $scalar
where
    $scalar: Clone + $trt<A, Output=A>,
    A: Clone,
    S: DataOwned<Elem=A> + DataMut,
    D: Dimension,
{
    type Output = ArrayBase<S, D>;
    fn $mth(self, mut rhs: ArrayBase<S, D>) -> ArrayBase<S, D> {
        rhs.unordered_foreach_mut(move |elt| {
            *elt = self.clone() $operator elt.clone();
        });
        rhs
    }
}

/// Perform elementwise
#[doc=$doc]
/// between the scalar `self` and array `rhs`,
/// and return the result as a new `Array`.
impl<'a, A, S, D, B> $trt<&'a ArrayBase<S, D>> for $scalar
where
      $scalar: Clone + $trt<A, Output=B>,
      A: Clone,
      S: Data<Elem=A>,
      D: Dimension,
{
    type Output = Array<B, D>;
    fn $mth(self, rhs: &ArrayBase<S, D>) -> Array<B, D> {
        rhs.map(move |elt| self.clone() $operator elt.clone())
    }
}
    );
}

mod arithmetic_ops {
    use super::*;
    use crate::imp_prelude::*;

    use num_complex::Complex;
    use std::ops::*;

    impl_binary_op!(Add, +, add, +=, "addition");
    impl_binary_op!(Sub, -, sub, -=, "subtraction");
    impl_binary_op!(Mul, *, mul, *=, "multiplication");
    impl_binary_op!(Div, /, div, /=, "division");
    impl_binary_op!(Rem, %, rem, %=, "remainder");
    impl_binary_op!(BitAnd, &, bitand, &=, "bit and");
    impl_binary_op!(BitOr, |, bitor, |=, "bit or");
    impl_binary_op!(BitXor, ^, bitxor, ^=, "bit xor");
    impl_binary_op!(Shl, <<, shl, <<=, "left shift");
    impl_binary_op!(Shr, >>, shr, >>=, "right shift");

    macro_rules! all_scalar_ops {
        ($int_scalar:ty) => (
            impl_scalar_lhs_op!($int_scalar, +, Add, add, "addition");
            impl_scalar_lhs_op!($int_scalar, -, Sub, sub, "subtraction");
            impl_scalar_lhs_op!($int_scalar, *, Mul, mul, "multiplication");
            impl_scalar_lhs_op!($int_scalar, /, Div, div, "division");
            impl_scalar_lhs_op!($int_scalar, %, Rem, rem, "remainder");
            impl_scalar_lhs_op!($int_scalar, &, BitAnd, bitand, "bit and");
            impl_scalar_lhs_op!($int_scalar, |, BitOr, bitor, "bit or");
            impl_scalar_lhs_op!($int_scalar, ^, BitXor, bitxor, "bit xor");
            impl_scalar_lhs_op!($int_scalar, <<, Shl, shl, "left shift");
            impl_scalar_lhs_op!($int_scalar, >>, Shr, shr, "right shift");
        );
    }
    all_scalar_ops!(i8);
    all_scalar_ops!(u8);
    all_scalar_ops!(i16);
    all_scalar_ops!(u16);
    all_scalar_ops!(i32);
    all_scalar_ops!(u32);
    all_scalar_ops!(i64);
    all_scalar_ops!(u64);
    all_scalar_ops!(i128);
    all_scalar_ops!(u128);

    impl_scalar_lhs_op!(bool, &, BitAnd, bitand, "bit and");
    impl_scalar_lhs_op!(bool, |, BitOr, bitor, "bit or");
    impl_scalar_lhs_op!(bool, ^, BitXor, bitxor, "bit xor");

    impl_scalar_lhs_op!(f32, +, Add, add, "addition");
    impl_scalar_lhs_op!(f32, -, Sub, sub, "subtraction");
    impl_scalar_lhs_op!(f32, *, Mul, mul, "multiplication");
    impl_scalar_lhs_op!(f32, /, Div, div, "division");
    impl_scalar_lhs_op!(f32, %, Rem, rem, "remainder");

    impl_scalar_lhs_op!(f64, +, Add, add, "addition");
    impl_scalar_lhs_op!(f64, -, Sub, sub, "subtraction");
    impl_scalar_lhs_op!(f64, *, Mul, mul, "multiplication");
    impl_scalar_lhs_op!(f64, /, Div, div, "division");
    impl_scalar_lhs_op!(f64, %, Rem, rem, "remainder");

    impl_scalar_lhs_op!(Complex<f32>, +, Add, add, "addition");
    impl_scalar_lhs_op!(Complex<f32>, -, Sub, sub, "subtraction");
    impl_scalar_lhs_op!(Complex<f32>, *, Mul, mul, "multiplication");
    impl_scalar_lhs_op!(Complex<f32>, /, Div, div, "division");

    impl_scalar_lhs_op!(Complex<f64>, +, Add, add, "addition");
    impl_scalar_lhs_op!(Complex<f64>, -, Sub, sub, "subtraction");
    impl_scalar_lhs_op!(Complex<f64>, *, Mul, mul, "multiplication");
    impl_scalar_lhs_op!(Complex<f64>, /, Div, div, "division");

    impl<A, S, D> Neg for ArrayBase<S, D>
    where
        A: Clone + Neg<Output = A>,
        S: DataOwned<Elem = A> + DataMut,
        D: Dimension,
    {
        type Output = Self;
        /// Perform an elementwise negation of `self` and return the result.
        fn neg(mut self) -> Self {
            self.unordered_foreach_mut(|elt| {
                *elt = -elt.clone();
            });
            self
        }
    }

    impl<'a, A, S, D> Neg for &'a ArrayBase<S, D>
    where
        &'a A: 'a + Neg<Output = A>,
        S: Data<Elem = A>,
        D: Dimension,
    {
        type Output = Array<A, D>;
        /// Perform an elementwise negation of reference `self` and return the
        /// result as a new `Array`.
        fn neg(self) -> Array<A, D> {
            self.map(Neg::neg)
        }
    }

    impl<A, S, D> Not for ArrayBase<S, D>
    where
        A: Clone + Not<Output = A>,
        S: DataOwned<Elem = A> + DataMut,
        D: Dimension,
    {
        type Output = Self;
        /// Perform an elementwise unary not of `self` and return the result.
        fn not(mut self) -> Self {
            self.unordered_foreach_mut(|elt| {
                *elt = !elt.clone();
            });
            self
        }
    }

    impl<'a, A, S, D> Not for &'a ArrayBase<S, D>
    where
        &'a A: 'a + Not<Output = A>,
        S: Data<Elem = A>,
        D: Dimension,
    {
        type Output = Array<A, D>;
        /// Perform an elementwise unary not of reference `self` and return the
        /// result as a new `Array`.
        fn not(self) -> Array<A, D> {
            self.map(Not::not)
        }
    }
}

mod assign_ops {
    use super::*;
    use crate::imp_prelude::*;

    macro_rules! impl_assign_op {
        ($trt:ident, $method:ident, $doc:expr) => {
            use std::ops::$trt;

            #[doc=$doc]
            /// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
            ///
            /// **Panics** if broadcasting isn’t possible.
            impl<'a, A, S, S2, D, E> $trt<&'a ArrayBase<S2, E>> for ArrayBase<S, D>
            where
                A: Clone + $trt<A>,
                S: DataMut<Elem = A>,
                S2: Data<Elem = A>,
                D: Dimension,
                E: Dimension,
            {
                fn $method(&mut self, rhs: &ArrayBase<S2, E>) {
                    self.zip_mut_with(rhs, |x, y| {
                        x.$method(y.clone());
                    });
                }
            }

            #[doc=$doc]
            impl<A, S, D> $trt<A> for ArrayBase<S, D>
            where
                A: ScalarOperand + $trt<A>,
                S: DataMut<Elem = A>,
                D: Dimension,
            {
                fn $method(&mut self, rhs: A) {
                    self.unordered_foreach_mut(move |elt| {
                        elt.$method(rhs.clone());
                    });
                }
            }
        };
    }

    impl_assign_op!(
        AddAssign,
        add_assign,
        "Perform `self += rhs` as elementwise addition (in place).\n"
    );
    impl_assign_op!(
        SubAssign,
        sub_assign,
        "Perform `self -= rhs` as elementwise subtraction (in place).\n"
    );
    impl_assign_op!(
        MulAssign,
        mul_assign,
        "Perform `self *= rhs` as elementwise multiplication (in place).\n"
    );
    impl_assign_op!(
        DivAssign,
        div_assign,
        "Perform `self /= rhs` as elementwise division (in place).\n"
    );
    impl_assign_op!(
        RemAssign,
        rem_assign,
        "Perform `self %= rhs` as elementwise remainder (in place).\n"
    );
    impl_assign_op!(
        BitAndAssign,
        bitand_assign,
        "Perform `self &= rhs` as elementwise bit and (in place).\n"
    );
    impl_assign_op!(
        BitOrAssign,
        bitor_assign,
        "Perform `self |= rhs` as elementwise bit or (in place).\n"
    );
    impl_assign_op!(
        BitXorAssign,
        bitxor_assign,
        "Perform `self ^= rhs` as elementwise bit xor (in place).\n"
    );
    impl_assign_op!(
        ShlAssign,
        shl_assign,
        "Perform `self <<= rhs` as elementwise left shift (in place).\n"
    );
    impl_assign_op!(
        ShrAssign,
        shr_assign,
        "Perform `self >>= rhs` as elementwise right shift (in place).\n"
    );
}