1use itertools::izip;
  2use num_traits::{WrappingAdd, WrappingMul, WrappingSub, Zero};
  3
  4use super::{ArithmeticOps, GetModulus, ModInit, Modulus, VectorOps};
  5
  6pub struct WordSizeModulus<T> {
  7    modulus: T,
  8}
  9
 10impl<T> ModInit for WordSizeModulus<T>
 11where
 12    T: Modulus,
 13{
 14    type M = T;
 15    fn new(modulus: T) -> Self {
 16        assert!(modulus.is_native());
 17        // For now assume ModulusOpsU64 is only used for u64
 18        Self { modulus: modulus }
 19    }
 20}
 21
 22impl<T> ArithmeticOps for WordSizeModulus<T>
 23where
 24    T: Modulus,
 25    T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
 26{
 27    type Element = T::Element;
 28    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
 29        T::Element::wrapping_add(a, b)
 30    }
 31
 32    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
 33        T::Element::wrapping_mul(a, b)
 34    }
 35
 36    fn neg(&self, a: &Self::Element) -> Self::Element {
 37        T::Element::wrapping_sub(&T::Element::zero(), a)
 38    }
 39
 40    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
 41        T::Element::wrapping_sub(a, b)
 42    }
 43}
 44
 45impl<T> VectorOps for WordSizeModulus<T>
 46where
 47    T: Modulus,
 48    T::Element: WrappingAdd + WrappingSub + WrappingMul + Zero,
 49{
 50    type Element = T::Element;
 51
 52    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
 53        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
 54            *ai = T::Element::wrapping_add(ai, bi);
 55        });
 56    }
 57
 58    fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
 59        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
 60            *ai = T::Element::wrapping_sub(ai, bi);
 61        });
 62    }
 63
 64    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
 65        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
 66            *ai = T::Element::wrapping_mul(ai, bi);
 67        });
 68    }
 69
 70    fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
 71        a.iter_mut()
 72            .for_each(|ai| *ai = T::Element::wrapping_sub(&T::Element::zero(), ai));
 73    }
 74
 75    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
 76        izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
 77            *oi = T::Element::wrapping_mul(ai, b);
 78        });
 79    }
 80
 81    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
 82        izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
 83            *oi = T::Element::wrapping_mul(ai, bi);
 84        });
 85    }
 86
 87    fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
 88        a.iter_mut().for_each(|ai| {
 89            *ai = T::Element::wrapping_mul(ai, b);
 90        });
 91    }
 92
 93    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
 94        izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
 95            *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, ci));
 96        });
 97    }
 98
 99    fn elwise_fma_scalar_mut(
100        &self,
101        a: &mut [Self::Element],
102        b: &[Self::Element],
103        c: &Self::Element,
104    ) {
105        izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
106            *ai = T::Element::wrapping_add(ai, &T::Element::wrapping_mul(bi, c));
107        });
108    }
109
110    // fn modulus(&self) -> &T {
111    //     &self.modulus
112    // }
113}
114
115impl<T> GetModulus for WordSizeModulus<T>
116where
117    T: Modulus,
118{
119    type Element = T::Element;
120    type M = T;
121    fn modulus(&self) -> &Self::M {
122        &self.modulus
123    }
124}