use p3_commit::{LagrangeSelectors, TwoAdicMultiplicativeCoset};
use p3_field::{AbstractField, TwoAdicField};
use sp1_recursion_compiler::prelude::*;
use super::types::FriConfigVariable;
use crate::commit::PolynomialSpaceVariable;
#[derive(DslVariable, Clone, Copy)]
pub struct TwoAdicMultiplicativeCosetVariable<C: Config> {
pub log_n: Var<C::N>,
pub size: Var<C::N>,
pub shift: Felt<C::F>,
pub g: Felt<C::F>,
}
impl<C: Config> TwoAdicMultiplicativeCosetVariable<C> {
pub const fn size(&self) -> Var<C::N> {
self.size
}
pub const fn first_point(&self) -> Felt<C::F> {
self.shift
}
pub const fn gen(&self) -> Felt<C::F> {
self.g
}
}
impl<C: Config> FromConstant<C> for TwoAdicMultiplicativeCosetVariable<C>
where
C::F: TwoAdicField,
{
type Constant = TwoAdicMultiplicativeCoset<C::F>;
fn constant(value: Self::Constant, builder: &mut Builder<C>) -> Self {
let log_d_val = value.log_n as u32;
let g_val = C::F::two_adic_generator(value.log_n);
TwoAdicMultiplicativeCosetVariable::<C> {
log_n: builder.eval::<Var<_>, _>(C::N::from_canonical_u32(log_d_val)),
size: builder.eval::<Var<_>, _>(C::N::from_canonical_u32(1 << (log_d_val))),
shift: builder.eval(value.shift),
g: builder.eval(g_val),
}
}
}
impl<C: Config> PolynomialSpaceVariable<C> for TwoAdicMultiplicativeCosetVariable<C>
where
C::F: TwoAdicField,
{
type Constant = p3_commit::TwoAdicMultiplicativeCoset<C::F>;
fn next_point(
&self,
builder: &mut Builder<C>,
point: Ext<<C as Config>::F, <C as Config>::EF>,
) -> Ext<<C as Config>::F, <C as Config>::EF> {
builder.eval(point * self.gen())
}
fn selectors_at_point(
&self,
builder: &mut Builder<C>,
point: Ext<<C as Config>::F, <C as Config>::EF>,
) -> LagrangeSelectors<Ext<<C as Config>::F, <C as Config>::EF>> {
let unshifted_point: Ext<_, _> = builder.eval(point * self.shift.inverse());
let z_h_expr = builder
.exp_power_of_2_v::<Ext<_, _>>(unshifted_point, Usize::Var(self.log_n))
- C::EF::one();
let z_h: Ext<_, _> = builder.eval(z_h_expr);
LagrangeSelectors {
is_first_row: builder.eval(z_h / (unshifted_point - C::EF::one())),
is_last_row: builder.eval(z_h / (unshifted_point - self.gen().inverse())),
is_transition: builder.eval(unshifted_point - self.gen().inverse()),
inv_zeroifier: builder.eval(z_h.inverse()),
}
}
fn zp_at_point(
&self,
builder: &mut Builder<C>,
point: Ext<<C as Config>::F, <C as Config>::EF>,
) -> Ext<<C as Config>::F, <C as Config>::EF> {
let unshifted_power = builder
.exp_power_of_2_v::<Ext<_, _>>(point * self.shift.inverse(), Usize::Var(self.log_n));
builder.eval(unshifted_power - C::EF::one())
}
fn split_domains(
&self,
builder: &mut Builder<C>,
log_num_chunks: impl Into<Usize<C::N>>,
num_chunks: impl Into<Usize<C::N>>,
) -> Array<C, Self> {
let log_num_chunks = log_num_chunks.into();
let num_chunks = num_chunks.into();
let log_n: Var<_> = builder.eval(self.log_n - log_num_chunks);
let size = builder.sll(C::N::one(), Usize::Var(log_n));
let g_dom = self.gen();
let g = builder.exp_power_of_2_v::<Felt<C::F>>(g_dom, log_num_chunks);
let domain_power: Felt<_> = builder.eval(C::F::one());
let mut domains = builder.dyn_array(num_chunks);
builder.range(0, num_chunks).for_each(|i, builder| {
let domain = TwoAdicMultiplicativeCosetVariable {
log_n,
size,
shift: builder.eval(self.shift * domain_power),
g,
};
builder.set(&mut domains, i, domain);
builder.assign(domain_power, domain_power * g_dom);
});
domains
}
fn split_domains_const(&self, builder: &mut Builder<C>, log_num_chunks: usize) -> Vec<Self> {
let num_chunks = 1 << log_num_chunks;
let log_n: Var<_> = builder.eval(self.log_n - C::N::from_canonical_usize(log_num_chunks));
let size = builder.sll(C::N::one(), Usize::Var(log_n));
let g_dom = self.gen();
let g = builder.exp_power_of_2_v::<Felt<C::F>>(g_dom, log_num_chunks);
let domain_power: Felt<_> = builder.eval(C::F::one());
let mut domains = vec![];
for _ in 0..num_chunks {
domains.push(TwoAdicMultiplicativeCosetVariable {
log_n,
size,
shift: builder.eval(self.shift * domain_power),
g,
});
builder.assign(domain_power, domain_power * g_dom);
}
domains
}
fn create_disjoint_domain(
&self,
builder: &mut Builder<C>,
log_degree: Usize<<C as Config>::N>,
config: Option<FriConfigVariable<C>>,
) -> Self {
let domain = config.unwrap().get_subgroup(builder, log_degree);
builder.assign(domain.shift, self.shift * C::F::generator());
domain
}
}
#[cfg(test)]
pub(crate) mod tests {
use sp1_recursion_compiler::asm::AsmBuilder;
use sp1_recursion_core::stark::utils::{run_test_recursion, TestConfig};
use sp1_stark::{
baby_bear_poseidon2::BabyBearPoseidon2, inner_fri_config, Dom, StarkGenericConfig,
};
use crate::utils::const_fri_config;
use super::*;
use p3_commit::{Pcs, PolynomialSpace};
use rand::{thread_rng, Rng};
pub(crate) fn domain_assertions<F: TwoAdicField, C: Config<N = F, F = F>>(
builder: &mut Builder<C>,
domain: &TwoAdicMultiplicativeCosetVariable<C>,
domain_val: &TwoAdicMultiplicativeCoset<F>,
zeta_val: C::EF,
) {
builder.assert_var_eq(domain.log_n, F::from_canonical_usize(domain_val.log_n));
builder.assert_var_eq(domain.size, F::from_canonical_usize(1 << domain_val.log_n));
builder.assert_felt_eq(domain.shift, domain_val.shift);
let zeta: Ext<_, _> = builder.eval(zeta_val.cons());
let sels_expected = domain_val.selectors_at_point(zeta_val);
let sels = domain.selectors_at_point(builder, zeta);
builder.assert_ext_eq(sels.is_first_row, sels_expected.is_first_row.cons());
builder.assert_ext_eq(sels.is_last_row, sels_expected.is_last_row.cons());
builder.assert_ext_eq(sels.is_transition, sels_expected.is_transition.cons());
let zp_val = domain_val.zp_at_point(zeta_val);
let zp = domain.zp_at_point(builder, zeta);
builder.assert_ext_eq(zp, zp_val.cons());
}
#[test]
fn test_domain() {
type SC = BabyBearPoseidon2;
type F = <SC as StarkGenericConfig>::Val;
type EF = <SC as StarkGenericConfig>::Challenge;
type Challenger = <SC as StarkGenericConfig>::Challenger;
type ScPcs = <SC as StarkGenericConfig>::Pcs;
let mut rng = thread_rng();
let config = SC::default();
let pcs = config.pcs();
let natural_domain_for_degree = |degree: usize| -> Dom<SC> {
<ScPcs as Pcs<EF, Challenger>>::natural_domain_for_degree(pcs, degree)
};
let mut builder = AsmBuilder::<F, EF>::default();
let config_var = const_fri_config(&mut builder, &inner_fri_config());
for i in 0..5 {
let log_d_val = 10 + i;
let log_quotient_degree = 2;
let domain_val = natural_domain_for_degree(1 << log_d_val);
let domain = builder.constant(domain_val);
let zeta_val = rng.gen::<EF>();
domain_assertions(&mut builder, &domain, &domain_val, zeta_val);
let disjoint_domain_val =
domain_val.create_disjoint_domain(1 << (log_d_val + log_quotient_degree));
let disjoint_domain = builder.constant(disjoint_domain_val);
domain_assertions(&mut builder, &disjoint_domain, &disjoint_domain_val, zeta_val);
let log_degree: Usize<_> = builder.eval(Usize::Const(log_d_val) + log_quotient_degree);
let disjoint_domain_gen =
domain.create_disjoint_domain(&mut builder, log_degree, Some(config_var.clone()));
domain_assertions(&mut builder, &disjoint_domain_gen, &disjoint_domain_val, zeta_val);
let qc_domains_val = disjoint_domain_val.split_domains(1 << log_quotient_degree);
for dom_val in qc_domains_val.iter() {
let dom = builder.constant(*dom_val);
domain_assertions(&mut builder, &dom, dom_val, zeta_val);
}
let quotient_size: Usize<_> = builder.eval(1 << log_quotient_degree);
let log_quotient_degree: Usize<_> = builder.eval(log_quotient_degree);
let qc_domains =
disjoint_domain.split_domains(&mut builder, log_quotient_degree, quotient_size);
for (i, dom_val) in qc_domains_val.iter().enumerate() {
let dom = builder.get(&qc_domains, i);
domain_assertions(&mut builder, &dom, dom_val, zeta_val);
}
}
builder.halt();
let program = builder.compile_program();
run_test_recursion(program, None, TestConfig::All);
}
}