1use itertools::izip;
  2
  3use crate::{ArithmeticOps, ModInit, VectorOps};
  4
  5use super::{GetModulus, Modulus};
  6
  7pub(crate) struct ModulusPowerOf2<T> {
  8    modulus: T,
  9    /// Modulus mask: (1 << q) - 1
 10    mask: u64,
 11}
 12
 13impl<T> ArithmeticOps for ModulusPowerOf2<T> {
 14    type Element = u64;
 15    #[inline]
 16    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
 17        (a.wrapping_add(*b)) & self.mask
 18    }
 19    #[inline]
 20    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
 21        (a.wrapping_sub(*b)) & self.mask
 22    }
 23    #[inline]
 24    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
 25        (a.wrapping_mul(*b)) & self.mask
 26    }
 27    #[inline]
 28    fn neg(&self, a: &Self::Element) -> Self::Element {
 29        (0u64.wrapping_sub(*a)) & self.mask
 30    }
 31}
 32
 33impl<T> VectorOps for ModulusPowerOf2<T> {
 34    type Element = u64;
 35
 36    #[inline]
 37    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
 38        izip!(a.iter_mut(), b.iter()).for_each(|(a0, b0)| *a0 = (a0.wrapping_add(*b0)) & self.mask);
 39    }
 40
 41    #[inline]
 42    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
 43        izip!(a.iter_mut(), b.iter()).for_each(|(a0, b0)| *a0 = (a0.wrapping_mul(*b0)) & self.mask);
 44    }
 45
 46    #[inline]
 47    fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
 48        a.iter_mut()
 49            .for_each(|a0| *a0 = 0u64.wrapping_sub(*a0) & self.mask);
 50    }
 51    #[inline]
 52    fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
 53        izip!(a.iter_mut(), b.iter()).for_each(|(a0, b0)| *a0 = (a0.wrapping_sub(*b0)) & self.mask);
 54    }
 55
 56    #[inline]
 57    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
 58        izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(a0, b0, c0)| {
 59            *a0 = a0.wrapping_add(b0.wrapping_mul(*c0)) & self.mask;
 60        });
 61    }
 62
 63    #[inline]
 64    fn elwise_fma_scalar_mut(
 65        &self,
 66        a: &mut [Self::Element],
 67        b: &[Self::Element],
 68        c: &Self::Element,
 69    ) {
 70        izip!(a.iter_mut(), b.iter()).for_each(|(a0, b0)| {
 71            *a0 = a0.wrapping_add(b0.wrapping_mul(*c)) & self.mask;
 72        });
 73    }
 74    #[inline]
 75    fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
 76        a.iter_mut()
 77            .for_each(|a0| *a0 = a0.wrapping_mul(*b) & self.mask)
 78    }
 79
 80    #[inline]
 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(|(o0, a0, b0)| {
 83            *o0 = a0.wrapping_mul(*b0) & self.mask;
 84        });
 85    }
 86
 87    #[inline]
 88    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
 89        izip!(out.iter_mut(), a.iter()).for_each(|(o0, a0)| {
 90            *o0 = a0.wrapping_mul(*b) & self.mask;
 91        });
 92    }
 93}
 94
 95impl<T: Modulus<Element = u64>> ModInit for ModulusPowerOf2<T> {
 96    type M = T;
 97    fn new(modulus: Self::M) -> Self {
 98        assert!(!modulus.is_native());
 99        assert!(modulus.q().unwrap().is_power_of_two());
100        let q = modulus.q().unwrap();
101        let mask = q - 1;
102        Self { modulus, mask }
103    }
104}
105
106impl<T: Modulus<Element = u64>> GetModulus for ModulusPowerOf2<T> {
107    type Element = u64;
108    type M = T;
109    fn modulus(&self) -> &Self::M {
110        &self.modulus
111    }
112}