From 3f83712fbbd223e2c472543a093d8a2437f45966 Mon Sep 17 00:00:00 2001
From: Arthur Meyre <arthur.meyre@zama.ai>
Date: Wed, 4 Sep 2024 17:32:53 +0200
Subject: [PATCH] feat(zk): manage D as an upper bound as in the report

- allows to prove less slots than what the CRS can handle
---
 tfhe-zk-pok/src/proofs/pke_v2.rs | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/tfhe-zk-pok/src/proofs/pke_v2.rs b/tfhe-zk-pok/src/proofs/pke_v2.rs
index 7fa6fa4009..20355cccd0 100644
--- a/tfhe-zk-pok/src/proofs/pke_v2.rs
+++ b/tfhe-zk-pok/src/proofs/pke_v2.rs
@@ -300,10 +300,10 @@ pub fn prove<G: Curve>(
     let (
         &PublicParams {
             ref g_lists,
-            D,
+            D: D_max,
             n,
             d,
-            k,
+            k: k_max,
             B,
             B_r: _,
             B_bound,
@@ -328,8 +328,14 @@ pub fn prove<G: Curve>(
 
     let PrivateCommit { r, e1, m, e2, .. } = private_commit;
 
-    assert!(c2.len() <= k);
-    let k = k.min(c2.len());
+    assert!(c2.len() <= k_max);
+    let k = k_max.min(c2.len());
+
+    // Recompute the D for our case if k is smaller than the k max
+    // formula in Prove_pp: 2.
+    let D = d + k * t_input.ilog2() as usize;
+
+    assert!(D <= D_max);
 
     // FIXME: div_round
     let delta = {
@@ -1456,10 +1462,10 @@ pub fn verify<G: Curve>(
 
     let &PublicParams {
         ref g_lists,
-        D,
+        D: D_max,
         n,
         d,
-        k,
+        k: k_max,
         B,
         B_r: _,
         B_bound: _,
@@ -1487,10 +1493,18 @@ pub fn verify<G: Curve>(
     };
 
     let PublicCommit { a, b, c1, c2, .. } = public.1;
-    if c2.len() > k {
+    if c2.len() > k_max {
+        return Err(());
+    }
+    let k = k_max.min(c2.len());
+
+    // Recompute the D for our case if k is smaller than the k max
+    // formula in Prove_pp: 2.
+    let D = d + k * t_input.ilog2() as usize;
+
+    if D > D_max {
         return Err(());
     }
-    let k = k.min(c2.len());
 
     let C_hat_h3_bytes = C_hat_h3.map(G::G2::to_bytes);
     let C_hat_w_bytes = C_hat_w.map(G::G2::to_bytes);
@@ -2128,7 +2142,7 @@ mod tests {
                 PublicParams::deserialize_with_mode(data.as_slice(), compress, Validate::No)
             };
 
-        let original_public_param = crs_gen_ghl::<Curve>(d, k, B, q, t, rng);
+        let original_public_param = crs_gen_ghl::<Curve>(d, 2 * k, B, q, t, rng);
         let public_param_that_was_compressed =
             serialize_then_deserialize(&original_public_param, Compress::No).unwrap();
         let public_param_that_was_not_compressed =