Add a new ScalarMap struct and use it to make witnesses more structured

Cargo.toml

@@ -19,26 +19,27 @@ description = "A toolkit for auto-generated implementations of Σ-protocols"
 exclude = [
     ".gitignore"
 ]
+rust-version = "1.81"
 
 [features]
 default = ["std"]
-std = ["thiserror", "rand", "num-bigint/std", "num-traits/std", "sha3/std", "rand_core/std"]
+std = ["rand", "num-bigint/std", "num-traits/std", "sha3/std", "rand_core/std"]
 
 [dependencies]
+ahash = { version = "0.8", default-features = false }
 ff = { version = "0.13", features = ["derive"] }
 group = "0.13.0"
+hashbrown = { version = "0.15", default-features = false }
+keccak = { version = "0.1.5", default-features = false }
 num-bigint = { version = "0.4.6", default-features = false }
 num-traits = { version = "0.2.19", default-features = false, features = ["libm"] }
 rand = { version = "0.8.5", optional = true }
 rand_core = { version = "0.6", default-features = false }
 sha3 = { version = "0.10.8", default-features = false }
 subtle = { version = "2.6.1", default-features = false }
-thiserror = { version = "1", optional = true }
-keccak = { version = "0.1.5", default-features = false }
+thiserror = { version = "2.0.16", default-features = false }
 zerocopy = { version = "0.8", default-features = false }
 zeroize = { version = "1.8.1", default-features = false, features = ["alloc"] }
-hashbrown = { version = "0.15", default-features = false }
-ahash = { version = "0.8", default-features = false }
 
 [dev-dependencies]
 bls12_381 = "0.8.0"

examples/schnorr.rs

@@ -6,64 +6,83 @@
 //!
 //! where $G$ is a generator of a prime-order group $\mathbb{G}$ and $P$ is a public group element.
 
+use std::process::ExitCode;
+
 use curve25519_dalek::scalar::Scalar;
 use curve25519_dalek::RistrettoPoint;
 use group::Group;
 use rand::rngs::OsRng;
 
 use sigma_proofs::errors::Error;
+use sigma_proofs::linear_relation::{Allocator, GroupVar, ScalarVar};
 use sigma_proofs::LinearRelation;
 
-type ProofResult<T> = Result<T, Error>;
-
 /// Create a discrete logarithm relation for the given public key P
 #[allow(non_snake_case)]
-fn create_relation(P: RistrettoPoint) -> LinearRelation<RistrettoPoint> {
+fn create_relation() -> (
+    LinearRelation<RistrettoPoint>,
+    ScalarVar<RistrettoPoint>,
+    GroupVar<RistrettoPoint>,
+) {
     let mut relation = LinearRelation::new();
 
     let x = relation.allocate_scalar();
-    let G = relation.allocate_element();
-    let P_var = relation.allocate_eq(x * G);
-    relation.set_element(G, RistrettoPoint::generator());
-    relation.set_element(P_var, P);
+    let G = relation.allocate_element_with(RistrettoPoint::generator());
+    let P = relation.allocate_eq(x * G);
 
-    relation
+    (relation, x, P)
 }
 
 /// Prove knowledge of the discrete logarithm: given witness x and public key P,
 /// generate a proof that P = x * G
 #[allow(non_snake_case)]
-fn prove(x: Scalar, P: RistrettoPoint) -> ProofResult<Vec<u8>> {
-    let nizk = create_relation(P).into_nizk(b"sigma-proofs-example");
-    nizk?.prove_batchable(&vec![x], &mut OsRng)
+fn prove(x: Scalar) -> Result<Vec<u8>, Error> {
+    let (mut relation, x_var, _) = create_relation();
+    let witness = [(x_var, x)];
+    relation.compute_image(witness)?;
+    relation
+        .into_nizk(b"sigma-proofs-example")?
+        .prove_batchable(witness, &mut OsRng)
 }
 
 /// Verify a proof of knowledge of discrete logarithm for the given public key P
 #[allow(non_snake_case)]
-fn verify(P: RistrettoPoint, proof: &[u8]) -> ProofResult<()> {
-    let nizk = create_relation(P).into_nizk(b"sigma-proofs-example");
-    nizk?.verify_batchable(proof)
+fn verify(P: RistrettoPoint, proof: &[u8]) -> Result<(), Error> {
+    let (mut relation, _, P_var) = create_relation();
+    relation.assign_element(P_var, P);
+    relation
+        .into_nizk(b"sigma-proofs-example")?
+        .verify_batchable(proof)
 }
 
 #[allow(non_snake_case)]
-fn main() {
+fn main() -> ExitCode {
     let x = Scalar::random(&mut OsRng); // Private key (witness)
     let P = RistrettoPoint::generator() * x; // Public key (statement)
 
     println!("Generated new key pair:");
     println!("Public key P: {:?}", hex::encode(P.compress().as_bytes()));
 
-    match prove(x, P) {
+    let proof = match prove(x) {
         Ok(proof) => {
             println!("Proof generated successfully:");
             println!("Proof (hex): {}", hex::encode(&proof));
+            proof
+        }
+        Err(e) => {
+            println!("✗ Failed to generate proof: {e:?}");
+            return ExitCode::FAILURE;
+        }
+    };
 
-            // Verify the proof
-            match verify(P, &proof) {
-                Ok(()) => println!("✓ Proof verified successfully!"),
-                Err(e) => println!("✗ Proof verification failed: {e:?}"),
-            }
+    // Verify the proof
+    match verify(P, &proof) {
+        Ok(()) => println!("✓ Proof verified successfully!"),
+        Err(e) => {
+            println!("✗ Proof verification failed: {e:?}");
+            return ExitCode::FAILURE;
         }
-        Err(e) => println!("✗ Failed to generate proof: {e:?}"),
     }
+
+    ExitCode::SUCCESS
 }

examples/simple_composition.rs

@@ -8,6 +8,7 @@ use sigma_proofs::{
     codec::Shake128DuplexSponge,
     composition::{ComposedRelation, ComposedWitness},
     errors::Error,
+    linear_relation::{Allocator, ScalarVar},
     LinearRelation, Nizk,
 };
 
@@ -18,29 +19,30 @@ type ProofResult<T> = Result<T, Error>;
 /// 1. Knowledge of discrete log: P1 = x1 * G
 /// 2. Knowledge of DLEQ: (P2 = x2 * G, Q = x2 * H)
 #[allow(non_snake_case)]
-fn create_relation(P1: G, P2: G, Q: G, H: G) -> ComposedRelation<G> {
+fn create_relation(P1: G, P2: G, Q: G, H: G) -> (ComposedRelation<G>, ScalarVar<G>, ScalarVar<G>) {
     // First relation: discrete logarithm P1 = x1 * G
     let mut rel1 = LinearRelation::<G>::new();
     let x1 = rel1.allocate_scalar();
-    let G1 = rel1.allocate_element();
+    let G1 = rel1.allocate_element_with(G::generator());
     let P1_var = rel1.allocate_eq(x1 * G1);
-    rel1.set_element(G1, G::generator());
-    rel1.set_element(P1_var, P1);
+    rel1.assign_element(P1_var, P1);
 
     // Second relation: DLEQ (P2 = x2 * G, Q = x2 * H)
     let mut rel2 = LinearRelation::<G>::new();
     let x2 = rel2.allocate_scalar();
-    let G2 = rel2.allocate_element();
-    let H_var = rel2.allocate_element();
+    let G2 = rel2.allocate_element_with(G::generator());
+    let H_var = rel2.allocate_element_with(H);
     let P2_var = rel2.allocate_eq(x2 * G2);
     let Q_var = rel2.allocate_eq(x2 * H_var);
-    rel2.set_element(G2, G::generator());
-    rel2.set_element(H_var, H);
-    rel2.set_element(P2_var, P2);
-    rel2.set_element(Q_var, Q);
+    rel2.assign_element(P2_var, P2);
+    rel2.assign_element(Q_var, Q);
 
     // Compose into OR protocol
-    ComposedRelation::or([rel1.canonical().unwrap(), rel2.canonical().unwrap()])
+    (
+        ComposedRelation::or([rel1.canonical().unwrap(), rel2.canonical().unwrap()]),
+        x1,
+        x2,
+    )
 }
 
 /// Prove knowledge of one of the witnesses (we know x2 for the DLEQ)
@@ -50,22 +52,18 @@ fn prove(P1: G, x2: Scalar, H: G) -> ProofResult<Vec<u8>> {
     let P2 = G::generator() * x2;
     let Q = H * x2;
 
-    let instance = create_relation(P1, P2, Q, H);
+    let (relation, var_x1, var_x2) = create_relation(P1, P2, Q, H);
     // Create OR witness with branch 1 being the real one (index 1)
-    let witness = ComposedWitness::Or(vec![
-        ComposedWitness::Simple(vec![Scalar::from(0u64)]),
-        ComposedWitness::Simple(vec![x2]),
-    ]);
-    let nizk = Nizk::<_, Shake128DuplexSponge<G>>::new(b"or_proof_example", instance);
-
-    nizk.prove_batchable(&witness, &mut OsRng)
+    let witness = ComposedWitness::or([[(var_x1, Scalar::from(0u64))], [(var_x2, x2)]]);
+    let nizk = Nizk::<_, Shake128DuplexSponge<G>>::new(b"or_proof_example", relation);
+    nizk.prove_batchable(witness, &mut OsRng)
 }
 
 /// Verify an OR proof given the public values
 #[allow(non_snake_case)]
 fn verify(P1: G, P2: G, Q: G, H: G, proof: &[u8]) -> ProofResult<()> {
-    let protocol = create_relation(P1, P2, Q, H);
-    let nizk = Nizk::<_, Shake128DuplexSponge<G>>::new(b"or_proof_example", protocol);
+    let (relation, _, _) = create_relation(P1, P2, Q, H);
+    let nizk = Nizk::<_, Shake128DuplexSponge<G>>::new(b"or_proof_example", relation);
 
     nizk.verify_batchable(proof)
 }

src/composition.rs

@@ -28,8 +28,10 @@ use rand_core::{CryptoRng, RngCore as Rng};
 use sha3::{Digest, Sha3_256};
 use subtle::{Choice, ConditionallySelectable, ConstantTimeEq};
 
+use crate::codec::Codec;
 use crate::errors::InvalidInstance;
 use crate::group::serialization::{deserialize_scalars, serialize_scalars};
+use crate::linear_relation::{Allocator, ScalarVar};
 use crate::{
     codec::Shake128DuplexSponge,
     errors::Error,
@@ -51,7 +53,7 @@ pub enum ComposedRelation<G: PrimeGroup> {
     Or(Vec<ComposedRelation<G>>),
 }
 
-impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> 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())
@@ -69,10 +71,10 @@ impl<G: PrimeGroup> From<CanonicalLinearRelation<G>> for ComposedRelation<G> {
     }
 }
 
-impl<G: PrimeGroup> TryFrom<LinearRelation<G>> for ComposedRelation<G> {
+impl<G: PrimeGroup, A: Allocator<G = G>> TryFrom<LinearRelation<G, A>> for ComposedRelation<G> {
     type Error = InvalidInstance;
 
-    fn try_from(value: LinearRelation<G>) -> Result<Self, Self::Error> {
+    fn try_from(value: LinearRelation<G, A>) -> Result<Self, Self::Error> {
         Ok(Self::Simple(CanonicalLinearRelation::try_from(value)?))
     }
 }
@@ -223,6 +225,18 @@ impl<G: PrimeGroup> ComposedWitness<G> {
     }
 }
 
+impl<G: PrimeGroup, const N: usize> From<[(ScalarVar<G>, G::Scalar); N]> for ComposedWitness<G> {
+    fn from(value: [(ScalarVar<G>, G::Scalar); N]) -> Self {
+        Self::Simple(value.into())
+    }
+}
+
+impl<G: PrimeGroup> From<Vec<(ScalarVar<G>, G::Scalar)>> for ComposedWitness<G> {
+    fn from(value: Vec<(ScalarVar<G>, G::Scalar)>) -> Self {
+        Self::Simple(value.into())
+    }
+}
+
 impl<G: PrimeGroup> From<<CanonicalLinearRelation<G> as SigmaProtocol>::Witness>
     for ComposedWitness<G>
 {
@@ -237,6 +251,29 @@ const fn composed_challenge_size<G: PrimeGroup>() -> usize {
     (G::Scalar::NUM_BITS as usize).div_ceil(8)
 }
 
+impl<G: PrimeGroup> ComposedRelation<G>
+where
+    Self: SigmaProtocol,
+{
+    /// Convert this Protocol into a non-interactive zero-knowledge proof
+    /// using the Shake128DuplexSponge codec and a specified session identifier.
+    ///
+    /// This method provides a convenient way to create a NIZK from a Protocol
+    /// without exposing the specific codec type to the API caller.
+    ///
+    /// # Parameters
+    /// - `session_identifier`: Domain separator bytes for the Fiat-Shamir transform
+    ///
+    /// # Returns
+    /// A `Nizk` instance ready for proving and verification
+    pub fn into_nizk(self, session_identifier: &[u8]) -> Nizk<Self, Shake128DuplexSponge<G>>
+    where
+        Shake128DuplexSponge<G>: Codec<Challenge = <Self as SigmaProtocol>::Challenge>,
+    {
+        Nizk::new(session_identifier, self)
+    }
+}
+
 impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<G> {
     fn is_witness_valid(&self, witness: &ComposedWitness<G>) -> Choice {
         match (self, witness) {
@@ -261,7 +298,7 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<
 
     fn prover_commit_simple(
         protocol: &CanonicalLinearRelation<G>,
-        witness: &<CanonicalLinearRelation<G> as SigmaProtocol>::Witness,
+        witness: <CanonicalLinearRelation<G> as SigmaProtocol>::Witness,
         rng: &mut (impl Rng + CryptoRng),
     ) -> Result<(ComposedCommitment<G>, ComposedProverState<G>), Error> {
         protocol.prover_commit(witness, rng).map(|(c, s)| {
@@ -284,7 +321,7 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<
 
     fn prover_commit_and(
         protocols: &[ComposedRelation<G>],
-        witnesses: &[ComposedWitness<G>],
+        witnesses: Vec<ComposedWitness<G>>,
         rng: &mut (impl Rng + CryptoRng),
     ) -> Result<(ComposedCommitment<G>, ComposedProverState<G>), Error> {
         if protocols.len() != witnesses.len() {
@@ -294,7 +331,7 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<
         let mut commitments = Vec::with_capacity(protocols.len());
         let mut prover_states = Vec::with_capacity(protocols.len());
 
-        for (p, w) in protocols.iter().zip(witnesses.iter()) {
+        for (p, w) in protocols.iter().zip(witnesses.into_iter()) {
             let (c, s) = p.prover_commit(w, rng)?;
             commitments.push(c);
             prover_states.push(s);
@@ -326,7 +363,7 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<
 
     fn prover_commit_or(
         instances: &[ComposedRelation<G>],
-        witnesses: &[ComposedWitness<G>],
+        witnesses: Vec<ComposedWitness<G>>,
         rng: &mut (impl Rng + CryptoRng),
     ) -> Result<(ComposedCommitment<G>, ComposedProverState<G>), Error>
     where
@@ -341,14 +378,15 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<
 
         // 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() {
+        for (i, w) in witnesses.into_iter().enumerate() {
+            // Determine whether or not to use the real witness for this relation. This rule uses
+            // the first valid witness found in the given list.
+            let select_witness = instances[i].is_witness_valid(&w) & !valid_witness_found;
+
             let (commitment, prover_state) = instances[i].prover_commit(w, rng)?;
             let (simulated_commitment, simulated_challenge, simulated_response) =
                 instances[i].simulate_transcript(rng)?;
 
-            let valid_witness = instances[i].is_witness_valid(w) & !valid_witness_found;
-            let select_witness = valid_witness;
-
             let commitment = ComposedCommitment::conditional_select(
                 &simulated_commitment,
                 &commitment,
@@ -363,7 +401,7 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<
                 simulated_response,
             ));
 
-            valid_witness_found |= valid_witness;
+            valid_witness_found |= select_witness;
         }
 
         if valid_witness_found.unwrap_u8() == 0 {
@@ -439,7 +477,7 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> SigmaProtocol
 
     fn prover_commit(
         &self,
-        witness: &Self::Witness,
+        witness: Self::Witness,
         rng: &mut (impl Rng + CryptoRng),
     ) -> Result<(Self::Commitment, Self::ProverState), Error> {
         match (self, witness) {
@@ -788,23 +826,3 @@ impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> SigmaProtocolSimu
         }
     }
 }
-
-impl<G: PrimeGroup + ConstantTimeEq + ConditionallySelectable> ComposedRelation<G> {
-    /// Convert this Protocol into a non-interactive zero-knowledge proof
-    /// using the Shake128DuplexSponge codec and a specified session identifier.
-    ///
-    /// This method provides a convenient way to create a NIZK from a Protocol
-    /// without exposing the specific codec type to the API caller.
-    ///
-    /// # Parameters
-    /// - `session_identifier`: Domain separator bytes for the Fiat-Shamir transform
-    ///
-    /// # Returns
-    /// A `Nizk` instance ready for proving and verification
-    pub fn into_nizk(
-        self,
-        session_identifier: &[u8],
-    ) -> Nizk<ComposedRelation<G>, Shake128DuplexSponge<G>> {
-        Nizk::new(session_identifier, self)
-    }
-}