phantom-zone/src/backend.rs
src/backend.rsRaw
1use itertools::izip;
2
3pub trait ModInit {
4 type Element;
5 fn new(q: Self::Element) -> Self;
6}
7
8pub trait VectorOps {
9 type Element;
10
11 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element);
12 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]);
13
14 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
15 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
16 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]);
17 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element);
18 fn elwise_neg_mut(&self, a: &mut [Self::Element]);
19 /// inplace mutates `a`: a = a + b*c
20 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]);
21 fn elwise_fma_scalar_mut(
22 &self,
23 a: &mut [Self::Element],
24 b: &[Self::Element],
25 c: &Self::Element,
26 );
27
28 fn modulus(&self) -> Self::Element;
29}
30
31pub trait ArithmeticOps {
32 type Element;
33
34 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
35 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
36 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element;
37 fn neg(&self, a: &Self::Element) -> Self::Element;
38
39 fn modulus(&self) -> Self::Element;
40}
41
42pub struct ModularOpsU64 {
43 q: u64,
44 logq: usize,
45 barrett_mu: u128,
46 barrett_alpha: usize,
47}
48
49impl ModInit for ModularOpsU64 {
50 type Element = u64;
51 fn new(q: u64) -> ModularOpsU64 {
52 let logq = 64 - q.leading_zeros();
53
54 // barrett calculation
55 let mu = (1u128 << (logq * 2 + 3)) / (q as u128);
56 let alpha = logq + 3;
57
58 ModularOpsU64 {
59 q,
60 logq: logq as usize,
61 barrett_alpha: alpha as usize,
62 barrett_mu: mu,
63 }
64 }
65}
66
67impl ModularOpsU64 {
68 fn add_mod_fast(&self, a: u64, b: u64) -> u64 {
69 debug_assert!(a < self.q);
70 debug_assert!(b < self.q);
71
72 let mut o = a + b;
73 if o >= self.q {
74 o -= self.q;
75 }
76 o
77 }
78
79 fn sub_mod_fast(&self, a: u64, b: u64) -> u64 {
80 debug_assert!(a < self.q);
81 debug_assert!(b < self.q);
82
83 if a > b {
84 a - b
85 } else {
86 (self.q + a) - b
87 }
88 }
89
90 /// returns (a * b) % q
91 ///
92 /// - both a and b must be in range [0, 2q)
93 /// - output is in range [0 , q)
94 fn mul_mod_fast(&self, a: u64, b: u64) -> u64 {
95 debug_assert!(a < 2 * self.q);
96 debug_assert!(b < 2 * self.q);
97
98 let ab = a as u128 * b as u128;
99
100 // ab / (2^{n + \beta})
101 // note: \beta is assumed to -2
102 let tmp = ab >> (self.logq - 2);
103
104 // k = ((ab / (2^{n + \beta})) * \mu) / 2^{\alpha - (-2)}
105 let k = (tmp * self.barrett_mu) >> (self.barrett_alpha + 2);
106
107 // ab - k*p
108 let tmp = k * (self.q as u128);
109
110 let mut out = (ab - tmp) as u64;
111
112 if out >= self.q {
113 out -= self.q;
114 }
115
116 return out;
117 }
118}
119
120impl ArithmeticOps for ModularOpsU64 {
121 type Element = u64;
122
123 fn add(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
124 self.add_mod_fast(*a, *b)
125 }
126
127 fn mul(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
128 self.mul_mod_fast(*a, *b)
129 }
130
131 fn sub(&self, a: &Self::Element, b: &Self::Element) -> Self::Element {
132 self.sub_mod_fast(*a, *b)
133 }
134
135 fn neg(&self, a: &Self::Element) -> Self::Element {
136 self.q - *a
137 }
138
139 fn modulus(&self) -> Self::Element {
140 self.q
141 }
142}
143
144impl VectorOps for ModularOpsU64 {
145 type Element = u64;
146
147 fn elwise_add_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
148 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
149 *ai = self.add_mod_fast(*ai, *bi);
150 });
151 }
152
153 fn elwise_sub_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
154 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
155 *ai = self.sub_mod_fast(*ai, *bi);
156 });
157 }
158
159 fn elwise_mul_mut(&self, a: &mut [Self::Element], b: &[Self::Element]) {
160 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
161 *ai = self.mul_mod_fast(*ai, *bi);
162 });
163 }
164
165 fn elwise_neg_mut(&self, a: &mut [Self::Element]) {
166 a.iter_mut().for_each(|ai| *ai = self.q - *ai);
167 }
168
169 fn elwise_scalar_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &Self::Element) {
170 izip!(out.iter_mut(), a.iter()).for_each(|(oi, ai)| {
171 *oi = self.mul_mod_fast(*ai, *b);
172 });
173 }
174
175 fn elwise_mul(&self, out: &mut [Self::Element], a: &[Self::Element], b: &[Self::Element]) {
176 izip!(out.iter_mut(), a.iter(), b.iter()).for_each(|(oi, ai, bi)| {
177 *oi = self.mul_mod_fast(*ai, *bi);
178 });
179 }
180
181 fn elwise_scalar_mul_mut(&self, a: &mut [Self::Element], b: &Self::Element) {
182 a.iter_mut().for_each(|ai| {
183 *ai = self.mul_mod_fast(*ai, *b);
184 });
185 }
186
187 fn elwise_fma_mut(&self, a: &mut [Self::Element], b: &[Self::Element], c: &[Self::Element]) {
188 izip!(a.iter_mut(), b.iter(), c.iter()).for_each(|(ai, bi, ci)| {
189 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *ci));
190 });
191 }
192
193 fn elwise_fma_scalar_mut(
194 &self,
195 a: &mut [Self::Element],
196 b: &[Self::Element],
197 c: &Self::Element,
198 ) {
199 izip!(a.iter_mut(), b.iter()).for_each(|(ai, bi)| {
200 *ai = self.add_mod_fast(*ai, self.mul_mod_fast(*bi, *c));
201 });
202 }
203
204 fn modulus(&self) -> Self::Element {
205 self.q
206 }
207}