phantom-zone/src/rgsw/runtime.rs
src/rgsw/runtime.rsRaw
1use itertools::izip;
2use num_traits::Zero;
3
4use crate::{
5 backend::{ArithmeticOps, GetModulus, ShoupMatrixFMA, VectorOps},
6 decomposer::{Decomposer, RlweDecomposer},
7 ntt::Ntt,
8 Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut,
9};
10
11/// Degree 1 RLWE ciphertext.
12///
13/// RLWE(m) = [a, b] s.t. m+e = b - as
14pub(crate) trait RlweCiphertext {
15 type R: RowMut;
16 /// Returns polynomial `a` of RLWE ciphertext as slice of elements
17 fn part_a(&self) -> &[<Self::R as Row>::Element];
18 /// Returns polynomial `a` of RLWE ciphertext as mutable slice of elements
19 fn part_a_mut(&mut self) -> &mut [<Self::R as Row>::Element];
20 /// Returns polynomial `b` of RLWE ciphertext as slice of elements
21 fn part_b(&self) -> &[<Self::R as Row>::Element];
22 /// Returns polynomial `b` of RLWE ciphertext as mut slice of elements
23 fn part_b_mut(&mut self) -> &mut [<Self::R as Row>::Element];
24 /// Returns ring size of polynomials
25 fn ring_size(&self) -> usize;
26}
27
28/// RGSW ciphertext
29///
30/// RGSW is a collection of RLWE' ciphertext which are collection degree 1 of
31/// RLWE ciphertexts
32///
33/// RGSW = [RLWE'(-sm) || RLWE'(m)]
34///
35/// As usual we refer to decomposition count for RLWE_A in RLWE x RGSW
36/// multiplicaiton as `d_a` and decomposition count for RLWE_B in RLWE x RGDW
37/// multiplication as `d_b`.
38pub(crate) trait RgswCiphertext {
39 type R: Row;
40
41 /// Splits RGSW ciphertext and returns references:
42 /// (RLWE'_A(-sm), RLWE'_B(-sm)), (RLWE'_A(m), RLWE'_B(m))
43 fn split(&self) -> ((&[Self::R], &[Self::R]), (&[Self::R], &[Self::R]));
44}
45
46pub(crate) trait RgswCiphertextMut: RgswCiphertext {
47 /// Splits RGSW ciphertext and returns mutable references:
48 /// (RLWE'_A(-sm), RLWE'_B(-sm)), (RLWE'_A(m), RLWE'_B(m))
49 fn split_mut(
50 &mut self,
51 ) -> (
52 (&mut [Self::R], &mut [Self::R]),
53 (&mut [Self::R], &mut [Self::R]),
54 );
55}
56
57/// RLWE Key switching Key
58///
59/// Key switching key from s' -> s consists of multiple RLWE cipheretxts.
60/// For gadget vector: [1, beta, ..., beta^{d-1}]
61/// RLWE'_{s}(-s'm) = [RWLE_{s}(-s'm), ..., RLWE_{s}(beta^{d-1} -s'm)]
62pub(crate) trait RlweKsk {
63 type R: Row;
64 /// Returns reference to RLWE'_A(-s'm) polynomials
65 fn ksk_part_a(&self) -> &[Self::R];
66 /// Returns reference to RLWE'_B(-s'm) polynomials
67 fn ksk_part_b(&self) -> &[Self::R];
68}
69
70/// Scratch matrix used in several rlwe/rgsw runtime operations
71pub(crate) trait RuntimeScratchMatrix {
72 type R: RowMut;
73 type Rgsw: RgswCiphertext;
74
75 /// Returns scratch matrix for RLWE automorphism (not trivial case)
76 ///
77 /// RLWE auto requires scratch matric to store decomposed polynomials + 1
78 /// rlwe ciphertext temporarily.
79 ///
80 /// For example, if Auto decomposer has decompostion count `d` then the
81 /// scratch matrix must have dimension (d + 2, N) where N is the ring size.
82 fn scratch_for_rlwe_auto_and_zero_rlwe_space(
83 &mut self,
84 decompostion_count: usize,
85 ) -> (&mut [Self::R], &mut [Self::R]);
86
87 /// Returns scratch matrix for RLWE automorphism (trivial case)
88 ///
89 /// We refer to cases where RLWE(m) = [0, b] s.t. m = b as trivial cases. In
90 /// such a case a single row of length N, N being the ring dimension, is
91 /// required as scratch buffer to store automorphism of polynomial `b`
92 /// temporarily.
93 fn scratch_for_rlwe_auto_trivial_case(&mut self) -> &mut Self::R;
94
95 /// Returns scratch matrix + zeroed RLWE ciphertext space for
96 /// RLWE x RGSW
97 ///
98 /// RLWE x RGSW product requires scratch space to store decomposed
99 /// polynomials for both cases: (1) SignedDecompose(RLWE_A) x RLWE'(-sm) and
100 /// (2) SignedDecompose(RLWE_B) x RLWE'(m). Hence, scratch space returned to
101 /// store decomposed polynomials must have MAX(d_a, d_b) rows.
102 ///
103 /// Additional scratch space is required to store 1 RLWE ciphertext
104 /// temporarily. The space must be zeroed.
105 fn scratch_for_rlwe_x_rgsw_and_zero_rlwe_space<D: RlweDecomposer>(
106 &mut self,
107 decomposer: &D,
108 ) -> (&mut [Self::R], &mut [Self::R]);
109
110 /// Returns scracth matrix + zeroed RGSW ciphertext space for RGSW0 x RGSW1
111 ///
112 /// RGSW0 x RGSW1 requires `d_{0,a} + d_{0,b}` RLWE x RGSW1 products where
113 /// d_{0, a/b} are decomposition counts corresponding to decmposer used for
114 /// RGSW0. Hence, scratch space required to store decomposed polynomial for
115 /// RLWE x RGSW1 product should have MAX(d_{1, a}, d_{1, b}) rows.
116 ///
117 /// Additional scravth space is required to store RGSW0 ciphertext
118 /// temporarily. The space must be zeroed.
119 fn scratch_for_rgsw_x_rgsw_and_zero_rgsw0_space<D: RlweDecomposer>(
120 &mut self,
121 d0: &D,
122 d1: &D,
123 ) -> (&mut [Self::R], &mut [Self::R]);
124}
125
126pub(crate) struct RlweCiphertextMutRef<'a, R> {
127 data: &'a mut [R],
128}
129
130impl<'a, R> RlweCiphertextMutRef<'a, R> {
131 pub(crate) fn new(data: &'a mut [R]) -> Self {
132 Self { data }
133 }
134}
135
136impl<'a, R: RowMut> RlweCiphertext for RlweCiphertextMutRef<'a, R> {
137 type R = R;
138 fn part_a(&self) -> &[<Self::R as Row>::Element] {
139 self.data[0].as_ref()
140 }
141 fn part_a_mut(&mut self) -> &mut [<Self::R as Row>::Element] {
142 self.data[0].as_mut()
143 }
144 fn part_b(&self) -> &[<Self::R as Row>::Element] {
145 self.data[1].as_ref()
146 }
147 fn part_b_mut(&mut self) -> &mut [<Self::R as Row>::Element] {
148 self.data[1].as_mut()
149 }
150 fn ring_size(&self) -> usize {
151 self.data[0].as_ref().len()
152 }
153}
154
155pub(crate) struct RgswCiphertextRef<'a, R> {
156 data: &'a [R],
157 d_a: usize,
158 d_b: usize,
159}
160
161impl<'a, R> RgswCiphertextRef<'a, R> {
162 pub(crate) fn new(data: &'a [R], d_a: usize, d_b: usize) -> Self {
163 RgswCiphertextRef { data, d_a, d_b }
164 }
165}
166
167impl<'a, R> RgswCiphertext for RgswCiphertextRef<'a, R>
168where
169 R: Row,
170{
171 type R = R;
172
173 fn split(&self) -> ((&[Self::R], &[Self::R]), (&[Self::R], &[Self::R])) {
174 let (rlwe_dash_nsm, rlwe_dash_m) = self.data.split_at(self.d_a * 2);
175 (
176 rlwe_dash_nsm.split_at(self.d_a),
177 rlwe_dash_m.split_at(self.d_b),
178 )
179 }
180}
181
182pub(crate) struct RgswCiphertextMutRef<'a, R> {
183 data: &'a mut [R],
184 d_a: usize,
185 d_b: usize,
186}
187
188impl<'a, R> RgswCiphertextMutRef<'a, R> {
189 pub(crate) fn new(data: &'a mut [R], d_a: usize, d_b: usize) -> Self {
190 RgswCiphertextMutRef { data, d_a, d_b }
191 }
192}
193
194impl<'a, R: RowMut> AsMut<[R]> for RgswCiphertextMutRef<'a, R> {
195 fn as_mut(&mut self) -> &mut [R] {
196 &mut self.data
197 }
198}
199
200impl<'a, R> RgswCiphertext for RgswCiphertextMutRef<'a, R>
201where
202 R: Row,
203{
204 type R = R;
205
206 fn split(&self) -> ((&[Self::R], &[Self::R]), (&[Self::R], &[Self::R])) {
207 let (rlwe_dash_nsm, rlwe_dash_m) = self.data.split_at(self.d_a * 2);
208 (
209 rlwe_dash_nsm.split_at(self.d_a),
210 rlwe_dash_m.split_at(self.d_b),
211 )
212 }
213}
214
215impl<'a, R> RgswCiphertextMut for RgswCiphertextMutRef<'a, R>
216where
217 R: RowMut,
218{
219 fn split_mut(
220 &mut self,
221 ) -> (
222 (&mut [Self::R], &mut [Self::R]),
223 (&mut [Self::R], &mut [Self::R]),
224 ) {
225 let (rlwe_dash_nsm, rlwe_dash_m) = self.data.split_at_mut(self.d_a * 2);
226 (
227 rlwe_dash_nsm.split_at_mut(self.d_a),
228 rlwe_dash_m.split_at_mut(self.d_b),
229 )
230 }
231}
232
233pub(crate) struct RlweKskRef<'a, R> {
234 data: &'a [R],
235 decomposition_count: usize,
236}
237impl<'a, R: Row> RlweKskRef<'a, R> {
238 pub(crate) fn new(ksk: &'a [R], decomposition_count: usize) -> Self {
239 Self {
240 data: ksk,
241 decomposition_count,
242 }
243 }
244}
245
246impl<'a, R: Row> RlweKsk for RlweKskRef<'a, R> {
247 type R = R;
248
249 fn ksk_part_a(&self) -> &[Self::R] {
250 &self.data[..self.decomposition_count]
251 }
252
253 fn ksk_part_b(&self) -> &[Self::R] {
254 &self.data[self.decomposition_count..]
255 }
256}
257
258pub(crate) struct RuntimeScratchMutRef<'a, R> {
259 data: &'a mut [R],
260}
261
262impl<'a, R> RuntimeScratchMutRef<'a, R> {
263 pub(crate) fn new(data: &'a mut [R]) -> Self {
264 Self { data }
265 }
266}
267
268impl<'a, R: RowMut> RuntimeScratchMatrix for RuntimeScratchMutRef<'a, R>
269where
270 R::Element: Zero + Clone,
271{
272 type R = R;
273 type Rgsw = RgswCiphertextRef<'a, R>;
274
275 fn scratch_for_rlwe_auto_and_zero_rlwe_space(
276 &mut self,
277 decompostion_count: usize,
278 ) -> (&mut [Self::R], &mut [Self::R]) {
279 let (decomp_poly, other) = self.data.split_at_mut(decompostion_count);
280 let (rlwe, _) = other.split_at_mut(2);
281
282 // zero fill rlwe
283 rlwe.iter_mut()
284 .for_each(|r| r.as_mut().fill(R::Element::zero()));
285
286 (decomp_poly, rlwe)
287 }
288
289 fn scratch_for_rlwe_auto_trivial_case(&mut self) -> &mut Self::R {
290 &mut self.data[0]
291 }
292
293 fn scratch_for_rgsw_x_rgsw_and_zero_rgsw0_space<D: RlweDecomposer>(
294 &mut self,
295 rgsw0_decoposer: &D,
296 rgsw1_decoposer: &D,
297 ) -> (&mut [Self::R], &mut [Self::R]) {
298 let (decomp_poly, other) = self.data.split_at_mut(std::cmp::max(
299 rgsw1_decoposer.a().decomposition_count(),
300 rgsw1_decoposer.b().decomposition_count(),
301 ));
302 let (rgsw, _) = other.split_at_mut(
303 rgsw0_decoposer.a().decomposition_count() * 2
304 + rgsw0_decoposer.b().decomposition_count() * 2,
305 );
306
307 // zero fill rgsw0
308 rgsw.iter_mut()
309 .for_each(|r| r.as_mut().fill(R::Element::zero()));
310
311 (decomp_poly, rgsw)
312 }
313
314 fn scratch_for_rlwe_x_rgsw_and_zero_rlwe_space<D: RlweDecomposer>(
315 &mut self,
316 decomposer: &D,
317 ) -> (&mut [Self::R], &mut [Self::R]) {
318 let (decomp_poly, other) = self.data.split_at_mut(std::cmp::max(
319 decomposer.a().decomposition_count(),
320 decomposer.b().decomposition_count(),
321 ));
322
323 let (rlwe, _) = other.split_at_mut(2);
324
325 // zero fill rlwe
326 rlwe.iter_mut()
327 .for_each(|r| r.as_mut().fill(R::Element::zero()));
328
329 (decomp_poly, rlwe)
330 }
331}
332
333/// Returns no. of rows in scratch space for RGSW0 x RGSW1 product
334pub(crate) fn rgsw_x_rgsw_scratch_rows<D: RlweDecomposer>(
335 rgsw0_decomposer: &D,
336 rgsw1_decomposer: &D,
337) -> usize {
338 std::cmp::max(
339 rgsw1_decomposer.a().decomposition_count(),
340 rgsw1_decomposer.b().decomposition_count(),
341 ) + rgsw0_decomposer.a().decomposition_count() * 2
342 + rgsw0_decomposer.b().decomposition_count() * 2
343}
344
345/// Returns no. of rows in scratch space for RLWE x RGSW product
346pub(crate) fn rlwe_x_rgsw_scratch_rows<D: RlweDecomposer>(rgsw_decomposer: &D) -> usize {
347 std::cmp::max(
348 rgsw_decomposer.a().decomposition_count(),
349 rgsw_decomposer.b().decomposition_count(),
350 ) + 2
351}
352/// Returns no. of rows in scratch space for RLWE auto
353pub(crate) fn rlwe_auto_scratch_rows<D: Decomposer>(decomposer: &D) -> usize {
354 decomposer.decomposition_count() + 2
355}
356
357pub(crate) fn poly_fma_routine<R: RowMut, ModOp: VectorOps<Element = R::Element>>(
358 write_to_row: &mut [R::Element],
359 matrix_a: &[R],
360 matrix_b: &[R],
361 mod_op: &ModOp,
362) {
363 izip!(matrix_a.iter(), matrix_b.iter()).for_each(|(a, b)| {
364 mod_op.elwise_fma_mut(write_to_row, a.as_ref(), b.as_ref());
365 });
366}
367
368/// Decomposes ring polynomial r(X) into d polynomials using decomposer into
369/// output matrix decomp_r
370///
371/// Note that decomposition of r(X) requires decomposition of each of
372/// coefficients.
373///
374/// - decomp_r: must have dimensions d x ring_size. i^th decomposed polynomial
375/// will be stored at i^th row.
376pub(crate) fn decompose_r<R: RowMut, D: Decomposer<Element = R::Element>>(
377 r: &[R::Element],
378 decomp_r: &mut [R],
379 decomposer: &D,
380) where
381 R::Element: Copy,
382{
383 let ring_size = r.len();
384
385 for ri in 0..ring_size {
386 decomposer
387 .decompose_iter(&r[ri])
388 .enumerate()
389 .for_each(|(index, el)| {
390 decomp_r[index].as_mut()[ri] = el;
391 });
392 }
393}
394
395/// Sends RLWE_{s(X)}(m(X)) -> RLWE_{s(X)}(m{X^k}) where k is some galois
396/// element
397///
398/// - rlwe_in: Input ciphertext RLWE_{s(X)}(m(X)).
399/// - ksk: Auto key switching key with polynomials in evaluation domain
400/// - auto_map_index: If automorphism sends i^th coefficient of m(X) to j^th
401/// coefficient of m(X^k) then auto_map_index[i] = j
402/// - auto_sign_index: With a = m(X)[i], if m(X^k)[auto_map_index[i]] = -a, then
403/// auto_sign_index[i] = false, else auto_sign_index[i] = true
404/// - scratch_matrix: must have dimension at-least d+2 x ring_size. `d` rows to
405/// store decomposed polynomials nad 2 rows to store out RLWE temporarily.
406pub(crate) fn rlwe_auto<
407 Rlwe: RlweCiphertext,
408 Ksk: RlweKsk<R = Rlwe::R>,
409 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
410 ModOp: ArithmeticOps<Element = <Rlwe::R as Row>::Element>
411 + VectorOps<Element = <Rlwe::R as Row>::Element>,
412 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
413 D: Decomposer<Element = <Rlwe::R as Row>::Element>,
414>(
415 rlwe_in: &mut Rlwe,
416 ksk: &Ksk,
417 scratch_matrix: &mut Sc,
418 auto_map_index: &[usize],
419 auto_map_sign: &[bool],
420 mod_op: &ModOp,
421 ntt_op: &NttOp,
422 decomposer: &D,
423 is_trivial: bool,
424) where
425 <Rlwe::R as Row>::Element: Copy + Zero,
426{
427 // let ring_size = rlwe_in.dimension().1;
428 // assert!(rlwe_in.dimension().0 == 2);
429 // assert!(scratch_matrix.fits(d + 2, ring_size));
430
431 if !is_trivial {
432 let (decomp_poly_scratch, tmp_rlwe) = scratch_matrix
433 .scratch_for_rlwe_auto_and_zero_rlwe_space(decomposer.decomposition_count());
434 let mut tmp_rlwe = RlweCiphertextMutRef::new(tmp_rlwe);
435
436 // send a(X) -> a(X^k) and decompose a(X^k)
437 izip!(
438 rlwe_in.part_a(),
439 auto_map_index.iter(),
440 auto_map_sign.iter()
441 )
442 .for_each(|(el_in, to_index, sign)| {
443 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
444
445 decomposer
446 .decompose_iter(&el_out)
447 .enumerate()
448 .for_each(|(index, el)| {
449 decomp_poly_scratch[index].as_mut()[*to_index] = el;
450 });
451 });
452
453 // transform decomposed a(X^k) to evaluation domain
454 decomp_poly_scratch.iter_mut().for_each(|r| {
455 ntt_op.forward(r.as_mut());
456 });
457
458 // RLWE(m^k) = a', b'; RLWE(m) = a, b
459 // key switch: (a * RLWE'(s(X^k)))
460 // a' = decomp<a> * RLWE'_A(s(X^k))
461 poly_fma_routine(
462 tmp_rlwe.part_a_mut(),
463 decomp_poly_scratch,
464 ksk.ksk_part_a(),
465 mod_op,
466 );
467
468 // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
469 poly_fma_routine(
470 tmp_rlwe.part_b_mut(),
471 decomp_poly_scratch,
472 ksk.ksk_part_b(),
473 mod_op,
474 );
475
476 // transform RLWE(m^k) to coefficient domain
477 ntt_op.backward(tmp_rlwe.part_a_mut());
478 ntt_op.backward(tmp_rlwe.part_b_mut());
479
480 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
481 izip!(
482 rlwe_in.part_b(),
483 auto_map_index.iter(),
484 auto_map_sign.iter()
485 )
486 .for_each(|(el_in, to_index, sign)| {
487 let row = tmp_rlwe.part_b_mut();
488 if !*sign {
489 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
490 } else {
491 row[*to_index] = mod_op.add(&row[*to_index], el_in);
492 }
493 });
494
495 // copy over A; Leave B for later
496 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe.part_a());
497 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe.part_b());
498 } else {
499 // RLWE is trivial, a(X) is 0.
500 // send b(X) -> b(X^k)
501 let tmp_row = scratch_matrix.scratch_for_rlwe_auto_trivial_case();
502 izip!(
503 rlwe_in.part_b(),
504 auto_map_index.iter(),
505 auto_map_sign.iter()
506 )
507 .for_each(|(el_in, to_index, sign)| {
508 if !*sign {
509 tmp_row.as_mut()[*to_index] = mod_op.neg(el_in);
510 } else {
511 tmp_row.as_mut()[*to_index] = *el_in;
512 }
513 });
514 rlwe_in.part_b_mut().copy_from_slice(tmp_row.as_ref());
515 }
516}
517
518/// Sends RLWE_{s(X)}(m(X)) -> RLWE_{s(X)}(m{X^k}) where k is some galois
519/// element
520///
521/// This is same as `galois_auto` with the difference that alongside `ksk` with
522/// key switching polynomials in evaluation domain, shoup representation,
523/// `ksk_shoup`, of the polynomials in evaluation domain is also supplied.
524pub(crate) fn rlwe_auto_shoup<
525 Rlwe: RlweCiphertext,
526 Ksk: RlweKsk<R = Rlwe::R>,
527 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
528 ModOp: ArithmeticOps<Element = <Rlwe::R as Row>::Element>
529 // + VectorOps<Element = MT::MatElement>
530 + ShoupMatrixFMA<Rlwe::R>,
531 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
532 D: Decomposer<Element = <Rlwe::R as Row>::Element>,
533>(
534 rlwe_in: &mut Rlwe,
535 ksk: &Ksk,
536 ksk_shoup: &Ksk,
537 scratch_matrix: &mut Sc,
538 auto_map_index: &[usize],
539 auto_map_sign: &[bool],
540 mod_op: &ModOp,
541 ntt_op: &NttOp,
542 decomposer: &D,
543 is_trivial: bool,
544) where
545 <Rlwe::R as Row>::Element: Copy + Zero,
546{
547 // let d = decomposer.decomposition_count();
548 // let ring_size = rlwe_in.dimension().1;
549 // assert!(rlwe_in.dimension().0 == 2);
550 // assert!(scratch_matrix.fits(d + 2, ring_size));
551
552 if !is_trivial {
553 let (decomp_poly_scratch, tmp_rlwe) = scratch_matrix
554 .scratch_for_rlwe_auto_and_zero_rlwe_space(decomposer.decomposition_count());
555 let mut tmp_rlwe = RlweCiphertextMutRef::new(tmp_rlwe);
556
557 // send a(X) -> a(X^k) and decompose a(X^k)
558 izip!(
559 rlwe_in.part_a(),
560 auto_map_index.iter(),
561 auto_map_sign.iter()
562 )
563 .for_each(|(el_in, to_index, sign)| {
564 let el_out = if !*sign { mod_op.neg(el_in) } else { *el_in };
565
566 decomposer
567 .decompose_iter(&el_out)
568 .enumerate()
569 .for_each(|(index, el)| {
570 decomp_poly_scratch[index].as_mut()[*to_index] = el;
571 });
572 });
573
574 // transform decomposed a(X^k) to evaluation domain
575 decomp_poly_scratch.iter_mut().for_each(|r| {
576 ntt_op.forward_lazy(r.as_mut());
577 });
578
579 // RLWE(m^k) = a', b'; RLWE(m) = a, b
580 // key switch: (a * RLWE'(s(X^k)))
581 // a' = decomp<a> * RLWE'_A(s(X^k))
582 mod_op.shoup_matrix_fma(
583 tmp_rlwe.part_a_mut(),
584 ksk.ksk_part_a(),
585 ksk_shoup.ksk_part_a(),
586 decomp_poly_scratch,
587 );
588
589 // b'= decomp<a(X^k)> * RLWE'_B(s(X^k))
590 mod_op.shoup_matrix_fma(
591 tmp_rlwe.part_b_mut(),
592 ksk.ksk_part_b(),
593 ksk_shoup.ksk_part_b(),
594 decomp_poly_scratch,
595 );
596
597 // transform RLWE(m^k) to coefficient domain
598 ntt_op.backward(tmp_rlwe.part_a_mut());
599 ntt_op.backward(tmp_rlwe.part_b_mut());
600
601 // send b(X) -> b(X^k) and then b'(X) += b(X^k)
602 let row = tmp_rlwe.part_b_mut();
603 izip!(
604 rlwe_in.part_b(),
605 auto_map_index.iter(),
606 auto_map_sign.iter()
607 )
608 .for_each(|(el_in, to_index, sign)| {
609 if !*sign {
610 row[*to_index] = mod_op.sub(&row[*to_index], el_in);
611 } else {
612 row[*to_index] = mod_op.add(&row[*to_index], el_in);
613 }
614 });
615
616 // copy over A, B
617 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe.part_a());
618 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe.part_b());
619 } else {
620 // RLWE is trivial, a(X) is 0.
621 // send b(X) -> b(X^k)
622 let row = scratch_matrix.scratch_for_rlwe_auto_trivial_case();
623 izip!(
624 rlwe_in.part_b(),
625 auto_map_index.iter(),
626 auto_map_sign.iter()
627 )
628 .for_each(|(el_in, to_index, sign)| {
629 if !*sign {
630 row.as_mut()[*to_index] = mod_op.neg(el_in);
631 } else {
632 row.as_mut()[*to_index] = *el_in;
633 }
634 });
635 rlwe_in.part_b_mut().copy_from_slice(row.as_ref());
636 }
637}
638
639/// Inplace mutates RLWE(m0) to equal RLWE(m0m1) = RLWE(m0) x RGSW(m1).
640///
641/// - rlwe_in: is RLWE(m0) with polynomials in coefficient domain
642/// - rgsw_in: is RGSW(m1) with polynomials in evaluation domain
643/// - scratch_matrix: with dimension (max(d_a, d_b) + 2) x ring_size columns.
644/// It's used to store decomposed polynomials and out RLWE temporarily
645pub(crate) fn rlwe_by_rgsw<
646 Rlwe: RlweCiphertext,
647 Rgsw: RgswCiphertext<R = Rlwe::R>,
648 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
649 D: RlweDecomposer<Element = <Rlwe::R as Row>::Element>,
650 ModOp: VectorOps<Element = <Rlwe::R as Row>::Element>,
651 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
652>(
653 rlwe_in: &mut Rlwe,
654 rgsw_in: &Rgsw,
655 scratch_matrix: &mut Sc,
656 decomposer: &D,
657 ntt_op: &NttOp,
658 mod_op: &ModOp,
659 is_trivial: bool,
660) where
661 <Rlwe::R as Row>::Element: Copy + Zero,
662{
663 let decomposer_a = decomposer.a();
664 let decomposer_b = decomposer.b();
665 let d_a = decomposer_a.decomposition_count();
666 let d_b = decomposer_b.decomposition_count();
667
668 let ((rlwe_dash_nsm_parta, rlwe_dash_nsm_partb), (rlwe_dash_m_parta, rlwe_dash_m_partb)) =
669 rgsw_in.split();
670
671 let (decomposed_poly_scratch, tmp_rlwe) =
672 scratch_matrix.scratch_for_rlwe_x_rgsw_and_zero_rlwe_space(decomposer);
673
674 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
675 if !is_trivial {
676 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
677 // decomp<a_in>
678 let mut decomposed_polys_of_rlwea = &mut decomposed_poly_scratch[..d_a];
679 decompose_r(
680 rlwe_in.part_a(),
681 &mut decomposed_polys_of_rlwea,
682 decomposer_a,
683 );
684
685 decomposed_polys_of_rlwea
686 .iter_mut()
687 .for_each(|r| ntt_op.forward(r.as_mut()));
688
689 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
690 poly_fma_routine(
691 tmp_rlwe[0].as_mut(),
692 &decomposed_polys_of_rlwea,
693 rlwe_dash_nsm_parta,
694 mod_op,
695 );
696 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
697 poly_fma_routine(
698 tmp_rlwe[1].as_mut(),
699 &decomposed_polys_of_rlwea,
700 &rlwe_dash_nsm_partb,
701 mod_op,
702 );
703 }
704
705 {
706 // decomp<b_in>
707 let mut decomposed_polys_of_rlweb = &mut decomposed_poly_scratch[..d_b];
708 decompose_r(
709 rlwe_in.part_b(),
710 &mut decomposed_polys_of_rlweb,
711 decomposer_b,
712 );
713
714 decomposed_polys_of_rlweb
715 .iter_mut()
716 .for_each(|r| ntt_op.forward(r.as_mut()));
717
718 // a_out += decomp<b_in> \cdot RLWE_A'(m)
719 poly_fma_routine(
720 tmp_rlwe[0].as_mut(),
721 &decomposed_polys_of_rlweb,
722 &rlwe_dash_m_parta,
723 mod_op,
724 );
725 // b_out += decomp<b_in> \cdot RLWE_B'(m)
726 poly_fma_routine(
727 tmp_rlwe[1].as_mut(),
728 &decomposed_polys_of_rlweb,
729 &rlwe_dash_m_partb,
730 mod_op,
731 );
732 }
733
734 // transform rlwe_out to coefficient domain
735 tmp_rlwe
736 .iter_mut()
737 .for_each(|r| ntt_op.backward(r.as_mut()));
738
739 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe[0].as_mut());
740 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe[1].as_mut());
741}
742
743/// Inplace mutates RLWE(m0) to equal RLWE(m0m1) = RLWE(m0) x RGSW(m1).
744///
745/// Same as `rlwe_by_rgsw` with the difference that alongside `rgsw_in` with
746/// polynomials in evaluation domain, shoup representation of polynomials in
747/// evaluation domain, `rgsw_in_shoup`, is also supplied.
748pub(crate) fn rlwe_by_rgsw_shoup<
749 Rlwe: RlweCiphertext,
750 Rgsw: RgswCiphertext<R = Rlwe::R>,
751 Sc: RuntimeScratchMatrix<R = Rlwe::R>,
752 D: RlweDecomposer<Element = <Rlwe::R as Row>::Element>,
753 ModOp: ShoupMatrixFMA<Rlwe::R>,
754 NttOp: Ntt<Element = <Rlwe::R as Row>::Element>,
755>(
756 rlwe_in: &mut Rlwe,
757 rgsw_in: &Rgsw,
758 rgsw_in_shoup: &Rgsw,
759 scratch_matrix: &mut Sc,
760 decomposer: &D,
761 ntt_op: &NttOp,
762 mod_op: &ModOp,
763 is_trivial: bool,
764) where
765 <Rlwe::R as Row>::Element: Copy + Zero,
766{
767 let decomposer_a = decomposer.a();
768 let decomposer_b = decomposer.b();
769 let d_a = decomposer_a.decomposition_count();
770 let d_b = decomposer_b.decomposition_count();
771
772 let ((rlwe_dash_nsm_parta, rlwe_dash_nsm_partb), (rlwe_dash_m_parta, rlwe_dash_m_partb)) =
773 rgsw_in.split();
774
775 let (
776 (rlwe_dash_nsm_parta_shoup, rlwe_dash_nsm_partb_shoup),
777 (rlwe_dash_m_parta_shoup, rlwe_dash_m_partb_shoup),
778 ) = rgsw_in_shoup.split();
779
780 let (decomposed_poly_scratch, tmp_rlwe) =
781 scratch_matrix.scratch_for_rlwe_x_rgsw_and_zero_rlwe_space(decomposer);
782
783 // RLWE_in = a_in, b_in; RLWE_out = a_out, b_out
784 if !is_trivial {
785 // a_in = 0 when RLWE_in is trivial RLWE ciphertext
786 // decomp<a_in>
787 let mut decomposed_polys_of_rlwea = &mut decomposed_poly_scratch[..d_a];
788 decompose_r(
789 rlwe_in.part_a(),
790 &mut decomposed_polys_of_rlwea,
791 decomposer_a,
792 );
793 decomposed_polys_of_rlwea
794 .iter_mut()
795 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
796
797 // a_out += decomp<a_in> \cdot RLWE_A'(-sm)
798 mod_op.shoup_matrix_fma(
799 tmp_rlwe[0].as_mut(),
800 &rlwe_dash_nsm_parta,
801 &rlwe_dash_nsm_parta_shoup,
802 &decomposed_polys_of_rlwea,
803 );
804
805 // b_out += decomp<a_in> \cdot RLWE_B'(-sm)
806 mod_op.shoup_matrix_fma(
807 tmp_rlwe[1].as_mut(),
808 &rlwe_dash_nsm_partb,
809 &rlwe_dash_nsm_partb_shoup,
810 &decomposed_polys_of_rlwea,
811 );
812 }
813 {
814 // decomp<b_in>
815 let mut decomposed_polys_of_rlweb = &mut decomposed_poly_scratch[..d_b];
816 decompose_r(
817 rlwe_in.part_b(),
818 &mut decomposed_polys_of_rlweb,
819 decomposer_b,
820 );
821 decomposed_polys_of_rlweb
822 .iter_mut()
823 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
824
825 // a_out += decomp<b_in> \cdot RLWE_A'(m)
826 mod_op.shoup_matrix_fma(
827 tmp_rlwe[0].as_mut(),
828 &rlwe_dash_m_parta,
829 &rlwe_dash_m_parta_shoup,
830 &decomposed_polys_of_rlweb,
831 );
832
833 // b_out += decomp<b_in> \cdot RLWE_B'(m)
834 mod_op.shoup_matrix_fma(
835 tmp_rlwe[1].as_mut(),
836 &rlwe_dash_m_partb,
837 &rlwe_dash_m_partb_shoup,
838 &decomposed_polys_of_rlweb,
839 );
840 }
841
842 // transform rlwe_out to coefficient domain
843 tmp_rlwe
844 .iter_mut()
845 .for_each(|r| ntt_op.backward(r.as_mut()));
846
847 rlwe_in.part_a_mut().copy_from_slice(tmp_rlwe[0].as_mut());
848 rlwe_in.part_b_mut().copy_from_slice(tmp_rlwe[1].as_mut());
849}
850
851/// Inplace mutates RGSW(m0) to equal RGSW(m0m1) = RGSW(m0)xRGSW(m1)
852///
853/// RGSW x RGSW product requires multiple RLWE x RGSW products. For example,
854/// Define
855///
856/// RGSW(m0) = [RLWE(-sm), RLWE(\beta -sm), ..., RLWE(\beta^{d-1} -sm)
857/// RLWE(m), RLWE(\beta m), ..., RLWE(\beta^{d-1} m)]
858/// And RGSW(m1)
859///
860/// Then RGSW(m0) x RGSW(m1) equals:
861/// RGSW(m0m1) = [
862/// rlwe_x_rgsw(RLWE(-sm), RGSW(m1)),
863/// ...,
864/// rlwe_x_rgsw(RLWE(\beta^{d-1} -sm), RGSW(m1)),
865/// rlwe_x_rgsw(RLWE(m), RGSW(m1)),
866/// ...,
867/// rlwe_x_rgsw(RLWE(\beta^{d-1} m), RGSW(m1)),
868/// ]
869///
870/// Since noise growth in RLWE x RGSW depends on noise in RGSW ciphertext, it is
871/// clear to observe from above that noise in resulting RGSW(m0m1) equals noise
872/// accumulated in a single RLWE x RGSW and depends on noise in RGSW(m1) (i.e.
873/// rgsw_1_eval)
874///
875/// - rgsw_0: RGSW(m0) in coefficient domain
876/// - rgsw_1_eval: RGSW(m1) in evaluation domain
877pub(crate) fn rgsw_by_rgsw_inplace<
878 Rgsw: RgswCiphertext,
879 RgswMut: RgswCiphertextMut<R = Rgsw::R>,
880 Sc: RuntimeScratchMatrix<R = Rgsw::R, Rgsw = Rgsw>,
881 D: RlweDecomposer<Element = <Rgsw::R as Row>::Element>,
882 ModOp: VectorOps<Element = <Rgsw::R as Row>::Element>,
883 NttOp: Ntt<Element = <Rgsw::R as Row>::Element>,
884>(
885 rgsw0: &mut RgswMut,
886 rgsw1_eval: &Rgsw,
887 rgsw0_decomposer: &D,
888 rgsw1_decomposer: &D,
889 scratch_matrix: &mut Sc,
890 ntt_op: &NttOp,
891 mod_op: &ModOp,
892) where
893 <Rgsw::R as Row>::Element: Copy + Zero,
894 RgswMut: AsMut<[Rgsw::R]>,
895 RgswMut::R: RowMut,
896 // Rgsw: AsRef<[Rgsw::R]>,
897{
898 let (decomp_r_space, rgsw_space) = scratch_matrix
899 .scratch_for_rgsw_x_rgsw_and_zero_rgsw0_space(rgsw0_decomposer, rgsw1_decomposer);
900
901 let mut rgsw_space = RgswCiphertextMutRef::new(
902 rgsw_space,
903 rgsw0_decomposer.a().decomposition_count(),
904 rgsw0_decomposer.b().decomposition_count(),
905 );
906 let (
907 (rlwe_dash_space_nsm_parta, rlwe_dash_space_nsm_partb),
908 (rlwe_dash_space_m_parta, rlwe_dash_space_m_partb),
909 ) = rgsw_space.split_mut();
910
911 let ((rgsw0_nsm_parta, rgsw0_nsm_partb), (rgsw0_m_parta, rgsw0_m_partb)) = rgsw0.split();
912 let ((rgsw1_nsm_parta, rgsw1_nsm_partb), (rgsw1_m_parta, rgsw1_m_partb)) = rgsw1_eval.split();
913
914 // RGSW x RGSW
915 izip!(
916 rgsw0_nsm_parta.iter().chain(rgsw0_m_parta),
917 rgsw0_nsm_partb.iter().chain(rgsw0_m_partb),
918 rlwe_dash_space_nsm_parta
919 .iter_mut()
920 .chain(rlwe_dash_space_m_parta.iter_mut()),
921 rlwe_dash_space_nsm_partb
922 .iter_mut()
923 .chain(rlwe_dash_space_m_partb.iter_mut()),
924 )
925 .for_each(|(rlwe_a, rlwe_b, rlwe_out_a, rlwe_out_b)| {
926 // RLWE(m0) x RGSW(m1)
927
928 // Part A: Decomp<RLWE(m0)[A]> \cdot RLWE'(-sm1)
929 {
930 let decomp_r_parta = &mut decomp_r_space[..rgsw1_decomposer.a().decomposition_count()];
931 decompose_r(
932 rlwe_a.as_ref(),
933 decomp_r_parta.as_mut(),
934 rgsw1_decomposer.a(),
935 );
936 decomp_r_parta
937 .iter_mut()
938 .for_each(|ri| ntt_op.forward(ri.as_mut()));
939 poly_fma_routine(
940 rlwe_out_a.as_mut(),
941 &decomp_r_parta,
942 &rgsw1_nsm_parta,
943 mod_op,
944 );
945 poly_fma_routine(
946 rlwe_out_b.as_mut(),
947 &decomp_r_parta,
948 &rgsw1_nsm_partb,
949 mod_op,
950 );
951 }
952
953 // Part B: Decompose<RLWE(m0)[B]> \cdot RLWE'(m1)
954 {
955 let decomp_r_partb = &mut decomp_r_space[..rgsw1_decomposer.b().decomposition_count()];
956 decompose_r(
957 rlwe_b.as_ref(),
958 decomp_r_partb.as_mut(),
959 rgsw1_decomposer.b(),
960 );
961 decomp_r_partb
962 .iter_mut()
963 .for_each(|ri| ntt_op.forward(ri.as_mut()));
964 poly_fma_routine(rlwe_out_a.as_mut(), &decomp_r_partb, &rgsw1_m_parta, mod_op);
965 poly_fma_routine(rlwe_out_b.as_mut(), &decomp_r_partb, &rgsw1_m_partb, mod_op);
966 }
967 });
968
969 // copy over RGSW(m0m1) to RGSW(m0)
970 // let d = rgsw0.as_mut();
971 izip!(rgsw0.as_mut().iter_mut(), rgsw_space.data.iter())
972 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
973
974 // send back to coefficient domain
975 rgsw0
976 .as_mut()
977 .iter_mut()
978 .for_each(|ri| ntt_op.backward(ri.as_mut()));
979}
980
981/// Key switches input RLWE_{s'}(m) -> RLWE_{s}(m)
982///
983/// Let RLWE_{s'}(m) = [a, b] s.t. m+e = b - as'
984///
985/// Given key switchin key Ksk(s' -> s) = RLWE'_{s}(s') = [RLWE_{s}(beta^i s')]
986/// = [a, a*s + e + beta^i s'] for i \in [0,d), key switching computes:
987/// 1. RLWE_{s}(-s'a) = \sum signed_decompose(-a)[i] RLWE_{s}(beta^i s')
988/// 2. RLWE_{s}(m) = (b, 0) + RLWE_{s}(-s'a)
989///
990/// - rlwe_in: Input rlwe ciphertext
991/// - ksk: Key switching key Ksk(s' -> s) with polynomials in evaluation domain
992/// - ksk_shoup: Key switching key Ksk(s' -> s) with polynomials in evaluation
993/// domain in shoup representation
994/// - decomposer: Decomposer used for key switching
995pub(crate) fn rlwe_key_switch<
996 M: MatrixMut + MatrixEntity,
997 ModOp: GetModulus<Element = M::MatElement> + ShoupMatrixFMA<M::R> + VectorOps<Element = M::MatElement>,
998 NttOp: Ntt<Element = M::MatElement>,
999 D: Decomposer<Element = M::MatElement>,
1000>(
1001 rlwe_in: &M,
1002 ksk: &M,
1003 ksk_shoup: &M,
1004 decomposer: &D,
1005 ntt_op: &NttOp,
1006 mod_op: &ModOp,
1007) -> M
1008where
1009 <M as Matrix>::R: RowMut + RowEntity,
1010 M::MatElement: Copy,
1011{
1012 let ring_size = rlwe_in.dimension().1;
1013 assert!(rlwe_in.dimension().0 == 2);
1014 assert!(ksk.dimension() == (decomposer.decomposition_count() * 2, ring_size));
1015
1016 let mut rlwe_out = M::zeros(2, ring_size);
1017
1018 let mut tmp = M::zeros(decomposer.decomposition_count(), ring_size);
1019 let mut tmp_row = M::R::zeros(ring_size);
1020
1021 // key switch RLWE part -A
1022 // negative A
1023 tmp_row.as_mut().copy_from_slice(rlwe_in.get_row_slice(0));
1024 mod_op.elwise_neg_mut(tmp_row.as_mut());
1025 // decompose -A and send to evaluation domain
1026 decompose_r(tmp_row.as_ref(), tmp.as_mut(), decomposer);
1027 tmp.iter_rows_mut()
1028 .for_each(|r| ntt_op.forward_lazy(r.as_mut()));
1029
1030 // RLWE_s(-A u) = B' + B, A' = (decomp(-A) * Ksk(u -> s)) + (B, 0)
1031 let (ksk_part_a, ksk_part_b) = ksk.split_at_row(decomposer.decomposition_count());
1032 let (ksk_part_a_shoup, ksk_part_b_shoup) =
1033 ksk_shoup.split_at_row(decomposer.decomposition_count());
1034 // Part A'
1035 mod_op.shoup_matrix_fma(
1036 rlwe_out.get_row_mut(0),
1037 &ksk_part_a,
1038 &ksk_part_a_shoup,
1039 tmp.as_ref(),
1040 );
1041 // Part B'
1042 mod_op.shoup_matrix_fma(
1043 rlwe_out.get_row_mut(1),
1044 &ksk_part_b,
1045 &ksk_part_b_shoup,
1046 tmp.as_ref(),
1047 );
1048 // back to coefficient domain
1049 rlwe_out
1050 .iter_rows_mut()
1051 .for_each(|r| ntt_op.backward(r.as_mut()));
1052
1053 // B' + B
1054 mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), rlwe_in.get_row_slice(1));
1055
1056 rlwe_out
1057}