refactor: ProtocolProverState and ProtocolWitness in OR composition (#65)

Author: nougzarm

README.md

@@ -48,12 +48,12 @@ let or_protocol = Protocol::Or(vec![
 ]);
 
 // If we know the second option, create witness for index 1
-let witness = ProtocolWitness::Or(1, vec![
-    ProtocolWitness::And(vec![
+let witness = ProtocolWitness::from((1, 
+    ProtocolWitness::from(vec![
         ProtocolWitness::Simple(vec![y]),
         ProtocolWitness::Simple(vec![z]),
     ])
-]);
+));
 ```
 
 ## Examples

examples/simple_composition.rs

@@ -53,7 +53,7 @@ fn prove(P1: G, x2: Scalar, H: G) -> ProofResult<Vec<u8>> {
     let Q = H * x2;
 
     let instance = create_relation(P1, P2, Q, H);
-    let witness = ComposedWitness::Or(1, vec![ComposedWitness::Simple(vec![x2])]);
+    let witness = ComposedWitness::from((1, ComposedWitness::Simple(vec![x2])));
     let nizk = Nizk::<_, Shake128DuplexSponge<G>>::new(b"or_proof_example", instance);
 
     nizk.prove_batchable(&witness, &mut OsRng)

src/composition.rs

@@ -82,7 +82,7 @@ pub enum ComposedProverState<G: PrimeGroup> {
     And(Vec<ComposedProverState<G>>),
     Or(
         usize,                                                 // real index
-        Vec<ComposedProverState<G>>,                           // real ProverState
+        Box<ComposedProverState<G>>,                           // real ProverState
         (Vec<ComposedChallenge<G>>, Vec<ComposedResponse<G>>), // simulated transcripts
     ),
 }
@@ -99,7 +99,19 @@ pub enum ComposedResponse<G: PrimeGroup> {
 pub enum ComposedWitness<G: PrimeGroup> {
     Simple(<SchnorrProof<G> as SigmaProtocol>::Witness),
     And(Vec<ComposedWitness<G>>),
-    Or(usize, Vec<ComposedWitness<G>>),
+    Or(usize, Box<ComposedWitness<G>>),
+}
+
+impl<G: PrimeGroup> From<Vec<ComposedWitness<G>>> for ComposedWitness<G> {
+    fn from(value: Vec<ComposedWitness<G>>) -> Self {
+        Self::And(value)
+    }
+}
+
+impl<G: PrimeGroup> From<(usize, ComposedWitness<G>)> for ComposedWitness<G> {
+    fn from((i, witness): (usize, ComposedWitness<G>)) -> Self {
+        Self::Or(i, Box::new(witness))
+    }
 }
 
 // Structure representing the Challenge type of Protocol as SigmaProtocol
@@ -149,7 +161,7 @@ impl<G: PrimeGroup> SigmaProtocol for ComposedRelation<G> {
                 let mut simulated_challenges = Vec::new();
                 let mut simulated_responses = Vec::new();
 
-                let (real_commitment, real_state) = ps[*w_index].prover_commit(&w[0], rng)?;
+                let (real_commitment, real_state) = ps[*w_index].prover_commit(w, rng)?;
 
                 for i in (0..ps.len()).filter(|i| i != w_index) {
                     let (commitment, challenge, response) = ps[i].simulate_transcript(rng)?;
@@ -163,7 +175,7 @@ impl<G: PrimeGroup> SigmaProtocol for ComposedRelation<G> {
                     ComposedCommitment::Or(commitments),
                     ComposedProverState::Or(
                         *w_index,
-                        vec![real_state],
+                        Box::new(real_state),
                         (simulated_challenges, simulated_responses),
                     ),
                 ))
@@ -208,8 +220,7 @@ impl<G: PrimeGroup> SigmaProtocol for ComposedRelation<G> {
                 for ch in &simulated_challenges {
                     real_challenge -= ch;
                 }
-                let real_response =
-                    ps[w_index].prover_response(real_state[0].clone(), &real_challenge)?;
+                let real_response = ps[w_index].prover_response(*real_state, &real_challenge)?;
 
                 for (i, _) in ps.iter().enumerate() {
                     if i == w_index {

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

@@ -351,6 +351,8 @@ pub struct CanonicalLinearRelation<G: PrimeGroup> {
     pub num_scalars: usize,
 }
 
+type WeightedCache<A, B> = HashMap<B, Vec<(A, B)>>;
+
 impl<G: PrimeGroup> CanonicalLinearRelation<G> {
     /// Create a new empty canonical linear relation
     pub fn new() -> Self {
@@ -368,7 +370,7 @@ impl<G: PrimeGroup> CanonicalLinearRelation<G> {
         group_var: GroupVar<G>,
         weight: &G::Scalar,
         original_group_elements: &GroupMap<G>,
-        weighted_group_cache: &mut HashMap<GroupVar<G>, Vec<(G::Scalar, GroupVar<G>)>>,
+        weighted_group_cache: &mut WeightedCache<G::Scalar, GroupVar<G>>,
     ) -> Result<GroupVar<G>, Error> {
         // Check if we already have this (weight, group_var) combination
         let entry = weighted_group_cache.entry(group_var).or_default();
@@ -397,7 +399,7 @@ impl<G: PrimeGroup> CanonicalLinearRelation<G> {
         image_var: GroupVar<G>,
         equation: &LinearCombination<G>,
         original_relation: &LinearRelation<G>,
-        weighted_group_cache: &mut HashMap<GroupVar<G>, Vec<(G::Scalar, GroupVar<G>)>>,
+        weighted_group_cache: &mut WeightedCache<G::Scalar, GroupVar<G>>,
     ) -> Result<(), Error> {
         let mut rhs_terms = Vec::new();
 
@@ -626,8 +628,7 @@ impl<G: PrimeGroup> CanonicalLinearRelation<G> {
 
             let elem = Option::<G>::from(G::from_bytes(&repr)).ok_or_else(|| {
                 Error::from(InvalidInstance::new(format!(
-                    "Invalid group element at index {}",
-                    i
+                    "Invalid group element at index {i}"
                 )))
             })?;
 
@@ -698,17 +699,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

@@ -745,9 +745,9 @@ mod tests {
 
         let diff = x - y;
         assert_eq!(diff.terms().len(), 2);
-        assert_eq!(diff.terms()[0].term, y.into());
+        assert_eq!(diff.terms()[0].term, y);
         assert_eq!(diff.terms()[0].weight, -Scalar::ONE);
-        assert_eq!(diff.terms()[1].term, x.into());
+        assert_eq!(diff.terms()[1].term, x);
         assert_eq!(diff.terms()[1].weight, Scalar::ONE);
     }
 

src/schnorr_protocol.rs

@@ -27,6 +27,8 @@ use rand::{CryptoRng, Rng, RngCore};
 #[derive(Clone, Default, Debug)]
 pub struct SchnorrProof<G: PrimeGroup>(pub CanonicalLinearRelation<G>);
 
+type CommitResult<Commit, Scalar> = (Commit, (Scalar, Scalar));
+
 impl<G: PrimeGroup> SchnorrProof<G> {
     pub fn witness_length(&self) -> usize {
         self.0.num_scalars
@@ -58,7 +60,7 @@ impl<G: PrimeGroup> SchnorrProof<G> {
         &self,
         witness: &[G::Scalar],
         nonces: &[G::Scalar],
-    ) -> Result<(Vec<G>, (Vec<G::Scalar>, Vec<G::Scalar>)), Error> {
+    ) -> Result<CommitResult<Vec<G>, Vec<G::Scalar>>, Error> {
         if witness.len() != self.witness_length() {
             return Err(Error::InvalidInstanceWitnessPair);
         }

src/tests/spec/test_vectors.rs

@@ -71,8 +71,7 @@ fn test_spec_testvectors() {
         assert_eq!(
             parsed_instance.label(),
             vector.statement,
-            "parsed statement doesn't match original for {}",
-            test_name
+            "parsed statement doesn't match original for {test_name}"
         );
 
         // Create NIZK with the session_id from the test vector
@@ -90,8 +89,7 @@ fn test_spec_testvectors() {
         assert_eq!(
             computed_iv,
             vector.iv.as_slice(),
-            "Computed IV doesn't match test vector IV for {}",
-            test_name
+            "Computed IV doesn't match test vector IV for {test_name}"
         );
 
         // Generate proof with the proof generation RNG
@@ -101,70 +99,66 @@ fn test_spec_testvectors() {
         // Verify the proof matches
         assert_eq!(
             proof_bytes, vector.proof,
-            "proof bytes for test vector {} do not match",
-            test_name
+            "proof bytes for test vector {test_name} do not match"
         );
 
         // Verify the proof is valid
         let verified = nizk.verify_batchable(&proof_bytes).is_ok();
         assert!(
             verified,
-            "Fiat-Shamir Schnorr proof verification failed for {}",
-            test_name
+            "Fiat-Shamir Schnorr proof verification failed for {test_name}"
         );
     }
 }
 
 fn extract_vectors_new(path: &str) -> Result<HashMap<String, TestVector>, String> {
     use std::collections::HashMap;
 
-    let content =
-        fs::read_to_string(path).map_err(|e| format!("Unable to read JSON file: {}", e))?;
-    let root: JsonValue =
-        json::parse(&content).map_err(|e| format!("JSON parsing error: {}", e))?;
+    let content = fs::read_to_string(path).map_err(|e| format!("Unable to read JSON file: {e}"))?;
+    let root: JsonValue = json::parse(&content).map_err(|e| format!("JSON parsing error: {e}"))?;
 
     let mut vectors = HashMap::new();
 
     for (name, obj) in root.entries() {
         let ciphersuite = obj["Ciphersuite"]
             .as_str()
-            .ok_or_else(|| format!("Ciphersuite field not found for {}", name))?
+            .ok_or_else(|| format!("Ciphersuite field not found for {name}"))?
             .to_string();
 
         let session_id = Vec::from_hex(
             obj["SessionId"]
                 .as_str()
-                .ok_or_else(|| format!("SessionId field not found for {}", name))?,
+                .ok_or_else(|| format!("SessionId field not found for {name}"))?,
         )
-        .map_err(|e| format!("Invalid hex in SessionId for {}: {}", name, e))?;
+        .map_err(|e| format!("Invalid hex in SessionId for {name}: {e}"))?;
 
         let statement = Vec::from_hex(
             obj["Statement"]
                 .as_str()
-                .ok_or_else(|| format!("Statement field not found for {}", name))?,
+                .ok_or_else(|| format!("Statement field not found for {name}"))?,
         )
-        .map_err(|e| format!("Invalid hex in Statement for {}: {}", name, e))?;
+        .map_err(|e| format!("Invalid hex in Statement for {name}: {e}"))?;
 
         let witness = Vec::from_hex(
             obj["Witness"]
                 .as_str()
-                .ok_or_else(|| format!("Witness field not found for {}", name))?,
+                .ok_or_else(|| format!("Witness field not found for {name}"))?,
         )
-        .map_err(|e| format!("Invalid hex in Witness for {}: {}", name, e))?;
+        .map_err(|e| format!("Invalid hex in Witness for {name}: {e}"))?;
 
         let iv = Vec::from_hex(
             obj["IV"]
                 .as_str()
-                .ok_or_else(|| format!("IV field not found for {}", name))?,
+                .ok_or_else(|| format!("IV field not found for {name}"))?,
         )
-        .map_err(|e| format!("Invalid hex in IV for {}: {}", name, e))?;
+        .map_err(|e| format!("Invalid hex in IV for {name}: {e}"))?;
 
         let proof = Vec::from_hex(
             obj["Proof"]
                 .as_str()
-                .ok_or_else(|| format!("Proof field not found for {}", name))?,
+                .ok_or_else(|| format!("Proof field not found for {name}"))?,
         )
-        .map_err(|e| format!("Invalid hex in Proof for {}: {}", name, e))?;
+        .map_err(|e| format!("Invalid hex in Proof for {name}: {e}"))?;
 
         vectors.insert(
             name.to_string(),

src/tests/test_composition.rs

@@ -34,7 +34,7 @@ fn composition_proof_correct() {
         ComposedRelation::Simple(SchnorrProof(relation1)),
         ComposedRelation::Simple(SchnorrProof(relation2)),
     ]);
-    let or_witness1 = ComposedWitness::Or(0, vec![ComposedWitness::Simple(witness1)]);
+    let or_witness1 = ComposedWitness::from((0, ComposedWitness::Simple(witness1)));
 
     let simple_protocol1 = ComposedRelation::Simple(SchnorrProof(relation3));
     let simple_witness1 = ComposedWitness::Simple(witness3);
@@ -43,14 +43,14 @@ fn composition_proof_correct() {
         ComposedRelation::Simple(SchnorrProof(relation4)),
         ComposedRelation::Simple(SchnorrProof(relation5)),
     ]);
-    let and_witness1 = ComposedWitness::And(vec![
+    let and_witness1 = ComposedWitness::from(vec![
         ComposedWitness::Simple(witness4),
         ComposedWitness::Simple(witness5),
     ]);
 
     // definition of the final protocol
     let protocol = ComposedRelation::And(vec![or_protocol1, simple_protocol1, and_protocol1]);
-    let witness = ComposedWitness::And(vec![or_witness1, simple_witness1, and_witness1]);
+    let witness = ComposedWitness::from(vec![or_witness1, simple_witness1, and_witness1]);
 
     let nizk = Nizk::<ComposedRelation<RistrettoPoint>, Shake128DuplexSponge<G>>::new(
         domain_sep, protocol,