sp1_recursion_program/
utils.rs

1use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
2use p3_commit::{ExtensionMmcs, TwoAdicMultiplicativeCoset};
3use p3_field::{extension::BinomialExtensionField, AbstractField, Field, TwoAdicField};
4use p3_fri::FriConfig;
5use p3_merkle_tree::FieldMerkleTreeMmcs;
6use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
7use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
8use sp1_recursion_compiler::{
9    asm::AsmConfig,
10    ir::{Array, Builder, Config, Felt, MemVariable, Var},
11};
12use sp1_recursion_core::{
13    air::ChallengerPublicValues,
14    runtime::{DIGEST_SIZE, PERMUTATION_WIDTH},
15};
16use sp1_stark::{
17    air::MachineAir, baby_bear_poseidon2::BabyBearPoseidon2, Dom, ShardProof, StarkGenericConfig,
18    StarkMachine, StarkVerifyingKey,
19};
20
21use crate::{
22    challenger::DuplexChallengerVariable,
23    fri::{types::FriConfigVariable, TwoAdicMultiplicativeCosetVariable},
24    stark::EMPTY,
25    types::{QuotientDataValues, VerifyingKeyVariable},
26};
27
28type SC = BabyBearPoseidon2;
29type F = <SC as StarkGenericConfig>::Val;
30type EF = <SC as StarkGenericConfig>::Challenge;
31type C = AsmConfig<F, EF>;
32type Val = BabyBear;
33type Challenge = BinomialExtensionField<Val, 4>;
34type Perm = Poseidon2<Val, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, 16, 7>;
35type Hash = PaddingFreeSponge<Perm, 16, 8, 8>;
36type Compress = TruncatedPermutation<Perm, 2, 8, 16>;
37type ValMmcs =
38    FieldMerkleTreeMmcs<<Val as Field>::Packing, <Val as Field>::Packing, Hash, Compress, 8>;
39type ChallengeMmcs = ExtensionMmcs<Val, Challenge, ValMmcs>;
40type RecursionConfig = AsmConfig<Val, Challenge>;
41type RecursionBuilder = Builder<RecursionConfig>;
42
43pub fn const_fri_config(
44    builder: &mut RecursionBuilder,
45    config: &FriConfig<ChallengeMmcs>,
46) -> FriConfigVariable<RecursionConfig> {
47    let two_addicity = Val::TWO_ADICITY;
48    let mut generators = builder.dyn_array(two_addicity);
49    let mut subgroups = builder.dyn_array(two_addicity);
50    for i in 0..two_addicity {
51        let constant_generator = Val::two_adic_generator(i);
52        builder.set(&mut generators, i, constant_generator);
53
54        let constant_domain = TwoAdicMultiplicativeCoset { log_n: i, shift: Val::one() };
55        let domain_value: TwoAdicMultiplicativeCosetVariable<_> = builder.constant(constant_domain);
56        builder.set(&mut subgroups, i, domain_value);
57    }
58    FriConfigVariable {
59        log_blowup: builder.eval(BabyBear::from_canonical_usize(config.log_blowup)),
60        blowup: builder.eval(BabyBear::from_canonical_usize(1 << config.log_blowup)),
61        num_queries: builder.eval(BabyBear::from_canonical_usize(config.num_queries)),
62        proof_of_work_bits: builder.eval(BabyBear::from_canonical_usize(config.proof_of_work_bits)),
63        subgroups,
64        generators,
65    }
66}
67
68pub fn clone<T: MemVariable<C>>(builder: &mut RecursionBuilder, var: &T) -> T {
69    let mut arr = builder.dyn_array(1);
70    builder.set(&mut arr, 0, var.clone());
71    builder.get(&arr, 0)
72}
73
74pub fn clone_array<T: MemVariable<C>>(
75    builder: &mut RecursionBuilder,
76    arr: &Array<C, T>,
77) -> Array<C, T> {
78    let mut new_arr = builder.dyn_array(arr.len());
79    builder.range(0, arr.len()).for_each(|i, builder| {
80        let var = builder.get(arr, i);
81        builder.set(&mut new_arr, i, var);
82    });
83    new_arr
84}
85
86// OPT: this can be done much more efficiently, but in the meantime this should work
87pub fn felt2var<C: Config>(builder: &mut Builder<C>, felt: Felt<C::F>) -> Var<C::N> {
88    let bits = builder.num2bits_f(felt);
89    builder.bits2num_v(&bits)
90}
91
92pub fn var2felt<C: Config>(builder: &mut Builder<C>, var: Var<C::N>) -> Felt<C::F> {
93    let bits = builder.num2bits_v(var);
94    builder.bits2num_f(&bits)
95}
96
97/// Asserts that the challenger variable is equal to a challenger in public values.
98pub fn assert_challenger_eq_pv<C: Config>(
99    builder: &mut Builder<C>,
100    var: &DuplexChallengerVariable<C>,
101    values: ChallengerPublicValues<Felt<C::F>>,
102) {
103    for i in 0..PERMUTATION_WIDTH {
104        let element = builder.get(&var.sponge_state, i);
105        builder.assert_felt_eq(element, values.sponge_state[i]);
106    }
107    let num_inputs_var = felt2var(builder, values.num_inputs);
108    builder.assert_var_eq(var.nb_inputs, num_inputs_var);
109    let mut input_buffer_array: Array<_, Felt<_>> = builder.dyn_array(PERMUTATION_WIDTH);
110    for i in 0..PERMUTATION_WIDTH {
111        builder.set(&mut input_buffer_array, i, values.input_buffer[i]);
112    }
113    builder.range(0, num_inputs_var).for_each(|i, builder| {
114        let element = builder.get(&var.input_buffer, i);
115        let values_element = builder.get(&input_buffer_array, i);
116        builder.assert_felt_eq(element, values_element);
117    });
118    let num_outputs_var = felt2var(builder, values.num_outputs);
119    builder.assert_var_eq(var.nb_outputs, num_outputs_var);
120    let mut output_buffer_array: Array<_, Felt<_>> = builder.dyn_array(PERMUTATION_WIDTH);
121    for i in 0..PERMUTATION_WIDTH {
122        builder.set(&mut output_buffer_array, i, values.output_buffer[i]);
123    }
124    builder.range(0, num_outputs_var).for_each(|i, builder| {
125        let element = builder.get(&var.output_buffer, i);
126        let values_element = builder.get(&output_buffer_array, i);
127        builder.assert_felt_eq(element, values_element);
128    });
129}
130
131/// Assigns a challenger variable from a challenger in public values.
132pub fn assign_challenger_from_pv<C: Config>(
133    builder: &mut Builder<C>,
134    dst: &mut DuplexChallengerVariable<C>,
135    values: ChallengerPublicValues<Felt<C::F>>,
136) {
137    for i in 0..PERMUTATION_WIDTH {
138        builder.set(&mut dst.sponge_state, i, values.sponge_state[i]);
139    }
140    let num_inputs_var = felt2var(builder, values.num_inputs);
141    builder.assign(dst.nb_inputs, num_inputs_var);
142    for i in 0..PERMUTATION_WIDTH {
143        builder.set(&mut dst.input_buffer, i, values.input_buffer[i]);
144    }
145    let num_outputs_var = felt2var(builder, values.num_outputs);
146    builder.assign(dst.nb_outputs, num_outputs_var);
147    for i in 0..PERMUTATION_WIDTH {
148        builder.set(&mut dst.output_buffer, i, values.output_buffer[i]);
149    }
150}
151
152pub fn get_challenger_public_values<C: Config>(
153    builder: &mut Builder<C>,
154    var: &DuplexChallengerVariable<C>,
155) -> ChallengerPublicValues<Felt<C::F>> {
156    let sponge_state = core::array::from_fn(|i| builder.get(&var.sponge_state, i));
157    let num_inputs = var2felt(builder, var.nb_inputs);
158    let input_buffer = core::array::from_fn(|i| builder.get(&var.input_buffer, i));
159    let num_outputs = var2felt(builder, var.nb_outputs);
160    let output_buffer = core::array::from_fn(|i| builder.get(&var.output_buffer, i));
161
162    ChallengerPublicValues { sponge_state, num_inputs, input_buffer, num_outputs, output_buffer }
163}
164
165/// Hash the verifying key + prep domains into a single digest.
166/// poseidon2( commit[0..8] || pc_start || prep_domains[N].{log_n, .size, .shift, .g})
167pub fn hash_vkey<C: Config>(
168    builder: &mut Builder<C>,
169    vk: &VerifyingKeyVariable<C>,
170) -> Array<C, Felt<C::F>> {
171    let domain_slots: Var<_> = builder.eval(vk.prep_domains.len() * 4);
172    let vkey_slots: Var<_> = builder.constant(C::N::from_canonical_usize(DIGEST_SIZE + 1));
173    let total_slots: Var<_> = builder.eval(vkey_slots + domain_slots);
174    let mut inputs = builder.dyn_array(total_slots);
175    builder.range(0, DIGEST_SIZE).for_each(|i, builder| {
176        let element = builder.get(&vk.commitment, i);
177        builder.set(&mut inputs, i, element);
178    });
179    builder.set(&mut inputs, DIGEST_SIZE, vk.pc_start);
180    let four: Var<_> = builder.constant(C::N::from_canonical_usize(4));
181    let one: Var<_> = builder.constant(C::N::one());
182    builder.range(0, vk.prep_domains.len()).for_each(|i, builder| {
183        let sorted_index = builder.get(&vk.preprocessed_sorted_idxs, i);
184        let domain = builder.get(&vk.prep_domains, i);
185        let log_n_index: Var<_> = builder.eval(vkey_slots + sorted_index * four);
186        let size_index: Var<_> = builder.eval(log_n_index + one);
187        let shift_index: Var<_> = builder.eval(size_index + one);
188        let g_index: Var<_> = builder.eval(shift_index + one);
189        let log_n_felt = var2felt(builder, domain.log_n);
190        let size_felt = var2felt(builder, domain.size);
191        builder.set(&mut inputs, log_n_index, log_n_felt);
192        builder.set(&mut inputs, size_index, size_felt);
193        builder.set(&mut inputs, shift_index, domain.shift);
194        builder.set(&mut inputs, g_index, domain.g);
195    });
196    builder.poseidon2_hash(&inputs)
197}
198
199pub(crate) fn get_sorted_indices<SC: StarkGenericConfig, A: MachineAir<SC::Val>>(
200    machine: &StarkMachine<SC, A>,
201    proof: &ShardProof<SC>,
202) -> Vec<usize> {
203    machine
204        .chips_sorted_indices(proof)
205        .into_iter()
206        .map(|x| match x {
207            Some(x) => x,
208            None => EMPTY,
209        })
210        .collect()
211}
212
213pub(crate) fn get_preprocessed_data<SC: StarkGenericConfig, A: MachineAir<SC::Val>>(
214    machine: &StarkMachine<SC, A>,
215    vk: &StarkVerifyingKey<SC>,
216) -> (Vec<usize>, Vec<Dom<SC>>) {
217    let chips = machine.chips();
218    let (prep_sorted_indices, prep_domains) = machine
219        .preprocessed_chip_ids()
220        .into_iter()
221        .map(|chip_idx| {
222            let name = chips[chip_idx].name().clone();
223            let prep_sorted_idx = vk.chip_ordering[&name];
224            (prep_sorted_idx, vk.chip_information[prep_sorted_idx].1)
225        })
226        .unzip();
227    (prep_sorted_indices, prep_domains)
228}
229
230pub(crate) fn get_chip_quotient_data<SC: StarkGenericConfig, A: MachineAir<SC::Val>>(
231    machine: &StarkMachine<SC, A>,
232    proof: &ShardProof<SC>,
233) -> Vec<QuotientDataValues> {
234    machine
235        .shard_chips_ordered(&proof.chip_ordering)
236        .map(|chip| {
237            let log_quotient_degree = chip.log_quotient_degree();
238            QuotientDataValues { log_quotient_degree, quotient_size: 1 << log_quotient_degree }
239        })
240        .collect()
241}
sp1_recursion_program/utils.rs

sp1_recursion_program/
utils.rs
