From 9de8c2c73b1e31146445809e3ac73d66bb2ad400 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Tue, 10 May 2022 15:08:08 +0200
Subject: [PATCH] Open lookup polys

---
 starky2/src/cross_table_lookup.rs |  4 ++++
 starky2/src/prover.rs             | 10 ++++------
 starky2/src/stark.rs              | 31 ++++++++++++++++++++-----------
 starky2/src/verifier.rs           |  2 ++
 4 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/starky2/src/cross_table_lookup.rs b/starky2/src/cross_table_lookup.rs
index 50ee8817..3d689b50 100644
--- a/starky2/src/cross_table_lookup.rs
+++ b/starky2/src/cross_table_lookup.rs
@@ -54,6 +54,10 @@ impl<F: Field> Default for LookupData<F> {
 }
 
 impl<F: Field> LookupData<F> {
+    pub fn len(&self) -> usize {
+        self.zs_beta_gammas.len()
+    }
+
     pub fn is_empty(&self) -> bool {
         self.zs_beta_gammas.is_empty()
     }
diff --git a/starky2/src/prover.rs b/starky2/src/prover.rs
index 02de8560..4ab821e4 100644
--- a/starky2/src/prover.rs
+++ b/starky2/src/prover.rs
@@ -159,14 +159,12 @@ where
     });
     let num_permutation_zs = permutation_zs.as_ref().map(|v| v.len()).unwrap_or(0);
 
-    let z_polys = match (permutation_zs, lookup_data.is_empty()) {
-        (None, true) => lookup_data.z_polys(),
-        (None, false) => vec![],
-        (Some(mut permutation_zs), true) => {
+    let z_polys = match permutation_zs {
+        None => lookup_data.z_polys(),
+        Some(mut permutation_zs) => {
             permutation_zs.extend(lookup_data.z_polys());
             permutation_zs
         }
-        (Some(permutation_zs), false) => permutation_zs,
     };
 
     let permutation_lookup_zs_commitment = (!z_polys.is_empty()).then(|| {
@@ -262,7 +260,7 @@ where
         timing,
         "compute openings proof",
         PolynomialBatch::prove_openings(
-            &stark.fri_instance(zeta, g, config),
+            &stark.fri_instance(zeta, g, degree_bits, lookup_data.len(), config),
             &initial_merkle_trees,
             challenger,
             &fri_params,
diff --git a/starky2/src/stark.rs b/starky2/src/stark.rs
index c5ec5405..35abb8ad 100644
--- a/starky2/src/stark.rs
+++ b/starky2/src/stark.rs
@@ -1,4 +1,5 @@
 use plonky2::field::extension_field::{Extendable, FieldExtension};
+use plonky2::field::field_types::Field;
 use plonky2::field::packed_field::PackedField;
 use plonky2::fri::structure::{
     FriBatchInfo, FriBatchInfoTarget, FriInstanceInfo, FriInstanceInfoTarget, FriOracleInfo,
@@ -78,6 +79,8 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
         &self,
         zeta: F::Extension,
         g: F,
+        degree_bits: usize,
+        num_ctl_zs: usize,
         config: &StarkConfig,
     ) -> FriInstanceInfo<F, D> {
         let no_blinding_oracle = FriOracleInfo { blinding: false };
@@ -86,14 +89,16 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
         let trace_info =
             FriPolynomialInfo::from_range(oracle_indices.next().unwrap(), 0..Self::COLUMNS);
 
-        let permutation_zs_info = if self.uses_permutation_args() {
-            FriPolynomialInfo::from_range(
-                oracle_indices.next().unwrap(),
-                0..self.num_permutation_batches(config),
-            )
-        } else {
-            vec![]
-        };
+        let num_permutation_batches = self.num_permutation_batches(config);
+        let permutation_lookup_index = oracle_indices.next().unwrap();
+        let permutation_lookup_zs_info = FriPolynomialInfo::from_range(
+            permutation_lookup_index,
+            0..num_permutation_batches + num_ctl_zs,
+        );
+        let lookup_zs_info = FriPolynomialInfo::from_range(
+            permutation_lookup_index,
+            num_permutation_batches..num_permutation_batches + num_ctl_zs,
+        );
 
         let quotient_info = FriPolynomialInfo::from_range(
             oracle_indices.next().unwrap(),
@@ -104,18 +109,22 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
             point: zeta,
             polynomials: [
                 trace_info.clone(),
-                permutation_zs_info.clone(),
+                permutation_lookup_zs_info.clone(),
                 quotient_info,
             ]
             .concat(),
         };
         let zeta_right_batch = FriBatchInfo {
             point: zeta.scalar_mul(g),
-            polynomials: [trace_info, permutation_zs_info].concat(),
+            polynomials: [trace_info, permutation_lookup_zs_info].concat(),
+        };
+        let lookup_batch = FriBatchInfo {
+            point: F::Extension::primitive_root_of_unity(degree_bits).inverse(),
+            polynomials: lookup_zs_info,
         };
         FriInstanceInfo {
             oracles: vec![no_blinding_oracle; oracle_indices.next().unwrap()],
-            batches: vec![zeta_batch, zeta_right_batch],
+            batches: vec![zeta_batch, zeta_right_batch, lookup_batch],
         }
     }
 
diff --git a/starky2/src/verifier.rs b/starky2/src/verifier.rs
index 05718ee0..886d7512 100644
--- a/starky2/src/verifier.rs
+++ b/starky2/src/verifier.rs
@@ -131,6 +131,8 @@ where
         &stark.fri_instance(
             challenges.stark_zeta,
             F::primitive_root_of_unity(degree_bits),
+            degree_bits,
+            todo!(),
             config,
         ),
         &proof.openings.to_fri_openings(),