primitives/types/heap_array/

ops.rs

use std::{
    iter::Sum,
    marker::PhantomData,
    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
};

use ff::Field;
use subtle::{Choice, ConstantTimeEq};

use super::HeapArray;
use crate::{
    errors::PrimitiveError,
    izip_eq,
    random::{CryptoRngCore, Random, RandomNonZero},
    types::{ConditionallySelectable, Element, Positive},
};

// ----------
// | Random |
// ----------

impl<T: Sized + Random, M: Positive> Random for HeapArray<T, M> {
    /// Generate a random array of length `M`.
    fn random(mut rng: impl CryptoRngCore) -> Self {
        Self {
            data: T::random_n::<Box<[_]>>(&mut rng, M::SIZE),
            _len: PhantomData,
        }
    }
}

impl<T: Sized + RandomNonZero, M: Positive> RandomNonZero for HeapArray<T, M> {
    /// Generate a random array of length `M`.
    fn random_non_zero(mut rng: impl CryptoRngCore) -> Result<Self, PrimitiveError> {
        Ok(Self {
            data: T::random_n_non_zero::<Box<[_]>>(&mut rng, M::SIZE)?,
            _len: PhantomData,
        })
    }
}

// --------------
// | Arithmetic |
// --------------

// === Addition === //

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<T: Sized, M: Positive> Add<&HeapArray<T, M>> for HeapArray<T, M>
where
    T: for<'b> Add<&'b T, Output = T>,
{
    type Output = HeapArray<T, M>;

    fn add(self, other: &HeapArray<T, M>) -> Self::Output {
        Self::Output {
            data: izip_eq!(self.data, &other.data)
                .map(|(lhs_value, rhs_value)| lhs_value + rhs_value)
                .collect(),
            _len: PhantomData,
        }
    }
}

#[macros::op_variants(owned)]
impl<T: Sized, M: Positive> Add<&T> for HeapArray<T, M>
where
    T: for<'a> Add<&'a T, Output = T>,
{
    type Output = HeapArray<T, M>;

    fn add(self, other: &T) -> Self::Output {
        Self::Output {
            data: IntoIterator::into_iter(self.data)
                .map(|value| value + other)
                .collect(),
            _len: PhantomData,
        }
    }
}

#[macros::op_variants(owned)]
impl<T: Sized, M: Positive> AddAssign<&HeapArray<T, M>> for HeapArray<T, M>
where
    T: for<'b> AddAssign<&'b T>,
{
    fn add_assign(&mut self, other: &Self) {
        izip_eq!(self, other).for_each(|(lhs_value, rhs_value)| lhs_value.add_assign(rhs_value));
    }
}

#[macros::op_variants(owned)]
impl<T: Sized, M: Positive> AddAssign<&T> for HeapArray<T, M>
where
    T: for<'a> AddAssign<&'a T>,
{
    fn add_assign(&mut self, other: &T) {
        self.iter_mut()
            .for_each(|lhs_value| lhs_value.add_assign(other));
    }
}

// === Subtraction === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<T: Sized, M: Positive> Sub<&HeapArray<T, M>> for HeapArray<T, M>
where
    T: for<'b> Sub<&'b T, Output = T>,
{
    type Output = HeapArray<T, M>;

    fn sub(self, other: &HeapArray<T, M>) -> Self::Output {
        Self::Output {
            data: izip_eq!(self, other)
                .map(|(lhs_value, rhs_value)| lhs_value - rhs_value)
                .collect(),
            _len: PhantomData,
        }
    }
}

#[macros::op_variants(owned, borrowed, flipped)]
impl<T: Sized, M: Positive> Sub<&T> for HeapArray<T, M>
where
    T: for<'a> Sub<&'a T, Output = T>,
{
    type Output = HeapArray<T, M>;

    fn sub(self, other: &T) -> Self::Output {
        Self::Output {
            data: IntoIterator::into_iter(self.data)
                .map(|value| value - other)
                .collect(),
            _len: PhantomData,
        }
    }
}

#[macros::op_variants(owned)]
impl<T: Sized, M: Positive> SubAssign<&HeapArray<T, M>> for HeapArray<T, M>
where
    T: for<'b> SubAssign<&'b T>,
{
    fn sub_assign(&mut self, other: &Self) {
        izip_eq!(self, other).for_each(|(lhs_value, rhs_value)| lhs_value.sub_assign(rhs_value));
    }
}

#[macros::op_variants(owned)]
impl<T: Sized, M: Positive> SubAssign<&T> for HeapArray<T, M>
where
    T: for<'a> SubAssign<&'a T>,
{
    fn sub_assign(&mut self, other: &T) {
        self.iter_mut()
            .for_each(|lhs_value| lhs_value.sub_assign(other));
    }
}

// === Multiplication === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<T1: Sized, T2: Sized, T3: Sized, M: Positive> Mul<&HeapArray<T2, M>> for HeapArray<T1, M>
where
    T1: for<'b> Mul<&'b T2, Output = T3>,
{
    type Output = HeapArray<T3, M>;
    fn mul(self, other: &HeapArray<T2, M>) -> Self::Output {
        Self::Output {
            data: izip_eq!(self, other)
                .map(|(lhs_value, rhs_value)| lhs_value * rhs_value)
                .collect(),
            _len: PhantomData,
        }
    }
}

#[macros::op_variants(borrowed)]
impl<T1: Sized, T2: Sized + Element, T3: Sized, M: Positive> Mul<&T2> for HeapArray<T1, M>
where
    T1: for<'a> Mul<&'a T2, Output = T3>,
{
    type Output = HeapArray<T3, M>;
    fn mul(self, other: &T2) -> Self::Output {
        Self::Output {
            data: IntoIterator::into_iter(self.data)
                .map(|value| value * other)
                .collect(),
            _len: PhantomData,
        }
    }
}

// === MulAssign === //

#[macros::op_variants(owned)]
impl<T1: Sized, T2: Sized, M: Positive> MulAssign<&HeapArray<T2, M>> for HeapArray<T1, M>
where
    T1: for<'b> MulAssign<&'b T2>,
{
    fn mul_assign(&mut self, other: &HeapArray<T2, M>) {
        izip_eq!(self, other).for_each(|(lhs, rhs)| *lhs *= rhs);
    }
}

impl<T: Sized, T2: Sized + Element, M: Positive> MulAssign<&T2> for HeapArray<T, M>
where
    T: for<'a> MulAssign<&'a T2>,
{
    fn mul_assign(&mut self, other: &T2) {
        self.iter_mut().for_each(|lhs| *lhs *= other);
    }
}

// === Negation === //

#[macros::op_variants(borrowed)]
impl<T: Sized, M: Positive> Neg for HeapArray<T, M>
where
    T: Neg<Output = T>,
{
    type Output = HeapArray<T, M>;

    fn neg(self) -> Self::Output {
        HeapArray::<T, M> {
            data: IntoIterator::into_iter(self.data)
                .map(|value| value.neg())
                .collect(),
            _len: PhantomData,
        }
    }
}

// === Field Ops === //

impl<T: Field, M: Positive> HeapArray<T, M> {
    /// Squares each element in the array.
    pub fn square(&self) -> Self {
        Self {
            data: self.iter().map(|value| value.square()).collect(),
            _len: PhantomData,
        }
    }

    /// Doubles each element in the array.
    pub fn double(&self) -> Self {
        Self {
            data: self.iter().map(|value| value.double()).collect(),
            _len: PhantomData,
        }
    }

    /// Inverts each element in the array.
    pub fn invert(&self) -> Option<Self> {
        let inverted_data: Option<Vec<T>> =
            self.iter().map(|value| value.invert().into()).collect();

        inverted_data.map(|data| HeapArray {
            data: data.into(),
            _len: PhantomData,
        })
    }

    /// Exponentiates `self` by `exp`, where `exp` is a little-endian order integer
    /// exponent.
    pub fn pow<S: AsRef<[u64]> + Clone>(&self, exp: S) -> Self {
        Self {
            data: self.iter().map(|value| value.pow(exp.clone())).collect(),
            _len: PhantomData,
        }
    }
}

// === ConditionallySelectable === //

impl<T: Sized + ConditionallySelectable, M: Positive> ConditionallySelectable for HeapArray<T, M> {
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self {
            data: izip_eq!(a, b)
                .map(|(lhs, rhs)| T::conditional_select(lhs, rhs, choice))
                .collect(),
            _len: PhantomData,
        }
    }
}

impl<T: Sized + ConditionallySelectable, M: Positive, L: Positive> HeapArray<HeapArray<T, M>, L> {
    /// Select `a` or `b` according to `choice`.
    ///
    /// # Returns
    ///
    /// * `a` if `choice == Choice(0)`;
    /// * `b` if `choice == Choice(1)`.
    pub fn conditional_select_2d(a: &Self, b: &Self, choice: Choice) -> Self {
        Self {
            data: izip_eq!(a, b)
                .map(|(lhs, rhs)| HeapArray::conditional_select(lhs, rhs, choice))
                .collect(),
            _len: PhantomData,
        }
    }
}

// === ConstantTimeEq === //

impl<T: Sized, M: Positive> ConstantTimeEq for HeapArray<T, M>
where
    T: ConstantTimeEq,
{
    fn ct_eq(&self, other: &Self) -> Choice {
        izip_eq!(self, other).fold(Choice::from(1u8), |acc, (lhs_value, rhs_value)| {
            acc & lhs_value.ct_eq(rhs_value)
        })
    }
}

// === PartialEq === //

impl<T: Sized, M: Positive> PartialEq for HeapArray<T, M>
where
    T: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        izip_eq!(self, other).all(|(lhs_value, rhs_value)| lhs_value == rhs_value)
    }
}

// === Sum === //

impl<T: Default, M: Positive> Sum for HeapArray<T, M>
where
    HeapArray<T, M>: Add<Output = HeapArray<T, M>>,
{
    fn sum<I: Iterator<Item = Self>>(mut iter: I) -> Self {
        let first = iter.next().unwrap_or_default();
        iter.fold(first, |acc, value| acc + value)
    }
}

impl<'a, T: Default + Clone, M: Positive> Sum<&'a HeapArray<T, M>> for HeapArray<T, M>
where
    HeapArray<T, M>: for<'b> Add<&'b HeapArray<T, M>, Output = HeapArray<T, M>>,
{
    fn sum<I: Iterator<Item = &'a Self>>(mut iter: I) -> Self {
        let first = iter.next().cloned().unwrap_or_default();
        iter.fold(first, |acc, value| acc + value)
    }
}

// === Zero and One === //
impl<T: Clone + num_traits::Zero + Eq + for<'b> Add<&'b T, Output = T>, M: Positive + Eq>
    num_traits::Zero for HeapArray<T, M>
{
    fn zero() -> Self {
        Self {
            data: (0..M::to_usize()).map(|_| T::zero()).collect(),
            _len: PhantomData,
        }
    }

    fn is_zero(&self) -> bool {
        let zero = T::zero();
        self.iter().all(|value| value == &zero)
    }
}

impl<
        T: Sized + Default + num_traits::One + Eq + for<'b> Mul<&'b T, Output = T>,
        M: Positive + Eq,
    > num_traits::One for HeapArray<T, M>
{
    fn one() -> Self {
        Self {
            data: (0..M::to_usize()).map(|_| T::one()).collect(),
            _len: PhantomData,
        }
    }

    fn is_one(&self) -> bool {
        let one = T::one();
        self.iter().all(|value| value == &one)
    }
}

#[cfg(test)]
mod tests {
    use typenum::U10;

    use super::*;
    use crate::types::heap_array::HeapArray;

    #[test]
    fn test_addition() {
        let a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = HeapArray::from_fn(|i| i + 1);
        let c = HeapArray::from_fn(|i| i + i + 1);

        assert_eq!(a.clone() + &b, c);
        assert_eq!(&a + b.clone(), c);
        assert_eq!(&a + &b, c);
        assert_eq!(a + b, c);
    }

    #[test]
    fn test_broadcast_addition() {
        let a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = 1;
        let c = HeapArray::from_fn(|i| i + 1);

        assert_eq!(a + b, c);
    }

    #[test]
    fn test_add_assign() {
        let mut a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = HeapArray::from_fn(|i| i + 1);
        let c = HeapArray::from_fn(|i| i + i + 1);

        a += b;
        assert_eq!(a, c);
    }

    #[test]
    fn test_broadcast_add_assign() {
        let mut a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = 1;
        let c = HeapArray::from_fn(|i| i + 1);

        a += b;
        assert_eq!(a, c);
    }

    #[test]
    fn test_subtraction() {
        let a = HeapArray::<usize, U10>::from_fn(|i| 2 * i + 1);
        let b = HeapArray::from_fn(|i| i);
        let c = HeapArray::from_fn(|i| i + 1);

        assert_eq!(a - b, c);
    }

    #[test]
    fn test_broadcast_subtraction() {
        let a = HeapArray::<usize, U10>::from_fn(|i| i + 1);
        let b = 1;
        let c = HeapArray::from_fn(|i| i);

        assert_eq!(a - b, c);
    }

    #[test]
    fn test_sub_assign() {
        let mut a = HeapArray::<usize, U10>::from_fn(|i| i + 1);
        let b = HeapArray::from_fn(|i| i);
        let c = HeapArray::from_fn(|_| 1);

        a -= b;
        assert_eq!(a, c);
    }

    #[test]
    fn test_broadcast_sub_assign() {
        let mut a = HeapArray::<usize, U10>::from_fn(|i| i + 1);
        let b = 1;
        let c = HeapArray::from_fn(|i| i);

        a -= b;
        assert_eq!(a, c);
    }

    #[test]
    fn test_multiplication() {
        let a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = HeapArray::from_fn(|i| i + 1);
        let c = HeapArray::from_fn(|i| i * (i + 1));

        assert_eq!(a * b, c);
    }

    #[test]
    fn test_broadcast_multiplication() {
        let a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = 1;
        let c = HeapArray::from_fn(|i| i);

        assert_eq!(a * &b, c);
    }

    #[test]
    fn test_mul_assign() {
        let mut a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = HeapArray::from_fn(|i| i + 1);
        let c = HeapArray::from_fn(|i| i * (i + 1));

        a *= b;
        assert_eq!(a, c);
    }

    #[test]
    fn test_broadcast_mul_assign() {
        let mut a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = 1;
        let c = HeapArray::from_fn(|i| i);

        a *= &b;
        assert_eq!(a, c);
    }

    #[test]
    fn test_negation() {
        let a = HeapArray::<i64, U10>::from_fn(|i| i as i64);
        let b = HeapArray::from_fn(|i| -(i as i64));

        assert_eq!(-a, b);
    }

    #[test]
    fn test_conditional_select() {
        let a = HeapArray::<u32, U10>::from_fn(|i| i as u32);
        let b = HeapArray::<u32, U10>::from_fn(|i| i as u32 + 1);
        let choice = Choice::from(1u8);

        let selected = HeapArray::conditional_select(&a, &b, choice);
        let non_selected = HeapArray::conditional_select(&a, &b, !choice);
        assert_eq!(selected, b);
        assert_eq!(non_selected, a);
    }

    #[test]
    fn test_constant_time_eq() {
        let a = HeapArray::<usize, U10>::from_fn(|i| i);
        let b = HeapArray::from_fn(|i| i);
        let c = HeapArray::from_fn(|i| i + 1);

        assert!(a.ct_eq(&b).unwrap_u8() == 1);
        assert!(a.ct_eq(&c).unwrap_u8() == 0);
    }
}