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