MultiRoundNIZK

Author: vishady721

src/compressed.rs

@@ -3,6 +3,7 @@
 use crate::errors::Error as ProofError;
 use crate::errors::Result as ProofResult;
 use crate::linear_relation;
+use crate::serialization::deserialize_elements;
 use crate::serialization::deserialize_scalars;
 use crate::traits::InteractiveProof;
 use ff::Field;
@@ -74,24 +75,18 @@ fn fold_scalars<F: Field>(left: &[F], right: &[F], x: &F, x_inv: &F) -> Vec<F> {
         .collect()
 }
 
-struct CompressedProofMessage<G: PrimeGroup> {
-    final_message: Option<G::Scalar>,
-    intermediate_message: Option<[G; 2]>,
+enum CompressedProofMessage<G: PrimeGroup> {
+    FinalMessage(G::Scalar),
+    IntermediateMessage([G; 2]),
 }
 
 impl<G: PrimeGroup> CompressedProofMessage<G> {
     fn new_from_intermediate_message(intermediate_message: [G; 2]) -> Self {
-        Self {
-            final_message: None,
-            intermediate_message: Some(intermediate_message),
-        }
+        Self::IntermediateMessage(intermediate_message)
     }
 
     fn new_from_final_message(final_message: G::Scalar) -> Self {
-        Self {
-            final_message: Some(final_message),
-            intermediate_message: None,
-        }
+        Self::FinalMessage(final_message)
     }
 }
 
@@ -104,6 +99,25 @@ impl<G: PrimeGroup> InteractiveProof for SquashedLinearRelation<G> {
 
     type Challenge = G::Scalar;
 
+    type Witness = Vec<G::Scalar>;
+
+    fn get_initial_prover_state(&self, witness: &Self::Witness) -> Self::ProverState {
+        (
+            witness.to_vec(),
+            SquashedLinearRelation {
+                generators: self.generators.clone(),
+                image: self.image,
+            },
+        )
+    }
+
+    fn get_initial_verifier_state(&self) -> Self::VerifierState {
+        SquashedLinearRelation {
+            generators: self.generators.clone(),
+            image: self.image,
+        }
+    }
+
     fn prover_message(
         &self,
         state: &mut Self::ProverState,
@@ -148,44 +162,71 @@ impl<G: PrimeGroup> InteractiveProof for SquashedLinearRelation<G> {
         challenge: &Self::Challenge,
     ) -> Result<(), ProofError> {
         if state.generators.len() == 1 {
-            if prover_message.final_message.is_none() {
-                return Err(ProofError::VerificationFailure);
+            match prover_message {
+                CompressedProofMessage::FinalMessage(witness) => {
+                    let computed = state.generators[0] * witness;
+                    if computed == state.image {
+                        return Ok(());
+                    } else {
+                        return Err(ProofError::VerificationFailure);
+                    }
+                }
+                CompressedProofMessage::IntermediateMessage(_) => {
+                    return Err(ProofError::VerificationFailure);
+                }
             }
-            let witness = prover_message.final_message.unwrap();
-            let computed = state.generators[0] * witness[0];
-            if computed == state.image {
-                return Ok(());
-            } else {
+        }
+        match prover_message {
+            CompressedProofMessage::FinalMessage(_) => {
                 return Err(ProofError::VerificationFailure);
             }
+            CompressedProofMessage::IntermediateMessage([A, B]) => {
+                let n = state.generators.len() / 2;
+                let (g_left, g_right) = state.generators.split_at(n);
+                let new_generators = fold_generators(g_left, g_right, &challenge, &G::Scalar::ONE);
+                let new_image = *A + state.image * challenge + *B * challenge.square();
+                state.generators = new_generators;
+                state.image = new_image;
+                Ok(())
+            }
         }
-        if prover_message.intermediate_message.is_none() {
-            return Err(ProofError::VerificationFailure);
-        }
-        let [A, B] = prover_message.intermediate_message.unwrap();
-        let n = state.generators.len() / 2;
-        let (g_left, g_right) = state.generators.split_at(n);
-        let new_generators = fold_generators(g_left, g_right, &challenge, &G::Scalar::ONE);
-        let new_image = A + state.image * challenge + B * challenge.square();
-        state.generators = new_generators;
-        state.image = new_image;
-        Ok(())
     }
 
     fn serialize_message(&self, prover_message: &Self::ProverMessage) -> Vec<u8> {
-        todo!()
+        match prover_message {
+            CompressedProofMessage::FinalMessage(witness) => serialize_scalars::<G>(&[*witness]),
+            CompressedProofMessage::IntermediateMessage(prover_message) => {
+                serialize_elements(prover_message)
+            }
+        }
     }
 
     fn serialize_challenge(&self, challenge: &Self::Challenge) -> Vec<u8> {
-        todo!()
+        serialize_scalars::<G>(&[*challenge])
     }
 
-    fn deserialize_message(&self, data: &[u8]) -> Result<Self::ProverMessage, ProofError> {
-        todo!()
+    fn deserialize_message(
+        &self,
+        data: &[u8],
+        is_final_message: bool,
+    ) -> Result<Self::ProverMessage, ProofError> {
+        if is_final_message {
+            let witness =
+                deserialize_scalars::<G>(data, 1).ok_or(ProofError::VerificationFailure)?;
+            Ok(CompressedProofMessage::new_from_final_message(witness[0]))
+        } else {
+            let elements =
+                deserialize_elements::<G>(data, 2).ok_or(ProofError::VerificationFailure)?;
+            let intermediate_message: [G; 2] = [elements[0], elements[1]];
+            Ok(CompressedProofMessage::IntermediateMessage(
+                intermediate_message,
+            ))
+        }
     }
 
     fn deserialize_challenge(&self, data: &[u8]) -> Result<Self::Challenge, ProofError> {
-        todo!()
+        let scalars = deserialize_scalars::<G>(data, 1).ok_or(ProofError::VerificationFailure)?;
+        Ok(scalars[0])
     }
 
     fn protocol_identifier(&self) -> impl AsRef<[u8]> {
@@ -195,6 +236,10 @@ impl<G: PrimeGroup> InteractiveProof for SquashedLinearRelation<G> {
     fn instance_label(&self) -> impl AsRef<[u8]> {
         "TODO"
     }
+
+    fn num_rounds(&self) -> usize {
+        self.generators.len().next_power_of_two().ilog2() as usize + 1
+    }
 }
 
 impl<G: PrimeGroup> SquashedLinearRelation<G> {

src/fiat_shamir.rs

@@ -13,7 +13,7 @@
 //! - `C`: the codec ([`Codec`] trait).
 
 use crate::errors::Error;
-use crate::traits::SigmaProtocol;
+use crate::traits::{InteractiveProof, SigmaProtocol};
 use crate::{codec::Codec, traits::SigmaProtocolSimulator};
 use alloc::vec::Vec;
 
@@ -28,6 +28,11 @@ type Transcript<P> = (
     <P as SigmaProtocol>::Response,
 );
 
+type MultiRoundTranscript<P> = (
+    Vec<<P as InteractiveProof>::ProverMessage>,
+    Vec<<P as InteractiveProof>::Challenge>,
+);
+
 /// A Fiat-Shamir transformation of a [`SigmaProtocol`] into a non-interactive proof.
 ///
 /// [`Nizk`] wraps an interactive Sigma protocol `P`
@@ -304,3 +309,84 @@ where
         self.verify(&commitment, &challenge, &response)
     }
 }
+
+pub struct MultiRoundNizk<P, C>
+where
+    P: InteractiveProof,
+    P::Challenge: PartialEq,
+    C: Codec<Challenge = P::Challenge>,
+{
+    /// Current codec state.
+    pub hash_state: C,
+    /// Underlying interactive proof.
+    pub interactive_proof: P,
+}
+
+impl<P, C> MultiRoundNizk<P, C>
+where
+    P: InteractiveProof,
+    P::Challenge: PartialEq,
+    C: Codec<Challenge = P::Challenge> + Clone,
+{
+    /// Constructs a new [`MultiRoundNizk`] instance.
+    ///
+    /// # Parameters
+    /// - `iv`: Domain separation tag for the hash function (e.g., protocol name or context).
+    /// - `instance`: An instance of the [`InteractiveProof`].
+    ///
+    /// # Returns
+    /// A new [`MultiRoundNizk`] that can generate and verify non-interactive proofs.
+    pub fn new(session_identifier: &[u8], interactive_proof: P) -> Self {
+        let hash_state = C::new(
+            interactive_proof.protocol_identifier().as_ref(),
+            session_identifier,
+            interactive_proof.instance_label().as_ref(),
+        );
+        Self {
+            hash_state,
+            interactive_proof,
+        }
+    }
+
+    pub fn from_iv(iv: [u8; 64], interactive_proof: P) -> Self {
+        let hash_state = C::from_iv(iv);
+        Self {
+            hash_state,
+            interactive_proof,
+        }
+    }
+
+    fn prove(&self, witness: &P::Witness) -> Result<MultiRoundTranscript<P>, Error> {
+        let mut hash_state = self.hash_state.clone();
+        let num_rounds = self.interactive_proof.num_rounds();
+        let mut statement = self.interactive_proof.get_initial_prover_state(witness);
+        let mut messages = vec![];
+        let mut challenges = vec![];
+        (0..num_rounds).for_each(|_| {
+            let challenge = hash_state.verifier_challenge();
+            let message = self
+                .interactive_proof
+                .prover_message(&mut statement, &challenge)
+                .unwrap();
+            let serialized_message = self.interactive_proof.serialize_message(&message);
+            hash_state.prover_message(&serialized_message);
+            messages.push(message);
+            challenges.push(challenge);
+        });
+        Ok((messages, challenges))
+    }
+
+    fn verify(&self, prover_messages: &[P::ProverMessage]) -> Result<(), Error> {
+        let mut hash_state = self.hash_state.clone();
+        let num_rounds = self.interactive_proof.num_rounds();
+        assert_eq!(prover_messages.len(), num_rounds);
+        let mut statement = self.interactive_proof.get_initial_verifier_state();
+        for message in prover_messages {
+            let challenge = hash_state.verifier_challenge();
+            P::update_verifier_state(message, &mut statement, &challenge)?;
+            let serialized_message = self.interactive_proof.serialize_message(&message);
+            hash_state.prover_message(&serialized_message);
+        }
+        Ok(())
+    }
+}

src/traits.rs

@@ -148,6 +148,11 @@ pub trait InteractiveProof {
     type ProverMessage;
     type VerifierState;
     type Challenge;
+    type Witness;
+
+    fn get_initial_prover_state(&self, witness: &Self::Witness) -> Self::ProverState;
+
+    fn get_initial_verifier_state(&self) -> Self::VerifierState;
 
     fn prover_message(
         &self,
@@ -166,12 +171,18 @@ pub trait InteractiveProof {
     /// Serializes a challenge to bytes.
     fn serialize_challenge(&self, challenge: &Self::Challenge) -> Vec<u8>;
 
-    fn deserialize_message(&self, data: &[u8]) -> Result<Self::ProverMessage, Error>;
+    fn deserialize_message(
+        &self,
+        data: &[u8],
+        is_final_message: bool,
+    ) -> Result<Self::ProverMessage, Error>;
 
     /// Deserializes a challenge from bytes.
     fn deserialize_challenge(&self, data: &[u8]) -> Result<Self::Challenge, Error>;
 
     fn protocol_identifier(&self) -> impl AsRef<[u8]>;
 
     fn instance_label(&self) -> impl AsRef<[u8]>;
+
+    fn num_rounds(&self) -> usize;
 }