1use std::{usize, vec};
  2
  3use itertools::{izip, Itertools};
  4use num_traits::{One, PrimInt, Signed};
  5
  6use crate::{
  7    backend::Modulus,
  8    decomposer::NumInfo,
  9    random::{RandomElementInModulus, RandomFill},
 10    Matrix, RowEntity, RowMut,
 11};
 12pub trait WithLocal {
 13    fn with_local<F, R>(func: F) -> R
 14    where
 15        F: Fn(&Self) -> R;
 16
 17    fn with_local_mut<F, R>(func: F) -> R
 18    where
 19        F: Fn(&mut Self) -> R;
 20
 21    fn with_local_mut_mut<F, R>(func: &mut F) -> R
 22    where
 23        F: FnMut(&mut Self) -> R;
 24}
 25
 26pub trait Global {
 27    fn global() -> &'static Self;
 28}
 29
 30pub(crate) trait ShoupMul {
 31    fn representation(value: Self, q: Self) -> Self;
 32    fn mul(a: Self, b: Self, b_shoup: Self, q: Self) -> Self;
 33}
 34
 35impl ShoupMul for u64 {
 36    #[inline]
 37    fn representation(value: Self, q: Self) -> Self {
 38        ((value as u128 * (1u128 << 64)) / q as u128) as u64
 39    }
 40
 41    #[inline]
 42    /// Returns a * b % q
 43    fn mul(a: Self, b: Self, b_shoup: Self, q: Self) -> Self {
 44        (b.wrapping_mul(a))
 45            .wrapping_sub(q.wrapping_mul(((b_shoup as u128 * a as u128) >> 64) as u64))
 46    }
 47}
 48
 49pub(crate) trait ToShoup {
 50    type Modulus;
 51    fn to_shoup(value: &Self, modulus: Self::Modulus) -> Self;
 52}
 53
 54impl ToShoup for u64 {
 55    type Modulus = u64;
 56    fn to_shoup(value: &Self, modulus: Self) -> Self {
 57        ((*value as u128 * (1u128 << 64)) / modulus as u128) as u64
 58    }
 59}
 60
 61impl ToShoup for Vec<Vec<u64>> {
 62    type Modulus = u64;
 63    fn to_shoup(value: &Self, modulus: Self::Modulus) -> Self {
 64        let (row, col) = value.dimension();
 65        let mut shoup_value = vec![vec![0u64; col]; row];
 66        izip!(shoup_value.iter_mut(), value.iter()).for_each(|(shoup_r, r)| {
 67            izip!(shoup_r.iter_mut(), r.iter()).for_each(|(s, e)| {
 68                *s = u64::to_shoup(e, modulus);
 69            })
 70        });
 71        shoup_value
 72    }
 73}
 74
 75pub fn fill_random_ternary_secret_with_hamming_weight<
 76    T: Signed,
 77    R: RandomFill<[u8]> + RandomElementInModulus<usize, usize>,
 78>(
 79    out: &mut [T],
 80    hamming_weight: usize,
 81    rng: &mut R,
 82) {
 83    let mut bytes = vec![0u8; hamming_weight.div_ceil(8)];
 84    RandomFill::<[u8]>::random_fill(rng, &mut bytes);
 85
 86    let size = out.len();
 87    let mut secret_indices = (0..size).into_iter().map(|i| i).collect_vec();
 88    let mut bit_index = 0;
 89    let mut byte_index = 0;
 90    for i in 0..hamming_weight {
 91        let s_index = RandomElementInModulus::<usize, usize>::random(rng, &secret_indices.len());
 92
 93        let curr_bit = (bytes[byte_index] >> bit_index) & 1;
 94        if curr_bit == 1 {
 95            out[secret_indices[s_index]] = T::one();
 96        } else {
 97            out[secret_indices[s_index]] = -T::one();
 98        }
 99
100        secret_indices[s_index] = *secret_indices.last().unwrap();
101        secret_indices.truncate(secret_indices.len() - 1);
102
103        if bit_index == 7 {
104            bit_index = 0;
105            byte_index += 1;
106        } else {
107            bit_index += 1;
108        }
109    }
110}
111
112// TODO (Jay): this is only a workaround. Add a propoer way to perform primality
113// tests.
114fn is_probably_prime(candidate: u64) -> bool {
115    num_bigint_dig::prime::probably_prime(&num_bigint_dig::BigUint::from(candidate), 0)
116}
117
118/// Finds prime that satisfy
119/// - $prime \lt upper_bound$
120/// - $\log{prime} = num_bits$
121/// - `prime % modulo == 1`
122pub(crate) fn generate_prime(num_bits: usize, modulo: u64, upper_bound: u64) -> Option<u64> {
123    let leading_zeros = (64 - num_bits) as u32;
124
125    let mut tentative_prime = upper_bound - 1;
126    while tentative_prime % modulo != 1 && tentative_prime.leading_zeros() == leading_zeros {
127        tentative_prime -= 1;
128    }
129
130    while !is_probably_prime(tentative_prime)
131        && tentative_prime.leading_zeros() == leading_zeros
132        && tentative_prime >= modulo
133    {
134        tentative_prime -= modulo;
135    }
136
137    if is_probably_prime(tentative_prime) && tentative_prime.leading_zeros() == leading_zeros {
138        Some(tentative_prime)
139    } else {
140        None
141    }
142}
143
144/// Returns a^b mod q
145pub fn mod_exponent(a: u64, mut b: u64, q: u64) -> u64 {
146    let mod_mul = |v1: &u64, v2: &u64| {
147        let tmp = *v1 as u128 * *v2 as u128;
148        (tmp % q as u128) as u64
149    };
150
151    let mut acc = a;
152    let mut out = 1;
153    while b != 0 {
154        let flag = b & 1;
155
156        if flag == 1 {
157            out = mod_mul(&acc, &out);
158        }
159        acc = mod_mul(&acc, &acc);
160
161        b >>= 1;
162    }
163
164    out
165}
166
167pub(crate) fn mod_inverse(a: u64, q: u64) -> u64 {
168    mod_exponent(a, q - 2, q)
169}
170
171pub(crate) fn negacyclic_mul<T: PrimInt, F: Fn(&T, &T) -> T>(
172    a: &[T],
173    b: &[T],
174    mul: F,
175    modulus: T,
176) -> Vec<T> {
177    let mut r = vec![T::zero(); a.len()];
178    for i in 0..a.len() {
179        for j in 0..i + 1 {
180            // println!("i: {j} {}", i - j);
181            r[i] = (r[i] + mul(&a[j], &b[i - j])) % modulus;
182        }
183
184        for j in i + 1..a.len() {
185            // println!("i: {j} {}", a.len() - j + i);
186            r[i] = (r[i] + modulus - mul(&a[j], &b[a.len() - j + i])) % modulus;
187        }
188        // println!("")
189    }
190
191    return r;
192}
193
194/// Returns a polynomial X^{emebedding_factor * si} \mod {Z_Q / X^{N}+1}
195pub(crate) fn encode_x_pow_si_with_emebedding_factor<
196    R: RowEntity + RowMut,
197    M: Modulus<Element = R::Element>,
198>(
199    si: i32,
200    embedding_factor: usize,
201    ring_size: usize,
202    modulus: &M,
203) -> R
204where
205    R::Element: One,
206{
207    assert!((si.abs() as usize) < ring_size);
208    let mut m = R::zeros(ring_size);
209    let si = si * (embedding_factor as i32);
210    if si < 0 {
211        // X^{-si} = X^{2N-si} = -X^{N-si}, assuming abs(si) < N
212        m.as_mut()[ring_size - (si.abs() as usize)] = modulus.neg_one();
213    } else {
214        m.as_mut()[si as usize] = R::Element::one();
215    }
216    m
217}
218
219pub(crate) fn puncture_p_rng<S: Default + Copy, R: RandomFill<S>>(
220    p_rng: &mut R,
221    times: usize,
222) -> S {
223    let mut out = S::default();
224    for _ in 0..times {
225        RandomFill::<S>::random_fill(p_rng, &mut out);
226    }
227    return out;
228}
229
230pub(crate) fn log2<T: PrimInt + NumInfo>(v: &T) -> usize {
231    if (*v & (*v - T::one())) == T::zero() {
232        // value is power of 2
233        (T::BITS - v.leading_zeros() - 1) as usize
234    } else {
235        (T::BITS - v.leading_zeros()) as usize
236    }
237}
238
239pub trait TryConvertFrom1<T: ?Sized, P> {
240    fn try_convert_from(value: &T, parameters: &P) -> Self;
241}
242
243impl<P: Modulus<Element = u64>> TryConvertFrom1<[i64], P> for Vec<u64> {
244    fn try_convert_from(value: &[i64], parameters: &P) -> Self {
245        value
246            .iter()
247            .map(|v| parameters.map_element_from_i64(*v))
248            .collect_vec()
249    }
250}
251
252impl<P: Modulus<Element = u64>> TryConvertFrom1<[i32], P> for Vec<u64> {
253    fn try_convert_from(value: &[i32], parameters: &P) -> Self {
254        value
255            .iter()
256            .map(|v| parameters.map_element_from_i64(*v as i64))
257            .collect_vec()
258    }
259}
260
261impl<P: Modulus> TryConvertFrom1<[P::Element], P> for Vec<i64> {
262    fn try_convert_from(value: &[P::Element], parameters: &P) -> Self {
263        value
264            .iter()
265            .map(|v| parameters.map_element_to_i64(v))
266            .collect_vec()
267    }
268}
269
270#[cfg(test)]
271pub(crate) mod tests {
272    use std::fmt::Debug;
273
274    use num_traits::ToPrimitive;
275
276    use crate::random::DefaultSecureRng;
277
278    use super::fill_random_ternary_secret_with_hamming_weight;
279
280    #[derive(Clone)]
281    pub(crate) struct Stats<T> {
282        pub(crate) samples: Vec<T>,
283    }
284
285    impl<T> Default for Stats<T> {
286        fn default() -> Self {
287            Stats { samples: vec![] }
288        }
289    }
290
291    impl<T: Copy + ToPrimitive + Debug> Stats<T>
292    where
293        // T: for<'a> Sum<&'a T>,
294        T: for<'a> std::iter::Sum<&'a T> + std::iter::Sum<T>,
295    {
296        pub(crate) fn new() -> Self {
297            Self { samples: vec![] }
298        }
299
300        pub(crate) fn mean(&self) -> f64 {
301            self.samples.iter().sum::<T>().to_f64().unwrap() / (self.samples.len() as f64)
302        }
303
304        pub(crate) fn variance(&self) -> f64 {
305            let mean = self.mean();
306
307            // diff
308            let diff_sq = self
309                .samples
310                .iter()
311                .map(|v| {
312                    let t = v.to_f64().unwrap() - mean;
313                    t * t
314                })
315                .into_iter()
316                .sum::<f64>();
317
318            diff_sq / (self.samples.len() as f64 - 1.0)
319        }
320
321        pub(crate) fn std_dev(&self) -> f64 {
322            self.variance().sqrt()
323        }
324
325        pub(crate) fn add_many_samples(&mut self, values: &[T]) {
326            self.samples.extend(values.iter());
327        }
328
329        pub(crate) fn add_sample(&mut self, value: T) {
330            self.samples.push(value)
331        }
332
333        pub(crate) fn merge_in(&mut self, other: &Self) {
334            self.samples.extend(other.samples.iter());
335        }
336    }
337
338    #[test]
339    fn ternary_secret_has_correct_hw() {
340        let mut rng = DefaultSecureRng::new();
341        for n in 4..15 {
342            let ring_size = 1 << n;
343            let mut out = vec![0i32; ring_size];
344            fill_random_ternary_secret_with_hamming_weight(&mut out, ring_size >> 1, &mut rng);
345
346            // check hamming weight of out equals ring_size/2
347            let mut non_zeros = 0;
348            out.iter().for_each(|i| {
349                if *i != 0 {
350                    non_zeros += 1;
351                }
352            });
353
354            assert_eq!(ring_size >> 1, non_zeros);
355        }
356    }
357}