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