chore: add implementation and tests for more linear relation operations.

Author: mmaker

src/linear_relation/convert.rs

@@ -153,5 +153,4 @@ impl<G: Group> From<Sum<Weighted<GroupVar<G>, G::Scalar>>> for Sum<Weighted<Term
         let sum = sum.0.into_iter().map(|x| x.into()).collect::<Vec<_>>();
         Self(sum)
     }
-
-}
+}

src/linear_relation/mod.rs

@@ -698,17 +698,13 @@ impl<G: PrimeGroup> TryFrom<&LinearRelation<G>> for CanonicalLinearRelation<G> {
         }
 
         // If any linear combination has no witness variables, the relation is invalid
-        if relation
-            .linear_map
-            .linear_combinations
-            .iter()
-            .any(|lc| lc.0.iter().all(|weighted| matches!(weighted.term.scalar, ScalarTerm::Unit)))
-        {
+        if relation.linear_map.linear_combinations.iter().any(|lc| {
+            lc.0.iter()
+                .all(|weighted| matches!(weighted.term.scalar, ScalarTerm::Unit))
+        }) {
             return Err(Error::InvalidInstanceWitnessPair);
         }
 
-
-
         let mut canonical = CanonicalLinearRelation::new();
         canonical.num_scalars = relation.linear_map.num_scalars;
 

src/linear_relation/ops.rs

@@ -268,6 +268,30 @@ mod add {
             rhs + self
         }
     }
+
+    impl<T: Field + Into<G::Scalar>, G: Group> Add<T> for Sum<ScalarVar<G>> {
+        type Output = Sum<Weighted<ScalarTerm<G>, G::Scalar>>;
+
+        fn add(self, rhs: T) -> Self::Output {
+            // Convert Sum<ScalarVar<G>> to Sum<Weighted<ScalarTerm<G>, G::Scalar>>
+            let mut weighted_terms = Vec::new();
+            for var in self.0 {
+                weighted_terms.push(Weighted {
+                    term: ScalarTerm::from(var),
+                    weight: G::Scalar::ONE,
+                });
+            }
+            let weighted_sum: Sum<Weighted<ScalarTerm<G>, G::Scalar>> = Sum(weighted_terms);
+
+            // Convert the scalar to a weighted term
+            let weighted_scalar = Weighted {
+                term: ScalarTerm::Unit,
+                weight: rhs.into(),
+            };
+
+            weighted_sum + weighted_scalar
+        }
+    }
 }
 
 mod mul {

src/tests/test_relations.rs

@@ -1,16 +1,12 @@
-use std::collections::HashMap;
-
-use ff::Field;
+use ff::{Field, PrimeField};
 use group::prime::PrimeGroup;
 use rand::rngs::OsRng;
 use rand::RngCore;
 
+use crate::codec::Shake128DuplexSponge;
 use crate::fiat_shamir::Nizk;
-use crate::{
-    codec::Shake128DuplexSponge, linear_relation::CanonicalLinearRelation,
-    schnorr_protocol::SchnorrProof,
-};
-use ff::PrimeField;
+use crate::linear_relation::CanonicalLinearRelation;
+use crate::schnorr_protocol::SchnorrProof;
 
 use crate::linear_relation::{msm_pr, LinearRelation};
 
@@ -155,7 +151,10 @@ pub fn twisted_pedersen_commitment<G: PrimeGroup, R: RngCore>(
     let [var_x, var_r] = relation.allocate_scalars();
     let [var_G, var_H] = relation.allocate_elements();
 
-    relation.allocate_eq((var_x * G::Scalar::from(3)) * var_G + (var_r * G::Scalar::from(2) + G::Scalar::from(3)) * var_H);
+    relation.allocate_eq(
+        (var_x * G::Scalar::from(3)) * var_G
+            + (var_r * G::Scalar::from(2) + G::Scalar::from(3)) * var_H,
+    );
 
     relation.set_elements([(var_H, H), (var_G, G::generator())]);
     relation.compute_image(&[x, r]).unwrap();
@@ -202,7 +201,6 @@ pub fn pedersen_commitment_dleq<G: PrimeGroup, R: RngCore>(
     (instance, witness_vec)
 }
 
-
 /// LinearMap for knowledge of an opening for use in a BBS commitment.
 // BBS message length is 3
 #[allow(non_snake_case)]
@@ -270,8 +268,7 @@ pub fn weird_linear_combination<G: PrimeGroup, R: RngCore>(
     let A = sigma__lr.allocate_element();
     let var_B = sigma__lr.allocate_element();
 
-    let sigma__eq1 =
-        sigma__lr.allocate_eq(A * G::Scalar::from(1) + gen__disj1_x_r_var * var_B);
+    let sigma__eq1 = sigma__lr.allocate_eq(A * G::Scalar::from(1) + gen__disj1_x_r_var * var_B);
 
     // Set the group elements
     sigma__lr.set_elements([(A, G::generator()), (var_B, B)]);
@@ -327,34 +324,92 @@ fn subtractions_with_shift<G: PrimeGroup, R: RngCore>(
     (instance, witness)
 }
 
+/// LinearMap for the failing relation from cmz wallet test: W.balance = N.balance + I.price + fee
+#[allow(non_snake_case)]
+fn cmz_wallet_spend_relation<G: PrimeGroup, R: RngCore>(
+    mut rng: &mut R,
+) -> (CanonicalLinearRelation<G>, Vec<G::Scalar>) {
+    // Simulate the wallet spend relation from cmz
+    let P_W = G::random(&mut rng);
+    let P_I = G::random(&mut rng);
+    let A = G::random(&mut rng);
+
+    // Secret values
+    let n_balance = G::Scalar::random(&mut rng);
+    let i_price = G::Scalar::random(&mut rng);
+    let fee = G::Scalar::from(5u64);
+    let z_w_balance = G::Scalar::random(&mut rng);
+
+    // W.balance = N.balance + I.price + fee
+    let w_balance = n_balance + i_price + fee;
+
+    let mut relation = LinearRelation::new();
+
+    // Allocate variables
+    let var_n_balance = relation.allocate_scalar();
+    let var_i_price = relation.allocate_scalar();
+    let var_z_w_balance = relation.allocate_scalar();
+
+    let var_P_W = relation.allocate_element();
+    let var_P_I = relation.allocate_element();
+    let var_A = relation.allocate_element();
+
+    // C_show_Hattr_W_balance = (N.balance + I.price + fee) * P_W + z_w_balance * A
+    // This is the problematic expression that fails in cmz
+    let var_C = relation.allocate_eq(
+        (var_n_balance + var_i_price + G::Scalar::from(5u64)) * var_P_W + var_z_w_balance * var_A,
+    );
+
+    relation.set_elements([(var_P_W, P_W), (var_P_I, P_I), (var_A, A)]);
+
+    relation
+        .compute_image(&[n_balance, i_price, z_w_balance])
+        .unwrap();
+
+    let C = relation.linear_map.group_elements.get(var_C).unwrap();
+    let expected = P_W * w_balance + A * z_w_balance;
+    assert_eq!(C, expected);
+
+    let witness = vec![n_balance, i_price, z_w_balance];
+    let instance = (&relation).try_into().unwrap();
+    (instance, witness)
+}
+
 /// Generic helper function to test both relation correctness and NIZK functionality
 #[test]
-fn test_common_relations() {
+fn test_relations() {
     use group::Group;
     type G = bls12_381::G1Projective;
 
-    let mut instance_generators = HashMap::<
+    let instance_generators: Vec<(
         &str,
         Box<dyn Fn(&mut OsRng) -> (CanonicalLinearRelation<G>, Vec<<G as Group>::Scalar>)>,
-    >::new();
-
-    instance_generators.insert("dlog", Box::new(discrete_logarithm));
-    instance_generators.insert("shifted_dlog", Box::new(shifted_dlog));
-    instance_generators.insert("dleq", Box::new(dleq));
-    instance_generators.insert("shifted_dleq", Box::new(shifted_dleq));
-    instance_generators.insert("pedersen_commitment", Box::new(pedersen_commitment));
-    instance_generators.insert("twisted_pedersen_commitment", Box::new(twisted_pedersen_commitment));
-    instance_generators.insert(
-        "pedersen_commitment_dleq",
-        Box::new(pedersen_commitment_dleq),
-    );
-    instance_generators.insert("bbs_blind_commitment", Box::new(bbs_blind_commitment));
-    instance_generators.insert(
-        "weird_linear_combination",
-        Box::new(weird_linear_combination),
-    );
-    instance_generators.insert("simple_subtractions", Box::new(simple_subtractions));
-    instance_generators.insert("subtractions_with_shift", Box::new(subtractions_with_shift));
+    )> = vec![
+        ("dlog", Box::new(discrete_logarithm)),
+        ("shifted_dlog", Box::new(shifted_dlog)),
+        ("dleq", Box::new(dleq)),
+        ("shifted_dleq", Box::new(shifted_dleq)),
+        ("pedersen_commitment", Box::new(pedersen_commitment)),
+        (
+            "twisted_pedersen_commitment",
+            Box::new(twisted_pedersen_commitment),
+        ),
+        (
+            "pedersen_commitment_dleq",
+            Box::new(pedersen_commitment_dleq),
+        ),
+        ("bbs_blind_commitment", Box::new(bbs_blind_commitment)),
+        (
+            "weird_linear_combination",
+            Box::new(weird_linear_combination),
+        ),
+        ("simple_subtractions", Box::new(simple_subtractions)),
+        ("subtractions_with_shift", Box::new(subtractions_with_shift)),
+        (
+            "cmz_wallet_spend_relation",
+            Box::new(cmz_wallet_spend_relation),
+        ),
+    ];
 
     for (relation_name, relation_sampler) in instance_generators.iter() {
         let mut rng = OsRng;

src/tests/test_validation_criteria.rs

@@ -163,7 +163,6 @@ mod instance_validation {
         let result = CanonicalLinearRelation::try_from(&linear_relation);
         assert!(result.is_err());
 
-
         // The following relation is for
         // X = B * x + B * pub_scalar + A * 3
         // and should be considered a valid instance.
@@ -172,7 +171,8 @@ mod instance_validation {
         let x_var = linear_relation.allocate_scalar();
         let B_var = linear_relation.allocate_element();
         let A_var = linear_relation.allocate_element();
-        let X_var = linear_relation.allocate_eq(B_var * x_var + B_var * pub_scalar + A_var * Scalar::from(3));
+        let X_var = linear_relation
+            .allocate_eq(B_var * x_var + B_var * pub_scalar + A_var * Scalar::from(3));
 
         linear_relation.set_element(B_var, B);
         linear_relation.set_element(A_var, A);
@@ -363,45 +363,45 @@ mod proof_validation {
     }
 
     #[test]
-    fn test_or_relation_wrong_branch() {
+    fn test_or_relation() {
         // This test reproduces the issue from sigma_compiler's simple_or test
         // where an OR relation fails verification when using the wrong branch
         let mut rng = thread_rng();
-        
+
         // Create generators
         // For this test, we'll use two different multiples of the generator
         let B = G::generator();
         let A = B * Scalar::from(42u64); // Different generator
-        
+
         // Create scalars
         let x = Scalar::random(&mut rng);
         let y = Scalar::random(&mut rng);
-        
+
         // Set C = y*B (so the second branch should be satisfied)
         let C = B * y;
-        
+
         // Create the first branch: C = x*A
         let mut lr1 = LinearRelation::<G>::new();
         let x_var = lr1.allocate_scalar();
         let A_var = lr1.allocate_element();
         let eq1 = lr1.allocate_eq(x_var * A_var);
         lr1.set_element(A_var, A);
         lr1.set_element(eq1, C);
-        
+
         // Create the second branch: C = y*B
         let mut lr2 = LinearRelation::<G>::new();
         let y_var = lr2.allocate_scalar();
         let B_var = lr2.allocate_element();
         let eq2 = lr2.allocate_eq(y_var * B_var);
         lr2.set_element(B_var, B);
         lr2.set_element(eq2, C);
-        
+
         // Create OR composition
         let or_relation = ComposedRelation::Or(vec![
             ComposedRelation::from(lr1),
             ComposedRelation::from(lr2),
         ]);
-        
+
         // The issue from sigma_compiler: the witness is always using branch 0
         // even when branch 1 should be used
         let witness_wrong = ComposedWitness::Or(
@@ -411,17 +411,20 @@ mod proof_validation {
                 ComposedWitness::Simple(vec![y]),
             ],
         );
-        
+
         // Test the bug: using branch 0 when C = y*B (branch 1 should be used)
         let nizk = Nizk::<_, KeccakByteSchnorrCodec<G>>::new(b"test_or_bug", or_relation);
         let proof_result = nizk.prove_batchable(&witness_wrong, &mut rng);
-        
+
         // This currently passes but should fail - this is the bug!
         // The prover is using branch 0 (C = x*A) but C actually equals y*B
         match proof_result {
             Ok(proof) => {
                 let verify_result = nizk.verify_batchable(&proof);
-                println!("Bug reproduced: Proof with wrong branch verified: {:?}", verify_result.is_ok());
+                println!(
+                    "Bug reproduced: Proof with wrong branch verified: {:?}",
+                    verify_result.is_ok()
+                );
                 assert!(
                     verify_result.is_err(),
                     "BUG: Proof should fail when using wrong branch in OR relation, but it passed!"