Rust implementation of extension field elements multiplication

hydra/garaga/hints/extf_mul.py

@@ -1,8 +1,10 @@
 import operator
 from functools import reduce
 
+from garaga import garaga_rs
 from garaga.algebra import ModuloCircuitElement, Polynomial, PyFelt
 from garaga.definitions import (
+    BLS12_381_ID,
     direct_to_tower,
     get_base_field,
     get_irreducible_poly,
@@ -19,29 +21,17 @@ def nondeterministic_extension_field_mul_divmod(
     curve_id: int,
     extension_degree: int,
 ) -> tuple[list[PyFelt], list[PyFelt]]:
-
-    Ps = [Polynomial(P) for P in Ps]
     field = get_base_field(curve_id)
-
-    P_irr = get_irreducible_poly(curve_id, extension_degree)
-
-    z_poly = reduce(operator.mul, Ps)  #  Π(Pi)
-    z_polyq, z_polyr = divmod(z_poly, P_irr)
-
-    z_polyr_coeffs = z_polyr.get_coeffs()
-    z_polyq_coeffs = z_polyq.get_coeffs()
-    # assert len(z_polyq_coeffs) <= (
-    #     extension_degree - 1
-    # ), f"len z_polyq_coeffs={len(z_polyq_coeffs)}, degree: {z_polyq.degree()}"
-    assert (
-        len(z_polyr_coeffs) <= extension_degree
-    ), f"len z_polyr_coeffs={len(z_polyr_coeffs)}, degree: {z_polyr.degree()}"
-
-    # Extend polynomials with 0 coefficients to match the expected lengths.
-    # TODO : pass exact expected max degree when len(Ps)>2.
-    z_polyq_coeffs += [field(0)] * (extension_degree - 1 - len(z_polyq_coeffs))
-    z_polyr_coeffs += [field(0)] * (extension_degree - len(z_polyr_coeffs))
-
+    if curve_id == BLS12_381_ID:
+        nb = 48
+    else:
+        nb = 32
+    ps = [[c.value.to_bytes(nb, "big") for c in P] for P in Ps]
+    q, r = garaga_rs.nondeterministic_extension_field_mul_divmod(
+        curve_id, extension_degree, ps
+    )
+    z_polyq_coeffs = [field(c) for c in q]
+    z_polyr_coeffs = [field(c) for c in r]
     return (z_polyq_coeffs, z_polyr_coeffs)
 
 

tools/garaga_rs/src/ecip/polynomial.rs

@@ -80,6 +80,10 @@ impl<F: IsPrimeField> Polynomial<F> {
         Polynomial::new(vec![FieldElement::<F>::zero()])
     }
 
+    pub fn one() -> Self {
+        Polynomial::new(vec![FieldElement::<F>::one()])
+    }
+
     pub fn mul_with_ref(&self, other: &Polynomial<F>) -> Polynomial<F> {
         if self.degree() == -1 || other.degree() == -1 {
             return Polynomial::zero();

tools/garaga_rs/src/extf_mul.rs

@@ -0,0 +1,155 @@
+use crate::ecip::polynomial::{pad_with_zero_coefficients_to_length, Polynomial};
+use lambdaworks_math::{
+    elliptic_curve::short_weierstrass::curves::{
+        bls12_381::field_extension::{BLS12381PrimeField, BLS12381_PRIME_FIELD_ORDER},
+        bn_254::field_extension::{BN254PrimeField, BN254_PRIME_FIELD_ORDER},
+    },
+    field::element::FieldElement,
+    traits::ByteConversion,
+    unsigned_integer::element::{U256, U384},
+};
+use num_bigint::BigUint;
+
+// irreducible polynomial coefficients
+const IPC: [[&[i8]; 2]; 2] = [
+    // curve 0: BN254
+    [
+        &[82, 0, 0, -18, 0, 0, 1],                   // 6
+        &[82, 0, 0, 0, 0, 0, -18, 0, 0, 0, 0, 0, 1], // 12
+    ],
+    // curve 1: BLS12-381
+    [
+        &[2, 0, 0, -2, 0, 0, 1],                   // 6
+        &[2, 0, 0, 0, 0, 0, -2, 0, 0, 0, 0, 0, 1], // 12
+    ],
+];
+
+// Returns (Q(X), R(X)) such that Π(Pi)(X) = Q(X) * P_irr(X) + R(X), for a given curve and extension degree.
+// R(X) is the result of the multiplication in the extension field.
+// Q(X) is used for verification.
+pub fn nondeterministic_extension_field_mul_divmod(
+    curve_id: usize,
+    ext_degree: usize,
+    list_coeffs: Vec<Vec<Vec<u8>>>,
+) -> Result<(Vec<BigUint>, Vec<BigUint>), String> {
+    assert!(6 <= ext_degree && ext_degree <= 12 && ext_degree % 6 == 0);
+    let ext_degree_index = ext_degree / 6 - 1;
+
+    if curve_id == 0 {
+        let mut ps = Vec::new();
+        for i in 0..list_coeffs.len() {
+            let coeffs = (&list_coeffs[i])
+                .into_iter()
+                .map(|x| {
+                    FieldElement::<BN254PrimeField>::from_bytes_be(&x)
+                        .map_err(|e| format!("Byte conversion error: {:?}", e))
+                })
+                .collect::<Result<Vec<FieldElement<BN254PrimeField>>, _>>()?;
+            ps.push(Polynomial::new(coeffs));
+        }
+
+        let coeffs = IPC[curve_id][ext_degree_index]
+            .into_iter()
+            .map(|x| {
+                if *x >= 0 {
+                    FieldElement::<BN254PrimeField>::from(*x as u64)
+                } else {
+                    FieldElement::<BN254PrimeField>::from(
+                        &(BN254_PRIME_FIELD_ORDER - U256::from_u64(-x as u64)),
+                    )
+                }
+            })
+            .collect::<Vec<FieldElement<BN254PrimeField>>>();
+        let p_irr = Polynomial::new(coeffs);
+
+        let mut z_poly = Polynomial::one();
+        for i in 0..ps.len() {
+            z_poly = z_poly.mul_with_ref(&ps[i]);
+        }
+
+        let (mut z_polyq, mut z_polyr) = z_poly.divmod(&p_irr);
+        assert!(z_polyr.coefficients.len() <= ext_degree);
+
+        // Extend polynomials with 0 coefficients to match the expected lengths.
+        // TODO : pass exact expected max degree when len(Ps)>2.
+        if z_polyq.coefficients.len() < ext_degree - 1 {
+            pad_with_zero_coefficients_to_length(&mut z_polyq, ext_degree - 1);
+        }
+        if z_polyr.coefficients.len() < ext_degree {
+            pad_with_zero_coefficients_to_length(&mut z_polyr, ext_degree);
+        }
+
+        let q = z_polyq
+            .coefficients
+            .into_iter()
+            .map(|x| BigUint::from_bytes_be(&x.to_bytes_be()))
+            .collect();
+        let r = z_polyr
+            .coefficients
+            .into_iter()
+            .map(|x| BigUint::from_bytes_be(&x.to_bytes_be()))
+            .collect();
+
+        return Ok((q, r));
+    }
+
+    if curve_id == 1 {
+        let mut ps = Vec::new();
+        for i in 0..list_coeffs.len() {
+            let coeffs = (&list_coeffs[i])
+                .into_iter()
+                .map(|x| {
+                    FieldElement::<BLS12381PrimeField>::from_bytes_be(&x)
+                        .map_err(|e| format!("Byte conversion error: {:?}", e))
+                })
+                .collect::<Result<Vec<FieldElement<BLS12381PrimeField>>, _>>()?;
+            ps.push(Polynomial::new(coeffs));
+        }
+
+        let coeffs = IPC[curve_id][ext_degree_index]
+            .into_iter()
+            .map(|x| {
+                if *x >= 0 {
+                    FieldElement::<BLS12381PrimeField>::from(*x as u64)
+                } else {
+                    FieldElement::<BLS12381PrimeField>::from(
+                        &(BLS12381_PRIME_FIELD_ORDER - U384::from_u64(-x as u64)),
+                    )
+                }
+            })
+            .collect::<Vec<FieldElement<BLS12381PrimeField>>>();
+        let p_irr = Polynomial::new(coeffs);
+
+        let mut z_poly = Polynomial::one();
+        for i in 0..ps.len() {
+            z_poly = z_poly.mul_with_ref(&ps[i]);
+        }
+
+        let (mut z_polyq, mut z_polyr) = z_poly.divmod(&p_irr);
+        assert!(z_polyr.coefficients.len() <= ext_degree);
+
+        // Extend polynomials with 0 coefficients to match the expected lengths.
+        // TODO : pass exact expected max degree when len(Ps)>2.
+        if z_polyq.coefficients.len() < ext_degree - 1 {
+            pad_with_zero_coefficients_to_length(&mut z_polyq, ext_degree - 1);
+        }
+        if z_polyr.coefficients.len() < ext_degree {
+            pad_with_zero_coefficients_to_length(&mut z_polyr, ext_degree);
+        }
+
+        let q = z_polyq
+            .coefficients
+            .into_iter()
+            .map(|x| BigUint::from_bytes_be(&x.to_bytes_be()))
+            .collect();
+        let r = z_polyr
+            .coefficients
+            .into_iter()
+            .map(|x| BigUint::from_bytes_be(&x.to_bytes_be()))
+            .collect();
+
+        return Ok((q, r));
+    }
+
+    panic!("Curve ID {} not supported", curve_id);
+}

tools/garaga_rs/src/lib.rs

@@ -1,6 +1,7 @@
 pub mod bls12_381_final_exp_witness;
 pub mod bn254_final_exp_witness;
 pub mod ecip;
+pub mod extf_mul;
 
 use ark_ec::{pairing::Pairing, AffineRepr};
 use ark_ff::PrimeField;
@@ -25,6 +26,10 @@ fn garaga_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(multi_miller_loop, m)?)?;
     m.add_function(wrap_pyfunction!(get_final_exp_witness, m)?)?;
     m.add_function(wrap_pyfunction!(hades_permutation, m)?)?;
+    m.add_function(wrap_pyfunction!(
+        nondeterministic_extension_field_mul_divmod,
+        m
+    )?)?;
     m.add_function(wrap_pyfunction!(zk_ecip_hint, m)?)?;
     Ok(())
 }
@@ -427,6 +432,23 @@ fn get_final_exp_witness(
     panic!("Curve ID {} not supported", curve_id);
 }
 
+#[pyfunction]
+fn nondeterministic_extension_field_mul_divmod(
+    py: Python,
+    curve_id: usize,
+    ext_degree: usize,
+    py_list: &Bound<'_, PyList>,
+) -> PyResult<PyObject> {
+    let ps = py_list
+        .into_iter()
+        .map(|x| x.extract())
+        .collect::<Result<Vec<Vec<Vec<u8>>>, _>>()?;
+    let (q, r) = extf_mul::nondeterministic_extension_field_mul_divmod(curve_id, ext_degree, ps)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
+    let py_tuple = PyTuple::new_bound(py, [PyList::new_bound(py, q), PyList::new_bound(py, r)]);
+    return Ok(py_tuple.into());
+}
+
 #[pyfunction]
 fn hades_permutation(
     py: Python,