use alloc::vec::Vec;
use p3_air::{AirBuilder, AirBuilderWithPublicValues};
use p3_field::AbstractField;
use p3_matrix::dense::RowMajorMatrixView;
use p3_matrix::stack::VerticalPair;
use crate::{PackedChallenge, PackedVal, StarkGenericConfig, Val};
#[derive(Debug)]
pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> {
pub main: RowMajorMatrixView<'a, PackedVal<SC>>,
pub public_values: &'a Vec<Val<SC>>,
pub is_first_row: PackedVal<SC>,
pub is_last_row: PackedVal<SC>,
pub is_transition: PackedVal<SC>,
pub alpha_powers: &'a [SC::Challenge],
pub accumulator: PackedChallenge<SC>,
pub constraint_index: usize,
}
type ViewPair<'a, T> = VerticalPair<RowMajorMatrixView<'a, T>, RowMajorMatrixView<'a, T>>;
#[derive(Debug)]
pub struct VerifierConstraintFolder<'a, SC: StarkGenericConfig> {
pub main: ViewPair<'a, SC::Challenge>,
pub public_values: &'a Vec<Val<SC>>,
pub is_first_row: SC::Challenge,
pub is_last_row: SC::Challenge,
pub is_transition: SC::Challenge,
pub alpha: SC::Challenge,
pub accumulator: SC::Challenge,
}
impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> {
type F = Val<SC>;
type Expr = PackedVal<SC>;
type Var = PackedVal<SC>;
type M = RowMajorMatrixView<'a, PackedVal<SC>>;
#[inline]
fn main(&self) -> Self::M {
self.main
}
#[inline]
fn is_first_row(&self) -> Self::Expr {
self.is_first_row
}
#[inline]
fn is_last_row(&self) -> Self::Expr {
self.is_last_row
}
#[inline]
fn is_transition_window(&self, size: usize) -> Self::Expr {
if size == 2 {
self.is_transition
} else {
panic!("uni-stark only supports a window size of 2")
}
}
#[inline]
fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
let x: PackedVal<SC> = x.into();
let alpha_power = self.alpha_powers[self.constraint_index];
self.accumulator += PackedChallenge::<SC>::from_f(alpha_power) * x;
self.constraint_index += 1;
}
}
impl<'a, SC: StarkGenericConfig> AirBuilderWithPublicValues for ProverConstraintFolder<'a, SC> {
type PublicVar = Self::F;
#[inline]
fn public_values(&self) -> &[Self::F] {
self.public_values
}
}
impl<'a, SC: StarkGenericConfig> AirBuilder for VerifierConstraintFolder<'a, SC> {
type F = Val<SC>;
type Expr = SC::Challenge;
type Var = SC::Challenge;
type M = ViewPair<'a, SC::Challenge>;
fn main(&self) -> Self::M {
self.main
}
fn is_first_row(&self) -> Self::Expr {
self.is_first_row
}
fn is_last_row(&self) -> Self::Expr {
self.is_last_row
}
fn is_transition_window(&self, size: usize) -> Self::Expr {
if size == 2 {
self.is_transition
} else {
panic!("uni-stark only supports a window size of 2")
}
}
fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
let x: SC::Challenge = x.into();
self.accumulator *= self.alpha;
self.accumulator += x;
}
}
impl<'a, SC: StarkGenericConfig> AirBuilderWithPublicValues for VerifierConstraintFolder<'a, SC> {
type PublicVar = Self::F;
fn public_values(&self) -> &[Self::F] {
self.public_values
}
}