phantom_zone/backend/
mod.rs

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