1use std::usize;
  2
  3use itertools::Itertools;
  4use num_traits::{PrimInt, Signed};
  5
  6use crate::RandomUniformDist;
  7pub trait WithLocal {
  8    fn with_local<F, R>(func: F) -> R
  9    where
 10        F: Fn(&Self) -> R;
 11
 12    fn with_local_mut<F, R>(func: F) -> R
 13    where
 14        F: Fn(&mut Self) -> R;
 15}
 16
 17pub fn fill_random_ternary_secret_with_hamming_weight<
 18    T: Signed,
 19    R: RandomUniformDist<[u8], Parameters = u8> + RandomUniformDist<usize, Parameters = usize>,
 20>(
 21    out: &mut [T],
 22    hamming_weight: usize,
 23    rng: &mut R,
 24) {
 25    let mut bytes = vec![0u8; hamming_weight.div_ceil(8)];
 26    RandomUniformDist::<[u8]>::random_fill(rng, &0, &mut bytes);
 27
 28    let size = out.len();
 29    let mut secret_indices = (0..size).into_iter().map(|i| i).collect_vec();
 30    let mut bit_index = 0;
 31    let mut byte_index = 0;
 32    for _ in 0..hamming_weight {
 33        let mut s_index = 0usize;
 34        RandomUniformDist::<usize>::random_fill(rng, &secret_indices.len(), &mut s_index);
 35
 36        let curr_bit = (bytes[byte_index] >> bit_index) & 1;
 37        if curr_bit == 1 {
 38            out[secret_indices[s_index]] = T::one();
 39        } else {
 40            out[secret_indices[s_index]] = -T::one();
 41        }
 42
 43        secret_indices[s_index] = *secret_indices.last().unwrap();
 44        secret_indices.truncate(secret_indices.len());
 45
 46        if bit_index == 7 {
 47            bit_index = 0;
 48            byte_index += 1;
 49        } else {
 50            bit_index += 1;
 51        }
 52    }
 53}
 54
 55// TODO (Jay): this is only a workaround. Add a propoer way to perform primality
 56// tests.
 57fn is_probably_prime(candidate: u64) -> bool {
 58    num_bigint_dig::prime::probably_prime(&num_bigint_dig::BigUint::from(candidate), 0)
 59}
 60
 61/// Finds prime that satisfy
 62/// - $prime \lt upper_bound$
 63/// - $\log{prime} = num_bits$
 64/// - `prime % modulo == 1`
 65pub fn generate_prime(num_bits: usize, modulo: u64, upper_bound: u64) -> Option<u64> {
 66    let leading_zeros = (64 - num_bits) as u32;
 67
 68    let mut tentative_prime = upper_bound - 1;
 69    while tentative_prime % modulo != 1 && tentative_prime.leading_zeros() == leading_zeros {
 70        tentative_prime -= 1;
 71    }
 72
 73    while !is_probably_prime(tentative_prime)
 74        && tentative_prime.leading_zeros() == leading_zeros
 75        && tentative_prime >= modulo
 76    {
 77        tentative_prime -= modulo;
 78    }
 79
 80    if is_probably_prime(tentative_prime) && tentative_prime.leading_zeros() == leading_zeros {
 81        Some(tentative_prime)
 82    } else {
 83        None
 84    }
 85}
 86
 87/// Returns a^b mod q
 88pub fn mod_exponent(a: u64, mut b: u64, q: u64) -> u64 {
 89    let mod_mul = |v1: &u64, v2: &u64| {
 90        let tmp = *v1 as u128 * *v2 as u128;
 91        (tmp % q as u128) as u64
 92    };
 93
 94    let mut acc = a;
 95    let mut out = 1;
 96    while b != 0 {
 97        let flag = b & 1;
 98
 99        if flag == 1 {
100            out = mod_mul(&acc, &out);
101        }
102        acc = mod_mul(&acc, &acc);
103
104        b >>= 1;
105    }
106
107    out
108}
109
110pub fn mod_inverse(a: u64, q: u64) -> u64 {
111    mod_exponent(a, q - 2, q)
112}
113
114pub fn shoup_representation_fq(v: u64, q: u64) -> u64 {
115    ((v as u128 * (1u128 << 64)) / q as u128) as u64
116}
117
118pub fn negacyclic_mul<T: PrimInt, F: Fn(&T, &T) -> T>(
119    a: &[T],
120    b: &[T],
121    mul: F,
122    modulus: T,
123) -> Vec<T> {
124    let mut r = vec![T::zero(); a.len()];
125    for i in 0..a.len() {
126        for j in 0..i + 1 {
127            // println!("i: {j} {}", i - j);
128            r[i] = (r[i] + mul(&a[j], &b[i - j])) % modulus;
129        }
130
131        for j in i + 1..a.len() {
132            // println!("i: {j} {}", a.len() - j + i);
133            r[i] = (r[i] + modulus - mul(&a[j], &b[a.len() - j + i])) % modulus;
134        }
135        // println!("")
136    }
137
138    return r;
139}
140
141pub trait TryConvertFrom<T: ?Sized> {
142    type Parameters: ?Sized;
143
144    fn try_convert_from(value: &T, parameters: &Self::Parameters) -> Self;
145}
146
147impl TryConvertFrom<[i32]> for Vec<Vec<u32>> {
148    type Parameters = u32;
149    fn try_convert_from(value: &[i32], parameters: &Self::Parameters) -> Self {
150        let row0 = value
151            .iter()
152            .map(|v| {
153                let is_neg = v.is_negative();
154                let v_u32 = v.abs() as u32;
155
156                assert!(v_u32 < *parameters);
157
158                if is_neg {
159                    parameters - v_u32
160                } else {
161                    v_u32
162                }
163            })
164            .collect_vec();
165
166        vec![row0]
167    }
168}
169
170impl TryConvertFrom<[i32]> for Vec<Vec<u64>> {
171    type Parameters = u64;
172    fn try_convert_from(value: &[i32], parameters: &Self::Parameters) -> Self {
173        let row0 = value
174            .iter()
175            .map(|v| {
176                let is_neg = v.is_negative();
177                let v_u64 = v.abs() as u64;
178
179                assert!(v_u64 < *parameters);
180
181                if is_neg {
182                    parameters - v_u64
183                } else {
184                    v_u64
185                }
186            })
187            .collect_vec();
188
189        vec![row0]
190    }
191}
192
193impl TryConvertFrom<[i32]> for Vec<u64> {
194    type Parameters = u64;
195    fn try_convert_from(value: &[i32], parameters: &Self::Parameters) -> Self {
196        value
197            .iter()
198            .map(|v| {
199                let is_neg = v.is_negative();
200                let v_u64 = v.abs() as u64;
201
202                assert!(v_u64 < *parameters);
203
204                if is_neg {
205                    parameters - v_u64
206                } else {
207                    v_u64
208                }
209            })
210            .collect_vec()
211    }
212}