arcis 0.6.0-alpha

use crate::{ArcisType, Reveal};
use rand::{thread_rng, Rng};
use std::cell::Cell;

thread_local! {
    static GENERATED_BOOLS: Cell<Vec<bool>> = const {Cell::new(Vec::new()) };
}

pub fn generated_bools() -> Vec<bool> {
    GENERATED_BOOLS.take()
}

/// This struct gives access to randomness operations.
/// ```
/// use arcis::*;
///
/// #[encrypted]
/// mod circuits {
///     use arcis::*;
///
///     #[instruction]
///     pub fn doc_arcis_rng(s: Shared) -> (Enc<Shared, bool>, [u8; 32], Enc<Mxe, [u128; 3]>, bool) {
///         let random_bool = ArcisRNG::bool();
///
///         // Generates an u128 with bit width 3.
///         let first_integer = ArcisRNG::gen_integer_from_width(3);
///         // 0 <= first_integer < 2^3 = 8
///
///         // Generates an integer between 17 and 27, both included.
///         // Will try 24 times with each attempt having > 1/2 chance of success.
///         // So chance of failure is below 2^-24.
///         let (second_integer, did_it_work) = ArcisRNG::gen_integer_in_range(17, 27, 24);
///
///         let mut u128_array = [first_integer, second_integer, 42];
///         // Shuffles a slice in-place.
///         ArcisRNG::shuffle(&mut u128_array);
///
///         // Generates uniformly an [u8; 32].
///         let u8_array = ArcisRNG::gen_uniform::<[u8;32]>(); // You need to put the type there.
///         // Our interpreter does not have type inference and needs to know the type right away.
///
///         // Does not work:
///         // let b: bool = ArcisRNG::gen_uniform(); // We do not do type inference.
///         // let random_float = ArcisRNG::gen_uniform::<f64>(); // floats cannot be generated uniformly.
///
///         // You can send the randomness to someone.
///         let enc_bool = s.from_arcis(random_bool);
///
///         // Or reveal it to the world.
///         let revealed_u8_array = u8_array.reveal();
///
///         // Or keep it for later.
///         let stored_shuffled_u128_array = Mxe::get().from_arcis(u128_array);
///
///         (enc_bool, revealed_u8_array, stored_shuffled_u128_array, did_it_work.reveal())
///     }
/// }
/// ```
pub struct ArcisRNG;

impl ArcisRNG {
    /// Generates a boolean, with 1/2 probability to be `true`, 1/2 probability to be `false`.
    pub fn bool() -> bool {
        let mut rng = thread_rng();
        let b = rng.gen_bool(0.5);
        let mut v = GENERATED_BOOLS.take();
        v.push(b);
        GENERATED_BOOLS.replace(v);
        b
    }
    /// Generates a secret integer between 0 and 2^width - 1 included.
    /// Note that `width` must be compile-time known.
    pub fn gen_integer_from_width(width: usize) -> u128 {
        assert!(width <= u128::BITS as usize);
        let mut integer: u128 = 0;
        for i in 0..width {
            integer += (Self::bool() as u128) << i;
        }
        integer
    }
    /// Generates a public integer between 0 and 2^width - 1 included.
    /// Note that `width` must be compile-time known.
    pub fn gen_public_integer_from_width(width: usize) -> u128 {
        Self::gen_integer_from_width(width).reveal()
    }
    /// Uniformly and secretly shuffles the array in argument.
    /// Complexity is in `O(n*log³(n) + n*log²(n)*sizeof(T))`.
    pub fn shuffle<T>(data: &mut [T]) {
        // We want to end up with the same shuffle as the MPC version
        // if we start from the same generated booleans.
        let n = data.len();
        // This needs to be the same as in arcis/src/core/circuits/sort.rs
        let randomness_width = if n < 2 {
            // No need for any randomness to shuffle 0 or 1 items.
            0
        } else {
            // We want P(two generated numbers are identical) ~<= 2^-64
            // P(two generated numbers are identical) < E(pair of identical numbers)
            // = n(n-1)/2 * 2^-randomness_width
            // With this value, we have from 2^-65 to 2^-63
            // depending on how close to a power of two n is.
            2 * (n.ilog2() as usize) + 64
        };
        let mut random_nums = Vec::with_capacity(n);
        for i in 0..n {
            random_nums.push((Self::gen_integer_from_width(randomness_width), i));
        }
        random_nums.sort();
        let final_perm = random_nums.into_iter().map(|(_, x)| x).collect::<Vec<_>>();
        apply_perm(data, &final_perm);
    }

    /// Generates uniformly a `T`.
    /// Works for every supported type except types containing public keys, floats, and slices.
    /// Our interpreter does not have type inference so you will need to explicit `T`.
    /// For instance: `ArcisRNG::gen_uniform::<[u8; 1024]>()`.
    pub fn gen_uniform<T: ArcisType>() -> T {
        let bools = (0..T::n_bools())
            .map(|_| ArcisRNG::bool())
            .collect::<Vec<_>>();
        T::from_bools(&bools)
    }
}

/// Applies perm, ie changes data
/// from data[i] = init_data[i]
/// to data[i] = init_data[perm[i]].
fn apply_perm<T>(data: &mut [T], perm: &[usize]) {
    let n = data.len();
    let mut current_perm: Vec<_> = (0..n).collect();
    let mut positions: Vec<_> = (0..n).collect();
    for i in 0..(n - 1) {
        let val_to_transfer = perm[i];
        let other_i = positions[val_to_transfer];
        if i == other_i {
            continue;
        }
        let i_val = current_perm[i];
        assert_eq!(current_perm[i], i_val);
        assert_eq!(current_perm[other_i], val_to_transfer);
        assert_eq!(positions[i_val], i);
        assert_eq!(positions[val_to_transfer], other_i);
        data.swap(i, other_i);
        current_perm.swap(i, other_i);
        positions.swap(i_val, val_to_transfer);
    }
    assert_eq!(current_perm.as_slice(), perm);
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_apply_perm() {
        let final_perm = [1, 2, 0, 5, 4, 3];
        let init_data = [10, 11, 12, 13, 14, 15];
        let mut data = init_data;
        apply_perm(&mut data, &final_perm);
        assert_eq!(data, [11, 12, 10, 15, 14, 13]);
        for i in 0..init_data.len() {
            assert_eq!(data[i], init_data[final_perm[i]]);
        }
    }
}