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