arcis-compiler 0.9.7

pub mod reed_solomon {
    use crate::{
        core::{
            actually_used_field::ActuallyUsedField,
            circuits::boolean::boolean_value::Boolean,
            expressions::field_expr::FieldExpr,
            global_value::value::FieldValue,
        },
        traits::{GreaterEqual, GreaterThan, Invert, Random, Reveal},
        utils::field::BaseField,
    };
    use num_traits::ToPrimitive;
    use std::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};

    #[derive(Debug)]
    #[allow(dead_code)]
    pub struct KeyRecoveryDesc {
        pub n: usize,
        pub k: usize,
        pub d: usize,
    }

    impl KeyRecoveryDesc {
        #[allow(dead_code)]
        pub fn new(n: usize) -> Self {
            let k = (n - 1) / 3 + 1;
            Self { n, k, d: n - k + 1 }
        }
    }

    // The notation is taken from [here](https://en.wikipedia.org/wiki/Reed%E2%80%93Solomon_error_correction#The_BCH_view:_The_codeword_as_a_sequence_of_coefficients)
    // and [here](https://en.wikipedia.org/wiki/Reed%E2%80%93Solomon_error_correction#Peterson%E2%80%93Gorenstein%E2%80%93Zierler_decoder).
    // N stands for an upper bound on the block length n, K for an upper bound
    // on the message length k, and D = N - K + 1 for the distance.
    // If t is the threshold of malicious nodes we want to support we set k = t + 1.
    #[derive(Clone, Debug, Default)]
    pub struct KeyRecoveryReedSolomonInit;

    #[allow(dead_code)]
    impl KeyRecoveryReedSolomonInit {
        // Compute the coefficients of the generator polynomial
        // g(x) = (x - alpha) * (x - alpha^2) * .. * (x - alpha^(d - 1))
        // and store them in an array of size D.
        // This is computed by each node.
        pub fn compute_generator_polynomial<const D: usize, F: ActuallyUsedField>(
            d: usize,
        ) -> [F; D] {
            let mut g = [F::from(0); D];
            let mut alpha_pow_j = F::MULTIPLICATIVE_GENERATOR;
            g[0] = -alpha_pow_j;
            g[1] = F::from(1);
            for _ in 0..d - 2 {
                alpha_pow_j *= F::MULTIPLICATIVE_GENERATOR;
                let mut tmp = g;
                tmp.rotate_right(1);
                for i in 0..D {
                    g[i] *= -alpha_pow_j;
                    g[i] += tmp[i];
                }
            }

            g
        }

        /// Encodes the value val using Reed-Solomon encoding.
        /// The message polynomial is any degree-k-1 polynomial
        /// whose sum of coefficients equals val (you can see this
        /// as k-secret-sharing val and then define a polynomial with
        /// coefficients the secret-shares). The message polynomial is
        /// then encoded by multiplying with the generator polynomial g.
        // Here, D is the number of coefficients of the generator polynomial
        // (i.e., it is of degree D - 1) and N is the length of the code.
        // The message is thus of length K = N - D + 1.
        // Typical values: N = 300, D = 201 and K = 100.
        // This guarantees that up to 99 corrupt peers cannot recover the message,
        // nor can they sabotage the recovery and would be identified if they tried.
        pub fn encode<
            const N: usize,
            const D: usize,
            B: Boolean,
            T: Random
                + Mul<T, Output = T>
                + Sub<T, Output = T>
                + AddAssign<T>
                + SubAssign<T>
                + GreaterEqual<T, Output = B>
                + From<B>
                + From<usize>
                + Copy,
        >(
            val: T,
            g: [T; D],
            k: T,
        ) -> [T; N] {
            // Secret-share val and view the secret-shares as the coefficients of
            // the message polynomial p(x). The first k coefficients are non-zero,
            // the remaining ones are 0.
            let mut p = vec![val];
            for i in 0..N - D {
                let r = T::from(T::from(i).lt(k - T::from(1))) * T::random();
                p[0] -= r;
                p.push(r);
            }

            // Perform the polynomial multiplication between
            // g(x) = g[0] + g[1] * x + .. + g[D - 1] * x^(D - 1) and
            // p(x) = p[0] + p[1] * x + .. + p[N - D] * x^(N - D),
            // which yields s(x) of degree N - 1.
            // Note: the actual degrees of g, p and s might be less (for p the degree is k - 1).
            let mut s = [T::from(0); N];
            for i in 0..N {
                for j in 0..=i {
                    if i - j < D && j < N - D + 1 {
                        s[i] += g[i - j] * p[j];
                    }
                }
            }

            s
        }
    }

    #[derive(Clone, Debug, Default)]
    pub struct KeyRecoveryReedSolomonFinal;

    #[allow(dead_code)]
    impl KeyRecoveryReedSolomonFinal {
        /// Given the N coefficients of the polynomial r, compute d - 1 syndromes
        /// r(alpha), r(alpha^2), .., r(alpha^(d - 1)).
        pub fn compute_syndromes<
            const N: usize,
            const DMINUSONE: usize,
            F: ActuallyUsedField,
            B: Boolean,
            T: Mul<T, Output = T>
                + AddAssign<T>
                + MulAssign<T>
                + GreaterEqual<T, Output = B>
                + From<B>
                + From<F>
                + From<usize>
                + Reveal
                + Copy,
        >(
            r: [T; N],
            d_minus_one: T,
        ) -> [T; DMINUSONE] {
            // Can be done more efficiently..
            let mut syndromes = [r[0]; DMINUSONE];
            let mut alpha_pow_j = F::MULTIPLICATIVE_GENERATOR;
            for (j, syndrome) in syndromes.iter_mut().enumerate() {
                let mut pow = alpha_pow_j;
                for r_i in r.iter().skip(1) {
                    *syndrome += *r_i * T::from(pow);
                    pow *= alpha_pow_j;
                }
                // Set syndrome to 0 for j >= d - 1.
                *syndrome *= T::from(T::from(j).lt(d_minus_one));
                alpha_pow_j *= F::MULTIPLICATIVE_GENERATOR;
            }

            syndromes.map(|syndrome| syndrome.reveal())
        }

        /// Compute the errors from plaintext syndromes, using Berlekamp-Massey, Chien search and
        /// Forney's algorithm.
        pub fn compute_errors<const N: usize, const DMINUSONE: usize>(
            d_minus_one: FieldValue<BaseField>,
            syndromes: [FieldValue<BaseField>; DMINUSONE],
        ) -> [FieldValue<BaseField>; N] {
            let mut errors = [FieldValue::from(0); N];
            for (i, e) in errors.iter_mut().enumerate() {
                *e = FieldValue::new(FieldExpr::KeyRecoveryComputeErrors(
                    d_minus_one,
                    syndromes.to_vec(),
                    i,
                ));
            }
            errors
        }

        /// Compute the errors from plaintext syndromes, using Berlekamp-Massey, Chien search and
        /// Forney's algorithm.
        #[allow(non_snake_case)]
        pub fn compute_errors_field<const N: usize, F: ActuallyUsedField>(
            d_minus_one: F,
            syndromes: Vec<F>,
        ) -> [F; N] {
            let d_minus_one = d_minus_one
                .to_unsigned_number()
                .to_usize()
                .expect("d_minus_one way too large");
            let syndromes = syndromes.into_iter().take(d_minus_one).collect::<Vec<F>>();

            // Compute the error locator polynomial Lambda using the Berlekamp-Massey algorithm.
            // Notation-wise we follow [this](https://en.wikipedia.org/wiki/Berlekamp%E2%80%93Massey_algorithm#Pseudocode).
            let mut C = vec![F::from(1)];
            let mut B = vec![F::from(1)];
            let mut L = 0usize;
            let mut m = 1usize;
            let mut b = F::from(1);
            for n in 0..syndromes.len() {
                let mut d = syndromes[n];
                for i in 1..=L {
                    d += C[i] * syndromes[n - i];
                }

                if d.eq(&F::from(0)) {
                    m += 1;
                } else if 2 * L <= n {
                    let T = C.clone();
                    C.resize((m + B.len()).max(C.len()), F::from(0));
                    // on plaintext values we simply set is_expected_non_zero = true
                    let db_inv = d * b.invert(true);
                    for i in 0..B.len() {
                        C[m + i] -= db_inv * B[i];
                    }
                    L = n + 1 - L;
                    B = T;
                    b = d;
                    m = 1usize;
                } else {
                    C.resize((m + B.len()).max(C.len()), F::from(0));
                    // on plaintext values we simply set is_expected_non_zero = true
                    let db_inv = d * b.invert(true);
                    for i in 0..B.len() {
                        C[m + i] -= db_inv * B[i];
                    }
                    m += 1;
                }
            }
            // L now corresponds to the number of errors and C corresponds to Lambda.

            // We perform the [Chien search](https://en.wikipedia.org/wiki/Chien_search)
            // to find the roots of the reversed Lambda.
            let mut alpha_pows = vec![F::from(1)];
            for i in 0..L {
                alpha_pows.push(alpha_pows[i] * F::MULTIPLICATIVE_GENERATOR);
            }
            let mut lambdas = C.iter().copied().rev().collect::<Vec<F>>();
            let mut error_locations = vec![lambdas
                .iter()
                .copied()
                .reduce(|a, b| a + b)
                .unwrap()
                .eq(&F::from(0))];
            for _ in 0..N - 1 {
                lambdas = lambdas
                    .iter()
                    .zip(alpha_pows.clone())
                    .map(|(lambda, pow)| *lambda * pow)
                    .collect::<Vec<F>>();
                error_locations.push(
                    lambdas
                        .iter()
                        .copied()
                        .reduce(|a, b| a + b)
                        .unwrap()
                        .eq(&F::from(0)),
                );
            }

            // [Forney's algorithm](https://en.wikipedia.org/wiki/Forney_algorithm).
            let Lambda_prime = C
                .iter()
                .enumerate()
                .skip(1)
                .map(|(i, lambda)| F::from(i as u64) * lambda)
                .collect::<Vec<F>>();
            let mut Omega = Vec::new();
            for i in 0..L {
                let mut omega_i = F::from(0);
                for j in 0..=i {
                    if i - j < syndromes.len() && j < L + 1 {
                        omega_i += syndromes[i - j] * C[j];
                    }
                }
                Omega.push(omega_i);
            }

            fn eval_poly<F: ActuallyUsedField>(poly: Vec<F>, x: F) -> F {
                let mut res = poly[0];
                let mut pow = F::from(1);
                for c in poly.into_iter().skip(1) {
                    pow *= x;
                    res += c * pow;
                }
                res
            }

            let mut errors = [F::from(0); N];
            let mut pow = F::from(1);
            for (i, e) in error_locations.into_iter().enumerate() {
                if e {
                    let X_inv = pow.invert(true);
                    let num = eval_poly(Omega.clone(), X_inv);
                    let denom = eval_poly(Lambda_prime.clone(), X_inv);
                    // on plaintext values we simply set is_expected_non_zero = true
                    errors[i] = -num * denom.invert(true);
                }
                pow *= F::MULTIPLICATIVE_GENERATOR;
            }

            errors
        }

        // We know that the generator polynomial g does not vanish at the consecutive
        // powers alpha^d, .., alpha^(d + k - 1). Observe that for i = 0, .., k - 1, we have
        // s(alpha^(d + i)) = p(alpha^(d + i)) * g(alpha^(d + i)).
        // It suffices to evaluate s(alpha^(d + i)), for i = 0, .., k - 1, to know the value
        // of p at k distinct points, more precisely,
        // p(alpha^(d + i)) = s(alpha^(d + i)) / g(alpha^(d + i)), for i = 0, .., k - 1.
        // Since we are interested in p(1) we compute the Lagrange polynomials evaluated at 1,
        // l_{alpha^(d + i)}(1), for i = 0, .., k - 1, and divide them by g(alpha^(d + i)).
        // This will be computed by each node.
        pub fn compute_scaled_polynomials<const K: usize, F: ActuallyUsedField>(
            d: usize,
            k: usize,
        ) -> ([F; K], [F; K]) {
            let alpha_pows = (0..d + k - 1)
                .scan(F::from(1), |acc, _| {
                    *acc *= F::MULTIPLICATIVE_GENERATOR;
                    Some(*acc)
                })
                .collect::<Vec<F>>();
            // Lagrange polynomials at 1 multiplied by the inverse of g(alpha^(d + i)).
            let mut scaled_polynomials = alpha_pows
                .iter()
                .copied()
                .skip(d - 1)
                .map(|x| {
                    // numerator of Lagrange polynomial
                    let num = alpha_pows
                        .iter()
                        .copied()
                        .skip(d - 1)
                        .filter(|val| *val != x)
                        .map(|val| F::from(1) - val)
                        .reduce(|a, b| a * b)
                        .unwrap();
                    // denominator of Lagrange polynomial
                    let denom = alpha_pows
                        .iter()
                        .copied()
                        .skip(d - 1)
                        .filter(|val| *val != x)
                        .map(|val| x - val)
                        .reduce(|a, b| a * b)
                        .unwrap();
                    let poly = num * denom.invert(true);
                    let g_at_alpha_pow = alpha_pows
                        .iter()
                        .copied()
                        .take(d - 1)
                        .map(|z| x - z)
                        .reduce(|a, b| a * b)
                        .unwrap();
                    poly * g_at_alpha_pow.invert(true)
                })
                .collect::<Vec<F>>();

            // We remove the first d - 1 powers of alpha and
            // fill up alpha_pows and scaled_polynomials with zeros.
            let mut alpha_pows = alpha_pows.iter().copied().skip(d - 1).collect::<Vec<F>>();
            alpha_pows.resize(K, F::from(0));
            scaled_polynomials.resize(K, F::from(0));

            (
                alpha_pows.try_into().unwrap_or_else(|v: Vec<F>| {
                    panic!("Expected a Vec of length {} (found {})", K, v.len())
                }),
                scaled_polynomials.try_into().unwrap_or_else(|v: Vec<F>| {
                    panic!("Expected a Vec of length {} (found {})", K, v.len())
                }),
            )
        }

        /// Once the error polynomial e is computed, subtract it from r.
        pub fn subtract_errors<const N: usize, T: Sub<T, Output = T>>(
            r: [T; N],
            e: [T; N],
        ) -> [T; N] {
            r.into_iter()
                .zip(e)
                .map(|(coeff_r, coeff_e)| coeff_r - coeff_e)
                .collect::<Vec<T>>()
                .try_into()
                .unwrap_or_else(|v: Vec<T>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                })
        }

        /// Recover the polynomial p from the corrected codeword s and evaluate at 1
        /// to obtain val (equivalently, sum the coefficients of p).
        pub fn decode<
            const N: usize,
            const K: usize,
            B: Boolean,
            T: Mul<T, Output = T>
                + MulAssign<T>
                + Add<T, Output = T>
                + AddAssign<T>
                + GreaterThan<T, Output = B>
                + From<usize>
                + Copy,
        >(
            s: [T; N],
            alpha_pows: [T; K],
            scaled_polynomials: [T; K],
        ) -> T {
            // compute s(alpha^d), .., s(alpha^(d + k - 1))
            let s_at_alpha_pows = alpha_pows
                .into_iter()
                .map(|alpha_pow_j| {
                    let mut pow = alpha_pow_j;
                    let mut s_at_alpha_pow = s[0];
                    for s_i in s.iter().skip(1) {
                        s_at_alpha_pow += *s_i * pow;
                        pow *= alpha_pow_j;
                    }
                    s_at_alpha_pow
                })
                .collect::<Vec<T>>();

            // scaled_polynomials are of the form l_{alpha^(d + i)}(1) / g(alpha^(d + i)),
            // for i = 0, .., k - 1. Multiplying the ith s_at_alpha_pow by the ith scaled_polynomial
            // yields s(alpha^(d + i)) * l_{alpha^(d + i)}(1) / g(alpha^(d + i))
            // = p(alpha^(d + i)) * l_{alpha^(d + i)}(1).
            // Summing all the above yields p(1).
            scaled_polynomials
                .into_iter()
                .zip(s_at_alpha_pows)
                .map(|(poly, s_at_alpha_pow)| poly * s_at_alpha_pow)
                .reduce(|a, b| a + b)
                .unwrap()
        }
    }
}

pub mod shamir {
    use crate::{
        core::{
            actually_used_field::ActuallyUsedField,
            bounds::FieldBounds,
            circuits::{
                boolean::boolean_value::{Boolean, BooleanValue},
                traits::arithmetic_circuit::ArithmeticCircuit,
            },
            expressions::expr::EvalFailure,
            global_value::value::FieldValue,
        },
        traits::{GreaterThan, Invert, Random},
        utils::number::Number,
    };
    use rand::{seq::SliceRandom, thread_rng};
    use std::ops::{Add, AddAssign, Mul, MulAssign};

    #[derive(Clone, Debug, Default)]
    pub struct KeyRecoveryShamirInit;

    impl KeyRecoveryShamirInit {
        /// Shamir secret-shares the value val as the constant term of a polynomial
        /// of degree deg and evaluates at the points 1, .., n (inclusive).
        /// The remaining N - n secret-shares are set to 0.
        pub fn shamir_secret_share<
            const N: usize,
            B: Boolean,
            T: Random
                + Mul<T, Output = T>
                + MulAssign<T>
                + AddAssign<T>
                + GreaterThan<T, Output = B>
                + From<B>
                + From<usize>
                + Copy,
        >(
            val: T,
            deg: T,
            n: T,
        ) -> [T; N] {
            let mut coeffs = vec![val];
            coeffs.extend((0..N - 1).map(|_| T::random()));
            (1..N + 1)
                .map(|i| {
                    let mut res = T::from(0);
                    // it is important to set pow = 0 as soon as i > n
                    // otherwise, whoever gets to decrypt sufficiently many of
                    // these secret-shares will be able to recover the key
                    let mut pow = T::from(T::from(i).le(n));
                    coeffs.iter().enumerate().for_each(|(j, c)| {
                        res += pow * *c * T::from(T::from(j).le(deg));
                        pow *= T::from(i);
                    });
                    res
                })
                .collect::<Vec<T>>()
                .try_into()
                .unwrap_or_else(|v: Vec<T>| {
                    panic!("Expected a Vec of length {} (found {})", N, v.len())
                })
        }
    }

    #[derive(Clone, Debug, Default)]
    pub struct KeyRecoveryShamirFinal;

    impl KeyRecoveryShamirFinal {
        pub fn lagrange_interpolate_at_zero<
            const N: usize,
            T: Mul<T, Output = T> + Add<T, Output = T> + From<usize> + Copy,
        >(
            inps: [(T, T); N],
        ) -> T {
            inps.into_iter()
                .fold(T::from(0), |acc, (poly, y)| acc + y * poly)
        }
    }

    #[derive(Clone, Debug)]
    #[allow(dead_code)]
    struct TestKeyRecovery {
        deg: usize,
        n: usize,
    }

    impl TestKeyRecovery {
        #[allow(dead_code)]
        pub fn new(deg: usize, n: usize) -> Self {
            Self { deg, n }
        }
    }

    const N: usize = 100;

    impl<F: ActuallyUsedField> ArithmeticCircuit<F> for TestKeyRecovery {
        fn eval(&self, x: Vec<F>) -> Result<Vec<F>, EvalFailure> {
            if x.len() != 1 {
                panic!("expected input of length 1 (found {})", x.len());
            }
            if self.n > N {
                return EvalFailure::err_ub("n too large");
            }
            Ok(x)
        }

        fn bounds(&self, _bounds: Vec<FieldBounds<F>>) -> Vec<FieldBounds<F>> {
            vec![FieldBounds::All]
        }

        fn run(&self, vals: Vec<FieldValue<F>>) -> Vec<FieldValue<F>>
        where
            F: ActuallyUsedField,
        {
            // computed by the source MXE M (we omit the encryption)
            let all_secret_shares =
                KeyRecoveryShamirInit::shamir_secret_share::<N, BooleanValue, FieldValue<F>>(
                    vals[0],
                    FieldValue::<F>::from(self.deg),
                    FieldValue::<F>::from(self.n),
                );

            // computed by the program/nodes of M and put on chain
            let mut secret_shares = all_secret_shares
                .into_iter()
                .take(self.n)
                .enumerate()
                .map(|(i, secret_share)| (i + 1, secret_share))
                .collect::<Vec<(usize, FieldValue<F>)>>();

            // simulates a random subset of deg + 1 shares for the recovery
            // these are authenticated via zkps and put on chain
            let mut rng = thread_rng();
            secret_shares.shuffle(&mut rng);
            let secret_shares = secret_shares
                .into_iter()
                .take(self.deg + 1)
                .collect::<Vec<(usize, FieldValue<F>)>>();

            // computed by the program/nodes of the target MXE M'
            let (xs, mut ys): (Vec<usize>, Vec<FieldValue<F>>) =
                secret_shares.iter().copied().unzip();
            let mut lagrange_basis_polynomials = xs
                .iter()
                .copied()
                .map(|x| {
                    let num = xs
                        .iter()
                        .copied()
                        .filter(|val| *val != x)
                        .map(|val| -Number::from(val))
                        .reduce(|a, b| a * b)
                        .unwrap();
                    let denom = xs
                        .iter()
                        .copied()
                        .filter(|val| *val != x)
                        .map(|val| Number::from(x) - Number::from(val))
                        .reduce(|a, b| a * b)
                        .unwrap();
                    FieldValue::from(F::from(num) * F::from(denom).invert(true))
                })
                .collect::<Vec<FieldValue<F>>>();
            lagrange_basis_polynomials.extend(std::iter::repeat_n(
                FieldValue::<F>::from(0),
                N - (self.deg + 1),
            ));
            ys.extend(std::iter::repeat_n(
                FieldValue::<F>::from(0),
                N - (self.deg + 1),
            ));

            // computed by M' (we omit the decryption)
            vec![KeyRecoveryShamirFinal::lagrange_interpolate_at_zero::<
                N,
                FieldValue<F>,
            >(
                lagrange_basis_polynomials
                    .into_iter()
                    .zip(ys)
                    .collect::<Vec<(FieldValue<F>, FieldValue<F>)>>()
                    .try_into()
                    .unwrap_or_else(|v: Vec<(FieldValue<F>, FieldValue<F>)>| {
                        panic!("Expected a Vec of length {} (found {})", N, v.len())
                    }),
            )]
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use crate::{
            core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit,
            utils::field::BaseField,
        };
        use rand::Rng;

        impl TestedArithmeticCircuit<BaseField> for TestKeyRecovery {
            fn gen_desc<R: Rng + ?Sized>(rng: &mut R) -> Self {
                let mut deg = 3;
                while rng.gen_bool(0.75) {
                    deg += 1;
                }
                let mut n = deg + 1;
                while rng.gen_bool(0.75) {
                    n += 1;
                }
                Self::new(deg as usize, n as usize)
            }

            fn gen_n_inputs<R: Rng + ?Sized>(&self, _rng: &mut R) -> usize {
                1
            }
        }

        #[test]
        fn tested_recovery() {
            TestKeyRecovery::test(2, 1)
        }
    }
}