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