p3_uni_stark/
folder.rs

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    /// # Panics
58    /// This function panics if `size` is not `2`.
59    #[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    /// # Panics
116    /// This function panics if `size` is not `2`.
117    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}