feat: emit public-scalar claims under disjunctions as constant-shape claim branches

sigma-compiler-core/src/sigma/codegen.rs

@@ -3,13 +3,33 @@
 //! If that crate gets its own macro interface, it can use this module
 //! directly.
 
-use super::combiners::StatementTree;
+use super::combiners::{PrivScalarMap, StatementTree};
 use super::types::{expr_type_tokens_id_closure, AExprType, VarDict};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote, ToTokens};
 use std::collections::HashSet;
+use syn::visit::Visit;
 use syn::{Expr, Ident};
 
+/// Test whether an expression is term-free: it involves no private
+/// `Scalar`s, so its truth can be evaluated directly from the public
+/// instance rather than being proven in zero knowledge.  How such an
+/// equation is lowered depends on where it sits; see
+/// [`CodeGen::linear_relation_codegen`].
+fn expr_is_term_free(vars: &VarDict, expr: &Expr) -> bool {
+    let mut has_priv = false;
+    let mut psmap = PrivScalarMap {
+        vars,
+        closure: &mut |_| {
+            has_priv = true;
+            Ok(())
+        },
+        result: Ok(()),
+    };
+    psmap.visit_expr(expr);
+    !has_priv
+}
+
 /// Names and types of fields that might end up in a generated struct
 #[derive(Clone)]
 pub enum StructField {
@@ -201,14 +221,127 @@ impl<'a> CodeGen<'a> {
         }
     }
 
+    /// Generate code that pushes constant-shape claim branches for a
+    /// term-free equation `left = right` onto the relation vector
+    /// `out_var` (or, with `witness` set, the matching
+    /// `ComposedWitness::Claim` entries onto the witness vector).  The
+    /// claim `left - right == identity` is evaluated by both parties
+    /// from the instance accessed through `acc`; a vector equation
+    /// contributes one claim per component.  The emission is identical
+    /// whether the claim is true or false, so the proof shape never
+    /// depends on the public values.
+    fn claim_push_tokens(
+        &self,
+        expr: &Expr,
+        acc: &Ident,
+        out_var: &Ident,
+        witness: bool,
+    ) -> TokenStream {
+        let Expr::Assign(syn::ExprAssign { left, right, .. }) = expr else {
+            let expr_str = quote! { #expr }.to_string();
+            panic!("Unrecognized expression: {expr_str}");
+        };
+        let mut closure = |id: &Ident, _: AExprType| Ok(quote! {#acc.#id});
+        let (left_type, left_tokens) =
+            expr_type_tokens_id_closure(self.vars, left, &mut closure, false).unwrap();
+        let AExprType::Point {
+            is_pub: true,
+            is_vec: left_is_vec,
+        } = left_type
+        else {
+            let expr_str = quote! { #expr }.to_string();
+            panic!("Left side of = does not evaluate to a public point: {expr_str}");
+        };
+        let Ok((right_type, right_tokens)) =
+            expr_type_tokens_id_closure(self.vars, right, &mut closure, false)
+        else {
+            let expr_str = quote! { #expr }.to_string();
+            panic!("Right side of = is not a valid arithmetic expression: {expr_str}");
+        };
+        let AExprType::Point {
+            is_vec: right_is_vec,
+            ..
+        } = right_type
+        else {
+            let expr_str = quote! { #expr }.to_string();
+            panic!("Right side of = does not evaluate to a Point: {expr_str}");
+        };
+        if left_is_vec != right_is_vec {
+            let expr_str = quote! { #expr }.to_string();
+            panic!("Only one side of = is a vector expression: {expr_str}");
+        }
+        let one = quote! { Scalar::from_u128(1u128) };
+        match (witness, left_is_vec) {
+            (false, false) => quote! {
+                #out_var.push(ComposedInstance::claim([
+                    (#one, #left_tokens),
+                    (#one.neg(), #right_tokens),
+                ])?);
+            },
+            (false, true) => {
+                let lvar = format_ident!("{}cl", self.unique_prefix);
+                let rvar = format_ident!("{}cr", self.unique_prefix);
+                let lelem = format_ident!("{}cle", self.unique_prefix);
+                let relem = format_ident!("{}cre", self.unique_prefix);
+                quote! {
+                    {
+                        let #lvar = &(#left_tokens);
+                        let #rvar = &(#right_tokens);
+                        if #lvar.len() != #rvar.len() {
+                            return Err(InvalidInstance::new(
+                                "the two sides of a public equality have different lengths",
+                            ));
+                        }
+                        for (#lelem, #relem) in #lvar.iter().zip(#rvar.iter()) {
+                            #out_var.push(ComposedInstance::claim([
+                                (#one, *#lelem),
+                                (#one.neg(), *#relem),
+                            ])?);
+                        }
+                    }
+                }
+            }
+            (true, false) => quote! {
+                #out_var.push(ComposedWitness::Claim);
+            },
+            (true, true) => quote! {
+                for _ in 0..(#left_tokens).len() {
+                    #out_var.push(ComposedWitness::Claim);
+                }
+            },
+        }
+    }
+
     /// Generate the code for the `protocol` and `protocol_witness`
     /// functions that create the `ComposedInstance` and `ComposedWitness`
     /// structs, respectively, given a slice of [`Expr`]s that will be
     /// bundled into a single `LinearRelation`.  The `protocol` code
     /// must evaluate to a `Result<ComposedInstance>` and the `protocol_witness`
     /// code must evaluate to a `Result<ComposedWitness>`.
-    fn linear_relation_codegen(&self, exprs: &[&Expr]) -> (TokenStream, TokenStream) {
+    ///
+    /// With `disjunctive` set (the bundle sits under an `Or` or a
+    /// threshold), term-free equations become native claim branches
+    /// conjoined with the relation: both parties evaluate their truth
+    /// from the public instance, and the emission (hence the proof
+    /// shape) is the same whether they are true or false, so a branch
+    /// whose claim is false stays simulatable.  Otherwise they stay in
+    /// the `LinearRelation` and `compile()` evaluates them: a true one
+    /// is stripped, a false one invalidates the instance.  That keeps
+    /// the claim leaf, which is outside the specification's wire
+    /// format, off statements that never disjoin.
+    fn linear_relation_codegen(
+        &self,
+        exprs: &[&Expr],
+        disjunctive: bool,
+    ) -> (TokenStream, TokenStream) {
         let instance_var = format_ident!("{}instance", self.unique_prefix);
+        let (claim_exprs, exprs): (Vec<&Expr>, Vec<&Expr>) = if disjunctive {
+            exprs
+                .iter()
+                .partition(|expr| expr_is_term_free(self.vars, expr))
+        } else {
+            (Vec::new(), exprs.to_vec())
+        };
         let lr_var = format_ident!("{}lr", self.unique_prefix);
         let mut allocated_vars: HashSet<Ident> = HashSet::new();
         let mut param_vec_code = quote! {};
@@ -442,9 +575,46 @@ impl<'a> CodeGen<'a> {
             }
         }
 
-        (
+        if claim_exprs.is_empty() {
+            return (
+                quote! {
+                    {
+                        let mut #lr_var = LinearRelation::<Point>::new();
+                        #param_vec_code
+                        #scalar_allocs
+                        #element_allocs
+                        #eq_code
+                        #element_assigns
+
+                        ComposedInstance::try_from(#lr_var)
+                    }
+                },
+                quote! {
+                    {
+                        #witness_vec_code
+                        let mut witnessvec = Vec::new();
+                        #witness_code
+                        ProverResult::Ok(ComposedWitness::Simple(witnessvec))
+                    }
+                },
+            );
+        }
+
+        let rels_var = format_ident!("{}rels", self.unique_prefix);
+        let wits_var = format_ident!("{}wits", self.unique_prefix);
+        let claims_proto: Vec<TokenStream> = claim_exprs
+            .iter()
+            .map(|expr| self.claim_push_tokens(expr, &instance_var, &rels_var, false))
+            .collect();
+        let claims_witness: Vec<TokenStream> = claim_exprs
+            .iter()
+            .map(|expr| self.claim_push_tokens(expr, &format_ident!("instance"), &wits_var, true))
+            .collect();
+        let core_push_proto = if exprs.is_empty() {
+            quote! {}
+        } else {
             quote! {
-                {
+                #rels_var.push({
                     let mut #lr_var = LinearRelation::<Point>::new();
                     #param_vec_code
                     #scalar_allocs
@@ -452,15 +622,45 @@ impl<'a> CodeGen<'a> {
                     #eq_code
                     #element_assigns
 
-                    ComposedInstance::try_from(#lr_var)
-                }
-            },
+                    ComposedInstance::try_from(#lr_var)?
+                });
+            }
+        };
+        let core_push_witness = if exprs.is_empty() {
+            quote! {}
+        } else {
             quote! {
-                {
+                #wits_var.push({
                     #witness_vec_code
                     let mut witnessvec = Vec::new();
                     #witness_code
-                    ProverResult::Ok(ComposedWitness::Simple(witnessvec))
+                    ComposedWitness::Simple(witnessvec)
+                });
+            }
+        };
+        (
+            quote! {
+                {
+                    let mut #rels_var: Vec<ComposedInstance<Point>> = Vec::new();
+                    #(#claims_proto)*
+                    #core_push_proto
+                    if #rels_var.len() == 1 {
+                        InstanceResult::Ok(#rels_var.pop().unwrap())
+                    } else {
+                        ComposedInstance::and(#rels_var)
+                    }
+                }
+            },
+            quote! {
+                {
+                    let mut #wits_var: Vec<ComposedWitness<Point>> = Vec::new();
+                    #(#claims_witness)*
+                    #core_push_witness
+                    if #wits_var.len() == 1 {
+                        ProverResult::Ok(#wits_var.pop().unwrap())
+                    } else {
+                        ProverResult::Ok(ComposedWitness::and(#wits_var))
+                    }
                 }
             },
         )
@@ -474,7 +674,17 @@ impl<'a> CodeGen<'a> {
     /// The `protocol` code must evaluate to a `Result<Protocol>` and
     /// the `protocol_witness` code must evaluate to a
     /// `Result<ComposedWitness>`.
-    fn proto_witness_codegen(&self, statement: &StatementTree) -> (TokenStream, TokenStream) {
+    ///
+    /// `disjunctive` records whether this subtree lies under an `Or` or
+    /// a threshold, and so whether term-free equations need claim
+    /// leaves ([`Self::linear_relation_codegen`]).  It is set on
+    /// descent into a disjunction and never cleared: an `And` nested in
+    /// an `Or` is still a branch that may have to be simulated.
+    fn proto_witness_codegen(
+        &self,
+        statement: &StatementTree,
+        disjunctive: bool,
+    ) -> (TokenStream, TokenStream) {
         match statement {
             // The StatementTree has no statements (it's just the single
             // leaf "true")
@@ -489,7 +699,7 @@ impl<'a> CodeGen<'a> {
             // The StatementTree is a single statement.  Generate a
             // single LinearRelation from it.
             StatementTree::Leaf(leafexpr) => {
-                self.linear_relation_codegen(std::slice::from_ref(&leafexpr))
+                self.linear_relation_codegen(std::slice::from_ref(&leafexpr), disjunctive)
             }
             // The StatementTree is an And.  Separate out the leaf
             // statements, and generate a single LinearRelation from
@@ -504,14 +714,15 @@ impl<'a> CodeGen<'a> {
                         _ => others.push(st),
                     }
                 }
-                let (proto_code, witness_code) = self.linear_relation_codegen(&leaves);
+                let (proto_code, witness_code) =
+                    self.linear_relation_codegen(&leaves, disjunctive);
                 if others.is_empty() {
                     (proto_code, witness_code)
                 } else {
                     let (others_proto, others_witness): (Vec<TokenStream>, Vec<TokenStream>) =
                         others
                             .iter()
-                            .map(|st| self.proto_witness_codegen(st))
+                            .map(|st| self.proto_witness_codegen(st, disjunctive))
                             .unzip();
                     (
                         quote! {
@@ -532,7 +743,7 @@ impl<'a> CodeGen<'a> {
             StatementTree::Or(stvec) => {
                 let (proto, witness): (Vec<TokenStream>, Vec<TokenStream>) = stvec
                     .iter()
-                    .map(|st| self.proto_witness_codegen(st))
+                    .map(|st| self.proto_witness_codegen(st, true))
                     .unzip();
                 (
                     quote! {
@@ -550,7 +761,7 @@ impl<'a> CodeGen<'a> {
             StatementTree::Thresh(thresh, stvec) => {
                 let (proto, witness): (Vec<TokenStream>, Vec<TokenStream>) = stvec
                     .iter()
-                    .map(|st| self.proto_witness_codegen(st))
+                    .map(|st| self.proto_witness_codegen(st, true))
                     .unzip();
                 (
                     quote! {
@@ -656,7 +867,7 @@ impl<'a> CodeGen<'a> {
             quote! {}
         };
 
-        let (protocol_code, witness_code) = self.proto_witness_codegen(self.statements);
+        let (protocol_code, witness_code) = self.proto_witness_codegen(self.statements, false);
 
         // Generate the function that creates the sigma-proofs Protocol
         let protocol_func = {

tests/claim_narrowing.rs

@@ -0,0 +1,130 @@
+#![allow(non_snake_case)]
+//! Term-free (public) equations are lowered one of two ways, and these
+//! tests pin down which: conjoined, they are evaluated by `compile()`
+//! and cost nothing on the wire; under a disjunction, they become claim
+//! leaves whose constant shape lets a false branch be simulated.
+use curve25519_dalek::ristretto::RistrettoPoint as G;
+use group::Group;
+use sha2::Sha512;
+use sigma_compiler::*;
+
+type Scalar = <G as Group>::Scalar;
+
+/// A conjoined term-free equation leaves no trace: `compile()` strips
+/// it, so the statement is the one written without it, down to the
+/// transcript.  A proof of the stripped statement therefore verifies
+/// here -- which it cannot if the equation became a claim leaf, since
+/// that leaf is part of the composition the transcript binds.
+#[test]
+fn conjoined_claim_leaves_no_trace() -> Result<(), Box<dyn std::error::Error>> {
+    sigma_compiler! { plain,
+        (x),
+        (C, const cind A),
+        C = x*A,
+    }
+    sigma_compiler! { claimed,
+        (x),
+        (C, D, E, const cind A),
+        C = x*A,
+        D = E,
+    }
+
+    let mut rng = rand::thread_rng();
+    let A = G::hash_from_bytes::<Sha512>(b"Generator A");
+    let D = G::hash_from_bytes::<Sha512>(b"Element D");
+    let x = Scalar::random(&mut rng);
+    let C = x * A;
+
+    let plain_proof = plain::prove(
+        &plain::Instance { C, A },
+        &plain::Witness { x },
+        b"claim_narrowing",
+        &mut rng,
+    )?;
+
+    // `D = E` holds, so the statement is exactly the one above.
+    let instance = claimed::Instance { C, D, E: D, A };
+    let claimed_proof = claimed::prove(
+        &instance,
+        &claimed::Witness { x },
+        b"claim_narrowing",
+        &mut rng,
+    )?;
+    claimed::verify(&instance, &claimed_proof, b"claim_narrowing")?;
+
+    assert_eq!(claimed_proof.len(), plain_proof.len());
+    Ok(claimed::verify(
+        &instance,
+        &plain_proof,
+        b"claim_narrowing",
+    )?)
+}
+
+/// A conjoined term-free equation that does not hold makes the instance
+/// invalid, for the prover and the verifier alike -- an error, not a
+/// panic and not a proof that fails to verify later.
+#[test]
+fn false_conjoined_claim_is_an_error() {
+    sigma_compiler! { claimed,
+        (x),
+        (C, D, E, const cind A),
+        C = x*A,
+        D = E,
+    }
+
+    let mut rng = rand::thread_rng();
+    let A = G::hash_from_bytes::<Sha512>(b"Generator A");
+    let D = G::hash_from_bytes::<Sha512>(b"Element D");
+    let E = G::hash_from_bytes::<Sha512>(b"Element E");
+    let x = Scalar::random(&mut rng);
+    let C = x * A;
+
+    let instance = claimed::Instance { C, D, E, A };
+    claimed::prove(
+        &instance,
+        &claimed::Witness { x },
+        b"claim_narrowing",
+        &mut rng,
+    )
+    .unwrap_err();
+    claimed::verify(&instance, &[], b"claim_narrowing").unwrap_err();
+}
+
+/// Under a disjunction the claim leaf stays: the branch has to be
+/// simulatable when its claim is false, and the proof must look the
+/// same either way.
+#[test]
+fn disjunctive_claim_keeps_constant_shape() -> Result<(), Box<dyn std::error::Error>> {
+    sigma_compiler! { proof,
+        (x),
+        (C, D, E, const cind A),
+        OR (
+            D = E,
+            C = x*A,
+        )
+    }
+
+    let mut rng = rand::thread_rng();
+    let A = G::hash_from_bytes::<Sha512>(b"Generator A");
+    let D = G::hash_from_bytes::<Sha512>(b"Element D");
+    let E = G::hash_from_bytes::<Sha512>(b"Element E");
+    let x = Scalar::random(&mut rng);
+    let C = x * A;
+
+    // The second disjunct is the real one in both runs; only the truth
+    // of the first disjunct's claim differs.
+    let mut sizes = Vec::new();
+    for E in [D, E] {
+        let instance = proof::Instance { C, D, E, A };
+        let proof = proof::prove(
+            &instance,
+            &proof::Witness { x },
+            b"claim_narrowing",
+            &mut rng,
+        )?;
+        proof::verify(&instance, &proof, b"claim_narrowing")?;
+        sizes.push(proof.len());
+    }
+    assert_eq!(sizes[0], sizes[1]);
+    Ok(())
+}

tests/pubstatements.rs

@@ -4,8 +4,7 @@ use group::ff::PrimeField;
 use group::Group;
 use sigma_compiler::*;
 
-#[test]
-fn pubstatements_test() -> Result<(), Box<dyn std::error::Error>> {
+fn pubstatements_with_a(a_val: u128) -> Result<(), Box<dyn std::error::Error>> {
     sigma_compiler! { proof,
         (x, pub a),
         (C, D, const cind B),
@@ -17,7 +16,7 @@ fn pubstatements_test() -> Result<(), Box<dyn std::error::Error>> {
     let mut rng = rand::thread_rng();
     let B = G::generator();
     let x = Scalar::from_u128(5);
-    let a = Scalar::from_u128(0);
+    let a = Scalar::from_u128(a_val);
     let C = a * x * B;
     let D = a * B;
 
@@ -27,3 +26,11 @@ fn pubstatements_test() -> Result<(), Box<dyn std::error::Error>> {
     let proof = proof::prove(&instance, &witness, b"pubstatements_test", &mut rng)?;
     Ok(proof::verify(&instance, &proof, b"pubstatements_test")?)
 }
+
+#[test]
+fn pubstatements_test() -> Result<(), Box<dyn std::error::Error>> {
+    // a = 0 sends both images to the identity, and sigma-proofs refuses an
+    // instance that carries an identity group element.
+    assert!(pubstatements_with_a(0).is_err());
+    pubstatements_with_a(7)
+}

tests/pubstatements_vec.rs

@@ -15,8 +15,14 @@ fn pubstatements_vec_test_vecsize(vecsize: usize) -> Result<(), Box<dyn std::err
     type Scalar = <G as Group>::Scalar;
     let mut rng = rand::thread_rng();
     let B = G::generator();
-    let a: Vec<Scalar> = (0..vecsize).map(|i| Scalar::from_u128(i as u128)).collect();
-    let x: Vec<Scalar> = (0..vecsize).map(|i| Scalar::from_u128(i as u128)).collect();
+    // Nonzero throughout: a zero in either vector sends an image point to the
+    // identity, and sigma-proofs refuses an instance that carries one.
+    let a: Vec<Scalar> = (0..vecsize)
+        .map(|i| Scalar::from_u128(i as u128 + 1))
+        .collect();
+    let x: Vec<Scalar> = (0..vecsize)
+        .map(|i| Scalar::from_u128(i as u128 + 1))
+        .collect();
     let C: Vec<G> = (0..vecsize).map(|i| a[i] * x[i] * B).collect();
     let D: Vec<G> = (0..vecsize).map(|i| a[i] * B).collect();
 
@@ -29,7 +35,8 @@ fn pubstatements_vec_test_vecsize(vecsize: usize) -> Result<(), Box<dyn std::err
 
 #[test]
 fn pubstatements_vec_test() {
-    pubstatements_vec_test_vecsize(0).unwrap();
+    // No equations at size 0: a relation that binds nothing is refused.
+    assert!(pubstatements_vec_test_vecsize(0).is_err());
     pubstatements_vec_test_vecsize(1).unwrap();
     pubstatements_vec_test_vecsize(2).unwrap();
     pubstatements_vec_test_vecsize(20).unwrap();