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.
114#[cfg(test)]
115fn is_probably_prime(candidate: u64) -> bool {
116    num_bigint_dig::prime::probably_prime(&num_bigint_dig::BigUint::from(candidate), 0)
117}
118
119/// Finds prime that satisfy
120/// - $prime \lt upper_bound$
121/// - $\log{prime} = num_bits$
122/// - `prime % modulo == 1`
123#[cfg(test)]
124pub(crate) fn generate_prime(num_bits: usize, modulo: u64, upper_bound: u64) -> Option<u64> {
125    let leading_zeros = (64 - num_bits) as u32;
126
127    let mut tentative_prime = upper_bound - 1;
128    while tentative_prime % modulo != 1 && tentative_prime.leading_zeros() == leading_zeros {
129        tentative_prime -= 1;
130    }
131
132    while !is_probably_prime(tentative_prime)
133        && tentative_prime.leading_zeros() == leading_zeros
134        && tentative_prime >= modulo
135    {
136        tentative_prime -= modulo;
137    }
138
139    if is_probably_prime(tentative_prime) && tentative_prime.leading_zeros() == leading_zeros {
140        Some(tentative_prime)
141    } else {
142        None
143    }
144}
145
146/// Returns a^b mod q
147pub fn mod_exponent(a: u64, mut b: u64, q: u64) -> u64 {
148    let mod_mul = |v1: &u64, v2: &u64| {
149        let tmp = *v1 as u128 * *v2 as u128;
150        (tmp % q as u128) as u64
151    };
152
153    let mut acc = a;
154    let mut out = 1;
155    while b != 0 {
156        let flag = b & 1;
157
158        if flag == 1 {
159            out = mod_mul(&acc, &out);
160        }
161        acc = mod_mul(&acc, &acc);
162
163        b >>= 1;
164    }
165
166    out
167}
168
169pub(crate) fn mod_inverse(a: u64, q: u64) -> u64 {
170    mod_exponent(a, q - 2, q)
171}
172
173pub(crate) fn negacyclic_mul<T: PrimInt, F: Fn(&T, &T) -> T>(
174    a: &[T],
175    b: &[T],
176    mul: F,
177    modulus: T,
178) -> Vec<T> {
179    let mut r = vec![T::zero(); a.len()];
180    for i in 0..a.len() {
181        for j in 0..i + 1 {
182            // println!("i: {j} {}", i - j);
183            r[i] = (r[i] + mul(&a[j], &b[i - j])) % modulus;
184        }
185
186        for j in i + 1..a.len() {
187            // println!("i: {j} {}", a.len() - j + i);
188            r[i] = (r[i] + modulus - mul(&a[j], &b[a.len() - j + i])) % modulus;
189        }
190        // println!("")
191    }
192
193    return r;
194}
195
196/// Returns a polynomial X^{emebedding_factor * si} \mod {Z_Q / X^{N}+1}
197pub(crate) fn encode_x_pow_si_with_emebedding_factor<
198    R: RowEntity + RowMut,
199    M: Modulus<Element = R::Element>,
200>(
201    si: i32,
202    embedding_factor: usize,
203    ring_size: usize,
204    modulus: &M,
205) -> R
206where
207    R::Element: One,
208{
209    assert!((si.abs() as usize) < ring_size);
210    let mut m = R::zeros(ring_size);
211    let si = si * (embedding_factor as i32);
212    if si < 0 {
213        // X^{-si} = X^{2N-si} = -X^{N-si}, assuming abs(si) < N
214        m.as_mut()[ring_size - (si.abs() as usize)] = modulus.neg_one();
215    } else {
216        m.as_mut()[si as usize] = R::Element::one();
217    }
218    m
219}
220
221pub(crate) fn puncture_p_rng<S: Default + Copy, R: RandomFill<S>>(
222    p_rng: &mut R,
223    times: usize,
224) -> S {
225    let mut out = S::default();
226    for _ in 0..times {
227        RandomFill::<S>::random_fill(p_rng, &mut out);
228    }
229    return out;
230}
231
232pub(crate) fn log2<T: PrimInt + NumInfo>(v: &T) -> usize {
233    if (*v & (*v - T::one())) == T::zero() {
234        // value is power of 2
235        (T::BITS - v.leading_zeros() - 1) as usize
236    } else {
237        (T::BITS - v.leading_zeros()) as usize
238    }
239}
240
241pub trait TryConvertFrom1<T: ?Sized, P> {
242    fn try_convert_from(value: &T, parameters: &P) -> Self;
243}
244
245impl<P: Modulus<Element = u64>> TryConvertFrom1<[i64], P> for Vec<u64> {
246    fn try_convert_from(value: &[i64], parameters: &P) -> Self {
247        value
248            .iter()
249            .map(|v| parameters.map_element_from_i64(*v))
250            .collect_vec()
251    }
252}
253
254impl<P: Modulus<Element = u64>> TryConvertFrom1<[i32], P> for Vec<u64> {
255    fn try_convert_from(value: &[i32], parameters: &P) -> Self {
256        value
257            .iter()
258            .map(|v| parameters.map_element_from_i64(*v as i64))
259            .collect_vec()
260    }
261}
262
263impl<P: Modulus> TryConvertFrom1<[P::Element], P> for Vec<i64> {
264    fn try_convert_from(value: &[P::Element], parameters: &P) -> Self {
265        value
266            .iter()
267            .map(|v| parameters.map_element_to_i64(v))
268            .collect_vec()
269    }
270}
271
272#[cfg(test)]
273pub(crate) mod tests {
274    use std::fmt::Debug;
275
276    use num_traits::ToPrimitive;
277
278    use crate::random::DefaultSecureRng;
279
280    use super::fill_random_ternary_secret_with_hamming_weight;
281
282    #[derive(Clone)]
283    pub(crate) struct Stats<T> {
284        pub(crate) samples: Vec<T>,
285    }
286
287    impl<T> Default for Stats<T> {
288        fn default() -> Self {
289            Stats { samples: vec![] }
290        }
291    }
292
293    impl<T: Copy + ToPrimitive + Debug> Stats<T>
294    where
295        // T: for<'a> Sum<&'a T>,
296        T: for<'a> std::iter::Sum<&'a T> + std::iter::Sum<T>,
297    {
298        pub(crate) fn new() -> Self {
299            Self { samples: vec![] }
300        }
301
302        pub(crate) fn mean(&self) -> f64 {
303            self.samples.iter().sum::<T>().to_f64().unwrap() / (self.samples.len() as f64)
304        }
305
306        pub(crate) fn variance(&self) -> f64 {
307            let mean = self.mean();
308
309            // diff
310            let diff_sq = self
311                .samples
312                .iter()
313                .map(|v| {
314                    let t = v.to_f64().unwrap() - mean;
315                    t * t
316                })
317                .into_iter()
318                .sum::<f64>();
319
320            diff_sq / (self.samples.len() as f64 - 1.0)
321        }
322
323        pub(crate) fn std_dev(&self) -> f64 {
324            self.variance().sqrt()
325        }
326
327        pub(crate) fn add_many_samples(&mut self, values: &[T]) {
328            self.samples.extend(values.iter());
329        }
330
331        pub(crate) fn add_sample(&mut self, value: T) {
332            self.samples.push(value)
333        }
334
335        pub(crate) fn merge_in(&mut self, other: &Self) {
336            self.samples.extend(other.samples.iter());
337        }
338    }
339
340    #[test]
341    fn ternary_secret_has_correct_hw() {
342        let mut rng = DefaultSecureRng::new();
343        for n in 4..15 {
344            let ring_size = 1 << n;
345            let mut out = vec![0i32; ring_size];
346            fill_random_ternary_secret_with_hamming_weight(&mut out, ring_size >> 1, &mut rng);
347
348            // check hamming weight of out equals ring_size/2
349            let mut non_zeros = 0;
350            out.iter().for_each(|i| {
351                if *i != 0 {
352                    non_zeros += 1;
353                }
354            });
355
356            assert_eq!(ring_size >> 1, non_zeros);
357        }
358    }
359}