openvm_native_recursion/fri/
domain.rs

1use openvm_native_compiler::prelude::*;
2use openvm_stark_backend::{
3    p3_commit::LagrangeSelectors,
4    p3_field::{coset::TwoAdicMultiplicativeCoset, Field, PrimeCharacteristicRing, TwoAdicField},
5};
6
7use super::types::FriConfigVariable;
8use crate::commit::PolynomialSpaceVariable;
9
10/// Reference: [`openvm_stark_backend::p3_commit::TwoAdicMultiplicativeCoset`]
11#[derive(DslVariable, Clone)]
12pub struct TwoAdicMultiplicativeCosetVariable<C: Config> {
13    pub log_n: Usize<C::N>,
14    pub shift: Felt<C::F>,
15    pub g: Felt<C::F>,
16}
17
18impl<C: Config> TwoAdicMultiplicativeCosetVariable<C> {
19    pub fn first_point(&self) -> Felt<C::F> {
20        self.shift
21    }
22
23    pub fn gen(&self) -> Felt<C::F> {
24        self.g
25    }
26}
27
28impl<C: Config> FromConstant<C> for TwoAdicMultiplicativeCosetVariable<C>
29where
30    C::F: TwoAdicField,
31{
32    type Constant = TwoAdicMultiplicativeCoset<C::F>;
33
34    fn constant(value: Self::Constant, builder: &mut Builder<C>) -> Self {
35        let g_val = C::F::two_adic_generator(value.log_size());
36        TwoAdicMultiplicativeCosetVariable::<C> {
37            // builder.eval is necessary to assign a variable in the dynamic mode.
38            log_n: builder.eval(RVar::from(value.log_size())),
39            shift: builder.eval(value.shift()),
40            g: builder.eval(g_val),
41        }
42    }
43}
44
45impl<C: Config> PolynomialSpaceVariable<C> for TwoAdicMultiplicativeCosetVariable<C>
46where
47    C::F: TwoAdicField,
48{
49    type Constant = TwoAdicMultiplicativeCoset<C::F>;
50
51    fn next_point(
52        &self,
53        builder: &mut Builder<C>,
54        point: Ext<<C as Config>::F, <C as Config>::EF>,
55    ) -> Ext<<C as Config>::F, <C as Config>::EF> {
56        builder.eval(point * self.gen())
57    }
58
59    fn selectors_at_point(
60        &self,
61        builder: &mut Builder<C>,
62        point: Ext<<C as Config>::F, <C as Config>::EF>,
63    ) -> LagrangeSelectors<Ext<<C as Config>::F, <C as Config>::EF>> {
64        let unshifted_point: Ext<_, _> = builder.eval(point * self.shift.inverse());
65        let z_h_expr =
66            builder.exp_power_of_2_v::<Ext<_, _>>(unshifted_point, self.log_n.clone()) - C::EF::ONE;
67        let z_h: Ext<_, _> = builder.eval(z_h_expr);
68
69        LagrangeSelectors {
70            is_first_row: builder.eval(z_h / (unshifted_point - C::EF::ONE)),
71            is_last_row: builder.eval(z_h / (unshifted_point - self.gen().inverse())),
72            is_transition: builder.eval(unshifted_point - self.gen().inverse()),
73            inv_vanishing: builder.eval(z_h.inverse()),
74        }
75    }
76
77    fn vanishing_poly_at_point(
78        &self,
79        builder: &mut Builder<C>,
80        point: Ext<<C as Config>::F, <C as Config>::EF>,
81    ) -> Ext<<C as Config>::F, <C as Config>::EF> {
82        let unshifted_power =
83            builder.exp_power_of_2_v::<Ext<_, _>>(point * self.shift.inverse(), self.log_n.clone());
84        builder.eval(unshifted_power - C::EF::ONE)
85    }
86
87    fn split_domains(
88        &self,
89        builder: &mut Builder<C>,
90        log_num_chunks: impl Into<RVar<C::N>>,
91        num_chunks: impl Into<RVar<C::N>>,
92    ) -> Array<C, Self> {
93        let log_num_chunks = log_num_chunks.into();
94        let num_chunks = num_chunks.into();
95        let log_n = builder.eval_expr(self.log_n.clone() - log_num_chunks);
96
97        let g_dom = self.gen();
98        let g = builder.exp_power_of_2_v::<Felt<C::F>>(g_dom, log_num_chunks);
99
100        let domain_power: Felt<_> = builder.eval(C::F::ONE);
101
102        let domains = builder.array(num_chunks);
103
104        builder.range(0, num_chunks).for_each(|i_vec, builder| {
105            let log_n = builder.eval(log_n);
106            let domain = TwoAdicMultiplicativeCosetVariable {
107                log_n,
108                shift: builder.eval(self.shift * domain_power),
109                g,
110            };
111            // ATTENTION: here must use `builder.set_value`. `builder.set` will convert
112            // `Usize::Const` to `Usize::Var` because it calls `builder.eval`.
113            builder.set_value(&domains, i_vec[0], domain);
114            builder.assign(&domain_power, domain_power * g_dom);
115        });
116
117        domains
118    }
119
120    fn split_domains_const(&self, builder: &mut Builder<C>, log_num_chunks: usize) -> Vec<Self> {
121        let num_chunks = 1 << log_num_chunks;
122        let log_n: Usize<_> = builder.eval(self.log_n.clone() - C::N::from_usize(log_num_chunks));
123
124        let g_dom = self.gen();
125        let g = builder.exp_power_of_2_v::<Felt<C::F>>(g_dom, log_num_chunks);
126
127        let domain_power: Felt<_> = builder.eval(C::F::ONE);
128        let mut domains = vec![];
129
130        for _ in 0..num_chunks {
131            domains.push(TwoAdicMultiplicativeCosetVariable {
132                log_n: log_n.clone(),
133                shift: builder.eval(self.shift * domain_power),
134                g,
135            });
136            builder.assign(&domain_power, domain_power * g_dom);
137        }
138        domains
139    }
140
141    fn create_disjoint_domain(
142        &self,
143        builder: &mut Builder<C>,
144        log_degree: RVar<<C as Config>::N>,
145        config: Option<FriConfigVariable<C>>,
146    ) -> Self {
147        let domain = config.unwrap().get_subgroup(builder, log_degree);
148        TwoAdicMultiplicativeCosetVariable {
149            log_n: domain.log_n,
150            shift: builder.eval(self.shift * C::F::GENERATOR),
151            g: domain.g,
152        }
153    }
154}
155
156#[cfg(test)]
157pub(crate) mod tests {
158    use openvm_native_circuit::execute_program;
159    use openvm_native_compiler::asm::AsmBuilder;
160    use openvm_stark_backend::{
161        config::{Domain, StarkGenericConfig, Val},
162        p3_commit::{Pcs, PolynomialSpace},
163        p3_field::PrimeField,
164    };
165    use openvm_stark_sdk::{
166        config::{
167            baby_bear_poseidon2::{config_from_perm, default_perm, BabyBearPoseidon2Config},
168            fri_params::SecurityParameters,
169        },
170        utils::create_seeded_rng,
171    };
172    use rand::Rng;
173
174    use super::*;
175    use crate::utils::const_fri_config;
176
177    pub(crate) fn domain_assertions<F: TwoAdicField + PrimeField, C: Config<N = F, F = F>>(
178        builder: &mut Builder<C>,
179        domain: &TwoAdicMultiplicativeCosetVariable<C>,
180        domain_val: &TwoAdicMultiplicativeCoset<F>,
181        zeta_val: C::EF,
182    ) {
183        // Assert the domain parameters are the same.
184        builder.assert_var_eq(domain.log_n.clone(), F::from_usize(domain_val.log_size()));
185        builder.assert_felt_eq(domain.shift, domain_val.shift());
186
187        // Get a random point.
188        let zeta: Ext<_, _> = builder.eval(zeta_val.cons());
189
190        // Compare the selector values of the reference and the builder.
191        let sels_expected = domain_val.selectors_at_point(zeta_val);
192        let sels = domain.selectors_at_point(builder, zeta);
193        builder.assert_ext_eq(sels.is_first_row, sels_expected.is_first_row.cons());
194        builder.assert_ext_eq(sels.is_last_row, sels_expected.is_last_row.cons());
195        builder.assert_ext_eq(sels.is_transition, sels_expected.is_transition.cons());
196
197        let zp_val = domain_val.vanishing_poly_at_point(zeta_val);
198        let zp = domain.vanishing_poly_at_point(builder, zeta);
199        builder.assert_ext_eq(zp, zp_val.cons());
200    }
201
202    fn test_domain_impl(static_only: bool) {
203        type SC = BabyBearPoseidon2Config;
204        type F = Val<SC>;
205        type EF = <SC as StarkGenericConfig>::Challenge;
206        type Challenger = <SC as StarkGenericConfig>::Challenger;
207        type ScPcs = <SC as StarkGenericConfig>::Pcs;
208
209        let mut rng = create_seeded_rng();
210        let security_params = SecurityParameters::standard_fast();
211        let config = config_from_perm(&default_perm(), security_params.clone());
212        let pcs = config.pcs();
213        let natural_domain_for_degree = |degree: usize| -> Domain<SC> {
214            <ScPcs as Pcs<EF, Challenger>>::natural_domain_for_degree(pcs, degree)
215        };
216
217        // Initialize a builder.
218        let mut builder = AsmBuilder::<F, EF>::default();
219        builder.flags.static_only = static_only;
220
221        let config_var = const_fri_config(&mut builder, &security_params.fri_params);
222        for i in 0..5 {
223            let log_d_val = 10 + i;
224
225            let log_quotient_degree = 2;
226
227            // Initialize a reference domain.
228            let domain_val = natural_domain_for_degree(1 << log_d_val);
229            let domain = builder.constant(domain_val);
230
231            // builder.assert_felt_eq(domain.shift, domain_val.shift);
232            let zeta_val = rng.random::<EF>();
233            domain_assertions(&mut builder, &domain, &domain_val, zeta_val);
234
235            // Try a shifted domain.
236            let disjoint_domain_val =
237                domain_val.create_disjoint_domain(1 << (log_d_val + log_quotient_degree));
238            let disjoint_domain = builder.constant(disjoint_domain_val);
239            domain_assertions(
240                &mut builder,
241                &disjoint_domain,
242                &disjoint_domain_val,
243                zeta_val,
244            );
245
246            let log_degree = log_d_val + log_quotient_degree;
247            let disjoint_domain_gen = domain.create_disjoint_domain(
248                &mut builder,
249                log_degree.into(),
250                Some(config_var.clone()),
251            );
252            domain_assertions(
253                &mut builder,
254                &disjoint_domain_gen,
255                &disjoint_domain_val,
256                zeta_val,
257            );
258
259            // Now try split domains
260            let qc_domains_val = disjoint_domain_val.split_domains(1 << log_quotient_degree);
261            for dom_val in qc_domains_val.iter() {
262                let dom = builder.constant(*dom_val);
263                domain_assertions(&mut builder, &dom, dom_val, zeta_val);
264            }
265
266            // Test the splitting of domains by the builder.
267            let quotient_size = 1 << log_quotient_degree;
268            let qc_domains =
269                disjoint_domain.split_domains(&mut builder, log_quotient_degree, quotient_size);
270            for (i, dom_val) in qc_domains_val.iter().enumerate() {
271                let dom = builder.get(&qc_domains, i);
272                domain_assertions(&mut builder, &dom, dom_val, zeta_val);
273            }
274        }
275        builder.halt();
276
277        let program = builder.compile_isa();
278        execute_program(program, vec![]);
279    }
280    #[test]
281    fn test_domain_static() {
282        test_domain_impl(true);
283    }
284
285    #[test]
286    fn test_domain_dynamic() {
287        test_domain_impl(false);
288    }
289}
openvm_native_recursion/fri/domain.rs

openvm_native_recursion/fri/
domain.rs
