arcis-compiler 0.12.0

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::{FieldBounds, IsBounds},
        circuits::{
            boolean::{boolean_value::BooleanValue, utils::subtraction_circuit},
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        compile_passes::new_eda_bit,
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::{FromLeBits, GetBit, Invert, Reveal},
    utils::{number::Number, used_field::UsedField},
};
use num_traits::{Signed, ToPrimitive};
use std::num::NonZeroU128;

/// Divides by a known number.
///
/// Idea:
/// Let's have x∈F. We want to find q and r such that x = q * b + r
/// Let's call f∈F close to p/b.
/// Then x * f in F is very close to p * r/b, and as p is very close to a power of two,
/// the MSBs of x * f will look like the binary decomposition of r/b.
/// We take k MSBs of x * f, multiply by b, right shift by k, and we have the remainder.
#[derive(Debug, Clone)]
pub struct EuclideanByKnown {
    b: NonZeroU128,
}

impl EuclideanByKnown {
    fn get_b<F: UsedField>(&self) -> F {
        let b = self.b.get();
        F::from((b >> 64) as u64) * F::power_of_two(64) + F::from(b as u64)
    }
    fn get_c_number<F: UsedField>(&self) -> Number {
        let modulo = F::modulus() % Number::from(self.b.get());
        if modulo.abs() > Number::from(self.b.get() / 2) {
            // f = ceil(p/b)
            modulo - Number::from(self.b.get())
        } else {
            // f = floor(p/b)
            modulo
        }
    }
    fn get_c<F: UsedField>(&self) -> F {
        F::from(self.get_c_number::<F>())
    }

    /// Returns true if there is a field element in `[-x, x]` for which the result could be wrong.
    fn could_fail_for_k<F: UsedField>(&self, x: F, k: usize) -> bool {
        let two_power_k = Number::power_of_two(k);
        let p = F::modulus();
        if &two_power_k * Number::from(self.b.get()) >= p {
            // There could be an overflow in the multiplication before the right shift.
            return true;
        }
        let exp = F::exponent_close_power_of_two();
        let two_power_exp = Number::power_of_two(exp);
        if &two_power_k * Number::from(self.b.get()) >= two_power_exp {
            // There could be an overflow in the multiplication before the right shift.
            return true;
        }
        let c_abs = self.get_c_number::<F>().abs();
        let b = self.get_b::<F>();
        // + 1 in case x is negative
        let q_abs = (x.abs().unsigned_euclidean_division(b) + F::ONE).to_unsigned_number();
        let b_number = b.to_unsigned_number();
        // The combined effect of q and the gap between p and two_power_exp
        // should not substantially mess with binary expansions.
        &b_number
            * (&c_abs * q_abs + c_abs + Number::power_of_two(exp - k) - 1
                + 2 * (&p - &two_power_exp).abs())
            >= two_power_exp / 2
    }
    fn div_bounds_inner<F: UsedField>(&self, b1: FieldBounds<F>) -> (F, F) {
        let (min1, max1) = b1.to_signed_number_pair();

        let div = Number::from(self.b.get().max(1));
        let offset = Number::from(min1 < 0);

        ((min1 / &div - offset).into(), (max1 / &div).into())
    }
    fn div_bounds<F: UsedField>(&self, b1: FieldBounds<F>) -> FieldBounds<F> {
        if let Some(c1) = b1.as_constant() {
            return self.eval(vec![c1]).unwrap()[0].into();
        }
        if self.b.get() == 1 {
            return b1;
        }
        let (min, max) = self.div_bounds_inner(b1);
        FieldBounds::new(min, max)
    }
    fn full_rem_bounds<F: UsedField>(&self) -> FieldBounds<F> {
        let b: F = self.get_b();
        FieldBounds::new(F::ZERO, b - F::ONE)
    }
    fn rem_bounds<F: UsedField>(&self, b1: FieldBounds<F>) -> FieldBounds<F> {
        if let Some(c1) = b1.as_constant() {
            return self.eval(vec![c1]).unwrap()[1].into();
        }
        self.full_rem_bounds()
    }

    /// Gives a whole range of possible `k`. Only needs to include all possible `k`,
    /// does not need to be smallest of such ranges.
    fn get_k_range<F: UsedField>(&self) -> core::ops::Range<usize> {
        let start = self.b.ilog2() as usize + 2;
        let end = F::exponent_close_power_of_two();
        start..end
    }

    /// Chooses lowest `k` such that and `y€[-x, x]` would work as input.
    fn choose_k<F: UsedField>(&self, x: F) -> Option<usize> {
        self.get_k_range::<F>()
            .find(|&k| !self.could_fail_for_k(x, k))
    }

    /// Returns true if there is no `k` guaranteeing 100% success rate between `-x` and `x`.
    fn could_fail<F: UsedField>(&self, x: F) -> bool {
        self.choose_k(x).is_none()
    }

    /// Chooses `k` with the biggest set of `x` that work.
    fn choose_safest_k<F: UsedField>(&self) -> Option<usize> {
        self.get_k_range::<F>()
            .rev()
            .find(|&k| !self.could_fail_for_k(F::ZERO, k))
    }

    /// Chooses a default `k` in case given the bounds we cannot guarantee success.
    fn default_k<F: UsedField>(&self) -> usize {
        self.choose_safest_k::<F>()
            .unwrap_or(self.b.ilog2() as usize + 3)
    }

    /// Tries to run `EuclideanByKnown` on the given inputs.
    pub fn try_perform<F: ActuallyUsedField>(
        x: FieldValue<F>,
        b: FieldValue<F>,
    ) -> Result<[FieldValue<F>; 2], &'static str> {
        let Some(b) = b.as_constant() else {
            return Err("Divisor should be constant.");
        };
        let b = b.to_signed_number();
        let Some(b) = b.to_u128() else {
            return Err("Divisor should be u128.");
        };
        let Some(b) = NonZeroU128::new(b) else {
            return Err("Divisor should not be zero.");
        };
        if b.get() == 1 << b.ilog2() {
            return Err("Divisor is a power of two. Do a normal division.");
        }
        let circ = EuclideanByKnown { b };
        if circ.could_fail(x.bounds().max_abs()) {
            return Err("Division could fail.");
        }
        let res = circ.run(vec![x]);
        // We could force the bounds,
        // but this has already been done by the caller of `SignedDivide`.
        Ok([res[0], res[1]])
    }

    /// Runs with provided randomness (and provided `k`).
    /// Extracting this is useful for testing.
    pub fn run_with_provided_randomness<F: ActuallyUsedField>(
        &self,
        x: FieldValue<F>,
        eda_bit: FieldValue<F>,
        eda_bit_bits: &[BooleanValue],
        k: usize,
    ) -> [FieldValue<F>; 2] {
        // See our proof.
        let b: F = self.get_b();
        let inv = b.invert(true);
        let c = self.get_c::<F>();
        let exp = F::exponent_close_power_of_two();
        // `revealed` is s_1 + o.
        // `-c * inv` is f (exercise to the reader)
        let revealed = (x * (-c * inv) + eda_bit).reveal();
        let b_number = b.to_unsigned_number();
        // `offset` is o
        let offset = F::from(Number::power_of_two(exp - 1) / b_number);
        // `revealed_with_offset` is s_3
        let revealed_with_offset = revealed + offset;
        let bit_range = exp - k..exp;
        let revealed_with_offset_bits = bit_range
            .clone()
            .map(|idx| revealed_with_offset.get_bit(idx, false))
            .collect::<Vec<_>>();
        let sub_bits = subtraction_circuit(
            revealed_with_offset_bits,
            eda_bit_bits[bit_range].to_vec(),
            Default::default(),
        );

        // `sub` is s_4
        let sub = FieldValue::from_le_bits(sub_bits, false);
        let remainder = (sub * b) >> k;

        [(x - remainder) * inv, remainder]
    }

    /// The function that we try to emulate.
    fn true_eval<F: UsedField>(&self, x: F) -> [F; 2] {
        let b = self.get_b::<F>();
        let q = x.unsigned_euclidean_division_better_bounds(b);
        let r = x - q * b;
        [q, r]
    }
}

impl<F: UsedField> ArithmeticCircuit<F> for EuclideanByKnown {
    fn eval(&self, x: Vec<F>) -> Result<Vec<F>, EvalFailure> {
        let &[x] = x.as_slice() else {
            return EvalFailure::err_imp("only 1 item allowed");
        };
        if self.could_fail(x) {
            return EvalFailure::err_ub("Too big.");
        }
        Ok(self.true_eval(x).to_vec())
    }

    fn bounds(&self, bounds: Vec<FieldBounds<F>>) -> Vec<FieldBounds<F>> {
        let bounds = bounds[0];
        let max_abs = bounds.max_abs();

        if self.could_fail(max_abs) {
            vec![FieldBounds::All, self.full_rem_bounds()]
        } else {
            vec![self.div_bounds(bounds), self.rem_bounds(bounds)]
        }
    }

    fn run(&self, vals: Vec<FieldValue<F>>) -> Vec<FieldValue<F>>
    where
        F: ActuallyUsedField,
    {
        let x = vals[0];
        let b: F = self.get_b();
        let inv = b.invert(true);
        if x.is_plaintext() {
            let sign = x.sign();
            let abs = x.abs();
            let modulo = abs % b;
            let remainder = (modulo * sign + b) % b;
            return vec![(x - remainder) * inv, remainder];
        }
        let exp = F::exponent_close_power_of_two();
        let (eda_bit, eda_bit_bits, ..) = new_eda_bit(exp, false);
        let k = self
            .choose_k(x.bounds().max_abs())
            .unwrap_or(self.default_k::<F>());
        self.run_with_provided_randomness(x, eda_bit, &eda_bit_bits, k)
            .to_vec()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::{
            circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit,
            expressions::{
                bit_expr::{BitExpr, BitInputInfo},
                expr::Expr,
                field_expr::{FieldExpr, InputInfo},
                InputKind,
            },
            global_value::global_expr_store::with_local_expr_store_as_global,
            ir_builder::IRBuilder,
        },
        utils::field::{BaseField, ScalarField},
        EvalValue,
    };
    use rand::Rng;
    use rustc_hash::FxHashMap;
    use std::{marker::PhantomData, rc::Rc};

    impl<F: ActuallyUsedField> TestedArithmeticCircuit<F> for EuclideanByKnown {
        fn gen_desc<R: Rng + ?Sized>(rng: &mut R) -> Self {
            let b = if rng.gen_bool(0.5) {
                // 2^n - 1 is a very interesting scenario, given our k.
                let n = rng.gen_range(1..128);
                NonZeroU128::new((1 << n) - 1).unwrap()
            } else if rng.gen_bool(0.5) {
                // Generate small integers sometimes.
                let n = rng.gen_range(1..128);
                let new_b = rng.gen_range(1..(1 << n));
                NonZeroU128::new(new_b).unwrap()
            } else {
                rng.gen()
            };

            Self { b }
        }

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

    #[test]
    fn tested_base() {
        EuclideanByKnown::test_with_marker(128, 1, PhantomData::<BaseField>)
    }
    #[test]
    fn tested_scalar() {
        EuclideanByKnown::test_with_marker(128, 1, PhantomData::<ScalarField>)
    }
    impl EuclideanByKnown {
        fn max_cannot_fail_for_k<F: UsedField>(&self, k: usize) -> F {
            let mut min = F::ZERO;
            let mut max = F::TWO_INV - F::ONE;
            if !self.could_fail_for_k(max, k) {
                return max;
            }
            assert!(!self.could_fail_for_k(min, k));
            while !(max - min - F::ONE).is_zero_vartime() {
                let mut mid_gap = (max - min) * F::TWO_INV;
                if !mid_gap.is_ge_zero() {
                    mid_gap -= F::TWO_INV;
                }
                assert!(mid_gap.is_ge_zero());
                let mid = min + mid_gap;
                if self.could_fail_for_k(mid, k) {
                    max = mid;
                } else {
                    min = mid;
                }
            }
            min
        }
    }
    fn test_divisor<F: ActuallyUsedField>(b: NonZeroU128) {
        let rng = &mut crate::utils::test_rng::AssertNoRandomness;
        let desc = EuclideanByKnown { b };
        let k = b.ilog2() as usize + 3;
        if desc.could_fail_for_k(F::ZERO, k) {
            // If no number can be divided, not even 0, it's better to not bother.
            return;
        }
        let max_cannot_fail = desc.max_cannot_fail_for_k::<F>(k);

        {
            // Testing that max_cannot_fail is really max_cannot_fail.
            assert!(!desc.could_fail_for_k(max_cannot_fail, k));
            if max_cannot_fail != F::TWO_INV - F::ONE {
                assert!(desc.could_fail_for_k(max_cannot_fail + F::ONE, k));
            }
        }
        let mut tested_v = vec![F::ZERO, F::ONE];
        for i in [F::ZERO, F::ONE, Number::from(b.get() - 1).into()] {
            if i > max_cannot_fail {
                break;
            }
            tested_v.push(max_cannot_fail - i);
            tested_v.push(i - max_cannot_fail);
        }
        let mut expr_store = IRBuilder::new(true);
        let exp = F::exponent_close_power_of_two();
        let bit_ids = (0..exp)
            .map(|idx| {
                expr_store.new_expr(Expr::Bit(BitExpr::Input(
                    idx,
                    Rc::new(BitInputInfo::default()),
                )))
            })
            .collect::<Vec<_>>();
        let output = with_local_expr_store_as_global(
            || {
                let unknown_ss = Rc::new(InputInfo::<F>::from(InputKind::Secret));
                let x = FieldValue::new(FieldExpr::Input(exp, unknown_ss.clone()));
                let eda_bit_bits = bit_ids
                    .into_iter()
                    .map(BooleanValue::new)
                    .collect::<Vec<_>>();
                let eda_bit = FieldValue::from_le_bits(eda_bit_bits.clone(), false);
                desc.run_with_provided_randomness(x, eda_bit, &eda_bit_bits, k)[1].get_id()
            },
            &mut expr_store,
        );
        let ir = expr_store.into_ir(vec![output]);
        // Testing a few values of the k MSBs.
        let msb_iter = [
            Number::from(0u128),
            Number::power_of_two(k) - 1,
            Number::power_of_two(k - 1),
        ];
        let mut inputs = FxHashMap::default();
        for eda_lsb_value in [false, true] {
            // Setting the LSBs of the edaBit to eda_lsb_value.
            for j in 0..(exp - k) {
                inputs.insert(j, EvalValue::Bit(eda_lsb_value));
            }
            for i in &msb_iter {
                // Setting the MSBs of the edaBit.
                for j in 0..k {
                    inputs.insert(j + exp - k, EvalValue::Bit(i.bit(j)));
                }
                // testing interesting values for x
                for x in tested_v.iter().copied() {
                    inputs.insert(exp, F::field_to_eval_value(x));
                    let truth = desc.true_eval(x)[1].to_signed_number();
                    assert_eq!(
                        ir.eval(rng, &mut inputs).unwrap()[0].to_signed_number(),
                        truth,
                        "b:{b:#x} x:{x:?} eda_bit:{i}-{eda_lsb_value} c:{:?}",
                        desc.get_c::<F>(),
                    );
                }
            }
        }
    }
    fn test_field_divisors<F: ActuallyUsedField>() {
        for b in 1..=62 {
            test_divisor::<F>(NonZeroU128::new(b).unwrap());
        }
        for pow in 7..=127 {
            // Divisors like this are the limit for `k`.
            // 6 seemed interesting during limit testing.
            for j in [1, 6] {
                test_divisor::<F>(NonZeroU128::new((1 << pow) - j).unwrap());
            }
        }
    }
    #[test]
    fn test_divisors() {
        test_field_divisors::<BaseField>();
        test_field_divisors::<ScalarField>();
    }
    #[test]
    fn test_base_field() {
        {
            // Test that the algorithm is actually usable for u64 % u64;
            let b = NonZeroU128::new(u64::MAX as u128).unwrap();
            let desc = EuclideanByKnown { b };
            assert!(!desc.could_fail::<BaseField>(u64::MAX.into()));
        }
        let test_divisor = test_divisor::<BaseField>;
        {
            // This check was useful once.
            // Checking multiple divisors of 2^255 - 18
            let mut b = (1u128 << 127) - 3;
            test_divisor(NonZeroU128::new(b).unwrap());
            for div in [5, 13, 41, 137, 53630713] {
                b /= div;
                if div > 63 {
                    test_divisor(NonZeroU128::new(div).unwrap());
                }
                test_divisor(NonZeroU128::new(b).unwrap());
            }
        }
        {
            // Checking multiple divisors of 2^255 - 20
            let mut b = 12 * 65147;
            test_divisor(NonZeroU128::new(b).unwrap());
            for div in [2, 2, 3] {
                b /= div;
                test_divisor(NonZeroU128::new(b).unwrap());
            }
        }
    }
}