1use std::fmt::Debug;
  2
  3use itertools::izip;
  4use num_traits::Zero;
  5
  6use crate::{
  7    backend::{GetModulus, Modulus, VectorOps},
  8    ntt::Ntt,
  9    random::{
 10        RandomFillGaussianInModulus, RandomFillUniformInModulus, RandomGaussianElementInModulus,
 11    },
 12    utils::TryConvertFrom1,
 13    ArithmeticOps, Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut,
 14};
 15
 16pub(crate) fn public_key_share<
 17    R: Row + RowMut + RowEntity,
 18    S,
 19    ModOp: VectorOps<Element = R::Element> + GetModulus<Element = R::Element>,
 20    NttOp: Ntt<Element = R::Element>,
 21    Rng: RandomFillGaussianInModulus<[R::Element], ModOp::M>,
 22    PRng: RandomFillUniformInModulus<[R::Element], ModOp::M>,
 23>(
 24    share_out: &mut R,
 25    s_i: &[S],
 26    modop: &ModOp,
 27    nttop: &NttOp,
 28    p_rng: &mut PRng,
 29    rng: &mut Rng,
 30) where
 31    R: TryConvertFrom1<[S], ModOp::M>,
 32{
 33    let ring_size = share_out.as_ref().len();
 34    assert!(s_i.len() == ring_size);
 35
 36    let q = modop.modulus();
 37
 38    // sample a
 39    let mut a = {
 40        let mut a = R::zeros(ring_size);
 41        RandomFillUniformInModulus::random_fill(p_rng, &q, a.as_mut());
 42        a
 43    };
 44
 45    // s*a
 46    nttop.forward(a.as_mut());
 47    let mut s = R::try_convert_from(s_i, &q);
 48    nttop.forward(s.as_mut());
 49    modop.elwise_mul_mut(s.as_mut(), a.as_ref());
 50    nttop.backward(s.as_mut());
 51
 52    RandomFillGaussianInModulus::random_fill(rng, &q, share_out.as_mut());
 53    modop.elwise_add_mut(share_out.as_mut(), s.as_ref()); // s*e + e
 54}
 55
 56/// Generate decryption share for LWE ciphertext `lwe_ct` with user's secret `s`
 57pub(crate) fn multi_party_decryption_share<
 58    R: RowMut + RowEntity,
 59    Mod: Modulus<Element = R::Element>,
 60    ModOp: ArithmeticOps<Element = R::Element> + VectorOps<Element = R::Element> + GetModulus<M = Mod>,
 61    Rng: RandomGaussianElementInModulus<R::Element, Mod>,
 62    S,
 63>(
 64    lwe_ct: &R,
 65    s: &[S],
 66    mod_op: &ModOp,
 67    rng: &mut Rng,
 68) -> R::Element
 69where
 70    R: TryConvertFrom1<[S], Mod>,
 71    R::Element: Zero,
 72{
 73    assert!(lwe_ct.as_ref().len() == s.len() + 1);
 74    let mut neg_s = R::try_convert_from(s, mod_op.modulus());
 75    mod_op.elwise_neg_mut(neg_s.as_mut());
 76
 77    // share =  (\sum -s_i * a_i) + e
 78    let mut share = R::Element::zero();
 79    izip!(neg_s.as_ref().iter(), lwe_ct.as_ref().iter().skip(1)).for_each(|(si, ai)| {
 80        share = mod_op.add(&share, &mod_op.mul(si, ai));
 81    });
 82
 83    let e = rng.random(mod_op.modulus());
 84    share = mod_op.add(&share, &e);
 85
 86    share
 87}
 88
 89/// Aggregate decryption shares for `lwe_ct` and return noisy decryption output
 90/// `m + e`
 91pub(crate) fn multi_party_aggregate_decryption_shares_and_decrypt<
 92    R: RowMut + RowEntity,
 93    ModOp: ArithmeticOps<Element = R::Element>,
 94>(
 95    lwe_ct: &R,
 96    shares: &[R::Element],
 97    mod_op: &ModOp,
 98) -> R::Element
 99where
100    R::Element: Zero,
101{
102    let mut sum_shares = R::Element::zero();
103    shares
104        .iter()
105        .for_each(|v| sum_shares = mod_op.add(&sum_shares, v));
106    mod_op.add(&lwe_ct.as_ref()[0], &sum_shares)
107}
108
109pub(crate) fn non_interactive_rgsw_ct<
110    M: MatrixMut + MatrixEntity,
111    S,
112    PRng: RandomFillUniformInModulus<[M::MatElement], ModOp::M>,
113    Rng: RandomFillGaussianInModulus<[M::MatElement], ModOp::M>,
114    NttOp: Ntt<Element = M::MatElement>,
115    ModOp: VectorOps<Element = M::MatElement> + GetModulus<Element = M::MatElement>,
116>(
117    s: &[S],
118    u: &[S],
119    m: &[M::MatElement],
120    gadget_vec: &[M::MatElement],
121    p_rng: &mut PRng,
122    rng: &mut Rng,
123    nttop: &NttOp,
124    modop: &ModOp,
125) -> (M, M)
126where
127    <M as Matrix>::R: RowMut + TryConvertFrom1<[S], ModOp::M> + RowEntity,
128    M::MatElement: Copy,
129{
130    assert_eq!(s.len(), u.len());
131    assert_eq!(s.len(), m.len());
132    let q = modop.modulus();
133    let d = gadget_vec.len();
134    let ring_size = s.len();
135
136    let mut s_poly_eval = M::R::try_convert_from(s, q);
137    let mut u_poly_eval = M::R::try_convert_from(u, q);
138    nttop.forward(s_poly_eval.as_mut());
139    nttop.forward(u_poly_eval.as_mut());
140
141    // encryptions of a_i*u + e + \beta m
142    let mut enc_beta_m = M::zeros(d, ring_size);
143    // zero encrypition: a_i*s + e'
144    let mut zero_encryptions = M::zeros(d, ring_size);
145
146    let mut scratch_space = M::R::zeros(ring_size);
147
148    izip!(
149        enc_beta_m.iter_rows_mut(),
150        zero_encryptions.iter_rows_mut(),
151        gadget_vec.iter()
152    )
153    .for_each(|(e_beta_m, e_zero, beta)| {
154        // sample a_i
155        RandomFillUniformInModulus::random_fill(p_rng, q, e_beta_m.as_mut());
156        e_zero.as_mut().copy_from_slice(e_beta_m.as_ref());
157
158        // a_i * u + \beta m + e //
159        // a_i * u
160        nttop.forward(e_beta_m.as_mut());
161        modop.elwise_mul_mut(e_beta_m.as_mut(), u_poly_eval.as_ref());
162        nttop.backward(e_beta_m.as_mut());
163        // sample error e
164        RandomFillGaussianInModulus::random_fill(rng, q, scratch_space.as_mut());
165        // a_i * u + e
166        modop.elwise_add_mut(e_beta_m.as_mut(), scratch_space.as_ref());
167        // beta * m
168        modop.elwise_scalar_mul(scratch_space.as_mut(), m.as_ref(), beta);
169        // a_i * u + e + \beta m
170        modop.elwise_add_mut(e_beta_m.as_mut(), scratch_space.as_ref());
171
172        // a_i * s + e //
173        // a_i * s
174        nttop.forward(e_zero.as_mut());
175        modop.elwise_mul_mut(e_zero.as_mut(), s_poly_eval.as_ref());
176        nttop.backward(e_zero.as_mut());
177        // sample error e
178        RandomFillGaussianInModulus::random_fill(rng, q, scratch_space.as_mut());
179        // a_i * s + e
180        modop.elwise_add_mut(e_zero.as_mut(), scratch_space.as_ref());
181    });
182
183    (enc_beta_m, zero_encryptions)
184}
185
186pub(crate) fn non_interactive_ksk_gen<
187    M: MatrixMut + MatrixEntity,
188    S,
189    PRng: RandomFillUniformInModulus<[M::MatElement], ModOp::M>,
190    Rng: RandomFillGaussianInModulus<[M::MatElement], ModOp::M>,
191    NttOp: Ntt<Element = M::MatElement>,
192    ModOp: VectorOps<Element = M::MatElement> + GetModulus<Element = M::MatElement>,
193>(
194    s: &[S],
195    u: &[S],
196    gadget_vec: &[M::MatElement],
197    p_rng: &mut PRng,
198    rng: &mut Rng,
199    nttop: &NttOp,
200    modop: &ModOp,
201) -> M
202where
203    <M as Matrix>::R: RowMut + TryConvertFrom1<[S], ModOp::M> + RowEntity,
204    M::MatElement: Copy + Debug,
205{
206    assert_eq!(s.len(), u.len());
207
208    let q = modop.modulus();
209    let d = gadget_vec.len();
210    let ring_size = s.len();
211
212    let mut s_poly_eval = M::R::try_convert_from(s, q);
213    nttop.forward(s_poly_eval.as_mut());
214    let u_poly = M::R::try_convert_from(u, q);
215    // a_i * s + \beta u + e
216    let mut ksk = M::zeros(d, ring_size);
217
218    let mut scratch_space = M::R::zeros(ring_size);
219
220    izip!(ksk.iter_rows_mut(), gadget_vec.iter()).for_each(|(e_ksk, beta)| {
221        // sample a_i
222        RandomFillUniformInModulus::random_fill(p_rng, q, e_ksk.as_mut());
223
224        // a_i * s + e + beta u
225        nttop.forward(e_ksk.as_mut());
226        modop.elwise_mul_mut(e_ksk.as_mut(), s_poly_eval.as_ref());
227        nttop.backward(e_ksk.as_mut());
228        // sample error e
229        RandomFillGaussianInModulus::random_fill(rng, q, scratch_space.as_mut());
230        // a_i * s + e
231        modop.elwise_add_mut(e_ksk.as_mut(), scratch_space.as_ref());
232        // \beta * u
233        modop.elwise_scalar_mul(scratch_space.as_mut(), u_poly.as_ref(), beta);
234        // a_i * s + e + \beta * u
235        modop.elwise_add_mut(e_ksk.as_mut(), scratch_space.as_ref());
236    });
237
238    ksk
239}
240
241pub(crate) fn non_interactive_ksk_zero_encryptions_for_other_party_i<
242    M: MatrixMut + MatrixEntity,
243    S,
244    PRng: RandomFillUniformInModulus<[M::MatElement], ModOp::M>,
245    Rng: RandomFillGaussianInModulus<[M::MatElement], ModOp::M>,
246    NttOp: Ntt<Element = M::MatElement>,
247    ModOp: VectorOps<Element = M::MatElement> + GetModulus<Element = M::MatElement>,
248>(
249    s: &[S],
250    gadget_vec: &[M::MatElement],
251    p_rng: &mut PRng,
252    rng: &mut Rng,
253    nttop: &NttOp,
254    modop: &ModOp,
255) -> M
256where
257    <M as Matrix>::R: RowMut + TryConvertFrom1<[S], ModOp::M> + RowEntity,
258    M::MatElement: Copy + Debug,
259{
260    let q = modop.modulus();
261    let d = gadget_vec.len();
262    let ring_size = s.len();
263
264    let mut s_poly_eval = M::R::try_convert_from(s, q);
265    nttop.forward(s_poly_eval.as_mut());
266
267    // a_i * s + e
268    let mut zero_encs = M::zeros(d, ring_size);
269
270    let mut scratch_space = M::R::zeros(ring_size);
271
272    izip!(zero_encs.iter_rows_mut()).for_each(|e_zero| {
273        // sample a_i
274        RandomFillUniformInModulus::random_fill(p_rng, q, e_zero.as_mut());
275
276        // a_i * s + e
277        nttop.forward(e_zero.as_mut());
278        modop.elwise_mul_mut(e_zero.as_mut(), s_poly_eval.as_ref());
279        nttop.backward(e_zero.as_mut());
280        // sample error e
281        RandomFillGaussianInModulus::random_fill(rng, q, scratch_space.as_mut());
282        modop.elwise_add_mut(e_zero.as_mut(), scratch_space.as_ref());
283    });
284
285    zero_encs
286}