arcis-compiler 0.9.4

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::FieldBounds,
        circuits::{
            boolean::{
                boolean_value::BooleanValue,
                utils::{decoder_circuit, shift_right},
            },
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::{FromLeBits, GetBit, Select},
    types::DOUBLE_PRECISION_MANTISSA,
    utils::{number::Number, used_field::UsedField},
};
use num_traits::ToPrimitive;

// This is the 53-bit number log2(e) * 2^DOUBLE_PRECISION_MANTISSA.
const LOG2_E: u64 = 6497320848556798;
// This is the 52-bit number ln(2) * 2^DOUBLE_PRECISION_MANTISSA.
const LN_2: u64 = 3121657384082679;
// Margin we add to the fractional part in order to reduce rounding noise.
const PRECISION_MARGIN: usize = 4;

// This is the array of the first 12 Chebyshev coefficients for the function exp2((z + 1)/2) with z
// in [-1, 1]. They are obtained by multiplying the Chebyshev coefficients for the
// function exp2(z / 2) with z in [-1, 1] by sqrt(2), and the former are computed as
// [I_0(log(2)/2), 2*I_1(log(2)/2), 2*I_2(log(2)/2), .., 2*I_11(log(2)/2)], where I_k(x) is the
// modified Bessel function of the first kind. For the below CHEBYSHEV_COEFFS, we multiply the
// coefficients by 2^(DOUBLE_PRECISION_MANTISSA + PRECISION_MARGIN) and round. Here's how to
// obtain the coefficients for exp2(z / 2):
// import numpy as np
// import scipy.special
// [scipy.special.iv(0, np.log(2)/2)] + [2 * scipy.special.iv(i, np.log(2)/2) for i in range(1, 12)]
const CHEBYSHEV_COEFFS: [usize; 12] = [
    104987905506995824,
    35850444948458576,
    3090774450775715,
    178085167335115,
    7703397532702,
    266712640298,
    7697462682,
    190450582,
    4123602,
    79370,
    1375,
    22,
];

/// An exponentiation algorithm. Computes exp2(x). The lowest DOUBLE_PRECISION_MANTISSA bits
/// represent the fractional part.
#[derive(Clone, Debug)]
pub struct Exp2 {
    // number of bits after the point of the input
    precision_in: usize,
    // number of bits after the point of the output
    precision_out: usize,
}

impl Exp2 {
    #[allow(unused)]
    pub const fn new(precision_in: usize) -> Self {
        if precision_in > DOUBLE_PRECISION_MANTISSA {
            panic!("precision_in must be at most 52");
        }
        Exp2 {
            precision_in,
            precision_out: DOUBLE_PRECISION_MANTISSA,
        }
    }

    /// Given an input x, this function computes 2^int(x) and frac(x), where int(x) = floor(x)
    /// and frac(x) = x - int(x) (this corresponds to the actual integer and fractional part
    /// if x >= 0 but is no longer true if x < 0).
    fn init_exp2<F: ActuallyUsedField>(&self, x: FieldValue<F>) -> (FieldValue<F>, FieldValue<F>) {
        let bounds = x.bounds();
        let bin_size = bounds.signed_bin_size();
        // TODO: the precision + 1 part is only for signed input -> one can make this better..
        let bits = (0..(self.precision_in + 1).max(bin_size))
            .map(|i| x.get_bit(i, true))
            .collect::<Vec<BooleanValue>>();
        let frac_x = FieldValue::<F>::from_le_bits(
            bits.iter()
                .copied()
                .take(self.precision_in)
                .collect::<Vec<BooleanValue>>(),
            false,
        );
        let (min, max) = bounds.to_signed_number_pair();
        // lower bound on the integer part of x
        let min_int = min >> self.precision_in;
        // upper bound on the integer part of x
        let max_int = max >> self.precision_in;
        let two_pow_int_x = if min_int == max_int {
            FieldValue::<F>::from(1)
        } else {
            // we apply a decoder circuit on an input of size l + 1
            // TODO: one can do better -> check the logic in speakeasy (input of size l - 1)
            let mut l = min_int.bits().max(max_int.bits());
            // TODO: make this cleaner
            if min_int == (-1 << (l - 1)) && max_int < (1 << (l - 1)) {
                l -= 1;
            }
            let mut int_x = bits[self.precision_in..self.precision_in + l].to_vec();
            let gap = (max_int.to_i32().unwrap() - min_int.to_i32().unwrap() + 1) as usize;
            if min_int >= 0 {
                // there are max_int - min_int + 1 integers in the interval [min_int, max_int],
                // hence the binary decomposition of 2^int(x) is of length max_int - min_int + 1
                let two_pow_int_x = decoder_circuit(int_x);
                let start = min_int.to_usize().unwrap();
                FieldValue::<F>::from_le_bits(two_pow_int_x[start..(start + gap)].to_vec(), false)
            } else if max_int < 0 {
                // there are max_int - min_int + 1 integers in the interval [min_int, max_int],
                // hence the binary decomposition of 2^int(x) is of length max_int - min_int + 1
                let two_pow_int_x = decoder_circuit(int_x);
                let start = (Number::power_of_two(l) + min_int).to_usize().unwrap();
                FieldValue::<F>::from_le_bits(two_pow_int_x[start..(start + gap)].to_vec(), false)
            } else {
                let sign = bits[self.precision_in + l];
                int_x.push(!sign);
                // there are max_int - min_int + 1 integers in the interval [min_int, max_int],
                // hence the binary decomposition of 2^int(x) is of length max_int - min_int + 1
                let two_pow_int_x = decoder_circuit(int_x);
                let start = (Number::power_of_two(l) + min_int).to_usize().unwrap();
                FieldValue::<F>::from_le_bits(two_pow_int_x[start..(start + gap)].to_vec(), false)
            }
        };
        (two_pow_int_x, frac_x)
    }

    fn exp2_approx<F: ActuallyUsedField>(&self, frac_x: FieldValue<F>) -> FieldValue<F> {
        // frac_x is in the interval [0, 2^precision_in - 1]
        let one =
            FieldValue::<F>::from(Number::power_of_two(self.precision_out + PRECISION_MARGIN));
        // affine change of variables frac_x -> z in [-2^(precision_out + PRECISION_MARGIN),
        // 2^(precision_out + PRECISION_MARGIN) - 1]
        let z = FieldValue::from(F::power_of_two(
            1 + (self.precision_out - self.precision_in) + PRECISION_MARGIN,
        )) * frac_x
            - one;
        let mut chebyshev_polynomials = vec![one, z];
        for i in 2..12 {
            let last = chebyshev_polynomials[i - 1];
            let second_last = chebyshev_polynomials[i - 2];
            chebyshev_polynomials.push(
                2 * shift_right(z * last, self.precision_out + PRECISION_MARGIN, true)
                    - second_last,
            );
        }

        CHEBYSHEV_COEFFS
            .into_iter()
            .zip(chebyshev_polynomials)
            .fold(FieldValue::<F>::from(0), |acc, (c, p)| {
                acc + FieldValue::from(F::from(Number::from(c))) * p
            })
            >> (self.precision_out + 2 * PRECISION_MARGIN)
    }

    pub fn exp2<F: ActuallyUsedField>(&self, x: FieldValue<F>) -> FieldValue<F> {
        let bounds = x.bounds();
        let negative_limit =
            F::from(self.precision_out as u64) * F::negative_power_of_two(self.precision_in);
        let x_clamped_left = x
            .signed_lt(FieldValue::from(negative_limit))
            .select(FieldValue::from(negative_limit), x);
        let positive_limit = self.exp2_positive_limit();
        let x_clamped = x_clamped_left
            .signed_gt(FieldValue::from(positive_limit))
            .select(FieldValue::from(positive_limit), x_clamped_left)
            .with_bounds(FieldBounds::new(negative_limit, positive_limit));
        let (two_pow_int_x, frac_x) = self.init_exp2(x_clamped);
        let approx_frac = self.exp2_approx(frac_x);
        let min = x_clamped.bounds().signed_min().to_signed_number();
        // lower bound on the integer part of x
        let min_int = min >> self.precision_in;
        let mut res = two_pow_int_x * approx_frac;
        if min_int > 0 {
            res *= F::power_of_two(min_int.to_usize().unwrap())
        } else {
            res = res >> (-min_int).to_usize().unwrap()
        }
        res.with_bounds(self.exp2_bounds(bounds))
    }

    fn exp2_public<F: UsedField>(&self, x: F) -> F {
        let x_float = f64::from(x.to_signed_number()) * 2f64.powi(-(self.precision_in as i32));
        F::from(x_float.exp2())
    }

    fn exp2_positive_limit_number<F: UsedField>(&self) -> Number {
        Number::from((F::NUM_BITS as usize - 2 * self.precision_out - 1) as u64)
            * Number::power_of_two(self.precision_in)
    }

    fn exp2_positive_limit<F: UsedField>(&self) -> F {
        F::from(self.exp2_positive_limit_number::<F>())
    }

    fn exp2_bounds<F: UsedField>(&self, bounds: FieldBounds<F>) -> FieldBounds<F> {
        let (min, max) = bounds.min_and_max(true);
        let negative_limit =
            F::from(self.precision_out as u64) * F::negative_power_of_two(self.precision_in);
        let clamped_min = (min - negative_limit)
            .is_lt_zero()
            .select(negative_limit, min);
        let positive_limit: F = self.exp2_positive_limit();
        let clamped_max = (max - positive_limit)
            .is_gt_zero()
            .select(positive_limit, max);
        FieldBounds::new(
            self.exp2_public(clamped_min) - self.eval_gap(&[clamped_min]),
            self.exp2_public(clamped_max) + self.eval_gap(&[clamped_max]),
        )
    }
}

impl<F: UsedField> ArithmeticCircuit<F> for Exp2 {
    fn eval(&self, x: Vec<F>) -> Result<Vec<F>, EvalFailure> {
        if x.len() != 1 {
            panic!("Exp2 requires one input")
        }
        let x = x[0];
        if x.is_ge_zero() && x > self.exp2_positive_limit() {
            return EvalFailure::err_ub("number too large for exp2");
        }
        Ok(vec![self.exp2_public(x)])
    }

    fn eval_gap(&self, x: &[F]) -> F {
        let x = x[0];
        if x.is_ge_zero() {
            self.exp2_public(x) >> 46
        } else {
            F::from(3)
        }
    }

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

    fn run(&self, vals: Vec<FieldValue<F>>) -> Vec<FieldValue<F>>
    where
        F: ActuallyUsedField,
    {
        if vals.len() != 1 {
            panic!("Exp2 requires one input")
        }
        vec![self.exp2(vals[0])]
    }
}

/// An exponentiation algorithm. Computes exp(x). The lowest DOUBLE_PRECISION_MANTISSA bits
/// represent the fractional part.
#[derive(Clone, Debug)]
pub struct Exp {
    // number of bits after the point of the input
    precision_in: usize,
    // number of bits after the point of the output
    precision_out: usize,
}

impl Exp {
    #[allow(unused)]
    pub const fn new(precision_in: usize) -> Self {
        if precision_in > DOUBLE_PRECISION_MANTISSA {
            panic!("precision_in must be at most 52");
        }
        Exp {
            precision_in,
            precision_out: DOUBLE_PRECISION_MANTISSA,
        }
    }

    pub fn exp<F: ActuallyUsedField>(&self, x: FieldValue<F>) -> FieldValue<F> {
        let bounds = x.bounds();
        // the precision of x is precision_in and the precision of LOG2_E is precision_out
        // we want the precision of LOG2_E * x to be precision_out
        let log2_e_x = shift_right(
            FieldValue::from(F::from(LOG2_E)) * x,
            self.precision_in,
            true,
        );
        Exp2::new(self.precision_out)
            .exp2(log2_e_x)
            .with_bounds(self.exp_bounds(bounds))
    }

    fn exp_public<F: UsedField>(&self, x: F) -> F {
        let x_float = f64::from(x.to_signed_number()) * 2f64.powi(-(self.precision_in as i32));
        F::from(x_float.exp())
    }

    fn exp_positive_limit<F: UsedField>(&self) -> F {
        F::from(
            (Number::from(LN_2) * Exp2::new(self.precision_in).exp2_positive_limit_number::<F>())
                >> self.precision_out,
        )
    }

    fn exp_bounds<F: UsedField>(&self, bounds: FieldBounds<F>) -> FieldBounds<F> {
        let (min, max) = bounds.min_and_max(true);
        let negative_limit: F =
            F::from(self.precision_out as u64) * F::negative_power_of_two(self.precision_in);
        let positive_limit: F = self.exp_positive_limit();
        let (clamped_min, clamped_max) =
            if min.signed_gt(positive_limit) || max.signed_lt(negative_limit) {
                (negative_limit, positive_limit)
            } else {
                (min.max(negative_limit, true), max.min(positive_limit, true))
            };
        FieldBounds::new(
            self.exp_public(clamped_min) - self.eval_gap(&[clamped_min]),
            self.exp_public(clamped_max) + self.eval_gap(&[clamped_max]),
        )
    }
}

impl<F: UsedField> ArithmeticCircuit<F> for Exp {
    fn eval(&self, x: Vec<F>) -> Result<Vec<F>, EvalFailure> {
        if x.len() != 1 {
            panic!("Exp requires one input")
        }
        let x = x[0];
        if x.is_ge_zero() && x > self.exp_positive_limit() {
            return EvalFailure::err_ub("number too large for exp");
        }
        if x.is_lt_zero()
            && x <= F::negative_power_of_two(F::NUM_BITS as usize - self.precision_out - 3)
        {
            return EvalFailure::err_ub("number too small for exp");
        }
        Ok(vec![self.exp_public(x)])
    }

    fn eval_gap(&self, x: &[F]) -> F {
        let x = x[0];
        if x.is_ge_zero() {
            self.exp_public(x) >> 44
        } else {
            F::from(3)
        }
    }

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

    fn run(&self, vals: Vec<FieldValue<F>>) -> Vec<FieldValue<F>>
    where
        F: ActuallyUsedField,
    {
        if vals.len() != 1 {
            panic!("Exp requires one input")
        }
        vec![self.exp(vals[0])]
    }
}

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

    impl TestedArithmeticCircuit<ScalarField> for Exp2 {
        fn gen_desc<R: Rng + ?Sized>(rng: &mut R) -> Self {
            let mut precision = 52;
            while rng.gen_bool(0.5) {
                precision -= 1;
            }
            Self::new(precision as usize)
        }

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

        fn input_precisions(&self) -> Vec<usize> {
            vec![self.precision_in]
        }
    }

    impl TestedArithmeticCircuit<ScalarField> for Exp {
        fn gen_desc<R: Rng + ?Sized>(rng: &mut R) -> Self {
            let mut precision = 52;
            while rng.gen_bool(0.5) {
                precision -= 1;
            }
            // we want a precision of at least 45 on the input, otherwise the result might get less
            // precise than the eval_gap
            Self::new(precision.max(45) as usize)
        }

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

        fn input_precisions(&self) -> Vec<usize> {
            vec![self.precision_in]
        }
    }

    #[test]
    fn tested_exp2() {
        Exp2::test(16, 4)
    }

    #[test]
    fn tested_exp() {
        Exp::test(16, 4)
    }
}