fix(tests): more testing of linear relations

Author: mmaker

src/linear_relation/mod.rs

@@ -160,9 +160,15 @@ impl<G: PrimeGroup> GroupMap<G> {
     ///
     /// Returns [`Error::UnassignedGroupVar`] if a value is not assigned.
     pub fn get(&self, var: GroupVar<G>) -> Result<G, Error> {
-        self.0[var.0].ok_or(Error::UnassignedGroupVar {
-            var_debug: format!("{var:?}"),
-        })
+        match self.0.get(var.0) {
+            Some(Some(elem)) => Ok(*elem),
+            Some(None) => Err(Error::UnassignedGroupVar {
+                var_debug: format!("{var:?}"),
+            }),
+            None => Err(Error::UnassignedGroupVar {
+                var_debug: format!("{var:?}"),
+            }),
+        }
     }
 
     /// Iterate over the assigned variable and group element pairs in this mapping.

src/tests/mod.rs

@@ -1,5 +1,5 @@
-mod composition;
-mod relations;
 mod spec;
-pub mod test_utils;
+
+mod test_composition;
+mod test_relations;
 mod test_validation_criteria;

src/tests/relations.rs

@@ -1,9 +1,7 @@
 use curve25519_dalek::ristretto::RistrettoPoint;
 use rand::rngs::OsRng;
 
-use super::test_utils::{
-    bbs_blind_commitment, discrete_logarithm, dleq, pedersen_commitment, pedersen_commitment_dleq,
-};
+use super::test_relations::*;
 use crate::codec::Shake128DuplexSponge;
 use crate::composition::{ComposedRelation, ProtocolWitness};
 use crate::fiat_shamir::Nizk;

src/tests/composition.rs renamed to src/tests/test_composition.rs

@@ -1,10 +1,17 @@
-//! Definitions used in tests for this crate.
+use std::collections::HashMap;
 
 use ff::Field;
 use group::prime::PrimeGroup;
+use rand::rngs::OsRng;
 use rand::RngCore;
 
-use crate::linear_relation::{msm_pr, CanonicalLinearRelation, LinearRelation};
+use crate::fiat_shamir::Nizk;
+use crate::{
+    codec::Shake128DuplexSponge, linear_relation::CanonicalLinearRelation,
+    schnorr_protocol::SchnorrProof,
+};
+
+use crate::linear_relation::{msm_pr, LinearRelation};
 
 /// LinearMap for knowledge of a discrete logarithm relative to a fixed basepoint.
 #[allow(non_snake_case)]
@@ -227,22 +234,10 @@ pub fn bbs_blind_commitment<G: PrimeGroup, R: RngCore>(
     (instance, witness)
 }
 
-/// Test function with the requested LinearRelation code
-#[allow(non_snake_case)]
-pub fn test_linear_relation_example<G: PrimeGroup>() -> LinearRelation<G> {
-    use ff::Field;
-
-    let mut sigma__lr = LinearRelation::<G>::new();
-    let x = sigma__lr.allocate_scalar();
-    let B = sigma__lr.allocate_element();
-    let _sigma__eq1 = sigma__lr.allocate_eq((x + (-<G::Scalar as Field>::ONE)) * B + (-B));
-
-    sigma__lr
-}
 
 /// LinearMap for the user's specific relation: A * 1 + gen__disj1_x_r * B
 #[allow(non_snake_case)]
-pub fn user_specific_linear_combination<G: PrimeGroup, R: RngCore>(
+pub fn weird_linear_combination<G: PrimeGroup, R: RngCore>(
     rng: &mut R,
 ) -> (CanonicalLinearRelation<G>, Vec<G::Scalar>) {
     let B = G::random(&mut *rng);
@@ -270,3 +265,64 @@ pub fn user_specific_linear_combination<G: PrimeGroup, R: RngCore>(
     let instance = (&sigma__lr).try_into().unwrap();
     (instance, witness)
 }
+
+/// Generic helper function to test both relation correctness and NIZK functionality
+#[test]
+fn test_common_relations() {
+    use group::Group;
+    type G = bls12_381::G1Projective;
+
+    let mut instance_generators = HashMap::<
+        &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_discrete_logarithm));
+    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(
+        "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),
+    );
+
+    for (relation_name, relation_sampler) in instance_generators.iter() {
+        let mut rng = OsRng;
+        let (canonical_relation, witness) = relation_sampler(&mut rng);
+
+        // Test the NIZK protocol
+        let protocol = SchnorrProof(canonical_relation);
+        let domain_sep = format!("test-fiat-shamir-{}", relation_name)
+            .as_bytes()
+            .to_vec();
+        let nizk = Nizk::<SchnorrProof<G>, Shake128DuplexSponge<G>>::new(&domain_sep, protocol);
+
+        // Test both proof types
+        let proof_batchable = nizk.prove_batchable(&witness, &mut OsRng).expect(&format!(
+            "Failed to create batchable proof for {}",
+            relation_name
+        ));
+        let proof_compact = nizk.prove_compact(&witness, &mut OsRng).expect(&format!(
+            "Failed to create compact proof for {}",
+            relation_name
+        ));
+
+        // Verify both proof types
+        assert!(
+            nizk.verify_batchable(&proof_batchable).is_ok(),
+            "Batchable proof verification failed for {}",
+            relation_name
+        );
+        assert!(
+            nizk.verify_compact(&proof_compact).is_ok(),
+            "Compact proof verification failed for {}",
+            relation_name
+        );
+    }
+}

src/tests/test_utils.rs renamed to src/tests/test_relations.rs

@@ -79,6 +79,30 @@ mod instance_validation {
         assert!(result.is_err());
     }
 
+    /// Test function with the requested LinearRelation code
+    #[test]
+    #[allow(non_snake_case)]
+    pub fn test_degenerate_equation() {
+        use ff::Field;
+
+        // This relation should fail for two reasons:
+        // 1. because B is not assigned
+        let mut relation = LinearRelation::<G>::new();
+        let x = relation.allocate_scalar();
+        let B = relation.allocate_element();
+        let _eq = relation.allocate_eq((x + (-Scalar::ONE)) * B + (-B));
+
+        assert!(CanonicalLinearRelation::try_from(&relation).is_err());
+
+        // 2. because the equation is void
+        let mut relation = LinearRelation::<G>::new();
+        let x = relation.allocate_scalar();
+        let B = relation.allocate_element();
+        let _eq = relation.allocate_eq((x + (-Scalar::ONE)) * B + (-B));
+        relation.set_element(B, G::generator());
+        assert!(CanonicalLinearRelation::try_from(&relation).is_err());
+    }
+
     #[test]
     fn test_inconsistent_equation_count() {
         // Create a relation with mismatched equations and image elements