arcis-compiler 0.9.0

//! Arcis implementation of <https://github.com/solana-program/zk-elgamal-proof/blob/main/zk-sdk/src/range_proof/util.rs>

use crate::{core::global_value::value::FieldValue, utils::field::ScalarField};

/// Byte length of a compressed Ristretto point or scalar in Curve255519
pub(crate) const UNIT_LEN: usize = 32;

/// Represents a degree-1 vector polynomial, such as `a + b*x`.
///
/// A vector polynomial is a vector of polynomials. For example:
/// `[a_1 + b_1*x, a_2 + b_2*x, ..., a_n + b_n*x]`
///
/// This struct stores the constant terms (the `a_i`'s) and the linear terms (the `b_i`'s) as two
/// separate vectors of scalars. In the Bulletproofs protocol, this is used to represent the `l(x)`
/// and `r(x)` vector polynomials.
pub struct VecPoly1(
    pub Vec<FieldValue<ScalarField>>,
    pub Vec<FieldValue<ScalarField>>,
);

impl VecPoly1 {
    /// Creates a new zero-filled `VecPoly1` of the given size.
    pub fn zero(n: usize) -> Self {
        VecPoly1(
            vec![FieldValue::<ScalarField>::from(0); n],
            vec![FieldValue::<ScalarField>::from(0); n],
        )
    }

    /// Computes the inner product of two vector polynomials.
    ///
    /// The result is a scalar polynomial of degree 2, `t(x) = t_0 + t_1*x + t_2*x^2`,
    /// which is the inner product `t(x) = <l(x), r(x)>` in the Bulletproofs protocol.
    ///
    /// This function uses Karatsuba's method for efficient polynomial multiplication.
    pub fn inner_product(&self, rhs: &VecPoly1) -> Option<Poly2> {
        // Uses Karatsuba's method
        let l = self;
        let r = rhs;

        let t0 = inner_product(&l.0, &r.0)?;
        let t2 = inner_product(&l.1, &r.1)?;

        let l0_plus_l1 = add_vec(&l.0, &l.1);
        let r0_plus_r1 = add_vec(&r.0, &r.1);

        let t1 = inner_product(&l0_plus_l1, &r0_plus_r1)? - t0 - t2;

        Some(Poly2(t0, t1, t2))
    }

    /// Evaluates the vector polynomial at a given scalar point `x`.
    ///
    /// The result is a vector of scalars, where each element is the evaluation
    /// of the corresponding polynomial at `x`.
    pub fn eval(&self, x: FieldValue<ScalarField>) -> Vec<FieldValue<ScalarField>> {
        let n = self.0.len();
        let mut out = vec![FieldValue::<ScalarField>::from(0); n];
        #[allow(clippy::needless_range_loop)]
        for i in 0..n {
            out[i] = self.0[i] + self.1[i] * x;
        }
        out
    }
}

/// Represents a degree-2 scalar polynomial `a + b*x + c*x^2`.
///
/// In the Bulletproofs protocol, this represents the `t(x)` polynomial
/// that results from the inner product of `l(x)` and `r(x)`.
pub struct Poly2(
    pub FieldValue<ScalarField>,
    pub FieldValue<ScalarField>,
    pub FieldValue<ScalarField>,
);

impl Poly2 {
    /// Evaluates the polynomial at a scalar point `x` using Horner's method.
    pub fn eval(&self, x: FieldValue<ScalarField>) -> FieldValue<ScalarField> {
        self.0 + x * (self.1 + x * self.2)
    }
}
/// Provides an iterator over the powers of a `FieldValue<ScalarField>` (e.g., `1, x, x^2, x^3,
/// ...`).
///
/// This struct is created by the `exp_iter` function.
pub struct ScalarExp {
    x: FieldValue<ScalarField>,
    next_exp_x: FieldValue<ScalarField>,
}

impl Iterator for ScalarExp {
    type Item = FieldValue<ScalarField>;

    fn next(&mut self) -> Option<FieldValue<ScalarField>> {
        let exp_x = self.next_exp_x;
        self.next_exp_x *= self.x;
        Some(exp_x)
    }
}

/// Returns an iterator of the powers of `x`, starting with `x^0 = 1`.
pub fn exp_iter(x: FieldValue<ScalarField>) -> ScalarExp {
    let next_exp_x = FieldValue::<ScalarField>::from(1);
    ScalarExp { x, next_exp_x }
}

/// Computes the component-wise sum of two vectors of scalars.
/// Panics if the lengths of the vectors two do not match.
pub fn add_vec(
    a: &[FieldValue<ScalarField>],
    b: &[FieldValue<ScalarField>],
) -> Vec<FieldValue<ScalarField>> {
    if a.len() != b.len() {
        panic!("lengths of vectors don't match for vector addition");
    }
    let mut out = vec![FieldValue::<ScalarField>::from(0); b.len()];
    for i in 0..a.len() {
        out[i] = a[i] + b[i];
    }
    out
}

/// Computes the inner product of two vectors of scalars.
///
/// The inner product is defined as:
/// `<a, b> = a_1*b_1 + a_2*b_2 + ... + a_n*b_n`
///
/// Returns `None` if the vectors have different lengths.
pub fn inner_product(
    a: &[FieldValue<ScalarField>],
    b: &[FieldValue<ScalarField>],
) -> Option<FieldValue<ScalarField>> {
    let mut out = FieldValue::<ScalarField>::from(0);
    if a.len() != b.len() {
        return None;
    }
    for i in 0..a.len() {
        out += a[i] * b[i];
    }
    Some(out)
}