1use alloc::vec::Vec;
2
3use p3_air::{AirBuilder, AirBuilderWithPublicValues};
4use p3_field::{BasedVectorSpace, PackedField};
5use p3_matrix::dense::RowMajorMatrixView;
6use p3_matrix::stack::VerticalPair;
7
8use crate::{PackedChallenge, PackedVal, StarkGenericConfig, Val};
9
10#[derive(Debug)]
11pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> {
12 pub main: RowMajorMatrixView<'a, PackedVal<SC>>,
13 pub public_values: &'a Vec<Val<SC>>,
14 pub is_first_row: PackedVal<SC>,
15 pub is_last_row: PackedVal<SC>,
16 pub is_transition: PackedVal<SC>,
17 pub alpha_powers: &'a [SC::Challenge],
18 pub decomposed_alpha_powers: &'a [Vec<Val<SC>>],
19 pub accumulator: PackedChallenge<SC>,
20 pub constraint_index: usize,
21}
22
23type ViewPair<'a, T> = VerticalPair<RowMajorMatrixView<'a, T>, RowMajorMatrixView<'a, T>>;
24
25#[derive(Debug)]
26pub struct VerifierConstraintFolder<'a, SC: StarkGenericConfig> {
27 pub main: ViewPair<'a, SC::Challenge>,
28 pub public_values: &'a Vec<Val<SC>>,
29 pub is_first_row: SC::Challenge,
30 pub is_last_row: SC::Challenge,
31 pub is_transition: SC::Challenge,
32 pub alpha: SC::Challenge,
33 pub accumulator: SC::Challenge,
34}
35
36impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> {
37 type F = Val<SC>;
38 type Expr = PackedVal<SC>;
39 type Var = PackedVal<SC>;
40 type M = RowMajorMatrixView<'a, PackedVal<SC>>;
41
42 #[inline]
43 fn main(&self) -> Self::M {
44 self.main
45 }
46
47 #[inline]
48 fn is_first_row(&self) -> Self::Expr {
49 self.is_first_row
50 }
51
52 #[inline]
53 fn is_last_row(&self) -> Self::Expr {
54 self.is_last_row
55 }
56
57 #[inline]
60 fn is_transition_window(&self, size: usize) -> Self::Expr {
61 if size == 2 {
62 self.is_transition
63 } else {
64 panic!("uni-stark only supports a window size of 2")
65 }
66 }
67
68 #[inline]
69 fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
70 let x: PackedVal<SC> = x.into();
71 let alpha_power = self.alpha_powers[self.constraint_index];
72 self.accumulator += Into::<PackedChallenge<SC>>::into(alpha_power) * x;
73 self.constraint_index += 1;
74 }
75
76 #[inline]
77 fn assert_zeros<const N: usize, I: Into<Self::Expr>>(&mut self, array: [I; N]) {
78 let expr_array: [Self::Expr; N] = array.map(Into::into);
79 self.accumulator += PackedChallenge::<SC>::from_basis_coefficients_fn(|i| {
80 let alpha_powers = &self.decomposed_alpha_powers[i]
81 [self.constraint_index..(self.constraint_index + N)];
82 PackedVal::<SC>::packed_linear_combination::<N>(alpha_powers, &expr_array)
83 });
84 self.constraint_index += N;
85 }
86}
87
88impl<SC: StarkGenericConfig> AirBuilderWithPublicValues for ProverConstraintFolder<'_, SC> {
89 type PublicVar = Self::F;
90
91 #[inline]
92 fn public_values(&self) -> &[Self::F] {
93 self.public_values
94 }
95}
96
97impl<'a, SC: StarkGenericConfig> AirBuilder for VerifierConstraintFolder<'a, SC> {
98 type F = Val<SC>;
99 type Expr = SC::Challenge;
100 type Var = SC::Challenge;
101 type M = ViewPair<'a, SC::Challenge>;
102
103 fn main(&self) -> Self::M {
104 self.main
105 }
106
107 fn is_first_row(&self) -> Self::Expr {
108 self.is_first_row
109 }
110
111 fn is_last_row(&self) -> Self::Expr {
112 self.is_last_row
113 }
114
115 fn is_transition_window(&self, size: usize) -> Self::Expr {
118 if size == 2 {
119 self.is_transition
120 } else {
121 panic!("uni-stark only supports a window size of 2")
122 }
123 }
124
125 fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
126 let x: SC::Challenge = x.into();
127 self.accumulator *= self.alpha;
128 self.accumulator += x;
129 }
130}
131
132impl<SC: StarkGenericConfig> AirBuilderWithPublicValues for VerifierConstraintFolder<'_, SC> {
133 type PublicVar = Self::F;
134
135 fn public_values(&self) -> &[Self::F] {
136 self.public_values
137 }
138}