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