Merge pull request #386 from imikushin/jolt-field-arithmetic-by-ref

Implement JoltField Arithmetic by Reference

jolt-core/src/field/ark.rs

@@ -1,6 +1,11 @@
 use ark_ff::{PrimeField, UniformRand};
+use ark_std::Zero;
 
-use super::JoltField;
+use super::{FieldOps, JoltField};
+
+impl FieldOps for ark_bn254::Fr {}
+impl<'a, 'b> FieldOps<&'b ark_bn254::Fr, ark_bn254::Fr> for &'a ark_bn254::Fr {}
+impl<'b> FieldOps<&'b ark_bn254::Fr, ark_bn254::Fr> for ark_bn254::Fr {}
 
 impl JoltField for ark_bn254::Fr {
     const NUM_BYTES: usize = 32;
@@ -9,22 +14,6 @@ impl JoltField for ark_bn254::Fr {
         <Self as UniformRand>::rand(rng)
     }
 
-    fn is_zero(&self) -> bool {
-        <Self as ark_std::Zero>::is_zero(self)
-    }
-
-    fn is_one(&self) -> bool {
-        <Self as ark_std::One>::is_one(self)
-    }
-
-    fn zero() -> Self {
-        <Self as ark_std::Zero>::zero()
-    }
-
-    fn one() -> Self {
-        <Self as ark_std::One>::one()
-    }
-
     fn from_u64(n: u64) -> Option<Self> {
         <Self as ark_ff::PrimeField>::from_u64(n)
     }

jolt-core/src/field/binius.rs

@@ -1,8 +1,9 @@
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use ark_std::{One, Zero};
 use binius_field::{BinaryField128b, BinaryField128bPolyval};
 use std::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
 
-use super::JoltField;
+use super::{FieldOps, JoltField};
 
 impl BiniusConstructable for BinaryField128b {
     fn new(n: u64) -> Self {
@@ -29,28 +30,84 @@ pub trait BiniusConstructable {
 #[derive(Default, Debug, Copy, Clone, Eq, PartialEq)]
 pub struct BiniusField<F: BiniusSpecific>(F);
 
-/// Wrapper for all generic BiniusField functionality.
-impl<F: BiniusSpecific> JoltField for BiniusField<F> {
-    const NUM_BYTES: usize = 16;
+impl<F: BiniusSpecific> FieldOps for BiniusField<F> {}
+impl<'a, 'b, F: BiniusSpecific> FieldOps<&'b BiniusField<F>, BiniusField<F>>
+    for &'a BiniusField<F>
+{
+}
+impl<'b, F: BiniusSpecific> FieldOps<&'b BiniusField<F>, BiniusField<F>> for BiniusField<F> {}
 
-    fn random<R: rand_core::RngCore>(rng: &mut R) -> Self {
-        Self(F::random(rng))
+impl<'a, 'b, F: BiniusSpecific> Add<&'b BiniusField<F>> for &'a BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn add(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        *self + other
     }
+}
 
-    fn is_zero(&self) -> bool {
-        self.0.is_zero()
+impl<'a, 'b, F: BiniusSpecific> Sub<&'b BiniusField<F>> for &'a BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn sub(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        *self - other
     }
+}
 
-    fn is_one(&self) -> bool {
-        self.0 == Self::one().0
+impl<'a, 'b, F: BiniusSpecific> Mul<&'b BiniusField<F>> for &'a BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn mul(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        *self * other
     }
+}
 
-    fn zero() -> Self {
-        Self(F::ZERO)
+impl<'a, 'b, F: BiniusSpecific> Div<&'b BiniusField<F>> for &'a BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn div(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        *self / other
     }
+}
 
-    fn one() -> Self {
-        Self(F::ONE)
+impl<'b, F: BiniusSpecific> Add<&'b BiniusField<F>> for BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn add(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        BiniusField(self.0 + other.0)
+    }
+}
+
+impl<'b, F: BiniusSpecific> Sub<&'b BiniusField<F>> for BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn sub(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        BiniusField(self.0 - other.0)
+    }
+}
+
+impl<'b, F: BiniusSpecific> Mul<&'b BiniusField<F>> for BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn mul(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        BiniusField(self.0 * other.0)
+    }
+}
+
+impl<'b, F: BiniusSpecific> Div<&'b BiniusField<F>> for BiniusField<F> {
+    type Output = BiniusField<F>;
+
+    fn div(self, other: &'b BiniusField<F>) -> BiniusField<F> {
+        #[allow(clippy::suspicious_arithmetic_impl)] // clippy doesn't know algebra
+        BiniusField(self.0 * other.0.invert().unwrap())
+    }
+}
+
+/// Wrapper for all generic BiniusField functionality.
+impl<F: BiniusSpecific> JoltField for BiniusField<F> {
+    const NUM_BYTES: usize = 16;
+
+    fn random<R: rand_core::RngCore>(rng: &mut R) -> Self {
+        Self(F::random(rng))
     }
 
     fn from_u64(n: u64) -> Option<Self> {
@@ -81,6 +138,26 @@ impl<F: BiniusSpecific> JoltField for BiniusField<F> {
     }
 }
 
+impl<F: BiniusSpecific> Zero for BiniusField<F> {
+    fn zero() -> Self {
+        Self(F::ZERO)
+    }
+
+    fn is_zero(&self) -> bool {
+        self.0.is_zero()
+    }
+}
+
+impl<F: BiniusSpecific> One for BiniusField<F> {
+    fn one() -> Self {
+        Self(F::ONE)
+    }
+
+    fn is_one(&self) -> bool {
+        self.0 == Self::one().0
+    }
+}
+
 impl<F: BiniusSpecific> Neg for BiniusField<F> {
     type Output = Self;
 

jolt-core/src/field/mod.rs

@@ -2,15 +2,24 @@ use std::fmt::Debug;
 use std::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
 
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use ark_std::{One, Zero};
+
+pub trait FieldOps<Rhs = Self, Output = Self>:
+    Add<Rhs, Output = Output>
+    + Sub<Rhs, Output = Output>
+    + Mul<Rhs, Output = Output>
+    + Div<Rhs, Output = Output>
+{
+}
 
 pub trait JoltField:
     'static
     + Sized
+    + Zero
+    + One
     + Neg<Output = Self>
-    + Add<Output = Self>
-    + Sub<Output = Self>
-    + Mul<Output = Self>
-    + Div<Output = Self>
+    + FieldOps<Self, Self>
+    + for<'a> FieldOps<&'a Self, Self>
     + AddAssign<Self>
     + SubAssign<Self>
     + MulAssign<Self>
@@ -30,18 +39,25 @@ pub trait JoltField:
     const NUM_BYTES: usize;
     fn random<R: rand_core::RngCore>(rng: &mut R) -> Self;
 
-    fn is_zero(&self) -> bool;
-    fn is_one(&self) -> bool;
-    fn zero() -> Self;
-    fn one() -> Self;
     fn from_u64(n: u64) -> Option<Self>;
     fn from_i64(val: i64) -> Self;
     fn square(&self) -> Self;
     fn from_bytes(bytes: &[u8]) -> Self;
     fn inverse(&self) -> Option<Self>;
+}
+
+pub trait OptimizedMul<Rhs, Output>: Sized + Mul<Rhs, Output = Output> {
+    fn mul_0_optimized(self, other: Rhs) -> Self::Output;
+    fn mul_1_optimized(self, other: Rhs) -> Self::Output;
+    fn mul_01_optimized(self, other: Rhs) -> Self::Output;
+}
 
+impl<T> OptimizedMul<T, T> for T
+where
+    T: JoltField,
+{
     #[inline(always)]
-    fn mul_0_optimized(self, other: Self) -> Self {
+    fn mul_0_optimized(self, other: T) -> T {
         if self.is_zero() || other.is_zero() {
             Self::zero()
         } else {
@@ -50,7 +66,7 @@ pub trait JoltField:
     }
 
     #[inline(always)]
-    fn mul_1_optimized(self, other: Self) -> Self {
+    fn mul_1_optimized(self, other: T) -> T {
         if self.is_one() {
             other
         } else if other.is_one() {
@@ -61,15 +77,11 @@ pub trait JoltField:
     }
 
     #[inline(always)]
-    fn mul_01_optimized(self, other: Self) -> Self {
+    fn mul_01_optimized(self, other: T) -> T {
         if self.is_zero() || other.is_zero() {
             Self::zero()
-        } else if self.is_one() {
-            other
-        } else if other.is_one() {
-            self
         } else {
-            self * other
+            self.mul_1_optimized(other)
         }
     }
 }

jolt-core/src/poly/commitment/hyperkzg.rs

@@ -216,7 +216,7 @@ where
         b"v",
         &v.iter().flatten().cloned().collect::<Vec<P::ScalarField>>(),
     );
-    let q_powers = transcript.challenge_scalar_powers(b"r", k);
+    let q_powers: Vec<P::ScalarField> = transcript.challenge_scalar_powers(b"r", k);
 
     transcript.append_points(
         b"W",
@@ -245,7 +245,7 @@ where
     // L0, L1, L2 can be replaced by single MSM of C with the powers of q multiplied by (1 + d_0 + d_1)
     // with additionally concatenated inputs for scalars/bases.
 
-    let q_power_multiplier = P::ScalarField::one() + d_0 + d_1;
+    let q_power_multiplier: P::ScalarField = P::ScalarField::one() + d_0 + d_1;
 
     let q_powers_multiplied: Vec<P::ScalarField> = q_powers
         .par_iter()

jolt-core/src/poly/commitment/kzg.rs

@@ -5,7 +5,7 @@ use crate::utils::errors::ProofVerifyError;
 use ark_ec::scalar_mul::fixed_base::FixedBase;
 use ark_ec::{pairing::Pairing, AffineRepr, CurveGroup};
 use ark_ff::PrimeField;
-use ark_std::UniformRand;
+use ark_std::{One, UniformRand};
 use rand_core::{CryptoRng, RngCore};
 use std::marker::PhantomData;
 use std::sync::Arc;

jolt-core/src/r1cs/builder.rs

@@ -1,11 +1,13 @@
 use crate::{
-    field::JoltField,
+    field::{JoltField, OptimizedMul},
     r1cs::key::{SparseConstraints, UniformR1CS},
     utils::{
         mul_0_1_optimized,
         thread::{drop_in_background_thread, unsafe_allocate_zero_vec},
     },
 };
+#[allow(unused_imports)] // clippy thinks these aren't needed lol
+use ark_std::{One, Zero};
 use rayon::prelude::*;
 use std::{collections::HashMap, fmt::Debug};
 

jolt-core/src/r1cs/jolt_constraints.rs

@@ -288,6 +288,7 @@ mod tests {
     };
 
     use ark_bn254::Fr;
+    use ark_std::Zero;
     use strum::EnumCount;
 
     #[test]

jolt-core/src/r1cs/key.rs

@@ -266,7 +266,7 @@ impl<F: JoltField> UniformSpartanKey<F> {
         let r_var = &r_rest[..var_bits];
 
         let r_var_eq = EqPolynomial::evals(r_var);
-        let eval_variables = (0..self.uniform_r1cs.num_vars)
+        let eval_variables: F = (0..self.uniform_r1cs.num_vars)
             .map(|var_index| r_var_eq[var_index] * segment_evals[var_index])
             .sum();
         let const_poly = SparsePolynomial::new(self.num_vars_total().log_2(), vec![(0, F::one())]);
@@ -397,6 +397,7 @@ impl<F: JoltField> UniformSpartanKey<F> {
 mod test {
     use super::*;
     use ark_bn254::Fr;
+    use ark_std::{One, Zero};
 
     use crate::{
         poly::dense_mlpoly::DensePolynomial,

jolt-core/src/r1cs/ops.rs

@@ -1,7 +1,10 @@
 //! Defines the Linear Combination (LC) object and associated operations.
 //! A LinearCombination is a vector of Terms, where each Term is a pair of a Variable and a coefficient.
 
-use crate::{field::JoltField, utils::thread::unsafe_allocate_zero_vec};
+use crate::{
+    field::{JoltField, OptimizedMul},
+    utils::thread::unsafe_allocate_zero_vec,
+};
 use rayon::prelude::*;
 use std::fmt::Debug;
 use std::hash::Hash;

jolt-core/src/r1cs/spartan.rs

@@ -411,6 +411,7 @@ impl<F: JoltField, C: CommitmentScheme<Field = F>> UniformSpartanProof<F, C> {
 #[cfg(test)]
 mod test {
     use ark_bn254::Fr;
+    use ark_std::One;
 
     use crate::{
         poly::commitment::{commitment_scheme::CommitShape, hyrax::HyraxScheme},

jolt-core/src/subprotocols/grand_product.rs

@@ -1,6 +1,6 @@
 use super::grand_product_quarks::QuarkGrandProductProof;
 use super::sumcheck::{BatchedCubicSumcheck, SumcheckInstanceProof};
-use crate::field::JoltField;
+use crate::field::{JoltField, OptimizedMul};
 use crate::poly::commitment::commitment_scheme::CommitmentScheme;
 use crate::poly::eq_poly::EqPolynomial;
 use crate::poly::{dense_mlpoly::DensePolynomial, unipoly::UniPoly};
@@ -1546,7 +1546,7 @@ mod grand_product_tests {
     use super::*;
     use crate::poly::commitment::zeromorph::Zeromorph;
     use ark_bn254::{Bn254, Fr};
-    use ark_std::test_rng;
+    use ark_std::{test_rng, One};
     use rand_core::RngCore;
 
     #[test]

jolt-core/src/subprotocols/grand_product_quarks.rs

@@ -10,6 +10,7 @@ use crate::poly::eq_poly::EqPolynomial;
 use crate::utils::math::Math;
 use crate::utils::transcript::{AppendToTranscript, ProofTranscript};
 use ark_serialize::*;
+use ark_std::{One, Zero};
 use itertools::Itertools;
 use rayon::prelude::*;
 use thiserror::Error;