Comments and cleanup

jolt-core/src/field/mod.rs

@@ -39,13 +39,20 @@ pub trait JoltField:
     + CanonicalDeserialize
     + Hash
 {
+    /// Number of bytes occupied by a single field element.
     const NUM_BYTES: usize;
+    /// An implementation of `JoltField` may use some precomputed lookup tables to speed up the
+    /// conversion of small primitive integers (e.g. `u16` values) into field elements. For example,
+    /// the arkworks BN254 scalar field requires a conversion into Montgomery form, which naively
+    /// requires a field multiplication, but can instead be looked up.
     type SmallValueLookupTables: Clone + Default + CanonicalSerialize + CanonicalDeserialize = ();
 
     fn random<R: rand_core::RngCore>(rng: &mut R) -> Self;
+    /// Computes the small-value lookup tables.
     fn compute_lookup_tables() -> Self::SmallValueLookupTables {
         unimplemented!("Small-value lookup tables are unimplemented")
     }
+    /// Initializes the static lookup tables using the provided values.
     fn initialize_lookup_tables(_init: Self::SmallValueLookupTables) {
         unimplemented!("Small-value lookup tables are unimplemented")
     }
@@ -70,6 +77,12 @@ pub trait JoltField:
     fn montgomery_r2() -> Option<Self> {
         None
     }
+
+    /// Does an "unchecked" field multiplication with a `u64`.
+    /// WARNING: For `x.mul_u64_unchecked(y)` to be equal to `x * F::from_u64(y)`,
+    /// which is presumably what you want, we need to correct for the fact that `y` is
+    /// not in Montgomery form. This is typically accomplished by multiplying the left
+    /// operand by an additional R^2 factor (see `JoltField::montgomery_r2`).
     #[inline(always)]
     fn mul_u64_unchecked(&self, n: u64) -> Self {
         *self * Self::from_u64(n)

jolt-core/src/poly/eq_poly.rs

@@ -31,6 +31,18 @@ impl<F: JoltField> EqPolynomial<F> {
         }
     }
 
+    /// When evaluating a multilinear polynomial on a point `r`, we first compute the EQ evaluation table
+    /// for `r`, then dot-product those values with the coefficients of the polynomial.
+    ///
+    /// However, if the polynomial in question is a `CompactPolynomial`, its coefficients are represented
+    /// by primitive integers while the dot product needs to be computed using Montgomery multiplication.
+    ///
+    /// To avoid converting every polynomial coefficient to Montgomery form, we can instead introduce an
+    /// additional R^2 factor to every element in the EQ evaluation table and performing the dot product
+    /// using `JoltField::mul_u64_unchecked`.
+    ///
+    /// We can efficiently compute the EQ table with this additional R^2 factor by initializing the root of
+    /// the dynamic programming tree to R^2 instead of 1.
     #[tracing::instrument(skip_all, name = "EqPolynomial::evals_with_r2")]
     pub fn evals_with_r2(r: &[F]) -> Vec<F> {
         let ell = r.len();

jolt-core/src/poly/spartan_interleaved_poly.rs

@@ -6,6 +6,8 @@ use super::{
 };
 #[cfg(test)]
 use crate::poly::dense_mlpoly::DensePolynomial;
+#[cfg(test)]
+use crate::r1cs::inputs::JoltR1CSInputs;
 use crate::{
     field::{JoltField, OptimizedMul},
     r1cs::builder::{eval_offset_lc, Constraint, OffsetEqConstraint},
@@ -64,7 +66,7 @@ impl<F: JoltField> SpartanInterleavedPolynomial<F> {
                         if !constraint.a.terms().is_empty() {
                             az_coeff = constraint
                                 .a
-                                .evaluate_row_i128(flattened_polynomials, step_index);
+                                .evaluate_row(flattened_polynomials, step_index);
                             if !az_coeff.is_zero() {
                                 coeffs.push((global_index, az_coeff).into());
                             }
@@ -74,7 +76,7 @@ impl<F: JoltField> SpartanInterleavedPolynomial<F> {
                         if !constraint.b.terms().is_empty() {
                             bz_coeff = constraint
                                 .b
-                                .evaluate_row_i128(flattened_polynomials, step_index);
+                                .evaluate_row(flattened_polynomials, step_index);
                             if !bz_coeff.is_zero() {
                                 coeffs.push((global_index + 1, bz_coeff).into());
                             }
@@ -84,12 +86,21 @@ impl<F: JoltField> SpartanInterleavedPolynomial<F> {
                             let cz_coeff = az_coeff * bz_coeff;
                             #[cfg(test)]
                             {
-                                assert_eq!(
-                                    cz_coeff,
-                                    constraint
-                                        .c
-                                        .evaluate_row_i128(flattened_polynomials, step_index)
-                                );
+                                if cz_coeff != constraint
+                                    .c
+                                    .evaluate_row(flattened_polynomials, step_index) {
+                                        let mut constraint_string = String::new();
+                                        let _ = constraint
+                                            .pretty_fmt::<4, JoltR1CSInputs, F>(
+                                                &mut constraint_string,
+                                                flattened_polynomials,
+                                                step_index,
+                                            );
+                                        println!("{constraint_string}");
+                                        panic!(
+                                            "Uniform constraint {constraint_index} violated at step {step_index}",
+                                        );
+                                    }
                             }
                             coeffs.push((global_index + 2, cz_coeff).into());
                         }
@@ -135,7 +146,7 @@ impl<F: JoltField> SpartanInterleavedPolynomial<F> {
                                     step_index,
                                     next_step_index,
                                 );
-                                assert_eq!(bz_coeff, 0);
+                                assert_eq!(bz_coeff, 0, "Non-uniform constraint {constraint_index} violated at step {step_index}");
                             }
                         } else {
                             // Bz