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