arnaucube
arnaucube/phantom-zonePublic Self-hosted repository
phantom-zone/src/rgsw/keygen.rs
src/rgsw/keygen.rsRaw
  1use std::{fmt::Debug, ops::Sub};
  2
  3use itertools::izip;
  4use num_traits::{PrimInt, Signed, ToPrimitive, Zero};
  5
  6use crate::{
  7    backend::{ArithmeticOps, GetModulus, Modulus, VectorOps},
  8    ntt::Ntt,
  9    random::{
 10        RandomElementInModulus, RandomFill, RandomFillGaussianInModulus, RandomFillUniformInModulus,
 11    },
 12    utils::{fill_random_ternary_secret_with_hamming_weight, TryConvertFrom1},
 13    Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut,
 14};
 15
 16pub(crate) fn generate_auto_map(ring_size: usize, k: isize) -> (Vec<usize>, Vec<bool>) {
 17    assert!(k & 1 == 1, "Auto {k} must be odd");
 18
 19    let k = if k < 0 {
 20        // k is -ve, return k%(2*N)
 21        (2 * ring_size) - (k.abs() as usize % (2 * ring_size))
 22    } else {
 23        k as usize
 24    };
 25    let (auto_map_index, auto_sign_index): (Vec<usize>, Vec<bool>) = (0..ring_size)
 26        .into_iter()
 27        .map(|i| {
 28            let mut to_index = (i * k) % (2 * ring_size);
 29            let mut sign = true;
 30
 31            // wrap around. false implies negative
 32            if to_index >= ring_size {
 33                to_index = to_index - ring_size;
 34                sign = false;
 35            }
 36
 37            (to_index, sign)
 38        })
 39        .unzip();
 40    (auto_map_index, auto_sign_index)
 41}
 42
 43/// Returns RGSW(m)
 44///
 45/// RGSW = [RLWE'(-sm) || RLWE(m)] = [RLWE'_A(-sm), RLWE'_B(-sm), RLWE'_A(m),
 46/// RLWE'_B(m)]
 47///
 48/// RGSW(m1) ciphertext is used for RLWE(m0) x RGSW(m1) multiplication.
 49/// Let RLWE(m) = [a, b] where b = as + e + m0.
 50/// For RLWExRGSW we calculate:
 51///       (\sum signed_decompose(a)[i] x RLWE(-s \beta^i' m1))
 52///     + (\sum signed_decompose(b)[i'] x RLWE(\beta^i' m1))
 53///     = RLWE(m0m1)
 54/// We denote decomposer count for signed_decompose(a)[i] with d_a and
 55/// corresponding gadget vector with `gadget_a`. We denote decomposer count for
 56/// signed_decompose(b)[i] with d_b and corresponding gadget vector with
 57/// `gadget_b`
 58///
 59/// In secret key RGSW encrypton RLWE'_A(m) can be seeded. Hence, we seed it
 60/// using the `p_rng` passed and the retured RGSW ciphertext has d_a * 2 + d_b
 61/// rows
 62///
 63/// - s: is the secret key
 64/// - m: message to encrypt
 65/// - gadget_a: Gadget vector for RLWE'(-sm)
 66/// - gadget_b: Gadget vector for RLWE'(m)
 67/// - p_rng: Seeded psuedo random generator used to sample RLWE'_A(m).
 68pub(crate) fn secret_key_encrypt_rgsw<
 69    Mmut: MatrixMut + MatrixEntity,
 70    S,
 71    R: RandomFillGaussianInModulus<[Mmut::MatElement], ModOp::M>
 72        + RandomFillUniformInModulus<[Mmut::MatElement], ModOp::M>,
 73    PR: RandomFillUniformInModulus<[Mmut::MatElement], ModOp::M>,
 74    ModOp: VectorOps<Element = Mmut::MatElement> + GetModulus<Element = Mmut::MatElement>,
 75    NttOp: Ntt<Element = Mmut::MatElement>,
 76>(
 77    out_rgsw: &mut Mmut,
 78    m: &[Mmut::MatElement],
 79    gadget_a: &[Mmut::MatElement],
 80    gadget_b: &[Mmut::MatElement],
 81    s: &[S],
 82    mod_op: &ModOp,
 83    ntt_op: &NttOp,
 84    p_rng: &mut PR,
 85    rng: &mut R,
 86) where
 87    <Mmut as Matrix>::R: RowMut + RowEntity + TryConvertFrom1<[S], ModOp::M> + Debug,
 88    Mmut::MatElement: Copy + Debug,
 89{
 90    let d_a = gadget_a.len();
 91    let d_b = gadget_b.len();
 92    let q = mod_op.modulus();
 93    let ring_size = s.len();
 94    assert!(out_rgsw.dimension() == (d_a * 2 + d_b, ring_size));
 95    assert!(m.as_ref().len() == ring_size);
 96
 97    // RLWE(-sm), RLWE(m)
 98    let (rlwe_dash_nsm, b_rlwe_dash_m) = out_rgsw.split_at_row_mut(d_a * 2);
 99
100    let mut s_eval = Mmut::R::try_convert_from(s, &q);
101    ntt_op.forward(s_eval.as_mut());
102
103    let mut scratch_space = Mmut::R::zeros(ring_size);
104
105    // RLWE'(-sm)
106    let (a_rlwe_dash_nsm, b_rlwe_dash_nsm) = rlwe_dash_nsm.split_at_mut(d_a);
107    izip!(
108        a_rlwe_dash_nsm.iter_mut(),
109        b_rlwe_dash_nsm.iter_mut(),
110        gadget_a.iter()
111    )
112    .for_each(|(ai, bi, beta_i)| {
113        // Sample a_i
114        RandomFillUniformInModulus::random_fill(rng, &q, ai.as_mut());
115
116        // a_i * s
117        scratch_space.as_mut().copy_from_slice(ai.as_ref());
118        ntt_op.forward(scratch_space.as_mut());
119        mod_op.elwise_mul_mut(scratch_space.as_mut(), s_eval.as_ref());
120        ntt_op.backward(scratch_space.as_mut());
121
122        // b_i = e_i + a_i * s
123        RandomFillGaussianInModulus::random_fill(rng, &q, bi.as_mut());
124        mod_op.elwise_add_mut(bi.as_mut(), scratch_space.as_ref());
125
126        // a_i + \beta_i * m
127        mod_op.elwise_scalar_mul(scratch_space.as_mut(), m.as_ref(), beta_i);
128        mod_op.elwise_add_mut(ai.as_mut(), scratch_space.as_ref());
129    });
130
131    // RLWE(m)
132    let mut a_rlwe_dash_m = {
133        // polynomials of part A of RLWE'(m) are sampled from seed
134        let mut a = Mmut::zeros(d_b, ring_size);
135        a.iter_rows_mut()
136            .for_each(|ai| RandomFillUniformInModulus::random_fill(p_rng, &q, ai.as_mut()));
137        a
138    };
139
140    izip!(
141        a_rlwe_dash_m.iter_rows_mut(),
142        b_rlwe_dash_m.iter_mut(),
143        gadget_b.iter()
144    )
145    .for_each(|(ai, bi, beta_i)| {
146        // ai * s
147        ntt_op.forward(ai.as_mut());
148        mod_op.elwise_mul_mut(ai.as_mut(), s_eval.as_ref());
149        ntt_op.backward(ai.as_mut());
150
151        // beta_i * m
152        mod_op.elwise_scalar_mul(scratch_space.as_mut(), m.as_ref(), beta_i);
153
154        // Sample e_i
155        RandomFillGaussianInModulus::random_fill(rng, &q, bi.as_mut());
156        // e_i + beta_i * m + ai*s
157        mod_op.elwise_add_mut(bi.as_mut(), scratch_space.as_ref());
158        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
159    });
160}
161
162/// Returns RGSW(m) encrypted with public key
163///
164/// Follows the same routine as `secret_key_encrypt_rgsw` but with the
165/// difference that each RLWE encryption uses public key instead of secret key.
166///
167/// Since public key encryption cannot be seeded `RLWE'_A(m)` is included in the
168/// ciphertext. Hence the returned RGSW ciphertext has d_a * 2 + d_b * 2 rows
169pub(crate) fn public_key_encrypt_rgsw<
170    Mmut: MatrixMut + MatrixEntity,
171    M: Matrix<MatElement = Mmut::MatElement>,
172    R: RandomFillGaussianInModulus<[Mmut::MatElement], ModOp::M>
173        + RandomFill<[u8]>
174        + RandomElementInModulus<usize, usize>,
175    ModOp: VectorOps<Element = Mmut::MatElement> + GetModulus<Element = Mmut::MatElement>,
176    NttOp: Ntt<Element = Mmut::MatElement>,
177>(
178    out_rgsw: &mut Mmut,
179    m: &[M::MatElement],
180    public_key: &M,
181    gadget_a: &[Mmut::MatElement],
182    gadget_b: &[Mmut::MatElement],
183    mod_op: &ModOp,
184    ntt_op: &NttOp,
185    rng: &mut R,
186) where
187    <Mmut as Matrix>::R: RowMut + RowEntity + TryConvertFrom1<[i32], ModOp::M>,
188    Mmut::MatElement: Copy,
189{
190    let ring_size = public_key.dimension().1;
191    let d_a = gadget_a.len();
192    let d_b = gadget_b.len();
193    assert!(public_key.dimension().0 == 2);
194    assert!(out_rgsw.dimension() == (d_a * 2 + d_b * 2, ring_size));
195
196    let mut pk_eval = Mmut::zeros(2, ring_size);
197    izip!(pk_eval.iter_rows_mut(), public_key.iter_rows()).for_each(|(to_i, from_i)| {
198        to_i.as_mut().copy_from_slice(from_i.as_ref());
199        ntt_op.forward(to_i.as_mut());
200    });
201    let p0 = pk_eval.get_row_slice(0);
202    let p1 = pk_eval.get_row_slice(1);
203
204    let q = mod_op.modulus();
205
206    // RGSW(m) = RLWE'(-sm), RLWE(m)
207    let (rlwe_dash_nsm, rlwe_dash_m) = out_rgsw.split_at_row_mut(d_a * 2);
208
209    // RLWE(-sm)
210    let (rlwe_dash_nsm_parta, rlwe_dash_nsm_partb) = rlwe_dash_nsm.split_at_mut(d_a);
211    izip!(
212        rlwe_dash_nsm_parta.iter_mut(),
213        rlwe_dash_nsm_partb.iter_mut(),
214        gadget_a.iter()
215    )
216    .for_each(|(ai, bi, beta_i)| {
217        // sample ephemeral secret u_i
218        let mut u = vec![0i32; ring_size];
219        fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
220        let mut u_eval = Mmut::R::try_convert_from(u.as_ref(), &q);
221        ntt_op.forward(u_eval.as_mut());
222
223        let mut u_eval_copy = Mmut::R::zeros(ring_size);
224        u_eval_copy.as_mut().copy_from_slice(u_eval.as_ref());
225
226        // p0 * u
227        mod_op.elwise_mul_mut(u_eval.as_mut(), p0.as_ref());
228        // p1 * u
229        mod_op.elwise_mul_mut(u_eval_copy.as_mut(), p1.as_ref());
230        ntt_op.backward(u_eval.as_mut());
231        ntt_op.backward(u_eval_copy.as_mut());
232
233        // sample error
234        RandomFillGaussianInModulus::random_fill(rng, &q, ai.as_mut());
235        RandomFillGaussianInModulus::random_fill(rng, &q, bi.as_mut());
236
237        // a = p0*u+e0
238        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
239        // b = p1*u+e1
240        mod_op.elwise_add_mut(bi.as_mut(), u_eval_copy.as_ref());
241
242        // a = p0*u + e0 + \beta*m
243        // use u_eval as scratch
244        mod_op.elwise_scalar_mul(u_eval.as_mut(), m.as_ref(), beta_i);
245        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
246    });
247
248    // RLWE(m)
249    let (rlwe_dash_m_parta, rlwe_dash_m_partb) = rlwe_dash_m.split_at_mut(d_b);
250    izip!(
251        rlwe_dash_m_parta.iter_mut(),
252        rlwe_dash_m_partb.iter_mut(),
253        gadget_b.iter()
254    )
255    .for_each(|(ai, bi, beta_i)| {
256        // sample ephemeral secret u_i
257        let mut u = vec![0i32; ring_size];
258        fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
259        let mut u_eval = Mmut::R::try_convert_from(u.as_ref(), &q);
260        ntt_op.forward(u_eval.as_mut());
261
262        let mut u_eval_copy = Mmut::R::zeros(ring_size);
263        u_eval_copy.as_mut().copy_from_slice(u_eval.as_ref());
264
265        // p0 * u
266        mod_op.elwise_mul_mut(u_eval.as_mut(), p0.as_ref());
267        // p1 * u
268        mod_op.elwise_mul_mut(u_eval_copy.as_mut(), p1.as_ref());
269        ntt_op.backward(u_eval.as_mut());
270        ntt_op.backward(u_eval_copy.as_mut());
271
272        // sample error
273        RandomFillGaussianInModulus::random_fill(rng, &q, ai.as_mut());
274        RandomFillGaussianInModulus::random_fill(rng, &q, bi.as_mut());
275
276        // a = p0*u+e0
277        mod_op.elwise_add_mut(ai.as_mut(), u_eval.as_ref());
278        // b = p1*u+e1
279        mod_op.elwise_add_mut(bi.as_mut(), u_eval_copy.as_ref());
280
281        // b = p1*u + e0 + \beta*m
282        // use u_eval as scratch
283        mod_op.elwise_scalar_mul(u_eval.as_mut(), m.as_ref(), beta_i);
284        mod_op.elwise_add_mut(bi.as_mut(), u_eval.as_ref());
285    });
286}
287
288/// Returns key switching key to key switch ciphertext RLWE_{from_s}(m)
289/// to RLWE_{to_s}(m).
290///
291/// Let key switching decomposer have `d` decompostion count with gadget vector:
292/// [1, \beta, ..., \beta^d-1]
293///
294/// Key switching key consists of `d` RLWE ciphertexts:
295/// RLWE'_{to_s}(-from_s) = [RLWE_{to_s}(\beta^i -from_s)]
296///
297/// In RLWE(m) s.t. b = as + e + m where s is the secret key, `a` can be seeded.
298/// And we seed all RLWE ciphertexts in key switchin key.
299///
300/// - neg_from_s: Negative of secret polynomial to key switch from (i.e.
301///   -from_s)
302/// - to_s: secret polynomial to key switch to.
303/// - gadget_vector: Gadget vector of decomposer used in key switch
304/// - p_rng: Seeded pseudo random generate used to generate `a` polynomials of
305///   key switching key RLWE ciphertexts
306fn seeded_rlwe_ksk_gen<
307    Mmut: MatrixMut + MatrixEntity,
308    ModOp: ArithmeticOps<Element = Mmut::MatElement>
309        + VectorOps<Element = Mmut::MatElement>
310        + GetModulus<Element = Mmut::MatElement>,
311    NttOp: Ntt<Element = Mmut::MatElement>,
312    R: RandomFillGaussianInModulus<[Mmut::MatElement], ModOp::M>,
313    PR: RandomFillUniformInModulus<[Mmut::MatElement], ModOp::M>,
314>(
315    ksk_out: &mut Mmut,
316    neg_from_s: Mmut::R,
317    mut to_s: Mmut::R,
318    gadget_vector: &[Mmut::MatElement],
319    mod_op: &ModOp,
320    ntt_op: &NttOp,
321    p_rng: &mut PR,
322    rng: &mut R,
323) where
324    <Mmut as Matrix>::R: RowMut,
325{
326    let ring_size = neg_from_s.as_ref().len();
327    let d = gadget_vector.len();
328    assert!(ksk_out.dimension() == (d, ring_size));
329
330    let q = mod_op.modulus();
331
332    ntt_op.forward(to_s.as_mut());
333
334    // RLWE'_{to_s}(-from_s)
335    let mut part_a = {
336        let mut a = Mmut::zeros(d, ring_size);
337        a.iter_rows_mut()
338            .for_each(|ai| RandomFillUniformInModulus::random_fill(p_rng, q, ai.as_mut()));
339        a
340    };
341    izip!(
342        part_a.iter_rows_mut(),
343        ksk_out.iter_rows_mut(),
344        gadget_vector.iter(),
345    )
346    .for_each(|(ai, bi, beta_i)| {
347        // si * ai
348        ntt_op.forward(ai.as_mut());
349        mod_op.elwise_mul_mut(ai.as_mut(), to_s.as_ref());
350        ntt_op.backward(ai.as_mut());
351
352        // ei + to_s*ai
353        RandomFillGaussianInModulus::random_fill(rng, &q, bi.as_mut());
354        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
355
356        // beta_i * -from_s
357        // use ai as scratch space
358        mod_op.elwise_scalar_mul(ai.as_mut(), neg_from_s.as_ref(), beta_i);
359
360        // bi = ei + to_s*ai + beta_i*-from_s
361        mod_op.elwise_add_mut(bi.as_mut(), ai.as_ref());
362    });
363}
364
365/// Returns auto key to send RLWE(m(X)) -> RLWE(m(X^k))
366///
367/// Auto key is key switchin key that key-switches RLWE_{s(X^k)}(m(X^k)) to
368/// RLWE_{s(X)}(m(X^k)).
369///
370/// - s: secret polynomial s(X)
371/// - auto_k: k used in for autmorphism X -> X^k
372/// - gadget_vector: Gadget vector corresponding to decomposer used in key
373///   switch
374/// - p_rng: pseudo random generator used to generate `a` polynomials of key
375///   switching key RLWE ciphertexts
376pub(crate) fn seeded_auto_key_gen<
377    Mmut: MatrixMut + MatrixEntity,
378    ModOp: ArithmeticOps<Element = Mmut::MatElement>
379        + VectorOps<Element = Mmut::MatElement>
380        + GetModulus<Element = Mmut::MatElement>,
381    NttOp: Ntt<Element = Mmut::MatElement>,
382    S,
383    R: RandomFillGaussianInModulus<[Mmut::MatElement], ModOp::M>,
384    PR: RandomFillUniformInModulus<[Mmut::MatElement], ModOp::M>,
385>(
386    ksk_out: &mut Mmut,
387    s: &[S],
388    auto_k: isize,
389    gadget_vector: &[Mmut::MatElement],
390    mod_op: &ModOp,
391    ntt_op: &NttOp,
392    p_rng: &mut PR,
393    rng: &mut R,
394) where
395    <Mmut as Matrix>::R: RowMut,
396    Mmut::R: TryConvertFrom1<[S], ModOp::M> + RowEntity,
397    Mmut::MatElement: Copy + Sub<Output = Mmut::MatElement>,
398{
399    let ring_size = s.len();
400    let (auto_map_index, auto_map_sign) = generate_auto_map(ring_size, auto_k);
401
402    let q = mod_op.modulus();
403
404    // s(X) -> -s(X^k)
405    let s = Mmut::R::try_convert_from(s, q);
406    let mut neg_s_auto = Mmut::R::zeros(s.as_ref().len());
407    izip!(s.as_ref(), auto_map_index.iter(), auto_map_sign.iter()).for_each(
408        |(el, to_index, sign)| {
409            // if sign is +ve (true), then negate because we need -s(X) (i.e. do the
410            // opposite than the usual case)
411            if *sign {
412                neg_s_auto.as_mut()[*to_index] = mod_op.neg(el);
413            } else {
414                neg_s_auto.as_mut()[*to_index] = *el;
415            }
416        },
417    );
418
419    // Ksk from -s(X^k) to s(X)
420    seeded_rlwe_ksk_gen(
421        ksk_out,
422        neg_s_auto,
423        s,
424        gadget_vector,
425        mod_op,
426        ntt_op,
427        p_rng,
428        rng,
429    );
430}
431
432/// Returns seeded RLWE(m(X))
433///
434/// RLWE(m(X)) = [a(X), b(X) = a(X)s(X) + e(X) + m(X)]
435///
436/// a(X) of RLWE encyrptions using secret key s(X) can be seeded. We use seeded
437/// pseudo random generator `p_rng` to sample a(X) and return seeded RLWE
438/// ciphertext (i.e. only b(X))
439pub(crate) fn seeded_secret_key_encrypt_rlwe<
440    Ro: Row + RowMut + RowEntity,
441    ModOp: VectorOps<Element = Ro::Element> + GetModulus<Element = Ro::Element>,
442    NttOp: Ntt<Element = Ro::Element>,
443    S,
444    R: RandomFillGaussianInModulus<[Ro::Element], ModOp::M>,
445    PR: RandomFillUniformInModulus<[Ro::Element], ModOp::M>,
446>(
447    m: &[Ro::Element],
448    b_rlwe_out: &mut Ro,
449    s: &[S],
450    mod_op: &ModOp,
451    ntt_op: &NttOp,
452    p_rng: &mut PR,
453    rng: &mut R,
454) where
455    Ro: TryConvertFrom1<[S], ModOp::M> + Debug,
456{
457    let ring_size = s.len();
458    assert!(m.as_ref().len() == ring_size);
459    assert!(b_rlwe_out.as_ref().len() == ring_size);
460
461    let q = mod_op.modulus();
462
463    // sample a
464    let mut a = {
465        let mut a = Ro::zeros(ring_size);
466        RandomFillUniformInModulus::random_fill(p_rng, q, a.as_mut());
467        a
468    };
469
470    // s * a
471    let mut sa = Ro::try_convert_from(s, q);
472    ntt_op.forward(sa.as_mut());
473    ntt_op.forward(a.as_mut());
474    mod_op.elwise_mul_mut(sa.as_mut(), a.as_ref());
475    ntt_op.backward(sa.as_mut());
476
477    // sample e
478    RandomFillGaussianInModulus::random_fill(rng, q, b_rlwe_out.as_mut());
479    mod_op.elwise_add_mut(b_rlwe_out.as_mut(), m.as_ref());
480    mod_op.elwise_add_mut(b_rlwe_out.as_mut(), sa.as_ref());
481}
482
483/// Returns RLWE(m(X)) encrypted using public key.
484///
485/// Unlike secret key encryption, public key encryption cannot be seeded
486pub(crate) fn public_key_encrypt_rlwe<
487    M: Matrix,
488    Mmut: MatrixMut<MatElement = M::MatElement>,
489    ModOp: VectorOps<Element = M::MatElement> + GetModulus<Element = M::MatElement>,
490    NttOp: Ntt<Element = M::MatElement>,
491    S,
492    R: RandomFillGaussianInModulus<[M::MatElement], ModOp::M>
493        + RandomFillUniformInModulus<[M::MatElement], ModOp::M>
494        + RandomFill<[u8]>
495        + RandomElementInModulus<usize, usize>,
496>(
497    rlwe_out: &mut Mmut,
498    pk: &M,
499    m: &[M::MatElement],
500    mod_op: &ModOp,
501    ntt_op: &NttOp,
502    rng: &mut R,
503) where
504    <Mmut as Matrix>::R: RowMut + TryConvertFrom1<[S], ModOp::M> + RowEntity,
505    M::MatElement: Copy,
506    S: Zero + Signed + Copy,
507{
508    let ring_size = m.len();
509    assert!(rlwe_out.dimension() == (2, ring_size));
510
511    let q = mod_op.modulus();
512
513    let mut u = vec![S::zero(); ring_size];
514    fill_random_ternary_secret_with_hamming_weight(u.as_mut(), ring_size >> 1, rng);
515    let mut u = Mmut::R::try_convert_from(&u, q);
516    ntt_op.forward(u.as_mut());
517
518    let mut ua = Mmut::R::zeros(ring_size);
519    ua.as_mut().copy_from_slice(pk.get_row_slice(0));
520    let mut ub = Mmut::R::zeros(ring_size);
521    ub.as_mut().copy_from_slice(pk.get_row_slice(1));
522
523    // a*u
524    ntt_op.forward(ua.as_mut());
525    mod_op.elwise_mul_mut(ua.as_mut(), u.as_ref());
526    ntt_op.backward(ua.as_mut());
527
528    // b*u
529    ntt_op.forward(ub.as_mut());
530    mod_op.elwise_mul_mut(ub.as_mut(), u.as_ref());
531    ntt_op.backward(ub.as_mut());
532
533    // sample error
534    rlwe_out.iter_rows_mut().for_each(|ri| {
535        RandomFillGaussianInModulus::random_fill(rng, &q, ri.as_mut());
536    });
537
538    // a*u + e0
539    mod_op.elwise_add_mut(rlwe_out.get_row_mut(0), ua.as_ref());
540    // b*u + e1
541    mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), ub.as_ref());
542
543    // b*u + e1 + m
544    mod_op.elwise_add_mut(rlwe_out.get_row_mut(1), m);
545}
546
547/// Returns RLWE public key generated using RLWE secret key
548pub(crate) fn rlwe_public_key<
549    Ro: RowMut + RowEntity,
550    S,
551    ModOp: VectorOps<Element = Ro::Element> + GetModulus<Element = Ro::Element>,
552    NttOp: Ntt<Element = Ro::Element>,
553    PRng: RandomFillUniformInModulus<[Ro::Element], ModOp::M>,
554    Rng: RandomFillGaussianInModulus<[Ro::Element], ModOp::M>,
555>(
556    part_b_out: &mut Ro,
557    s: &[S],
558    ntt_op: &NttOp,
559    mod_op: &ModOp,
560    p_rng: &mut PRng,
561    rng: &mut Rng,
562) where
563    Ro: TryConvertFrom1<[S], ModOp::M>,
564{
565    let ring_size = s.len();
566    assert!(part_b_out.as_ref().len() == ring_size);
567
568    let q = mod_op.modulus();
569
570    // sample a
571    let mut a = {
572        let mut tmp = Ro::zeros(ring_size);
573        RandomFillUniformInModulus::random_fill(p_rng, &q, tmp.as_mut());
574        tmp
575    };
576    ntt_op.forward(a.as_mut());
577
578    // s*a
579    let mut sa = Ro::try_convert_from(s, &q);
580    ntt_op.forward(sa.as_mut());
581    mod_op.elwise_mul_mut(sa.as_mut(), a.as_ref());
582    ntt_op.backward(sa.as_mut());
583
584    // s*a + e
585    RandomFillGaussianInModulus::random_fill(rng, &q, part_b_out.as_mut());
586    mod_op.elwise_add_mut(part_b_out.as_mut(), sa.as_ref());
587}
588
589/// Decrypts ciphertext RLWE(m) and returns noisy m
590///
591/// We assume RLWE(m) = [a, b] is a degree 1 ciphertext s.t. b - sa = e + m
592pub(crate) fn decrypt_rlwe<
593    R: RowMut,
594    M: Matrix<MatElement = R::Element>,
595    ModOp: VectorOps<Element = R::Element> + GetModulus<Element = R::Element>,
596    NttOp: Ntt<Element = R::Element>,
597    S,
598>(
599    rlwe_ct: &M,
600    s: &[S],
601    m_out: &mut R,
602    ntt_op: &NttOp,
603    mod_op: &ModOp,
604) where
605    R: TryConvertFrom1<[S], ModOp::M>,
606    R::Element: Copy,
607{
608    let ring_size = s.len();
609    assert!(rlwe_ct.dimension() == (2, ring_size));
610    assert!(m_out.as_ref().len() == ring_size);
611
612    // transform a to evluation form
613    m_out.as_mut().copy_from_slice(rlwe_ct.get_row_slice(0));
614    ntt_op.forward(m_out.as_mut());
615
616    // -s*a
617    let mut s = R::try_convert_from(&s, mod_op.modulus());
618    ntt_op.forward(s.as_mut());
619    mod_op.elwise_mul_mut(m_out.as_mut(), s.as_ref());
620    mod_op.elwise_neg_mut(m_out.as_mut());
621    ntt_op.backward(m_out.as_mut());
622
623    // m+e = b - s*a
624    mod_op.elwise_add_mut(m_out.as_mut(), rlwe_ct.get_row_slice(1));
625}
626
627// Measures maximum noise in degree 1 RLWE ciphertext against message `want_m`
628pub(crate) fn measure_max_noise<
629    Mmut: MatrixMut + Matrix,
630    ModOp: VectorOps<Element = Mmut::MatElement> + GetModulus<Element = Mmut::MatElement>,
631    NttOp: Ntt<Element = Mmut::MatElement>,
632    S,
633>(
634    rlwe_ct: &Mmut,
635    want_m: &Mmut::R,
636    ntt_op: &NttOp,
637    mod_op: &ModOp,
638    s: &[S],
639) -> f64
640where
641    <Mmut as Matrix>::R: RowMut,
642    Mmut::R: RowEntity + TryConvertFrom1<[S], ModOp::M>,
643    Mmut::MatElement: PrimInt + ToPrimitive + Debug,
644{
645    let ring_size = s.len();
646    assert!(rlwe_ct.dimension() == (2, ring_size));
647    assert!(want_m.as_ref().len() == ring_size);
648
649    // -(s * a)
650    let q = mod_op.modulus();
651    let mut s = Mmut::R::try_convert_from(s, &q);
652    ntt_op.forward(s.as_mut());
653    let mut a = Mmut::R::zeros(ring_size);
654    a.as_mut().copy_from_slice(rlwe_ct.get_row_slice(0));
655    ntt_op.forward(a.as_mut());
656    mod_op.elwise_mul_mut(s.as_mut(), a.as_ref());
657    mod_op.elwise_neg_mut(s.as_mut());
658    ntt_op.backward(s.as_mut());
659
660    // m+e = b - s*a
661    let mut m_plus_e = s;
662    mod_op.elwise_add_mut(m_plus_e.as_mut(), rlwe_ct.get_row_slice(1));
663
664    // difference
665    mod_op.elwise_sub_mut(m_plus_e.as_mut(), want_m.as_ref());
666
667    let mut max_diff_bits = f64::MIN;
668    m_plus_e.as_ref().iter().for_each(|v| {
669        let bits = (q.map_element_to_i64(v).to_f64().unwrap().abs()).log2();
670
671        if max_diff_bits < bits {
672            max_diff_bits = bits;
673        }
674    });
675
676    return max_diff_bits;
677}