1use std::{fmt::Debug, usize};
  2
  3use itertools::Itertools;
  4use num_traits::{FromPrimitive, PrimInt, Signed, Unsigned};
  5
  6use crate::{
  7    backend::Modulus,
  8    random::{RandomElement, RandomElementInModulus, RandomFill},
  9};
 10pub trait WithLocal {
 11    fn with_local<F, R>(func: F) -> R
 12    where
 13        F: Fn(&Self) -> R;
 14
 15    fn with_local_mut<F, R>(func: F) -> R
 16    where
 17        F: Fn(&mut Self) -> R;
 18
 19    fn with_local_mut_mut<F, R>(func: &mut F) -> R
 20    where
 21        F: FnMut(&mut Self) -> R;
 22}
 23
 24pub trait Global {
 25    fn global() -> &'static Self;
 26}
 27
 28pub trait ShoupMul {
 29    fn representation(value: Self, q: Self) -> Self;
 30    fn mul(a: Self, b: Self, b_shoup: Self, q: Self) -> Self;
 31}
 32
 33impl ShoupMul for u64 {
 34    #[inline]
 35    fn representation(value: Self, q: Self) -> Self {
 36        ((value as u128 * (1u128 << 64)) / q as u128) as u64
 37    }
 38
 39    #[inline]
 40    /// Returns a * b % q
 41    fn mul(a: Self, b: Self, b_shoup: Self, q: Self) -> Self {
 42        (b.wrapping_mul(a))
 43            .wrapping_sub(q.wrapping_mul(((b_shoup as u128 * a as u128) >> 64) as u64))
 44    }
 45}
 46
 47pub fn fill_random_ternary_secret_with_hamming_weight<
 48    T: Signed,
 49    R: RandomFill<[u8]> + RandomElementInModulus<usize, usize>,
 50>(
 51    out: &mut [T],
 52    hamming_weight: usize,
 53    rng: &mut R,
 54) {
 55    let mut bytes = vec![0u8; hamming_weight.div_ceil(8)];
 56    RandomFill::<[u8]>::random_fill(rng, &mut bytes);
 57
 58    let size = out.len();
 59    let mut secret_indices = (0..size).into_iter().map(|i| i).collect_vec();
 60    let mut bit_index = 0;
 61    let mut byte_index = 0;
 62    for _ in 0..hamming_weight {
 63        let s_index = RandomElementInModulus::<usize, usize>::random(rng, &secret_indices.len());
 64
 65        let curr_bit = (bytes[byte_index] >> bit_index) & 1;
 66        if curr_bit == 1 {
 67            out[secret_indices[s_index]] = T::one();
 68        } else {
 69            out[secret_indices[s_index]] = -T::one();
 70        }
 71
 72        secret_indices[s_index] = *secret_indices.last().unwrap();
 73        secret_indices.truncate(secret_indices.len());
 74
 75        if bit_index == 7 {
 76            bit_index = 0;
 77            byte_index += 1;
 78        } else {
 79            bit_index += 1;
 80        }
 81    }
 82}
 83
 84// TODO (Jay): this is only a workaround. Add a propoer way to perform primality
 85// tests.
 86fn is_probably_prime(candidate: u64) -> bool {
 87    num_bigint_dig::prime::probably_prime(&num_bigint_dig::BigUint::from(candidate), 0)
 88}
 89
 90/// Finds prime that satisfy
 91/// - $prime \lt upper_bound$
 92/// - $\log{prime} = num_bits$
 93/// - `prime % modulo == 1`
 94pub fn generate_prime(num_bits: usize, modulo: u64, upper_bound: u64) -> Option<u64> {
 95    let leading_zeros = (64 - num_bits) as u32;
 96
 97    let mut tentative_prime = upper_bound - 1;
 98    while tentative_prime % modulo != 1 && tentative_prime.leading_zeros() == leading_zeros {
 99        tentative_prime -= 1;
100    }
101
102    while !is_probably_prime(tentative_prime)
103        && tentative_prime.leading_zeros() == leading_zeros
104        && tentative_prime >= modulo
105    {
106        tentative_prime -= modulo;
107    }
108
109    if is_probably_prime(tentative_prime) && tentative_prime.leading_zeros() == leading_zeros {
110        Some(tentative_prime)
111    } else {
112        None
113    }
114}
115
116/// Returns a^b mod q
117pub fn mod_exponent(a: u64, mut b: u64, q: u64) -> u64 {
118    let mod_mul = |v1: &u64, v2: &u64| {
119        let tmp = *v1 as u128 * *v2 as u128;
120        (tmp % q as u128) as u64
121    };
122
123    let mut acc = a;
124    let mut out = 1;
125    while b != 0 {
126        let flag = b & 1;
127
128        if flag == 1 {
129            out = mod_mul(&acc, &out);
130        }
131        acc = mod_mul(&acc, &acc);
132
133        b >>= 1;
134    }
135
136    out
137}
138
139pub fn mod_inverse(a: u64, q: u64) -> u64 {
140    mod_exponent(a, q - 2, q)
141}
142
143pub fn negacyclic_mul<T: PrimInt, F: Fn(&T, &T) -> T>(
144    a: &[T],
145    b: &[T],
146    mul: F,
147    modulus: T,
148) -> Vec<T> {
149    let mut r = vec![T::zero(); a.len()];
150    for i in 0..a.len() {
151        for j in 0..i + 1 {
152            // println!("i: {j} {}", i - j);
153            r[i] = (r[i] + mul(&a[j], &b[i - j])) % modulus;
154        }
155
156        for j in i + 1..a.len() {
157            // println!("i: {j} {}", a.len() - j + i);
158            r[i] = (r[i] + modulus - mul(&a[j], &b[a.len() - j + i])) % modulus;
159        }
160        // println!("")
161    }
162
163    return r;
164}
165
166pub trait TryConvertFrom1<T: ?Sized, P> {
167    fn try_convert_from(value: &T, parameters: &P) -> Self;
168}
169
170impl<P: Modulus<Element = u64>> TryConvertFrom1<[i64], P> for Vec<u64> {
171    fn try_convert_from(value: &[i64], parameters: &P) -> Self {
172        value
173            .iter()
174            .map(|v| parameters.map_element_from_i64(*v))
175            .collect_vec()
176    }
177}
178
179impl<P: Modulus<Element = u64>> TryConvertFrom1<[i32], P> for Vec<u64> {
180    fn try_convert_from(value: &[i32], parameters: &P) -> Self {
181        value
182            .iter()
183            .map(|v| parameters.map_element_from_i64(*v as i64))
184            .collect_vec()
185    }
186}
187
188impl<P: Modulus> TryConvertFrom1<[P::Element], P> for Vec<i64> {
189    fn try_convert_from(value: &[P::Element], parameters: &P) -> Self {
190        value
191            .iter()
192            .map(|v| parameters.map_element_to_i64(v))
193            .collect_vec()
194    }
195}
196
197pub(crate) struct Stats<T> {
198    pub(crate) samples: Vec<T>,
199}
200
201impl<T: PrimInt + FromPrimitive + Debug> Stats<T>
202where
203    // T: for<'a> Sum<&'a T>,
204    T: for<'a> std::iter::Sum<&'a T> + std::iter::Sum<T>,
205{
206    pub(crate) fn new() -> Self {
207        Self { samples: vec![] }
208    }
209
210    pub(crate) fn mean(&self) -> f64 {
211        self.samples.iter().sum::<T>().to_f64().unwrap() / (self.samples.len() as f64)
212    }
213
214    pub(crate) fn std_dev(&self) -> f64 {
215        let mean = self.mean();
216
217        // diff
218        let diff_sq = self
219            .samples
220            .iter()
221            .map(|v| {
222                let t = v.to_f64().unwrap() - mean;
223                t * t
224            })
225            .into_iter()
226            .sum::<f64>();
227
228        (diff_sq / (self.samples.len() as f64)).sqrt()
229    }
230
231    pub(crate) fn add_more(&mut self, values: &[T]) {
232        self.samples.extend(values.iter());
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use rand::{thread_rng, Rng};
239
240    use super::ShoupMul;
241
242    #[test]
243    fn gg() {
244        let mut rng = thread_rng();
245        let p = 36028797018820609;
246
247        let a = rng.gen_range(0..p);
248        let b = rng.gen_range(0..p);
249        let a_shoup = ShoupMul::representation(a, p);
250
251        // let c = ShoupMul::mul(b, a, a_shoup, p);
252        // assert!(c == ((a as u128 * b as u128) % p as u128) as u64);
253
254        let mut quotient = ((a_shoup as u128 * b as u128) >> 64) as u64;
255        quotient -= 1;
256        let c = (b.wrapping_mul(a)).wrapping_sub(p.wrapping_mul(quotient));
257        assert!(c - p == ((a as u128 * b as u128) % p as u128) as u64);
258    }
259}