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}