Skip to main content

p3_lookup/
folder.rs

1use p3_air::{AirBuilder, ExtensionBuilder, PermutationAirBuilder, RowWindow};
2use p3_matrix::dense::RowMajorMatrixView;
3use p3_matrix::stack::ViewPair;
4use p3_uni_stark::{
5    PackedChallenge, PackedVal, ProverConstraintFolder, StarkGenericConfig, Val,
6    VerifierConstraintFolder,
7};
8
9pub struct ProverConstraintFolderWithLookups<'a, SC: StarkGenericConfig> {
10    pub inner: ProverConstraintFolder<'a, SC>,
11    pub permutation: RowMajorMatrixView<'a, PackedChallenge<SC>>,
12    pub permutation_challenges: &'a [PackedChallenge<SC>],
13    pub permutation_values: &'a [PackedChallenge<SC>],
14}
15
16impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolderWithLookups<'a, SC> {
17    type F = Val<SC>;
18    type Expr = PackedVal<SC>;
19    type Var = PackedVal<SC>;
20    type PreprocessedWindow = RowWindow<'a, PackedVal<SC>>;
21    type MainWindow = RowWindow<'a, PackedVal<SC>>;
22    type PublicVar = Val<SC>;
23
24    fn main(&self) -> Self::MainWindow {
25        RowWindow::from_view(&self.inner.main)
26    }
27
28    fn preprocessed(&self) -> &Self::PreprocessedWindow {
29        self.inner.preprocessed()
30    }
31
32    #[inline]
33    fn is_first_row(&self) -> Self::Expr {
34        self.inner.is_first_row
35    }
36
37    #[inline]
38    fn is_last_row(&self) -> Self::Expr {
39        self.inner.is_last_row
40    }
41
42    #[inline]
43    fn is_transition_window(&self, size: usize) -> Self::Expr {
44        assert!(size <= 2, "only two-row windows are supported, got {size}");
45        self.inner.is_transition
46    }
47
48    #[inline]
49    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
50        self.inner.assert_zero(x);
51    }
52
53    #[inline]
54    fn assert_zeros<const N: usize, I: Into<Self::Expr>>(&mut self, array: [I; N]) {
55        self.inner.assert_zeros(array);
56    }
57
58    #[inline]
59    fn public_values(&self) -> &[Self::PublicVar] {
60        self.inner.public_values
61    }
62}
63
64impl<SC: StarkGenericConfig> ExtensionBuilder for ProverConstraintFolderWithLookups<'_, SC> {
65    type EF = SC::Challenge;
66    type ExprEF = PackedChallenge<SC>;
67    type VarEF = PackedChallenge<SC>;
68
69    fn assert_zero_ext<I>(&mut self, x: I)
70    where
71        I: Into<Self::ExprEF>,
72    {
73        self.inner.assert_zero_ext(x);
74    }
75}
76
77impl<'a, SC: StarkGenericConfig> PermutationAirBuilder
78    for ProverConstraintFolderWithLookups<'a, SC>
79{
80    type RandomVar = PackedChallenge<SC>;
81    type MP = RowWindow<'a, PackedChallenge<SC>>;
82
83    type PermutationVar = PackedChallenge<SC>;
84
85    fn permutation(&self) -> Self::MP {
86        RowWindow::from_view(&self.permutation)
87    }
88
89    fn permutation_randomness(&self) -> &[PackedChallenge<SC>] {
90        self.permutation_challenges
91    }
92
93    fn permutation_values(&self) -> &[PackedChallenge<SC>] {
94        self.permutation_values
95    }
96}
97
98pub struct VerifierConstraintFolderWithLookups<'a, SC: StarkGenericConfig> {
99    pub inner: VerifierConstraintFolder<'a, SC>,
100    pub permutation: ViewPair<'a, SC::Challenge>,
101    pub permutation_challenges: &'a [SC::Challenge],
102    pub permutation_values: &'a [SC::Challenge],
103}
104
105impl<'a, SC: StarkGenericConfig> AirBuilder for VerifierConstraintFolderWithLookups<'a, SC> {
106    type F = Val<SC>;
107    type Expr = SC::Challenge;
108    type Var = SC::Challenge;
109    type PublicVar = Val<SC>;
110    type PreprocessedWindow = RowWindow<'a, SC::Challenge>;
111    type MainWindow = RowWindow<'a, SC::Challenge>;
112
113    fn main(&self) -> Self::MainWindow {
114        RowWindow::from_two_rows(self.inner.main.top.values, self.inner.main.bottom.values)
115    }
116
117    fn preprocessed(&self) -> &Self::PreprocessedWindow {
118        self.inner.preprocessed()
119    }
120
121    #[inline]
122    fn public_values(&self) -> &[Self::PublicVar] {
123        self.inner.public_values
124    }
125
126    #[inline]
127    fn is_first_row(&self) -> Self::Expr {
128        self.inner.is_first_row
129    }
130
131    #[inline]
132    fn is_last_row(&self) -> Self::Expr {
133        self.inner.is_last_row
134    }
135
136    #[inline]
137    fn is_transition_window(&self, size: usize) -> Self::Expr {
138        assert!(size <= 2, "only two-row windows are supported, got {size}");
139        self.inner.is_transition
140    }
141
142    #[inline]
143    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
144        self.inner.assert_zero(x);
145    }
146
147    #[inline]
148    fn assert_zeros<const N: usize, I: Into<Self::Expr>>(&mut self, array: [I; N]) {
149        self.inner.assert_zeros(array);
150    }
151}
152
153impl<SC: StarkGenericConfig> ExtensionBuilder for VerifierConstraintFolderWithLookups<'_, SC> {
154    type EF = SC::Challenge;
155    type ExprEF = SC::Challenge;
156    type VarEF = SC::Challenge;
157
158    fn assert_zero_ext<I>(&mut self, x: I)
159    where
160        I: Into<Self::ExprEF>,
161    {
162        self.inner.accumulator *= self.inner.alpha;
163        self.inner.accumulator += x.into();
164    }
165}
166
167impl<'a, SC: StarkGenericConfig> PermutationAirBuilder
168    for VerifierConstraintFolderWithLookups<'a, SC>
169{
170    type RandomVar = SC::Challenge;
171    type MP = RowWindow<'a, SC::Challenge>;
172
173    type PermutationVar = SC::Challenge;
174
175    fn permutation(&self) -> Self::MP {
176        RowWindow::from_two_rows(self.permutation.top.values, self.permutation.bottom.values)
177    }
178
179    fn permutation_randomness(&self) -> &[SC::Challenge] {
180        self.permutation_challenges
181    }
182
183    fn permutation_values(&self) -> &[SC::Challenge] {
184        self.permutation_values
185    }
186}