From 5d4d81c29f752dfea04b8c52b832d651298543e8 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Mon, 19 Sep 2022 21:41:24 -0700
Subject: [PATCH] Shape check in starky

---
 evm/src/stark.rs       |  4 +--
 evm/src/verifier.rs    |  3 +-
 starky/src/stark.rs    |  4 +++
 starky/src/verifier.rs | 62 ++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/evm/src/stark.rs b/evm/src/stark.rs
index 00818bbe..49c5b70b 100644
--- a/evm/src/stark.rs
+++ b/evm/src/stark.rs
@@ -111,7 +111,7 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
             num_permutation_batches..num_permutation_batches + num_ctl_zs,
         );
 
-        let num_quotient_polys = self.quotient_degree_factor() * config.num_challenges;
+        let num_quotient_polys = self.num_quotient_polys(config);
         let quotient_oracle = FriOracleInfo {
             num_polys: num_quotient_polys,
             blinding: false,
@@ -174,7 +174,7 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
             num_permutation_batches..num_permutation_batches + num_ctl_zs,
         );
 
-        let num_quotient_polys = self.quotient_degree_factor() * inner_config.num_challenges;
+        let num_quotient_polys = self.num_quotient_polys(inner_config);
         let quotient_oracle = FriOracleInfo {
             num_polys: num_quotient_polys,
             blinding: false,
diff --git a/evm/src/verifier.rs b/evm/src/verifier.rs
index ad919f72..8c238e93 100644
--- a/evm/src/verifier.rs
+++ b/evm/src/verifier.rs
@@ -241,8 +241,7 @@ where
     ensure!(openings.permutation_ctl_zs.len() == num_zs);
     ensure!(openings.permutation_ctl_zs_next.len() == num_zs);
     ensure!(openings.ctl_zs_last.len() == num_ctl_zs);
-    let num_quotient_polys = stark.num_quotient_polys(config);
-    ensure!(openings.quotient_polys.len() == num_quotient_polys);
+    ensure!(openings.quotient_polys.len() == stark.num_quotient_polys(config));
 
     Ok(())
 }
diff --git a/starky/src/stark.rs b/starky/src/stark.rs
index 7f0df197..8ebca87c 100644
--- a/starky/src/stark.rs
+++ b/starky/src/stark.rs
@@ -78,6 +78,10 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
         1.max(self.constraint_degree() - 1)
     }
 
+    fn num_quotient_polys(&self, config: &StarkConfig) -> usize {
+        self.quotient_degree_factor() * config.num_challenges
+    }
+
     /// Computes the FRI instance used to prove this Stark.
     fn fri_instance(
         &self,
diff --git a/starky/src/verifier.rs b/starky/src/verifier.rs
index efb3d29c..d4847859 100644
--- a/starky/src/verifier.rs
+++ b/starky/src/verifier.rs
@@ -1,6 +1,6 @@
 use std::iter::once;
 
-use anyhow::{ensure, Result};
+use anyhow::{anyhow, ensure, Result};
 use itertools::Itertools;
 use plonky2::field::extension::{Extendable, FieldExtension};
 use plonky2::field::types::Field;
@@ -12,7 +12,7 @@ use plonky2::plonk::plonk_common::reduce_with_powers;
 use crate::config::StarkConfig;
 use crate::constraint_consumer::ConstraintConsumer;
 use crate::permutation::PermutationCheckVars;
-use crate::proof::{StarkOpeningSet, StarkProofChallenges, StarkProofWithPublicInputs};
+use crate::proof::{StarkOpeningSet, StarkProof, StarkProofChallenges, StarkProofWithPublicInputs};
 use crate::stark::Stark;
 use crate::vanishing_poly::eval_vanishing_poly;
 use crate::vars::StarkEvaluationVars;
@@ -144,6 +144,64 @@ where
     Ok(())
 }
 
+fn validate_proof_shape<F, C, S, const D: usize>(
+    stark: &S,
+    proof_with_pis: StarkProofWithPublicInputs<F, C, D>,
+    config: &StarkConfig,
+) -> anyhow::Result<()>
+where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    [(); S::COLUMNS]:,
+    [(); C::Hasher::HASH_SIZE]:,
+{
+    let degree_bits = proof_with_pis.proof.recover_degree_bits(config);
+
+    let StarkProofWithPublicInputs {
+        proof:
+            StarkProof {
+                trace_cap,
+                permutation_zs_cap,
+                quotient_polys_cap,
+                openings,
+                // The shape of the opening proof will be checked in the FRI verifier (see
+                // validate_fri_proof_shape), so we ignore it here.
+                opening_proof: _,
+            },
+        public_inputs,
+    } = proof_with_pis;
+
+    ensure!(public_inputs.len() == S::PUBLIC_INPUTS);
+
+    let fri_params = config.fri_params(degree_bits);
+    let cap_height = fri_params.config.cap_height;
+    let num_zs = stark.num_permutation_batches(config);
+
+    ensure!(trace_cap.height() == cap_height);
+    ensure!(quotient_polys_cap.height() == cap_height);
+
+    ensure!(openings.local_values.len() == S::COLUMNS);
+    ensure!(openings.next_values.len() == S::COLUMNS);
+    ensure!(openings.quotient_polys.len() == stark.num_quotient_polys(config));
+
+    if stark.uses_permutation_args() {
+        let permutation_zs_cap = permutation_zs_cap.ok_or_else(|| anyhow!("Missing Zs cap"))?;
+        let permutation_zs = openings
+            .permutation_zs
+            .ok_or_else(|| anyhow!("Missing permutation_zs"))?;
+        let permutation_zs_next = openings
+            .permutation_zs_next
+            .ok_or_else(|| anyhow!("Missing permutation_zs_next"))?;
+
+        ensure!(permutation_zs_cap.height() == cap_height);
+        ensure!(permutation_zs.len() == num_zs);
+        ensure!(permutation_zs_next.len() == num_zs);
+    }
+
+    Ok(())
+}
+
 /// Evaluate the Lagrange polynomials `L_0` and `L_(n-1)` at a point `x`.
 /// `L_0(x) = (x^n - 1)/(n * (x - 1))`
 /// `L_(n-1)(x) = (x^n - 1)/(n * (g * x - 1))`, with `g` the first element of the subgroup.