Merge pull request #981 from private-attribution/final_proof

Added the different logic for the last ZKP

ipa-core/src/protocol/ipa_prf/malicious_security/prover.rs

@@ -17,12 +17,14 @@ pub struct ZeroKnowledgeProof<F: PrimeField, N: ArrayLength> {
     g: GenericArray<F, N>,
 }
 
+#[derive(Debug)]
 pub struct ProofGenerator<F: PrimeField> {
     u: Vec<F>,
     v: Vec<F>,
 }
 
 type TwoNMinusOne<N> = Diff<Sum<N, N>, U1>;
+type TwoNPlusOne<N> = Sum<Sum<N, N>, U1>;
 
 ///
 /// Distributed Zero Knowledge Proofs algorithm drawn from
@@ -91,53 +93,126 @@ where
         };
         (proof, next_proof_generator)
     }
+
+    pub fn compute_final_proof<λ: ArrayLength>(
+        &self,
+        p_0: F,
+        q_0: F,
+    ) -> ZeroKnowledgeProof<F, TwoNPlusOne<λ>>
+    where
+        λ: ArrayLength + Add + Add<U1>,
+        <λ as Add>::Output: Add<U1>,
+        <<λ as Add>::Output as Add<U1>>::Output: ArrayLength,
+        <λ as Add<U1>>::Output: ArrayLength,
+    {
+        assert_eq!(self.u.len(), λ::USIZE); // We should pad with zeroes eventually
+
+        // We need a table of size `λ + 1` since we add a random point at x=0
+        let denominator = CanonicalLagrangeDenominator::<F, Sum<λ, U1>>::new();
+        let lagrange_table = LagrangeTable::<F, Sum<λ, U1>, λ>::from(denominator);
+
+        let mut p = vec![p_0];
+        p.extend_from_slice(&self.u);
+        let mut q = vec![q_0];
+        q.extend_from_slice(&self.v);
+        let p_extrapolated = lagrange_table.eval(&p);
+        let q_extrapolated = lagrange_table.eval(&q);
+
+        ZeroKnowledgeProof {
+            g: zip(
+                p.into_iter().chain(p_extrapolated),
+                q.into_iter().chain(q_extrapolated),
+            )
+            .map(|(a, b)| a * b)
+            .collect(),
+        }
+    }
+}
+
+impl<F> PartialEq<(&[u128], &[u128])> for ProofGenerator<F>
+where
+    F: PrimeField + std::cmp::PartialEq<u128>,
+{
+    fn eq(&self, other: &(&[u128], &[u128])) -> bool {
+        let (cmp_a, cmp_b) = other;
+        for (i, elem) in cmp_a.iter().enumerate() {
+            if !self.u[i].eq(elem) {
+                return false;
+            }
+        }
+        for (i, elem) in cmp_b.iter().enumerate() {
+            if !self.v[i].eq(elem) {
+                return false;
+            }
+        }
+        true
+    }
 }
 
 #[cfg(all(test, unit_test))]
 mod test {
-    use typenum::U4;
+    use typenum::{U2, U4};
 
     use super::ProofGenerator;
     use crate::ff::{Fp31, U128Conversions};
 
     #[test]
     fn sample_proof() {
-        const U: [u128; 32] = [
+        const U_1: [u128; 32] = [
             0, 30, 0, 16, 0, 1, 0, 15, 0, 0, 0, 16, 0, 30, 0, 16, 29, 1, 1, 15, 0, 0, 1, 15, 2, 30,
             30, 16, 0, 0, 30, 16,
         ];
-        const V: [u128; 32] = [
+        const V_1: [u128; 32] = [
             0, 0, 0, 30, 0, 0, 0, 1, 30, 30, 30, 30, 0, 0, 30, 30, 0, 30, 0, 30, 0, 0, 0, 1, 0, 0,
             1, 1, 0, 0, 1, 1,
         ];
-        const EXPECTED: [u128; 7] = [0, 30, 29, 30, 5, 28, 13];
-        const R1: u128 = 22;
-        const EXPECTED_NEXT_U: [u128; 8] = [0, 0, 26, 0, 7, 18, 24, 13];
-        const EXPECTED_NEXT_V: [u128; 8] = [10, 21, 30, 28, 15, 21, 3, 3];
+        const EXPECTED_1: [u128; 7] = [0, 30, 29, 30, 5, 28, 13];
+        const R_1: u128 = 22;
+        const U_2: [u128; 8] = [0, 0, 26, 0, 7, 18, 24, 13];
+        const V_2: [u128; 8] = [10, 21, 30, 28, 15, 21, 3, 3];
+
+        const EXPECTED_2: [u128; 7] = [12, 6, 15, 8, 29, 30, 6];
+        const R_2: u128 = 17;
+        const U_3: [u128; 2] = [3, 3];
+        const V_3: [u128; 2] = [5, 24];
+
+        const EXPECTED_3: [u128; 5] = [12, 15, 10, 14, 17];
+        const P_RANDOM_WEIGHT: u128 = 12;
+        const Q_RANDOM_WEIGHT: u128 = 1;
+
         let pg: ProofGenerator<Fp31> = ProofGenerator::new(
-            U.into_iter().map(|x| Fp31::try_from(x).unwrap()).collect(),
-            V.into_iter().map(|x| Fp31::try_from(x).unwrap()).collect(),
+            U_1.into_iter()
+                .map(|x| Fp31::try_from(x).unwrap())
+                .collect(),
+            V_1.into_iter()
+                .map(|x| Fp31::try_from(x).unwrap())
+                .collect(),
         );
-        let (proof, next_proof_generator) = pg.compute_proof::<U4>(Fp31::try_from(R1).unwrap());
+
+        // first iteration
+        let (proof, pg_2) = pg.compute_proof::<U4>(Fp31::try_from(R_1).unwrap());
         assert_eq!(
-            proof.g.into_iter().map(|x| x.as_u128()).collect::<Vec<_>>(),
-            EXPECTED,
+            proof.g.iter().map(Fp31::as_u128).collect::<Vec<_>>(),
+            EXPECTED_1,
         );
+        assert_eq!(pg_2, (&U_2[..], &V_2[..]));
+
+        // next iteration
+        let (proof_2, pg_3) = pg_2.compute_proof::<U4>(Fp31::try_from(R_2).unwrap());
         assert_eq!(
-            next_proof_generator
-                .u
-                .into_iter()
-                .map(|x| x.as_u128())
-                .collect::<Vec<_>>(),
-            EXPECTED_NEXT_U,
+            proof_2.g.iter().map(Fp31::as_u128).collect::<Vec<_>>(),
+            EXPECTED_2,
+        );
+        assert_eq!(pg_3, (&U_3[..], &V_3[..]));
+
+        // final iteration
+        let proof_3 = pg_3.compute_final_proof::<U2>(
+            Fp31::try_from(P_RANDOM_WEIGHT).unwrap(),
+            Fp31::try_from(Q_RANDOM_WEIGHT).unwrap(),
         );
         assert_eq!(
-            next_proof_generator
-                .v
-                .into_iter()
-                .map(|x| x.as_u128())
-                .collect::<Vec<_>>(),
-            EXPECTED_NEXT_V,
+            proof_3.g.iter().map(Fp31::as_u128).collect::<Vec<_>>(),
+            EXPECTED_3,
         );
     }
 }