o1vm/pickles/
lookup_columns.rs

1use ark_ff::{FftField, Field, PrimeField};
2use ark_poly::{Evaluations, Radix2EvaluationDomain as D};
3use core::{iter::Once, ops::Index};
4use kimchi::{
5    circuits::{
6        domains::{Domain, EvaluationDomains},
7        expr::{
8            AlphaChallengeTerm, ColumnEnvironment, ColumnEvaluations, ConstantExpr, Constants,
9            Expr, ExprError,
10        },
11        gate::CurrOrNext,
12    },
13    curve::KimchiCurve,
14    proof::PointEvaluations,
15};
16use poly_commitment::ipa::OpeningProof;
17use serde::{Deserialize, Serialize};
18use std::iter::Chain;
19
20// This file contains the associated types and methods for the lookup prover.
21// It defines the columns, the proof, proof input, and constraint expressions.
22
23#[derive(Clone, PartialEq, Copy)]
24pub enum LookupColumns {
25    Wires(usize),
26    Inverses(usize),
27    Acc,
28}
29
30#[derive(Clone)]
31pub struct ColumnEnv<X> {
32    pub wires: Vec<X>,
33    pub inverses: Vec<X>,
34    pub acc: X,
35}
36
37impl<X> IntoIterator for ColumnEnv<X> {
38    type Item = X;
39    type IntoIter = Chain<
40        Chain<<Vec<X> as IntoIterator>::IntoIter, <Vec<X> as IntoIterator>::IntoIter>,
41        <Once<X> as IntoIterator>::IntoIter,
42    >;
43    fn into_iter(self) -> Self::IntoIter {
44        self.wires
45            .into_iter()
46            .chain(self.inverses)
47            .chain(std::iter::once(self.acc))
48    }
49}
50
51// TODO: I could not find a more elegant solution to map over this struct
52impl<X> ColumnEnv<X> {
53    pub fn my_map<Y, F>(self, f: F) -> ColumnEnv<Y>
54    where
55        F: FnMut(X) -> Y,
56        Self: Sized,
57    {
58        let nb_wires = self.wires.len();
59        let nb_inverses = self.inverses.len();
60        let mut iterator = self.into_iter().map(f);
61        let mut new_wires = vec![];
62        let mut new_inverses = vec![];
63        for _ in 0..nb_wires {
64            new_wires.push(iterator.next().unwrap());
65        }
66        for _ in 0..nb_inverses {
67            new_inverses.push(iterator.next().unwrap());
68        }
69        let new_acc = iterator.next().unwrap();
70        assert!(iterator.next().is_none());
71        ColumnEnv {
72            wires: new_wires,
73            inverses: new_inverses,
74            acc: new_acc,
75        }
76    }
77}
78
79pub struct LookupProofInput<F: PrimeField> {
80    pub wires: Vec<Vec<F>>,
81    pub arity: Vec<Vec<usize>>,
82    pub beta_challenge: F,
83    pub gamma_challenge: F,
84}
85#[derive(Clone)]
86pub struct AllColumns<X> {
87    pub cols: ColumnEnv<X>,
88    pub t_shares: Vec<X>,
89}
90
91impl<X> IntoIterator for AllColumns<X> {
92    type Item = X;
93    type IntoIter =
94        Chain<<ColumnEnv<X> as IntoIterator>::IntoIter, <Vec<X> as IntoIterator>::IntoIter>;
95    fn into_iter(self) -> Self::IntoIter {
96        self.cols.into_iter().chain(self.t_shares)
97    }
98}
99
100#[derive(Clone)]
101pub struct Eval<F: PrimeField> {
102    pub zeta: AllColumns<F>,
103    pub zeta_omega: AllColumns<F>,
104}
105
106impl<F: PrimeField> IntoIterator for Eval<F> {
107    type Item = F;
108    type IntoIter =
109        Chain<<AllColumns<F> as IntoIterator>::IntoIter, <AllColumns<F> as IntoIterator>::IntoIter>;
110    fn into_iter(self) -> Self::IntoIter {
111        self.zeta.into_iter().chain(self.zeta_omega)
112    }
113}
114
115pub struct Proof<G: KimchiCurve> {
116    pub commitments: AllColumns<G>,
117    pub evaluations: Eval<G::ScalarField>,
118    pub ipa_proof: OpeningProof<G>,
119}
120
121#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
122pub enum LookupChallengeTerm {
123    /// The challenge to compute 1/(beta + lookupvalue)
124    Beta,
125    /// The challenge to combine tuple sum gamma^i lookupvalue_i
126    Gamma,
127    /// The challenge to combine constraints
128    Alpha,
129}
130
131pub struct LookupChallenges<F> {
132    pub beta: F,
133    pub gamma: F,
134    pub alpha: F,
135}
136
137impl<F: Field> Index<LookupChallengeTerm> for LookupChallenges<F> {
138    type Output = F;
139
140    fn index(&self, term: LookupChallengeTerm) -> &Self::Output {
141        match term {
142            LookupChallengeTerm::Alpha => &self.alpha,
143            LookupChallengeTerm::Beta => &self.beta,
144            LookupChallengeTerm::Gamma => &self.gamma,
145        }
146    }
147}
148
149impl std::fmt::Display for LookupChallengeTerm {
150    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151        use LookupChallengeTerm::*;
152        let str = match self {
153            Alpha => "alpha".to_string(),
154            Beta => "beta".to_string(),
155            Gamma => "gamma".to_string(),
156        };
157        write!(f, "{}", str)
158    }
159}
160
161impl<'a> AlphaChallengeTerm<'a> for LookupChallengeTerm {
162    const ALPHA: Self = Self::Alpha;
163}
164
165/// The collection of polynomials (all in evaluation form) and constants
166/// required to evaluate an expression as a polynomial.
167/// All are evaluations.
168pub struct LookupEvalEnvironment<'a, F: FftField> {
169    pub columns: &'a ColumnEnv<Evaluations<F, D<F>>>,
170    pub challenges: LookupChallenges<F>,
171    pub domain: &'a EvaluationDomains<F>,
172    pub l0_1: F,
173}
174
175// Necessary trait to evaluate the numerator of T in the prover
176impl<'a, F: FftField> ColumnEnvironment<'a, F, LookupChallengeTerm, LookupChallenges<F>>
177    for LookupEvalEnvironment<'a, F>
178{
179    type Column = LookupColumns;
180
181    fn get_column(&self, col: &Self::Column) -> Option<&'a Evaluations<F, D<F>>> {
182        use LookupColumns::*;
183        match col {
184            Wires(i) => Some(&self.columns.wires[*i]),
185            Inverses(i) => Some(&self.columns.inverses[*i]),
186            Acc => Some(&self.columns.acc),
187        }
188    }
189
190    fn get_domain(&self, d: Domain) -> D<F> {
191        match d {
192            Domain::D1 => self.domain.d1,
193            Domain::D2 => self.domain.d2,
194            Domain::D4 => self.domain.d4,
195            Domain::D8 => self.domain.d8,
196        }
197    }
198    // TODO verify this
199    fn column_domain(&self, _col: &Self::Column) -> Domain {
200        Domain::D8
201    }
202    // We do not have constants here
203    fn get_constants(&self) -> &Constants<F> {
204        panic!("no constants are supposed to be used in this protocol")
205    }
206
207    fn get_challenges(&self) -> &LookupChallenges<F> {
208        &self.challenges
209    }
210
211    fn vanishes_on_zero_knowledge_and_previous_rows(&self) -> &'a Evaluations<F, D<F>> {
212        panic!("no zk is supposed to be used in this protocol")
213    }
214
215    fn l0_1(&self) -> F {
216        self.l0_1
217    }
218}
219
220// helper to implement the next trait
221impl<F> ColumnEnv<F> {
222    pub fn get_column(&self, col: &LookupColumns) -> Option<&F> {
223        match *col {
224            LookupColumns::Wires(i) => self.wires.get(i),
225            LookupColumns::Inverses(i) => self.inverses.get(i),
226            LookupColumns::Acc => Some(&self.acc),
227        }
228    }
229}
230// Necessary trait to evaluate the numerator of T at zeta in the verifier
231impl<F: PrimeField> ColumnEvaluations<F> for Eval<F> {
232    type Column = LookupColumns;
233    fn evaluate(&self, col: Self::Column) -> Result<PointEvaluations<F>, ExprError<Self::Column>> {
234        if let Some(&zeta) = self.zeta.cols.get_column(&col) {
235            if let Some(&zeta_omega) = self.zeta_omega.cols.get_column(&col) {
236                Ok(PointEvaluations { zeta, zeta_omega })
237            } else {
238                Err(ExprError::MissingEvaluation(col, CurrOrNext::Next))
239            }
240        } else {
241            Err(ExprError::MissingEvaluation(col, CurrOrNext::Curr))
242        }
243    }
244}
245
246pub type ELookup<F> = Expr<ConstantExpr<F, LookupChallengeTerm>, LookupColumns>;