Merge pull request #1279 from o1-labs/feature/num-zk-rows-variable2

Capture zk_rows as a variable in the expression framework

kimchi/src/alphas.rs

@@ -323,11 +323,8 @@ mod tests {
     fn get_alphas_for_spec() {
         let gates = vec![CircuitGate::<Fp>::zero(Wire::for_row(0)); 2];
         let index = new_index_for_test::<Vesta>(gates, 0);
-        let (_linearization, powers_of_alpha) = expr_linearization::<Fp>(
-            Some(&index.cs.feature_flags),
-            true,
-            index.cs.zk_rows as usize,
-        );
+        let (_linearization, powers_of_alpha) =
+            expr_linearization::<Fp>(Some(&index.cs.feature_flags), true);
         // make sure this is present in the specification
         let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
         let spec_path = Path::new(&manifest_dir)

kimchi/src/circuits/expr.rs

@@ -21,11 +21,11 @@ use o1_utils::{foreign_field::ForeignFieldHelpers, FieldHelpers};
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
 use std::ops::{Add, AddAssign, Mul, Neg, Sub};
+use std::{cmp::Ordering, fmt, iter::FromIterator};
 use std::{
     collections::{HashMap, HashSet},
     ops::MulAssign,
 };
-use std::{fmt, iter::FromIterator};
 use thiserror::Error;
 use CurrOrNext::{Curr, Next};
 
@@ -459,6 +459,12 @@ impl FeatureFlag {
     }
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RowOffset {
+    pub zk_rows: bool,
+    pub offset: i32,
+}
+
 /// An multi-variate polynomial over the base ring `C` with
 /// variables
 ///
@@ -479,7 +485,7 @@ pub enum Expr<C> {
     VanishesOnZeroKnowledgeAndPreviousRows,
     /// UnnormalizedLagrangeBasis(i) is
     /// (x^n - 1) / (x - omega^i)
-    UnnormalizedLagrangeBasis(i32),
+    UnnormalizedLagrangeBasis(RowOffset),
     Pow(Box<Expr<C>>, u64),
     Cache(CacheId, Box<Expr<C>>),
     /// If the feature flag is enabled, return the first expression; otherwise, return the second.
@@ -649,7 +655,7 @@ pub enum PolishToken<F> {
     Mul,
     Sub,
     VanishesOnZeroKnowledgeAndPreviousRows,
-    UnnormalizedLagrangeBasis(i32),
+    UnnormalizedLagrangeBasis(RowOffset),
     Store,
     Load(usize),
     /// Skip the given number of tokens if the feature is enabled.
@@ -764,7 +770,12 @@ impl<F: FftField> PolishToken<F> {
                     stack.push(eval_vanishes_on_last_n_rows(d, c.zk_rows + 1, pt))
                 }
                 UnnormalizedLagrangeBasis(i) => {
-                    stack.push(unnormalized_lagrange_basis(&d, *i, &pt))
+                    let offset = if i.zk_rows {
+                        -(c.zk_rows as i32) + i.offset
+                    } else {
+                        i.offset
+                    };
+                    stack.push(unnormalized_lagrange_basis(&d, offset, &pt))
                 }
                 Literal(x) => stack.push(*x),
                 Dup => stack.push(stack[stack.len() - 1]),
@@ -1587,7 +1598,14 @@ impl<F: FftField> Expr<ConstantExpr<F>> {
             VanishesOnZeroKnowledgeAndPreviousRows => {
                 Ok(eval_vanishes_on_last_n_rows(d, c.zk_rows + 1, pt))
             }
-            UnnormalizedLagrangeBasis(i) => Ok(unnormalized_lagrange_basis(&d, *i, &pt)),
+            UnnormalizedLagrangeBasis(i) => {
+                let offset = if i.zk_rows {
+                    -(c.zk_rows as i32) + i.offset
+                } else {
+                    i.offset
+                };
+                Ok(unnormalized_lagrange_basis(&d, offset, &pt))
+            }
             Cell(v) => v.evaluate(evals),
             Cache(_, e) => e.evaluate_(d, pt, evals, c),
             IfFeature(feature, e1, e2) => {
@@ -1649,7 +1667,14 @@ impl<F: FftField> Expr<F> {
             VanishesOnZeroKnowledgeAndPreviousRows => {
                 Ok(eval_vanishes_on_last_n_rows(d, zk_rows + 1, pt))
             }
-            UnnormalizedLagrangeBasis(i) => Ok(unnormalized_lagrange_basis(&d, *i, &pt)),
+            UnnormalizedLagrangeBasis(i) => {
+                let offset = if i.zk_rows {
+                    -(zk_rows as i32) + i.offset
+                } else {
+                    i.offset
+                };
+                Ok(unnormalized_lagrange_basis(&d, offset, &pt))
+            }
             Cell(v) => v.evaluate(evals),
             Cache(_, e) => e.evaluate(d, pt, zk_rows, evals),
             IfFeature(feature, e1, e2) => {
@@ -1787,10 +1812,17 @@ impl<F: FftField> Expr<F> {
                 evals: env.vanishes_on_zero_knowledge_and_previous_rows,
             },
             Expr::Constant(x) => EvalResult::Constant(*x),
-            Expr::UnnormalizedLagrangeBasis(i) => EvalResult::Evals {
-                domain: d,
-                evals: unnormalized_lagrange_evals(env.l0_1, *i, d, env),
-            },
+            Expr::UnnormalizedLagrangeBasis(i) => {
+                let offset = if i.zk_rows {
+                    -(env.constants.zk_rows as i32) + i.offset
+                } else {
+                    i.offset
+                };
+                EvalResult::Evals {
+                    domain: d,
+                    evals: unnormalized_lagrange_evals(env.l0_1, offset, d, env),
+                }
+            }
             Expr::Cell(Variable { col, row }) => {
                 let evals: &'a Evaluations<F, D<F>> = {
                     match env.get_column(col) {
@@ -2481,7 +2513,9 @@ where
             Double(x) => format!("double({})", x.ocaml(cache)),
             Constant(x) => x.ocaml(),
             Cell(v) => format!("cell({})", v.ocaml()),
-            UnnormalizedLagrangeBasis(i) => format!("unnormalized_lagrange_basis({})", *i),
+            UnnormalizedLagrangeBasis(i) => {
+                format!("unnormalized_lagrange_basis({}, {})", i.zk_rows, i.offset)
+            }
             VanishesOnZeroKnowledgeAndPreviousRows => {
                 "vanishes_on_zero_knowledge_and_previous_rows".to_string()
             }
@@ -2532,7 +2566,18 @@ where
             Double(x) => format!("2 ({})", x.latex(cache)),
             Constant(x) => x.latex(),
             Cell(v) => v.latex(),
-            UnnormalizedLagrangeBasis(i) => format!("unnormalized\\_lagrange\\_basis({})", *i),
+            UnnormalizedLagrangeBasis(RowOffset {
+                zk_rows: true,
+                offset: i,
+            }) => {
+                format!("unnormalized\\_lagrange\\_basis(zk\\_rows + {})", *i)
+            }
+            UnnormalizedLagrangeBasis(RowOffset {
+                zk_rows: false,
+                offset: i,
+            }) => {
+                format!("unnormalized\\_lagrange\\_basis({})", *i)
+            }
             VanishesOnZeroKnowledgeAndPreviousRows => {
                 "vanishes\\_on\\_zero\\_knowledge\\_and\\_previous\\_row".to_string()
             }
@@ -2557,7 +2602,20 @@ where
             Double(x) => format!("double({})", x.text(cache)),
             Constant(x) => x.text(),
             Cell(v) => v.text(),
-            UnnormalizedLagrangeBasis(i) => format!("unnormalized_lagrange_basis({})", *i),
+            UnnormalizedLagrangeBasis(RowOffset {
+                zk_rows: true,
+                offset: i,
+            }) => match i.cmp(&0) {
+                Ordering::Greater => format!("unnormalized_lagrange_basis(zk_rows + {})", *i),
+                Ordering::Equal => "unnormalized_lagrange_basis(zk_rows)".to_string(),
+                Ordering::Less => format!("unnormalized_lagrange_basis(zk_rows - {})", (-*i)),
+            },
+            UnnormalizedLagrangeBasis(RowOffset {
+                zk_rows: false,
+                offset: i,
+            }) => {
+                format!("unnormalized_lagrange_basis({})", *i)
+            }
             VanishesOnZeroKnowledgeAndPreviousRows => {
                 "vanishes_on_zero_knowledge_and_previous_rows".to_string()
             }

kimchi/src/circuits/lookup/constraints.rs

@@ -1,6 +1,6 @@
 use crate::{
     circuits::{
-        expr::{prologue::*, Column, ConstantExpr},
+        expr::{prologue::*, Column, ConstantExpr, RowOffset},
         gate::{CircuitGate, CurrOrNext},
         lookup::lookups::{
             JointLookup, JointLookupSpec, JointLookupValue, LocalPosition, LookupInfo,
@@ -371,7 +371,6 @@ impl<F: Zero> LookupConfiguration<F> {
 pub fn constraints<F: FftField>(
     configuration: &LookupConfiguration<F>,
     generate_feature_flags: bool,
-    zk_rows: usize,
 ) -> Vec<E<F>> {
     // Something important to keep in mind is that the last 2 rows of
     // all columns will have random values in them to maintain zero-knowledge.
@@ -601,14 +600,20 @@ pub fn constraints<F: FftField>(
     let aggreg_equation = E::cell(Column::LookupAggreg, Next) * denominator
         - E::cell(Column::LookupAggreg, Curr) * numerator;
 
-    let final_lookup_row: i32 = -(zk_rows as i32) - 1;
+    let final_lookup_row = RowOffset {
+        zk_rows: true,
+        offset: -1,
+    };
 
     let mut res = vec![
         // the accumulator except for the last zk_rows+1 rows
         // (contains the zk-rows and the last value of the accumulator)
         E::VanishesOnZeroKnowledgeAndPreviousRows * aggreg_equation,
         // the initial value of the accumulator
-        E::UnnormalizedLagrangeBasis(0) * (E::cell(Column::LookupAggreg, Curr) - E::one()),
+        E::UnnormalizedLagrangeBasis(RowOffset {
+            zk_rows: false,
+            offset: 0,
+        }) * (E::cell(Column::LookupAggreg, Curr) - E::one()),
         // Check that the final value of the accumulator is 1
         E::UnnormalizedLagrangeBasis(final_lookup_row)
             * (E::cell(Column::LookupAggreg, Curr) - E::one()),
@@ -622,7 +627,10 @@ pub fn constraints<F: FftField>(
                 final_lookup_row
             } else {
                 // Check compatibility of the first elements
-                0
+                RowOffset {
+                    zk_rows: false,
+                    offset: 0,
+                }
             };
             let mut expr = E::UnnormalizedLagrangeBasis(first_or_last)
                 * (column(Column::LookupSorted(i)) - column(Column::LookupSorted(i + 1)));

kimchi/src/linearization.rs

@@ -38,7 +38,6 @@ use ark_ff::{FftField, PrimeField, SquareRootField, Zero};
 pub fn constraints_expr<F: PrimeField + SquareRootField>(
     feature_flags: Option<&FeatureFlags>,
     generic: bool,
-    zk_rows: usize,
 ) -> (Expr<ConstantExpr<F>>, Alphas<F>) {
     // register powers of alpha so that we don't reuse them across mutually inclusive constraints
     let mut powers_of_alpha = Alphas::<F>::default();
@@ -166,8 +165,7 @@ pub fn constraints_expr<F: PrimeField + SquareRootField>(
         if feature_flags.lookup_features.patterns != LookupPatterns::default() {
             let lookup_configuration =
                 LookupConfiguration::new(LookupInfo::create(feature_flags.lookup_features));
-            let constraints =
-                lookup::constraints::constraints(&lookup_configuration, false, zk_rows);
+            let constraints = lookup::constraints::constraints(&lookup_configuration, false);
 
             // note: the number of constraints depends on the lookup configuration,
             // specifically the presence of runtime tables.
@@ -193,7 +191,7 @@ pub fn constraints_expr<F: PrimeField + SquareRootField>(
             joint_lookup_used: true,
         };
         let lookup_configuration = LookupConfiguration::new(LookupInfo::create(all_features));
-        let constraints = lookup::constraints::constraints(&lookup_configuration, true, zk_rows);
+        let constraints = lookup::constraints::constraints(&lookup_configuration, true);
 
         // note: the number of constraints depends on the lookup configuration,
         // specifically the presence of runtime tables.
@@ -224,7 +222,7 @@ pub fn constraints_expr<F: PrimeField + SquareRootField>(
     // flags.
     if cfg!(feature = "check_feature_flags") {
         if let Some(feature_flags) = feature_flags {
-            let (feature_flagged_expr, _) = constraints_expr(None, generic, zk_rows);
+            let (feature_flagged_expr, _) = constraints_expr(None, generic);
             let feature_flagged_expr = feature_flagged_expr.apply_feature_flags(feature_flags);
             assert_eq!(expr, feature_flagged_expr);
         }
@@ -341,11 +339,10 @@ pub fn linearization_columns<F: FftField + SquareRootField>(
 pub fn expr_linearization<F: PrimeField + SquareRootField>(
     feature_flags: Option<&FeatureFlags>,
     generic: bool,
-    zk_rows: usize,
 ) -> (Linearization<Vec<PolishToken<F>>>, Alphas<F>) {
     let evaluated_cols = linearization_columns::<F>(feature_flags);
 
-    let (expr, powers_of_alpha) = constraints_expr(feature_flags, generic, zk_rows);
+    let (expr, powers_of_alpha) = constraints_expr(feature_flags, generic);
 
     let linearization = expr
         .linearize(evaluated_cols)

kimchi/src/prover.rs

@@ -818,11 +818,7 @@ where
             // lookup
             {
                 if let Some(lcs) = index.cs.lookup_constraint_system.as_ref() {
-                    let constraints = lookup::constraints::constraints(
-                        &lcs.configuration,
-                        false,
-                        index.cs.zk_rows as usize,
-                    );
+                    let constraints = lookup::constraints::constraints(&lcs.configuration, false);
                     let constraints_len = u32::try_from(constraints.len())
                         .expect("not expecting a large amount of constraints");
                     let lookup_alphas =

kimchi/src/prover_index.rs

@@ -69,8 +69,7 @@ where
         cs.endo = endo_q;
 
         // pre-compute the linearization
-        let (linearization, powers_of_alpha) =
-            expr_linearization(Some(&cs.feature_flags), true, cs.zk_rows as usize);
+        let (linearization, powers_of_alpha) = expr_linearization(Some(&cs.feature_flags), true);
 
         let evaluated_column_coefficients = cs.evaluated_column_coefficients();