openvm_native_recursion/
helper.rs

1use itertools::Itertools;
2use openvm_native_compiler::prelude::*;
3
4use crate::vars::{MultiStarkVerificationAdviceVariable, StarkProofVariable};
5
6impl<C: Config> StarkProofVariable<C> {
7    pub fn get_air_ids(&self, builder: &mut Builder<C>) -> Array<C, Usize<C::N>> {
8        if builder.flags.static_only {
9            builder.vec(
10                (0..self.per_air.len().value())
11                    .map(Usize::from)
12                    .collect_vec(),
13            )
14        } else {
15            let air_ids = builder.array(self.per_air.len());
16            builder
17                .range(0, self.per_air.len())
18                .for_each(|i_vec, builder| {
19                    let i = i_vec[0];
20                    let air_proof_data = builder.get(&self.per_air, i);
21                    builder.set_value(&air_ids, i, air_proof_data.air_id);
22                });
23            air_ids
24        }
25    }
26}
27
28impl<C: Config> MultiStarkVerificationAdviceVariable<C> {
29    /// Assumption: at most 1 phase is supported.
30    pub fn num_challenges_to_sample(&self, builder: &mut Builder<C>) -> Array<C, Usize<C::N>> {
31        if self.num_challenges_to_sample_mask.is_empty() {
32            return builder.array(0);
33        }
34        // If all 0s, no phase 1.
35        let num_phases: Usize<_> = builder.eval(self.num_challenges_to_sample_mask[0][0].clone());
36        let ret = builder.array(num_phases.clone());
37        // If phase 1 exists:
38        builder
39            .if_eq(num_phases.clone(), RVar::one())
40            .then(|builder| {
41                // Find the biggest index where the mask is 1.
42                let num_challenges: Usize<_> = builder.eval(RVar::zero());
43                for i in 1..self.num_challenges_to_sample_mask[0].len() {
44                    builder
45                        .if_eq(
46                            self.num_challenges_to_sample_mask[0][i].clone(),
47                            RVar::one(),
48                        )
49                        .then(|builder| {
50                            builder.assign(&num_challenges, RVar::from(i));
51                        });
52                }
53                builder.set(&ret, RVar::zero(), num_challenges + RVar::one());
54            });
55        ret
56    }
57}