sp1_stark/
folder.rs

1use std::{
2    marker::PhantomData,
3    ops::{Add, Mul, MulAssign, Sub},
4};
5
6use p3_field::{AbstractField, ExtensionField, Field};
7use p3_matrix::{dense::RowMajorMatrixView, stack::VerticalPair};
8
9use super::{Challenge, PackedChallenge, PackedVal, StarkGenericConfig, Val};
10use crate::{
11    air::{EmptyMessageBuilder, MultiTableAirBuilder},
12    septic_digest::SepticDigest,
13};
14use p3_air::{
15    AirBuilder, AirBuilderWithPublicValues, ExtensionBuilder, PairBuilder, PermutationAirBuilder,
16};
17
18/// A folder for prover constraints.
19pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> {
20    /// The preprocessed trace.
21    pub preprocessed:
22        VerticalPair<RowMajorMatrixView<'a, PackedVal<SC>>, RowMajorMatrixView<'a, PackedVal<SC>>>,
23    /// The main trace.
24    pub main:
25        VerticalPair<RowMajorMatrixView<'a, PackedVal<SC>>, RowMajorMatrixView<'a, PackedVal<SC>>>,
26    /// The permutation trace.
27    pub perm: VerticalPair<
28        RowMajorMatrixView<'a, PackedChallenge<SC>>,
29        RowMajorMatrixView<'a, PackedChallenge<SC>>,
30    >,
31    /// The challenges for the permutation.
32    pub perm_challenges: &'a [PackedChallenge<SC>],
33    /// The local cumulative sum for the permutation.
34    pub local_cumulative_sum: &'a PackedChallenge<SC>,
35    /// The global cumulative sum for the permutation.
36    pub global_cumulative_sum: &'a SepticDigest<Val<SC>>,
37    /// The selector for the first row.
38    pub is_first_row: PackedVal<SC>,
39    /// The selector for the last row.
40    pub is_last_row: PackedVal<SC>,
41    /// The selector for the transition.
42    pub is_transition: PackedVal<SC>,
43    /// The powers of the constraint folding challenge.
44    pub powers_of_alpha: &'a Vec<SC::Challenge>,
45    /// The accumulator for the constraint folding.
46    pub accumulator: PackedChallenge<SC>,
47    /// The public values.
48    pub public_values: &'a [Val<SC>],
49    /// The constraint index.
50    pub constraint_index: usize,
51}
52
53impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> {
54    type F = Val<SC>;
55    type Expr = PackedVal<SC>;
56    type Var = PackedVal<SC>;
57    type M =
58        VerticalPair<RowMajorMatrixView<'a, PackedVal<SC>>, RowMajorMatrixView<'a, PackedVal<SC>>>;
59
60    fn main(&self) -> Self::M {
61        self.main
62    }
63
64    fn is_first_row(&self) -> Self::Expr {
65        self.is_first_row
66    }
67
68    fn is_last_row(&self) -> Self::Expr {
69        self.is_last_row
70    }
71
72    fn is_transition_window(&self, size: usize) -> Self::Expr {
73        if size == 2 {
74            self.is_transition
75        } else {
76            panic!("uni-stark only supports a window size of 2")
77        }
78    }
79
80    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
81        let x: PackedVal<SC> = x.into();
82        self.accumulator +=
83            PackedChallenge::<SC>::from_f(self.powers_of_alpha[self.constraint_index]) * x;
84        self.constraint_index += 1;
85    }
86}
87
88impl<SC: StarkGenericConfig> ExtensionBuilder for ProverConstraintFolder<'_, SC> {
89    type EF = SC::Challenge;
90
91    type ExprEF = PackedChallenge<SC>;
92
93    type VarEF = PackedChallenge<SC>;
94
95    fn assert_zero_ext<I>(&mut self, x: I)
96    where
97        I: Into<Self::ExprEF>,
98    {
99        let x: PackedChallenge<SC> = x.into();
100        self.accumulator +=
101            PackedChallenge::<SC>::from_f(self.powers_of_alpha[self.constraint_index]) * x;
102        self.constraint_index += 1;
103    }
104}
105
106impl<'a, SC: StarkGenericConfig> PermutationAirBuilder for ProverConstraintFolder<'a, SC> {
107    type MP = VerticalPair<
108        RowMajorMatrixView<'a, PackedChallenge<SC>>,
109        RowMajorMatrixView<'a, PackedChallenge<SC>>,
110    >;
111
112    type RandomVar = PackedChallenge<SC>;
113
114    fn permutation(&self) -> Self::MP {
115        self.perm
116    }
117
118    fn permutation_randomness(&self) -> &[Self::RandomVar] {
119        self.perm_challenges
120    }
121}
122
123impl<'a, SC: StarkGenericConfig> MultiTableAirBuilder<'a> for ProverConstraintFolder<'a, SC> {
124    type LocalSum = PackedChallenge<SC>;
125    type GlobalSum = Val<SC>;
126
127    fn local_cumulative_sum(&self) -> &'a Self::LocalSum {
128        self.local_cumulative_sum
129    }
130
131    fn global_cumulative_sum(&self) -> &'a SepticDigest<Self::GlobalSum> {
132        self.global_cumulative_sum
133    }
134}
135
136impl<SC: StarkGenericConfig> PairBuilder for ProverConstraintFolder<'_, SC> {
137    fn preprocessed(&self) -> Self::M {
138        self.preprocessed
139    }
140}
141
142impl<SC: StarkGenericConfig> EmptyMessageBuilder for ProverConstraintFolder<'_, SC> {}
143
144impl<SC: StarkGenericConfig> AirBuilderWithPublicValues for ProverConstraintFolder<'_, SC> {
145    type PublicVar = Self::F;
146
147    fn public_values(&self) -> &[Self::PublicVar] {
148        self.public_values
149    }
150}
151
152/// A folder for verifier constraints.
153pub type VerifierConstraintFolder<'a, SC> = GenericVerifierConstraintFolder<
154    'a,
155    Val<SC>,
156    Challenge<SC>,
157    Val<SC>,
158    Challenge<SC>,
159    Challenge<SC>,
160>;
161
162/// A folder for verifier constraints.
163pub struct GenericVerifierConstraintFolder<'a, F, EF, PubVar, Var, Expr> {
164    /// The preprocessed trace.
165    pub preprocessed: VerticalPair<RowMajorMatrixView<'a, Var>, RowMajorMatrixView<'a, Var>>,
166    /// The main trace.
167    pub main: VerticalPair<RowMajorMatrixView<'a, Var>, RowMajorMatrixView<'a, Var>>,
168    /// The permutation trace.
169    pub perm: VerticalPair<RowMajorMatrixView<'a, Var>, RowMajorMatrixView<'a, Var>>,
170    /// The challenges for the permutation.
171    pub perm_challenges: &'a [Var],
172    /// The local cumulative sum of the permutation.
173    pub local_cumulative_sum: &'a Var,
174    /// The global cumulative sum of the permutation.
175    pub global_cumulative_sum: &'a SepticDigest<PubVar>,
176    /// The selector for the first row.
177    pub is_first_row: Var,
178    /// The selector for the last row.
179    pub is_last_row: Var,
180    /// The selector for the transition.
181    pub is_transition: Var,
182    /// The constraint folding challenge.
183    pub alpha: Var,
184    /// The accumulator for the constraint folding.
185    pub accumulator: Expr,
186    /// The public values.
187    pub public_values: &'a [PubVar],
188    /// The marker type.
189    pub _marker: PhantomData<(F, EF)>,
190}
191
192impl<'a, F, EF, PubVar, Var, Expr> AirBuilder
193    for GenericVerifierConstraintFolder<'a, F, EF, PubVar, Var, Expr>
194where
195    F: Field,
196    EF: ExtensionField<F>,
197    Expr: AbstractField
198        + From<F>
199        + Add<Var, Output = Expr>
200        + Add<F, Output = Expr>
201        + Sub<Var, Output = Expr>
202        + Sub<F, Output = Expr>
203        + Mul<Var, Output = Expr>
204        + Mul<F, Output = Expr>
205        + MulAssign<EF>,
206    Var: Into<Expr>
207        + Copy
208        + Add<F, Output = Expr>
209        + Add<Var, Output = Expr>
210        + Add<Expr, Output = Expr>
211        + Sub<F, Output = Expr>
212        + Sub<Var, Output = Expr>
213        + Sub<Expr, Output = Expr>
214        + Mul<F, Output = Expr>
215        + Mul<Var, Output = Expr>
216        + Mul<Expr, Output = Expr>
217        + Send
218        + Sync,
219    PubVar: Into<Expr> + Copy,
220{
221    type F = F;
222    type Expr = Expr;
223    type Var = Var;
224    type M = VerticalPair<RowMajorMatrixView<'a, Var>, RowMajorMatrixView<'a, Var>>;
225
226    fn main(&self) -> Self::M {
227        self.main
228    }
229
230    fn is_first_row(&self) -> Self::Expr {
231        self.is_first_row.into()
232    }
233
234    fn is_last_row(&self) -> Self::Expr {
235        self.is_last_row.into()
236    }
237
238    fn is_transition_window(&self, size: usize) -> Self::Expr {
239        if size == 2 {
240            self.is_transition.into()
241        } else {
242            panic!("uni-stark only supports a window size of 2")
243        }
244    }
245
246    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
247        let x: Expr = x.into();
248        self.accumulator *= self.alpha.into();
249        self.accumulator += x;
250    }
251}
252
253impl<F, EF, PubVar, Var, Expr> ExtensionBuilder
254    for GenericVerifierConstraintFolder<'_, F, EF, PubVar, Var, Expr>
255where
256    F: Field,
257    EF: ExtensionField<F>,
258    Expr: AbstractField<F = EF>
259        + From<F>
260        + Add<Var, Output = Expr>
261        + Add<F, Output = Expr>
262        + Sub<Var, Output = Expr>
263        + Sub<F, Output = Expr>
264        + Mul<Var, Output = Expr>
265        + Mul<F, Output = Expr>
266        + MulAssign<EF>,
267    Var: Into<Expr>
268        + Copy
269        + Add<F, Output = Expr>
270        + Add<Var, Output = Expr>
271        + Add<Expr, Output = Expr>
272        + Sub<F, Output = Expr>
273        + Sub<Var, Output = Expr>
274        + Sub<Expr, Output = Expr>
275        + Mul<F, Output = Expr>
276        + Mul<Var, Output = Expr>
277        + Mul<Expr, Output = Expr>
278        + Send
279        + Sync,
280    PubVar: Into<Expr> + Copy,
281{
282    type EF = EF;
283    type ExprEF = Expr;
284    type VarEF = Var;
285
286    fn assert_zero_ext<I>(&mut self, x: I)
287    where
288        I: Into<Self::ExprEF>,
289    {
290        self.assert_zero(x);
291    }
292}
293
294impl<'a, F, EF, PubVar, Var, Expr> PermutationAirBuilder
295    for GenericVerifierConstraintFolder<'a, F, EF, PubVar, Var, Expr>
296where
297    F: Field,
298    EF: ExtensionField<F>,
299    Expr: AbstractField<F = EF>
300        + From<F>
301        + Add<Var, Output = Expr>
302        + Add<F, Output = Expr>
303        + Sub<Var, Output = Expr>
304        + Sub<F, Output = Expr>
305        + Mul<Var, Output = Expr>
306        + Mul<F, Output = Expr>
307        + MulAssign<EF>,
308    Var: Into<Expr>
309        + Copy
310        + Add<F, Output = Expr>
311        + Add<Var, Output = Expr>
312        + Add<Expr, Output = Expr>
313        + Sub<F, Output = Expr>
314        + Sub<Var, Output = Expr>
315        + Sub<Expr, Output = Expr>
316        + Mul<F, Output = Expr>
317        + Mul<Var, Output = Expr>
318        + Mul<Expr, Output = Expr>
319        + Send
320        + Sync,
321    PubVar: Into<Expr> + Copy,
322{
323    type MP = VerticalPair<RowMajorMatrixView<'a, Var>, RowMajorMatrixView<'a, Var>>;
324    type RandomVar = Var;
325
326    fn permutation(&self) -> Self::MP {
327        self.perm
328    }
329
330    fn permutation_randomness(&self) -> &[Self::Var] {
331        self.perm_challenges
332    }
333}
334
335impl<'a, F, EF, PubVar, Var, Expr> MultiTableAirBuilder<'a>
336    for GenericVerifierConstraintFolder<'a, F, EF, PubVar, Var, Expr>
337where
338    F: Field,
339    EF: ExtensionField<F>,
340    Expr: AbstractField<F = EF>
341        + From<F>
342        + Add<Var, Output = Expr>
343        + Add<F, Output = Expr>
344        + Sub<Var, Output = Expr>
345        + Sub<F, Output = Expr>
346        + Mul<Var, Output = Expr>
347        + Mul<F, Output = Expr>
348        + MulAssign<EF>,
349    Var: Into<Expr>
350        + Copy
351        + Add<F, Output = Expr>
352        + Add<Var, Output = Expr>
353        + Add<Expr, Output = Expr>
354        + Sub<F, Output = Expr>
355        + Sub<Var, Output = Expr>
356        + Sub<Expr, Output = Expr>
357        + Mul<F, Output = Expr>
358        + Mul<Var, Output = Expr>
359        + Mul<Expr, Output = Expr>
360        + Send
361        + Sync,
362    PubVar: Into<Expr> + Copy,
363{
364    type LocalSum = Var;
365    type GlobalSum = PubVar;
366
367    fn local_cumulative_sum(&self) -> &'a Self::LocalSum {
368        self.local_cumulative_sum
369    }
370
371    fn global_cumulative_sum(&self) -> &'a SepticDigest<Self::GlobalSum> {
372        self.global_cumulative_sum
373    }
374}
375
376impl<F, EF, PubVar, Var, Expr> PairBuilder
377    for GenericVerifierConstraintFolder<'_, F, EF, PubVar, Var, Expr>
378where
379    F: Field,
380    EF: ExtensionField<F>,
381    Expr: AbstractField<F = EF>
382        + From<F>
383        + Add<Var, Output = Expr>
384        + Add<F, Output = Expr>
385        + Sub<Var, Output = Expr>
386        + Sub<F, Output = Expr>
387        + Mul<Var, Output = Expr>
388        + Mul<F, Output = Expr>
389        + MulAssign<EF>,
390    Var: Into<Expr>
391        + Copy
392        + Add<F, Output = Expr>
393        + Add<Var, Output = Expr>
394        + Add<Expr, Output = Expr>
395        + Sub<F, Output = Expr>
396        + Sub<Var, Output = Expr>
397        + Sub<Expr, Output = Expr>
398        + Mul<F, Output = Expr>
399        + Mul<Var, Output = Expr>
400        + Mul<Expr, Output = Expr>
401        + Send
402        + Sync,
403    PubVar: Into<Expr> + Copy,
404{
405    fn preprocessed(&self) -> Self::M {
406        self.preprocessed
407    }
408}
409
410impl<F, EF, PubVar, Var, Expr> EmptyMessageBuilder
411    for GenericVerifierConstraintFolder<'_, F, EF, PubVar, Var, Expr>
412where
413    F: Field,
414    EF: ExtensionField<F>,
415    Expr: AbstractField<F = EF>
416        + From<F>
417        + Add<Var, Output = Expr>
418        + Add<F, Output = Expr>
419        + Sub<Var, Output = Expr>
420        + Sub<F, Output = Expr>
421        + Mul<Var, Output = Expr>
422        + Mul<F, Output = Expr>
423        + MulAssign<EF>,
424    Var: Into<Expr>
425        + Copy
426        + Add<F, Output = Expr>
427        + Add<Var, Output = Expr>
428        + Add<Expr, Output = Expr>
429        + Sub<F, Output = Expr>
430        + Sub<Var, Output = Expr>
431        + Sub<Expr, Output = Expr>
432        + Mul<F, Output = Expr>
433        + Mul<Var, Output = Expr>
434        + Mul<Expr, Output = Expr>
435        + Send
436        + Sync,
437    PubVar: Into<Expr> + Copy,
438{
439}
440
441impl<F, EF, PubVar, Var, Expr> AirBuilderWithPublicValues
442    for GenericVerifierConstraintFolder<'_, F, EF, PubVar, Var, Expr>
443where
444    F: Field,
445    EF: ExtensionField<F>,
446    Expr: AbstractField<F = EF>
447        + From<F>
448        + Add<Var, Output = Expr>
449        + Add<F, Output = Expr>
450        + Sub<Var, Output = Expr>
451        + Sub<F, Output = Expr>
452        + Mul<Var, Output = Expr>
453        + Mul<F, Output = Expr>
454        + MulAssign<EF>,
455    Var: Into<Expr>
456        + Copy
457        + Add<F, Output = Expr>
458        + Add<Var, Output = Expr>
459        + Add<Expr, Output = Expr>
460        + Sub<F, Output = Expr>
461        + Sub<Var, Output = Expr>
462        + Sub<Expr, Output = Expr>
463        + Mul<F, Output = Expr>
464        + Mul<Var, Output = Expr>
465        + Mul<Expr, Output = Expr>
466        + Send
467        + Sync,
468    PubVar: Into<Expr> + Copy,
469{
470    type PublicVar = PubVar;
471
472    fn public_values(&self) -> &[Self::PublicVar] {
473        self.public_values
474    }
475}