openvm_native_recursion/
view.rs

Help
1use itertools::Itertools;
2use openvm_native_compiler::{
3    ir::{Builder, Config},
4    prelude::*,
5};
6use openvm_stark_backend::{
7    keygen::types::TraceWidth,
8    p3_field::{FieldAlgebra, PrimeField32},
9    p3_util::log2_strict_usize,
10};
11
12use crate::{
13    types::{MultiStarkVerificationAdvice, StarkVerificationAdvice},
14    vars::{
15        LinearConstraintVariable, MultiStarkVerificationAdviceVariable, OptionalVar,
16        StarkVerificationAdviceVariable, TraceHeightConstraintSystem, TraceWidthVariable,
17    },
18};
19
20pub fn get_advice_per_air<C: Config>(
21    builder: &mut Builder<C>,
22    m_advice: &MultiStarkVerificationAdvice<C>,
23    air_ids: &Array<C, Usize<C::N>>,
24) -> MultiStarkVerificationAdviceVariable<C> {
25    let num_challenges_to_sample_mask = m_advice
26        .num_challenges_to_sample
27        .iter()
28        .map(|&num_challenges_to_sample| vec![builder.eval(RVar::zero()); num_challenges_to_sample])
29        .collect_vec();
30    let advice_per_air = builder.array(air_ids.len());
31
32    let idx: Usize<_> = builder.eval(RVar::zero());
33    for (air_id, advice) in m_advice.per_air.iter().enumerate() {
34        builder.if_ne(idx.clone(), air_ids.len()).then(|builder| {
35            let curr_air_id = builder.get(air_ids, idx.clone());
36            let air_id = Usize::from(air_id);
37            builder.if_eq(air_id, curr_air_id).then(|builder| {
38                let advice_var = constant_advice_and_update_mask(
39                    builder,
40                    advice,
41                    &num_challenges_to_sample_mask,
42                );
43                builder.set_value(&advice_per_air, idx.clone(), advice_var);
44                builder.inc(&idx);
45            });
46        });
47    }
48    // Assert that all AIRs in air_ids are covered.
49    // This will ensure that
50    // - `air_ids` are in increasing order and that
51    // - `air_ids.len() <= m_advice.per_air.len()`.
52    builder.assert_var_eq(idx, air_ids.len());
53
54    let trace_height_constraints = m_advice
55        .trace_height_constraints
56        .iter()
57        .map(|constraint| {
58            let coefficients = builder.array(constraint.coefficients.len());
59            for (i, coeff) in constraint.coefficients.iter().enumerate() {
60                let coefficient: Var<_> = builder.constant(C::N::from_canonical_u32(*coeff));
61                builder.set(&coefficients, i, coefficient);
62            }
63            assert!(constraint.threshold <= C::F::ORDER_U32);
64            let threshold: Var<_> = builder.constant(C::N::from_wrapped_u32(constraint.threshold));
65            let is_threshold_at_p = constraint.threshold == C::F::ORDER_U32;
66            LinearConstraintVariable {
67                coefficients,
68                threshold,
69                is_threshold_at_p,
70            }
71        })
72        .collect();
73
74    let height_maxes = builder.array(m_advice.per_air.len());
75    for i in 0..m_advice.per_air.len() {
76        let max_coefficient = m_advice
77            .trace_height_constraints
78            .iter()
79            .map(|constraint| constraint.coefficients[i])
80            .max()
81            .unwrap();
82        let height_max = if max_coefficient <= 1 {
83            OptionalVar {
84                is_some: Usize::from(0),
85                value: builder.constant(C::N::ZERO),
86            }
87        } else {
88            OptionalVar {
89                is_some: Usize::from(1),
90                // Because `C::F::ORDER_U32` is prime and `max_coefficient > 1`, `floor(C::F::ORDER_U32 / max_coefficient) * max_coefficient < C::F::ORDER_U32`,
91                // `height * max_coefficient` cannot overflow `C::F`.
92                value: builder.constant(C::N::from_canonical_u32(
93                    C::F::ORDER_U32 / max_coefficient + 1,
94                )),
95            }
96        };
97        builder.set(&height_maxes, i, height_max);
98    }
99
100    MultiStarkVerificationAdviceVariable {
101        per_air: advice_per_air,
102        num_challenges_to_sample_mask,
103        trace_height_constraint_system: TraceHeightConstraintSystem {
104            height_constraints: trace_height_constraints,
105            height_maxes,
106        },
107    }
108}
109
110fn constant_advice_and_update_mask<C: Config>(
111    builder: &mut Builder<C>,
112    advice: &StarkVerificationAdvice<C>,
113    num_challenges_to_sample_mask: &[Vec<Usize<C::N>>],
114) -> StarkVerificationAdviceVariable<C> {
115    let preprocessed_data = if let Some(preprocessed_data) = advice.preprocessed_data.as_ref() {
116        let commit = builder.constant(preprocessed_data.commit.clone());
117        let arr = builder.array(1);
118        builder.set_value(&arr, 0, commit);
119        arr
120    } else {
121        builder.array(0)
122    };
123    for (phase, &num_challenges) in advice.num_challenges_to_sample.iter().enumerate() {
124        for i in 0..num_challenges {
125            builder.assign(&num_challenges_to_sample_mask[phase][i], RVar::one());
126        }
127    }
128    StarkVerificationAdviceVariable {
129        preprocessed_data,
130        width: constant_trace_width(builder, &advice.width),
131        log_quotient_degree: builder.eval(RVar::from(log2_strict_usize(advice.quotient_degree))),
132        num_public_values: builder.eval(RVar::from(advice.num_public_values)),
133        num_challenges_to_sample: constant_usize_array(builder, &advice.num_challenges_to_sample),
134        num_exposed_values_after_challenge: constant_usize_array(
135            builder,
136            &advice.num_exposed_values_after_challenge,
137        ),
138    }
139}
140
141fn constant_trace_width<C: Config>(
142    builder: &mut Builder<C>,
143    trace_width: &TraceWidth,
144) -> TraceWidthVariable<C> {
145    TraceWidthVariable {
146        preprocessed: constant_usize_option(builder, &trace_width.preprocessed),
147        cached_mains: constant_usize_array(builder, &trace_width.cached_mains),
148        common_main: builder.eval(RVar::from(trace_width.common_main)),
149        after_challenge: constant_usize_array(builder, &trace_width.after_challenge),
150    }
151}
152
153fn constant_usize_option<C: Config>(
154    builder: &mut Builder<C>,
155    opt: &Option<usize>,
156) -> Array<C, Usize<C::N>> {
157    match opt {
158        Some(val) => {
159            let arr = builder.array(1);
160            let v: Usize<_> = builder.eval(RVar::from(*val));
161            builder.set_value(&arr, 0, v);
162            arr
163        }
164        None => builder.array(0),
165    }
166}
167
168fn constant_usize_array<C: Config>(
169    builder: &mut Builder<C>,
170    arr: &[usize],
171) -> Array<C, Usize<C::N>> {
172    let carr = builder.array(arr.len());
173    for (i, val) in arr.iter().enumerate() {
174        let v: Usize<_> = builder.eval(RVar::from(*val));
175        builder.set(&carr, i, v);
176    }
177    carr
178}
openvm_native_recursion/view.rs

openvm_native_recursion/
view.rs
