1use num_traits::ToPrimitive;
  2
  3use crate::Row;
  4
  5mod modulus_u64;
  6mod power_of_2;
  7mod word_size;
  8
  9pub use modulus_u64::ModularOpsU64;
 10pub(crate) use power_of_2::ModulusPowerOf2;
 11
 12pub trait Modulus {
 13    type Element;
 14    /// Modulus value if it fits in Element
 15    fn q(&self) -> Option<Self::Element>;
 16    /// Modulus value as f64 if it fits in f64
 17    fn q_as_f64(&self) -> Option<f64>;
 18    /// Is modulus native?
 19    fn is_native(&self) -> bool;
 20    /// -1 in signed representaiton
 21    fn neg_one(&self) -> Self::Element;
 22    /// Largest unsigned value that fits in the modulus. That is, q - 1.
 23    fn largest_unsigned_value(&self) -> Self::Element;
 24    /// Smallest unsigned value that fits in the modulus
 25    /// Always assmed to be 0.
 26    fn smallest_unsigned_value(&self) -> Self::Element;
 27    /// Convert unsigned value in signed represetation to i64
 28    fn map_element_to_i64(&self, v: &Self::Element) -> i64;
 29    /// Convert f64 to signed represented in modulus
 30    fn map_element_from_f64(&self, v: f64) -> Self::Element;
 31    /// Convert i64 to signed represented in modulus
 32    fn map_element_from_i64(&self, v: i64) -> Self::Element;
 33}
 34
 35impl Modulus for u64 {
 36    type Element = u64;
 37    fn is_native(&self) -> bool {
 38        // q of size u64 can never be a naitve modulus
 39        false
 40    }
 41    fn largest_unsigned_value(&self) -> Self::Element {
 42        self - 1
 43    }
 44    fn neg_one(&self) -> Self::Element {
 45        self - 1
 46    }
 47    fn smallest_unsigned_value(&self) -> Self::Element {
 48        0
 49    }
 50    fn map_element_to_i64(&self, v: &Self::Element) -> i64 {
 51        assert!(v <= self, "{v} must be <= {self}");
 52        if *v >= (self >> 1) {
 53            -ToPrimitive::to_i64(&(self - v)).unwrap()
 54        } else {
 55            ToPrimitive::to_i64(v).unwrap()
 56        }
 57    }
 58    fn map_element_from_f64(&self, v: f64) -> Self::Element {
 59        //FIXME (Jay): Before I check whether v is smaller than 0 with `let is_neg =
 60        // o.is_sign_negative() && o != 0.0; I'm ocnfused why didn't I simply check <
 61        // 0.0?
 62        let v = v.round();
 63        if v < 0.0 {
 64            self - v.abs().to_u64().unwrap()
 65        } else {
 66            v.to_u64().unwrap()
 67        }
 68    }
 69    fn map_element_from_i64(&self, v: i64) -> Self::Element {
 70        if v < 0 {
 71            self - v.abs().to_u64().unwrap()
 72        } else {
 73            v.to_u64().unwrap()
 74        }
 75    }
 76    fn q(&self) -> Option<Self::Element> {
 77        Some(*self)
 78    }
 79    fn q_as_f64(&self) -> Option<f64> {
 80        self.to_f64()
 81    }
 82}
 83
 84pub trait ModInit {
 85    type M;
 86    fn new(modulus: Self::M) -> Self;
 87}
 88
 89pub trait GetModulus {
 90    type Element;
 91    type M: Modulus<Element = Self::Element>;
 92    fn modulus(&self) -> &Self::M;
 93}
 94
 95pub trait VectorOps {
 96    type Element;
 97
 98    fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element);
 99    fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]);
100
101    fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
102    fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
103    fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
104    fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element);
105    fn elwise_neg_mut(&self, a: &mut [Self::Element]);
106    /// inplace mutates `a`: a = a + b*c
107    fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]);
108    fn elwise_fma_scalar_mut(
109        &self,
110        a: &mut [Self::Element],
111        b: &[Self::Element],
112        c: &Self::Element,
113    );
114}
115
116pub trait ArithmeticOps {
117    type Element;
118    fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
119    fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
120    fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
121    fn neg(&self, a: &Self::Element) -> Self::Element;
122}
123
124pub trait ArithmeticLazyOps {
125    type Element;
126    fn mul_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
127    fn add_lazy(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
128}
129
130pub trait ShoupMatrixFMA<R: Row> {
131    /// Returns summation of `row-wise product of matrix a and b` + out where
132    /// each element is in range [0, 2q)
133    fn shoup_matrix_fma(&self, out: &mut [R::Element], a: &[R], a_shoup: &[R], b: &[R]);
134}