refactor a bit more

Author: tgeoghegan

src/sumcheck/constraints.rs

@@ -149,19 +149,22 @@ impl<FE: CodecFieldElement + LagrangePolynomialFieldElement> ProofConstraints<FE
 
             let mut layer_claim = SymbolicExpression::new(layer_index);
 
-            // Known portion of layer's initial claim.
-            layer_claim += claims[0] + alpha * claims[1];
+            let mut claim_0 = SymbolicTerm::from_known(claims[0]);
+            let mut claim_1 = SymbolicTerm::from_known(claims[1]) * alpha;
 
             if layer_index > 0 {
                 // For layers past the first, claims is computed from the previous layer's vl and
                 // vr, so we need linear constraint terms for that symbolic manipulation.
                 let (vl_witness, vr_witness, _) =
                     witness_layout.wire_witness_indices(layer_index - 1);
 
-                layer_claim += SymbolicTerm::new(vl_witness);
-                layer_claim += SymbolicTerm::new(vr_witness) * alpha;
+                claim_0.with_witness(vl_witness);
+                claim_1.with_witness(vr_witness);
             }
 
+            layer_claim += claim_0;
+            layer_claim += claim_1;
+
             for (round, polynomial_pair) in proof_layer.polynomials.iter().enumerate() {
                 for (hand, polynomial) in polynomial_pair.iter().enumerate() {
                     transcript.write_polynomial(polynomial)?;
@@ -179,34 +182,26 @@ impl<FE: CodecFieldElement + LagrangePolynomialFieldElement> ProofConstraints<FE
                     //
                     // Compute the current claim:
                     //
-                    //   claim = lag_0(challenge) * p0
-                    //     + lag_1(challenge) * p1
-                    //     + lag_2(challenge) * p2
+                    //   claim = p0 * lag_0(challenge)
+                    //     + p1 * lag_1(challenge)
+                    //     + p2 * lag_2(challenge)
                     //
                     // Expanding p1 = prev_claim - p0 and rearranging:
                     //
-                    //   claim = lag_1(challenge) * prev_claim
-                    //     + (lag_0(challenge) - lag_1(challenge)) * p0
-                    //     + lag_2(challenge) * p2
+                    //   claim = prev_claim * lag_1(challenge)
+                    //     + p0 * (lag_0(challenge) - lag_1(challenge))
+                    //     + p2 * lag_2(challenge)
 
                     // lag_1(challenge) * prev_claim
                     layer_claim *= FE::lagrange_basis_polynomial_1(challenge[0]);
 
-                    // (lag_0(challenge) - lag_1(challenge)) * p0, known part:
-                    layer_claim += polynomial.p0
-                        * (FE::lagrange_basis_polynomial_0(challenge[0])
-                            - FE::lagrange_basis_polynomial_1(challenge[0]));
-
-                    // (lag_0(challenge) - lag_1(challenge)) * p0, symbolic part:
-                    layer_claim += SymbolicTerm::new(p0_witness)
+                    // p0 * (lag_0(challenge) - lag_1(challenge)):
+                    layer_claim += (SymbolicTerm::new(p0_witness) + polynomial.p0)
                         * (FE::lagrange_basis_polynomial_0(challenge[0])
                             - FE::lagrange_basis_polynomial_1(challenge[0]));
 
-                    // lag_2(challenge) * p2, known part:
-                    layer_claim += polynomial.p2 * FE::lagrange_basis_polynomial_2(challenge[0]);
-
-                    // lag_2(challenge) * p2, symbolic part:
-                    layer_claim += SymbolicTerm::new(p2_witness)
+                    // p2 * lag_2(challenge):
+                    layer_claim += (SymbolicTerm::new(p2_witness) + polynomial.p2)
                         * FE::lagrange_basis_polynomial_2(challenge[0]);
 
                     bound_quad = bound_quad.bind(&challenge).transpose();
@@ -296,13 +291,15 @@ impl<FE: CodecFieldElement + LagrangePolynomialFieldElement> ProofConstraints<FE
         final_claim += SymbolicTerm::new(vr_witness) * -gamma;
 
         // Linear constraint RHS
-        final_claim += claims[0] + gamma * claims[1]
-            - public_inputs
-                .iter()
-                .zip(eq2.iter())
-                .fold(FE::ZERO, |sum, (public_input_i, eq2_i)| {
-                    sum + *public_input_i * eq2_i
-                });
+        final_claim += SymbolicTerm::from_known(
+            claims[0] + gamma * claims[1]
+                - public_inputs
+                    .iter()
+                    .zip(eq2.iter())
+                    .fold(FE::ZERO, |sum, (public_input_i, eq2_i)| {
+                        sum + *public_input_i * eq2_i
+                    }),
+        );
 
         constraints
             .linear_constraint_lhs

src/sumcheck/symbolic.rs

@@ -1,13 +1,12 @@
 use crate::{fields::FieldElement, sumcheck::constraints::LinearConstraintLhsTerm};
-use std::ops::{AddAssign, Mul, MulAssign, Sub};
+use std::ops::{Add, AddAssign, Mul, MulAssign};
 
 /// A symbolic expression, used to accumulate symbolic terms that contribute to a circuit layer's
 /// linear constraint.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct SymbolicExpression<FieldElement> {
     constraint_number: usize,
-    known: FieldElement,
-    symbolic_terms: Vec<SymbolicTerm<FieldElement>>,
+    terms: Vec<SymbolicTerm<FieldElement>>,
 }
 
 impl<FE: FieldElement> SymbolicExpression<FE> {
@@ -16,76 +15,90 @@ impl<FE: FieldElement> SymbolicExpression<FE> {
     pub fn new(layer_index: usize) -> Self {
         Self {
             constraint_number: layer_index,
-            known: FE::ZERO,
-            symbolic_terms: Vec::new(),
+            terms: Vec::new(),
         }
     }
 
     /// The known portion of this expression. Its contribution to the linear constraint's right hand
     /// side.
     pub fn known(&self) -> FE {
-        self.known
+        self.terms
+            .iter()
+            .fold(FE::ZERO, |sum, term| sum + term.known)
     }
 
     /// The linear constraint LHS terms for this expression.
     pub fn lhs_terms(&self) -> Vec<LinearConstraintLhsTerm<FE>> {
-        self.symbolic_terms
+        self.terms
             .iter()
-            .map(|term| LinearConstraintLhsTerm {
-                constraint_number: self.constraint_number,
-                witness_index: term.witness_index,
-                constant_factor: term.constant_factor,
+            // Terms with no witness index do not contribute to LHS
+            .filter_map(|term| {
+                term.witness_index
+                    .map(|witness_index| LinearConstraintLhsTerm {
+                        constraint_number: self.constraint_number,
+                        witness_index,
+                        constant_factor: term.constant_factor,
+                    })
             })
             .collect()
     }
 }
 
-impl<FE: FieldElement> AddAssign<FE> for SymbolicExpression<FE> {
-    fn add_assign(&mut self, rhs: FE) {
-        self.known += rhs;
-    }
-}
-
 impl<FE: FieldElement> AddAssign<SymbolicTerm<FE>> for SymbolicExpression<FE> {
     fn add_assign(&mut self, rhs: SymbolicTerm<FE>) {
-        self.symbolic_terms.push(rhs);
-    }
-}
-
-impl<FE: FieldElement> Sub<FE> for SymbolicExpression<FE> {
-    type Output = Self;
-
-    fn sub(self, rhs: FE) -> Self::Output {
-        Self {
-            known: self.known - rhs,
-            ..self
-        }
+        self.terms.push(rhs);
     }
 }
 
 impl<FE: FieldElement> MulAssign<FE> for SymbolicExpression<FE> {
     fn mul_assign(&mut self, rhs: FE) {
-        self.known *= rhs;
-        self.symbolic_terms.iter_mut().for_each(|term| *term *= rhs);
+        self.terms.iter_mut().for_each(|term| *term *= rhs);
     }
 }
 
-/// A symbolic term in a symbolic expression.
+/// A symbolic term in a symbolic expression, consisting of
+/// `known + constant_factor * W[witness_index]`.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct SymbolicTerm<FieldElement> {
+    /// The known portion of the expression.
+    pub known: FieldElement,
     /// The index into the witness vector W. This is `j` in the specification.
-    pub witness_index: usize,
+    pub witness_index: Option<usize>,
     /// The constant factor `k`.
     pub constant_factor: FieldElement,
 }
 
 impl<FE: FieldElement> SymbolicTerm<FE> {
     pub fn new(witness_index: usize) -> Self {
         Self {
-            witness_index,
+            known: FE::ZERO,
+            witness_index: Some(witness_index),
+            constant_factor: FE::ONE,
+        }
+    }
+
+    pub fn from_known(known: FE) -> Self {
+        Self {
+            known,
+            witness_index: None,
             constant_factor: FE::ONE,
         }
     }
+
+    pub fn with_witness(&mut self, index: usize) {
+        self.witness_index = Some(index);
+    }
+}
+
+impl<FE: FieldElement> Add<FE> for SymbolicTerm<FE> {
+    type Output = Self;
+
+    fn add(self, rhs: FE) -> Self::Output {
+        Self {
+            known: self.known + rhs,
+            ..self
+        }
+    }
 }
 
 impl<FE: FieldElement> Mul<FE> for SymbolicTerm<FE> {
@@ -94,13 +107,111 @@ impl<FE: FieldElement> Mul<FE> for SymbolicTerm<FE> {
     fn mul(self, rhs: FE) -> Self::Output {
         Self {
             constant_factor: self.constant_factor * rhs,
+            known: self.known * rhs,
             ..self
         }
     }
 }
 
 impl<FE: FieldElement> MulAssign<FE> for SymbolicTerm<FE> {
     fn mul_assign(&mut self, rhs: FE) {
+        self.known *= rhs;
         self.constant_factor *= rhs;
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::fields::fieldp256::FieldP256;
+
+    #[test]
+    fn term_ops() {
+        let term = SymbolicTerm::new(1);
+
+        let term = term + FieldP256::from_u128(2);
+
+        assert_eq!(
+            term,
+            SymbolicTerm {
+                known: FieldP256::from_u128(2),
+                witness_index: Some(1),
+                constant_factor: FieldP256::ONE
+            }
+        );
+
+        let mut term = term * FieldP256::from_u128(5);
+
+        assert_eq!(
+            term,
+            SymbolicTerm {
+                known: FieldP256::from_u128(10),
+                witness_index: Some(1),
+                constant_factor: FieldP256::from_u128(5)
+            }
+        );
+
+        term *= FieldP256::from_u128(6);
+
+        assert_eq!(
+            term,
+            SymbolicTerm {
+                known: FieldP256::from_u128(60),
+                witness_index: Some(1),
+                constant_factor: FieldP256::from_u128(30)
+            }
+        );
+    }
+
+    #[test]
+    fn expression_ops() {
+        let mut expression = SymbolicExpression::new(11);
+        assert_eq!(expression.lhs_terms(), vec![]);
+        assert_eq!(expression.known(), FieldP256::ZERO);
+
+        // Term with both known and symbolic part
+        expression += SymbolicTerm::new(22) + FieldP256::from_u128(11);
+
+        // Term with only symbolic part
+        expression += SymbolicTerm::new(33);
+
+        // Term with only known part
+        expression += SymbolicTerm::from_known(FieldP256::from_u128(3));
+
+        assert_eq!(
+            expression.lhs_terms(),
+            vec![
+                LinearConstraintLhsTerm {
+                    constraint_number: 11,
+                    witness_index: 22,
+                    constant_factor: FieldP256::ONE,
+                },
+                LinearConstraintLhsTerm {
+                    constraint_number: 11,
+                    witness_index: 33,
+                    constant_factor: FieldP256::ONE,
+                },
+            ]
+        );
+        assert_eq!(expression.known(), FieldP256::from_u128(14));
+
+        expression *= FieldP256::from_u128(6);
+
+        assert_eq!(
+            expression.lhs_terms(),
+            vec![
+                LinearConstraintLhsTerm {
+                    constraint_number: 11,
+                    witness_index: 22,
+                    constant_factor: FieldP256::from_u128(6),
+                },
+                LinearConstraintLhsTerm {
+                    constraint_number: 11,
+                    witness_index: 33,
+                    constant_factor: FieldP256::from_u128(6),
+                },
+            ]
+        );
+        assert_eq!(expression.known(), FieldP256::from_u128(14 * 6));
+    }
+}