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