fix: support OR composition when all supplied branches are true (#77)

Author: nategraf

examples/simple_composition.rs

@@ -40,9 +40,7 @@ fn create_relation(P1: G, P2: G, Q: G, H: G) -> ComposedRelation<G> {
     rel2.set_element(Q_var, Q);
 
     // Compose into OR protocol
-    let proto1 = ComposedRelation::from(rel1);
-    let proto2 = ComposedRelation::from(rel2);
-    ComposedRelation::Or(vec![proto1, proto2])
+    ComposedRelation::or([rel1.canonical().unwrap(), rel2.canonical().unwrap()])
 }
 
 /// Prove knowledge of one of the witnesses (we know x2 for the DLEQ)

src/composition.rs

@@ -20,12 +20,10 @@
 
 use ff::{Field, PrimeField};
 use group::prime::PrimeGroup;
-use sha3::Digest;
-use sha3::Sha3_256;
-use subtle::Choice;
-use subtle::ConditionallySelectable;
-use subtle::ConstantTimeEq;
+use sha3::{Digest, Sha3_256};
+use subtle::{Choice, ConditionallySelectable, ConstantTimeEq};
 
+use crate::errors::InvalidInstance;
 use crate::{
     codec::Shake128DuplexSponge,
     errors::Error,
@@ -48,18 +46,29 @@ pub enum ComposedRelation<G: PrimeGroup> {
     Or(Vec<ComposedRelation<G>>),
 }
 
+impl<G: PrimeGroup> ComposedRelation<G> {
+    /// Create a [ComposedRelation] for an AND relation from the given list of relations.
+    pub fn and<T: Into<ComposedRelation<G>>>(witness: impl IntoIterator<Item = T>) -> Self {
+        Self::And(witness.into_iter().map(|x| x.into()).collect())
+    }
+
+    /// Create a [ComposedRelation] for an OR relation from the given list of relations.
+    pub fn or<T: Into<ComposedRelation<G>>>(witness: impl IntoIterator<Item = T>) -> Self {
+        Self::Or(witness.into_iter().map(|x| x.into()).collect())
+    }
+}
+
 impl<G: PrimeGroup> From<CanonicalLinearRelation<G>> for ComposedRelation<G> {
     fn from(value: CanonicalLinearRelation<G>) -> Self {
         ComposedRelation::Simple(value)
     }
 }
 
-impl<G: PrimeGroup> From<LinearRelation<G>> for ComposedRelation<G> {
-    fn from(value: LinearRelation<G>) -> Self {
-        Self::Simple(
-            CanonicalLinearRelation::try_from(value)
-                .expect("Failed to convert LinearRelation to CanonicalLinearRelation"),
-        )
+impl<G: PrimeGroup> TryFrom<LinearRelation<G>> for ComposedRelation<G> {
+    type Error = InvalidInstance;
+
+    fn try_from(value: LinearRelation<G>) -> Result<Self, Self::Error> {
+        Ok(Self::Simple(CanonicalLinearRelation::try_from(value)?))
     }
 }
 
@@ -102,6 +111,26 @@ pub enum ComposedWitness<G: PrimeGroup> {
     Or(Vec<ComposedWitness<G>>),
 }
 
+impl<G: PrimeGroup> ComposedWitness<G> {
+    /// Create a [ComposedWitness] for an AND relation from the given list of witnesses.
+    pub fn and<T: Into<ComposedWitness<G>>>(witness: impl IntoIterator<Item = T>) -> Self {
+        Self::And(witness.into_iter().map(|x| x.into()).collect())
+    }
+
+    /// Create a [ComposedWitness] for an OR relation from the given list of witnesses.
+    pub fn or<T: Into<ComposedWitness<G>>>(witness: impl IntoIterator<Item = T>) -> Self {
+        Self::Or(witness.into_iter().map(|x| x.into()).collect())
+    }
+}
+
+impl<G: PrimeGroup> From<<CanonicalLinearRelation<G> as SigmaProtocol>::Witness>
+    for ComposedWitness<G>
+{
+    fn from(value: <CanonicalLinearRelation<G> as SigmaProtocol>::Witness) -> Self {
+        Self::Simple(value)
+    }
+}
+
 type ComposedChallenge<G> = <CanonicalLinearRelation<G> as SigmaProtocol>::Challenge;
 
 const fn composed_challenge_size<G: PrimeGroup>() -> usize {
@@ -207,37 +236,37 @@ impl<G: PrimeGroup + ConstantTimeEq> ComposedRelation<G> {
         let mut commitments = Vec::new();
         let mut prover_states = Vec::new();
 
+        // Selector value set when the first valid witness is found.
+        let mut valid_witness_found = Choice::from(0);
         for (i, w) in witnesses.iter().enumerate() {
             let (commitment, prover_state) = instances[i].prover_commit(w, rng)?;
             let (simulated_commitment, simulated_challenge, simulated_response) =
                 instances[i].simulate_transcript(rng)?;
 
+            // TODO: Implement and use ConditionallySelectable here
             let valid_witness = instances[i].is_witness_valid(w);
-            commitments.push(if valid_witness.unwrap_u8() == 1 {
+            let select_witness = valid_witness & !valid_witness_found;
+            commitments.push(if select_witness.unwrap_u8() == 1 {
                 commitment
             } else {
                 simulated_commitment.clone()
             });
             prover_states.push(ComposedOrProverStateEntry(
-                valid_witness,
+                select_witness,
                 prover_state,
                 simulated_challenge,
                 simulated_response,
             ));
+
+            valid_witness_found |= valid_witness;
         }
-        // check that we have only one witness set
-        let witnesses_found = prover_states
-            .iter()
-            .map(|x| x.0.unwrap_u8() as usize)
-            .sum::<usize>();
-        let prover_state = prover_states;
 
-        if witnesses_found != 1 {
+        if valid_witness_found.unwrap_u8() == 0 {
             Err(Error::InvalidInstanceWitnessPair)
         } else {
             Ok((
                 ComposedCommitment::Or(commitments),
-                ComposedProverState::Or(prover_state),
+                ComposedProverState::Or(prover_states),
             ))
         }
     }
@@ -250,15 +279,13 @@ impl<G: PrimeGroup + ConstantTimeEq> ComposedRelation<G> {
         let mut result_challenges = Vec::with_capacity(instances.len());
         let mut result_responses = Vec::with_capacity(instances.len());
 
-        let prover_states = prover_state;
-
         let mut witness_challenge = *challenge;
         for ComposedOrProverStateEntry(
             valid_witness,
             _prover_state,
             simulated_challenge,
             _simulated_response,
-        ) in &prover_states
+        ) in &prover_state
         {
             let c = G::Scalar::conditional_select(
                 simulated_challenge,
@@ -275,7 +302,7 @@ impl<G: PrimeGroup + ConstantTimeEq> ComposedRelation<G> {
                 simulated_challenge,
                 simulated_response,
             ),
-        ) in instances.iter().zip(prover_states)
+        ) in instances.iter().zip(prover_state)
         {
             let challenge_i = G::Scalar::conditional_select(
                 &simulated_challenge,

src/linear_relation/mod.rs

@@ -541,4 +541,11 @@ impl<G: PrimeGroup> LinearRelation<G> {
     ) -> Result<Nizk<CanonicalLinearRelation<G>, Shake128DuplexSponge<G>>, InvalidInstance> {
         Ok(Nizk::new(session_identifier, self.try_into()?))
     }
+
+    /// Construct a [CanonicalLinearRelation] from this generalized linear relation.
+    ///
+    /// The construction may fail if the linear relation is malformed, unsatisfiable, or trivial.
+    pub fn canonical(&self) -> Result<CanonicalLinearRelation<G>, InvalidInstance> {
+        self.try_into()
+    }
 }

src/tests/test_composition.rs

@@ -9,7 +9,7 @@ type G = RistrettoPoint;
 
 #[allow(non_snake_case)]
 #[test]
-fn test_composition_correctness() {
+fn test_composition_example() {
     // Composition and verification of proof for the following protocol :
     //
     // And(
@@ -31,30 +31,15 @@ fn test_composition_correctness() {
         .map(|_| <G as Group>::Scalar::random(&mut rng))
         .collect::<Vec<_>>();
     // second layer protocol definitions
-    let or_protocol1 = ComposedRelation::<G>::Or(vec![
-        ComposedRelation::Simple(relation1),
-        ComposedRelation::Simple(relation2),
-    ]);
-    let or_witness1 = ComposedWitness::Or(vec![
-        ComposedWitness::Simple(witness1),
-        ComposedWitness::Simple(wrong_witness2),
-    ]);
-
-    let simple_protocol1 = ComposedRelation::Simple(relation3);
-    let simple_witness1 = ComposedWitness::Simple(witness3);
-
-    let and_protocol1 = ComposedRelation::And(vec![
-        ComposedRelation::Simple(relation4),
-        ComposedRelation::Simple(relation5),
-    ]);
-    let and_witness1 = ComposedWitness::And(vec![
-        ComposedWitness::Simple(witness4),
-        ComposedWitness::Simple(witness5),
-    ]);
+    let or_protocol1 = ComposedRelation::<G>::or([relation1, relation2]);
+    let or_witness1 = ComposedWitness::or([witness1, wrong_witness2]);
+
+    let and_protocol1 = ComposedRelation::and([relation4, relation5]);
+    let and_witness1 = ComposedWitness::and([witness4, witness5]);
 
     // definition of the final protocol
-    let instance = ComposedRelation::And(vec![or_protocol1, simple_protocol1, and_protocol1]);
-    let witness = ComposedWitness::And(vec![or_witness1, simple_witness1, and_witness1]);
+    let instance = ComposedRelation::and([or_protocol1, relation3.into(), and_protocol1]);
+    let witness = ComposedWitness::and([or_witness1, witness3.into(), and_witness1]);
 
     let nizk = instance.into_nizk(domain_sep);
 
@@ -65,3 +50,61 @@ fn test_composition_correctness() {
     assert!(nizk.verify_batchable(&proof_batchable_bytes).is_ok());
     assert!(nizk.verify_compact(&proof_compact_bytes).is_ok());
 }
+
+#[allow(non_snake_case)]
+#[test]
+fn test_or_one_true() {
+    // Test composition of a basic OR protocol, with one of the two witnesses being valid.
+
+    // definitions of the underlying protocols
+    let mut rng = OsRng;
+    let (relation1, witness1) = dleq::<G, _>(&mut rng);
+    let (relation2, witness2) = dleq::<G, _>(&mut rng);
+
+    let wrong_witness1 = (0..witness1.len())
+        .map(|_| <G as Group>::Scalar::random(&mut rng))
+        .collect::<Vec<_>>();
+    let wrong_witness2 = (0..witness2.len())
+        .map(|_| <G as Group>::Scalar::random(&mut rng))
+        .collect::<Vec<_>>();
+
+    let or_protocol = ComposedRelation::or([relation1, relation2]);
+
+    // Construct two witnesses to the protocol, the first and then the second as the true branch.
+    let witness_or_1 = ComposedWitness::or([witness1, wrong_witness2]);
+    let witness_or_2 = ComposedWitness::or([wrong_witness1, witness2]);
+
+    let nizk = or_protocol.into_nizk(b"test_or_one_true");
+
+    for witness in [witness_or_1, witness_or_2] {
+        // Batchable and compact proofs
+        let proof_batchable_bytes = nizk.prove_batchable(&witness, &mut OsRng).unwrap();
+        let proof_compact_bytes = nizk.prove_compact(&witness, &mut OsRng).unwrap();
+        // Verify proofs
+        assert!(nizk.verify_batchable(&proof_batchable_bytes).is_ok());
+        assert!(nizk.verify_compact(&proof_compact_bytes).is_ok());
+    }
+}
+
+#[allow(non_snake_case)]
+#[test]
+fn test_or_both_true() {
+    // Test composition of a basic OR protocol, with both of the two witnesses being valid.
+
+    // definitions of the underlying protocols
+    let mut rng = OsRng;
+    let (relation1, witness1) = dleq::<G, _>(&mut rng);
+    let (relation2, witness2) = dleq::<G, _>(&mut rng);
+
+    let or_protocol = ComposedRelation::or([relation1, relation2]);
+
+    let witness = ComposedWitness::or([witness1, witness2]);
+    let nizk = or_protocol.into_nizk(b"test_or_one_true");
+
+    // Batchable and compact proofs
+    let proof_batchable_bytes = nizk.prove_batchable(&witness, &mut OsRng).unwrap();
+    let proof_compact_bytes = nizk.prove_compact(&witness, &mut OsRng).unwrap();
+    // Verify proofs
+    assert!(nizk.verify_batchable(&proof_batchable_bytes).is_ok());
+    assert!(nizk.verify_compact(&proof_compact_bytes).is_ok());
+}

src/tests/test_validation_criteria.rs

@@ -406,10 +406,8 @@ mod proof_validation {
         lr2.set_element(eq2, C);
 
         // Create OR composition
-        let or_relation = ComposedRelation::Or(vec![
-            ComposedRelation::from(lr1),
-            ComposedRelation::from(lr2),
-        ]);
+        let or_relation =
+            ComposedRelation::or([lr1.canonical().unwrap(), lr2.canonical().unwrap()]);
         let nizk = or_relation.into_nizk(b"test_or_relation");
 
         // Create a correct witness for branch 1 (C = y*B)