openvm_recursion_circuit/batch_constraint/
mod.rs

1use core::iter::zip;
2use std::sync::Arc;
3
4use itertools::Itertools;
5use openvm_cpu_backend::CpuBackend;
6use openvm_stark_backend::{
7    air_builders::symbolic::{symbolic_variable::Entry, SymbolicExpressionNode},
8    keygen::types::MultiStarkVerifyingKey,
9    poly_common::{eval_eq_sharp_uni, eval_eq_uni, eval_eq_uni_at_one},
10    proof::{column_openings_by_rot, BatchConstraintProof, Proof},
11    prover::{
12        AirProvingContext, ColMajorMatrix, CommittedTraceData, DeviceDataTransporter,
13        TraceCommitter,
14    },
15    AirRef, FiatShamirTranscript, StarkEngine, StarkProtocolConfig, TranscriptHistory,
16};
17use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, EF, F};
18use p3_baby_bear::BabyBear;
19use p3_field::{Field, PrimeCharacteristicRing, TwoAdicField};
20use p3_matrix::dense::RowMajorMatrix;
21use p3_maybe_rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
22use strum::{EnumCount, EnumDiscriminants};
23
24use crate::{
25    batch_constraint::{
26        bus::{
27            BatchConstraintConductorBus, ConstraintsFoldingBus, Eq3bBus, EqNOuterBus,
28            EqNegInternalBus, EqSharpUniBus, EqZeroNBus, ExpressionClaimBus,
29            InteractionsFoldingBus, SumcheckClaimBus, SymbolicExpressionBus,
30            UnivariateSumcheckInputBus,
31        },
32        eq_airs::{
33            generate_eq_sharp_uni_blob, Eq3bAir, Eq3bBlob, EqNegAir, EqNegTraceGenerator, EqNsAir,
34            EqSharpUniAir, EqSharpUniBlob, EqSharpUniReceiverAir, EqUniAir,
35        },
36        expr_eval::{
37            CachedTraceRecord, ConstraintsFoldingAir, ConstraintsFoldingBlob, DagCommitSubAir,
38            InteractionsFoldingAir, InteractionsFoldingBlob, SymbolicExpressionAir,
39        },
40        expression_claim::{
41            generate_expression_claim_blob, ExpressionClaimAir, ExpressionClaimBlob,
42            ExpressionClaimTraceGenerator,
43        },
44        fractions_folder::{FractionsFolderAir, FractionsFolderTraceGenerator},
45        sumcheck::{
46            multilinear::MultilinearSumcheckTraceGenerator,
47            univariate::UnivariateSumcheckTraceGenerator, MultilinearSumcheckAir,
48            UnivariateSumcheckAir,
49        },
50    },
51    bus::{
52        AirPresenceBus, AirShapeBus, BatchConstraintModuleBus, ColumnClaimsBus,
53        ConstraintSumcheckRandomnessBus, ConstraintsFoldingInputBus, Eq3bShapeBus,
54        EqNegBaseRandBus, EqNegResultBus, EqNsNLogupMaxBus, ExpressionClaimNMaxBus,
55        FractionFolderInputBus, HyperdimBus, InteractionsFoldingInputBus, NLiftBus,
56        PublicValuesBus, SelHypercubeBus, SelUniBus, StackingModuleBus, TranscriptBus,
57        XiRandomnessBus,
58    },
59    primitives::{bus::PowerCheckerBus, pow::PowerCheckerCpuTraceGenerator},
60    system::{
61        AirModule, BatchConstraintPreflight, BusIndexManager, BusInventory, GlobalCtxCpu,
62        Preflight, TraceGenModule, POW_CHECKER_HEIGHT,
63    },
64    tracegen::{ModuleChip, RowMajorChip, StandardTracegenCtx},
65    utils::MultiVecWithBounds,
66};
67
68pub mod bus;
69pub mod eq_airs;
70pub mod expr_eval;
71pub mod expression_claim;
72pub mod fractions_folder;
73pub mod sumcheck;
74
75#[cfg(feature = "cuda")]
76mod cuda_abi;
77#[cfg(feature = "cuda")]
78mod cuda_utils;
79
80/// AIR index within the BatchConstraintModule
81pub(crate) const LOCAL_SYMBOLIC_EXPRESSION_AIR_IDX: usize = 0;
82
83pub fn commit_child_vk<E>(
84    engine: &E,
85    child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
86    has_cached: bool,
87) -> CommittedTraceData<E::PB>
88where
89    E: StarkEngine,
90    E::SC: StarkProtocolConfig<F = F>,
91{
92    let cached_trace_record = expr_eval::build_cached_trace_record(child_vk, has_cached);
93    let cached_trace = expr_eval::generate_symbolic_expr_cached_trace(&cached_trace_record);
94    let cached_trace = ColMajorMatrix::from_row_major(&cached_trace);
95    let cached_trace = engine.device().transport_matrix_to_device(&cached_trace);
96    let (commitment, data) = engine.device().commit(&[&cached_trace]).unwrap();
97    CommittedTraceData {
98        commitment,
99        data: Arc::new(data),
100        trace: cached_trace,
101    }
102}
103
104pub struct BatchConstraintModule {
105    transcript_bus: TranscriptBus,
106    constraint_sumcheck_randomness_bus: ConstraintSumcheckRandomnessBus,
107    xi_randomness_bus: XiRandomnessBus,
108    gkr_claim_bus: BatchConstraintModuleBus,
109    constraints_folding_input_bus: ConstraintsFoldingInputBus,
110    interactions_folding_input_bus: InteractionsFoldingInputBus,
111    fraction_folder_input_bus: FractionFolderInputBus,
112    univariate_sumcheck_input_bus: UnivariateSumcheckInputBus,
113    stacking_module_bus: StackingModuleBus,
114    column_opening_bus: ColumnClaimsBus,
115    air_shape_bus: AirShapeBus,
116    air_presence_bus: AirPresenceBus,
117    hyperdim_bus: HyperdimBus,
118    public_values_bus: PublicValuesBus,
119    sel_uni_bus: SelUniBus,
120    eq_n_outer_bus: EqNOuterBus,
121
122    batch_constraint_conductor_bus: BatchConstraintConductorBus,
123    sumcheck_bus: SumcheckClaimBus,
124    expression_claim_n_max_bus: ExpressionClaimNMaxBus,
125    n_lift_bus: NLiftBus,
126    eq_n_logup_n_max_bus: EqNsNLogupMaxBus,
127    eq_3b_shape_bus: Eq3bShapeBus,
128
129    zero_n_bus: EqZeroNBus,
130    eq_sharp_uni_bus: EqSharpUniBus,
131    eq_3b_bus: Eq3bBus,
132    sel_hypercube_bus: SelHypercubeBus,
133    eq_neg_result_bus: EqNegResultBus,
134    eq_neg_base_rand_bus: EqNegBaseRandBus,
135    eq_neg_internal_bus: EqNegInternalBus,
136
137    symbolic_expression_bus: SymbolicExpressionBus,
138    expression_claim_bus: ExpressionClaimBus,
139    interactions_folding_bus: InteractionsFoldingBus,
140    constraints_folding_bus: ConstraintsFoldingBus,
141    power_checker_bus: PowerCheckerBus,
142
143    l_skip: usize,
144    max_constraint_degree: usize,
145
146    max_num_proofs: usize,
147    pub(crate) has_cached: bool,
148}
149
150impl BatchConstraintModule {
151    pub fn new(
152        child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
153        b: &mut BusIndexManager,
154        bus_inventory: BusInventory,
155        max_num_proofs: usize,
156        has_cached: bool,
157    ) -> Self {
158        let l_skip = child_vk.inner.params.l_skip;
159        let max_constraint_degree = child_vk.max_constraint_degree();
160        BatchConstraintModule {
161            transcript_bus: bus_inventory.transcript_bus,
162            constraint_sumcheck_randomness_bus: bus_inventory.constraint_randomness_bus,
163            xi_randomness_bus: bus_inventory.xi_randomness_bus,
164            gkr_claim_bus: bus_inventory.bc_module_bus,
165            constraints_folding_input_bus: bus_inventory.constraints_folding_input_bus,
166            interactions_folding_input_bus: bus_inventory.interactions_folding_input_bus,
167            fraction_folder_input_bus: bus_inventory.fraction_folder_input_bus,
168            stacking_module_bus: bus_inventory.stacking_module_bus,
169            column_opening_bus: bus_inventory.column_claims_bus,
170            air_shape_bus: bus_inventory.air_shape_bus,
171            air_presence_bus: bus_inventory.air_presence_bus,
172            hyperdim_bus: bus_inventory.hyperdim_bus,
173            public_values_bus: bus_inventory.public_values_bus,
174            sel_uni_bus: bus_inventory.sel_uni_bus,
175            eq_neg_base_rand_bus: bus_inventory.eq_neg_base_rand_bus,
176            eq_neg_result_bus: bus_inventory.eq_neg_result_bus,
177            expression_claim_n_max_bus: bus_inventory.expression_claim_n_max_bus,
178            n_lift_bus: bus_inventory.n_lift_bus,
179            eq_n_logup_n_max_bus: bus_inventory.eq_n_logup_n_max_bus,
180            eq_3b_shape_bus: bus_inventory.eq_3b_shape_bus,
181            batch_constraint_conductor_bus: BatchConstraintConductorBus::new(b.new_bus_idx()),
182            univariate_sumcheck_input_bus: UnivariateSumcheckInputBus::new(b.new_bus_idx()),
183            sumcheck_bus: SumcheckClaimBus::new(b.new_bus_idx()),
184
185            zero_n_bus: EqZeroNBus::new(b.new_bus_idx()),
186            eq_sharp_uni_bus: EqSharpUniBus::new(b.new_bus_idx()),
187            eq_3b_bus: Eq3bBus::new(b.new_bus_idx()),
188            eq_neg_internal_bus: EqNegInternalBus::new(b.new_bus_idx()),
189            sel_hypercube_bus: SelHypercubeBus::new(b.new_bus_idx()),
190            eq_n_outer_bus: EqNOuterBus::new(b.new_bus_idx()),
191            // sel_uni bus is shared via inventory
192            symbolic_expression_bus: SymbolicExpressionBus::new(b.new_bus_idx()),
193            expression_claim_bus: ExpressionClaimBus::new(b.new_bus_idx()),
194            interactions_folding_bus: InteractionsFoldingBus::new(b.new_bus_idx()),
195            constraints_folding_bus: ConstraintsFoldingBus::new(b.new_bus_idx()),
196            power_checker_bus: bus_inventory.power_checker_bus,
197            l_skip,
198            max_constraint_degree,
199            max_num_proofs,
200            has_cached,
201        }
202    }
203
204    #[tracing::instrument(level = "trace", skip_all)]
205    pub fn run_preflight<TS>(
206        &self,
207        child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
208        proof: &Proof<BabyBearPoseidon2Config>,
209        preflight: &mut Preflight,
210        ts: &mut TS,
211    ) where
212        TS: FiatShamirTranscript<BabyBearPoseidon2Config> + TranscriptHistory,
213    {
214        let BatchConstraintProof {
215            numerator_term_per_air,
216            denominator_term_per_air,
217            univariate_round_coeffs,
218            sumcheck_round_polys,
219            column_openings,
220        } = &proof.batch_constraint_proof;
221
222        let mut sumcheck_rnd = vec![];
223
224        let mut xi = preflight.gkr.xi.iter().map(|(_, x)| *x).collect_vec();
225        let l_skip = preflight.proof_shape.l_skip;
226        let n_global = preflight.proof_shape.n_global();
227        for _ in xi.len()..(l_skip + n_global) {
228            xi.push(ts.sample_ext());
229        }
230
231        // Constraint batching
232        let lambda_tidx = ts.len();
233        let _lambda = ts.sample_ext();
234
235        for (sum_claim_p, sum_claim_q) in zip(numerator_term_per_air, denominator_term_per_air) {
236            ts.observe_ext(*sum_claim_p);
237            ts.observe_ext(*sum_claim_q);
238        }
239        let _mu = ts.sample_ext();
240
241        let tidx_before_univariate = ts.len();
242
243        // univariate round
244        for coef in univariate_round_coeffs {
245            ts.observe_ext(*coef);
246        }
247        let r0 = ts.sample_ext();
248        sumcheck_rnd.push(r0);
249
250        let tidx_before_multilinear = ts.len();
251
252        for polys in sumcheck_round_polys {
253            for eval in polys {
254                ts.observe_ext(*eval);
255            }
256            let ri = ts.sample_ext();
257            sumcheck_rnd.push(ri);
258        }
259
260        let tidx_before_column_openings = ts.len();
261
262        // Common main
263        for (sort_idx, (air_id, _)) in preflight.proof_shape.sorted_trace_vdata.iter().enumerate() {
264            let need_rot = child_vk.inner.per_air[*air_id].params.need_rot;
265            for (col_opening, rot_opening) in
266                column_openings_by_rot(&column_openings[sort_idx][0], need_rot)
267            {
268                ts.observe_ext(col_opening);
269                ts.observe_ext(rot_opening);
270            }
271        }
272
273        for (sort_idx, (air_id, _)) in preflight.proof_shape.sorted_trace_vdata.iter().enumerate() {
274            let need_rot = child_vk.inner.per_air[*air_id].params.need_rot;
275            for part in column_openings[sort_idx].iter().skip(1) {
276                for (col_opening, rot_opening) in column_openings_by_rot(part, need_rot) {
277                    ts.observe_ext(col_opening);
278                    ts.observe_ext(rot_opening);
279                }
280            }
281        }
282
283        let omega_skip_pows = F::two_adic_generator(l_skip)
284            .powers()
285            .take(1 << l_skip)
286            .collect_vec();
287
288        let mut eq_ns = Vec::with_capacity(preflight.proof_shape.n_max + 1);
289        let mut eq_sharp_ns = Vec::with_capacity(preflight.proof_shape.n_max + 1);
290        let mut eq = eval_eq_uni(l_skip, xi[0], sumcheck_rnd[0]);
291        let mut eq_sharp = eval_eq_sharp_uni(&omega_skip_pows, &xi[..l_skip], sumcheck_rnd[0]);
292        eq_ns.push(eq);
293        eq_sharp_ns.push(eq_sharp);
294        for i in 0..preflight.proof_shape.n_max {
295            let mult = EF::ONE - xi[l_skip + i] - sumcheck_rnd[1 + i]
296                + (xi[l_skip + i] * sumcheck_rnd[1 + i]).double();
297            eq *= mult;
298            eq_sharp *= mult;
299            eq_ns.push(eq);
300            eq_sharp_ns.push(eq_sharp);
301        }
302
303        let mut r_rev_prod = sumcheck_rnd[preflight.proof_shape.n_max];
304        let mut eq_ns_frontloaded = Vec::with_capacity(preflight.proof_shape.n_max + 1);
305        let mut eq_sharp_ns_frontloaded = Vec::with_capacity(preflight.proof_shape.n_max + 1);
306        // Product with r_i's to account for \hat{f} vs \tilde{f} for different n's in front-loaded
307        // batch sumcheck.
308        for i in (0..preflight.proof_shape.n_max).rev() {
309            eq_ns_frontloaded.push(eq_ns[i] * r_rev_prod);
310            eq_sharp_ns_frontloaded.push(eq_sharp_ns[i] * r_rev_prod);
311            r_rev_prod *= sumcheck_rnd[i];
312        }
313        eq_ns_frontloaded.reverse();
314        eq_sharp_ns_frontloaded.reverse();
315        eq_ns_frontloaded.push(eq_ns[preflight.proof_shape.n_max]);
316        eq_sharp_ns_frontloaded.push(eq_sharp_ns[preflight.proof_shape.n_max]);
317
318        preflight.batch_constraint = BatchConstraintPreflight {
319            lambda_tidx,
320            tidx_before_univariate,
321            tidx_before_multilinear,
322            tidx_before_column_openings,
323            post_tidx: ts.len(),
324            xi,
325            sumcheck_rnd,
326            eq_ns,
327            eq_sharp_ns,
328            eq_ns_frontloaded,
329            eq_sharp_ns_frontloaded,
330        }
331    }
332}
333
334impl AirModule for BatchConstraintModule {
335    fn num_airs(&self) -> usize {
336        BatchConstraintModuleChipDiscriminants::COUNT
337    }
338
339    fn airs<SC: StarkProtocolConfig<F = BabyBear>>(&self) -> Vec<AirRef<SC>> {
340        let l_skip = self.l_skip;
341
342        let symbolic_expression_air = SymbolicExpressionAir::<BabyBear> {
343            expr_bus: self.symbolic_expression_bus,
344            air_shape_bus: self.air_shape_bus,
345            air_presence_bus: self.air_presence_bus,
346            column_claims_bus: self.column_opening_bus,
347            interactions_folding_bus: self.interactions_folding_bus,
348            constraints_folding_bus: self.constraints_folding_bus,
349            hyperdim_bus: self.hyperdim_bus,
350            public_values_bus: self.public_values_bus,
351            sel_hypercube_bus: self.sel_hypercube_bus,
352            sel_uni_bus: self.sel_uni_bus,
353            cnt_proofs: self.max_num_proofs,
354            dag_commit_subair: (!self.has_cached).then_some(Arc::new(DagCommitSubAir::new())),
355        };
356        let fraction_folder_air = FractionsFolderAir {
357            transcript_bus: self.transcript_bus,
358            univariate_sumcheck_input_bus: self.univariate_sumcheck_input_bus,
359            fraction_folder_input_bus: self.fraction_folder_input_bus,
360            sumcheck_bus: self.sumcheck_bus,
361            mu_bus: self.batch_constraint_conductor_bus,
362            gkr_claim_bus: self.gkr_claim_bus,
363        };
364        let sumcheck_uni_air = UnivariateSumcheckAir {
365            l_skip,
366            univariate_deg: (self.max_constraint_degree + 1) * ((1 << l_skip) - 1),
367            univariate_sumcheck_input_bus: self.univariate_sumcheck_input_bus,
368            stacking_module_bus: self.stacking_module_bus,
369            claim_bus: self.sumcheck_bus,
370            transcript_bus: self.transcript_bus,
371            randomness_bus: self.constraint_sumcheck_randomness_bus,
372            batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
373        };
374        let sumcheck_lin_air = MultilinearSumcheckAir {
375            max_constraint_degree: self.max_constraint_degree,
376            claim_bus: self.sumcheck_bus,
377            transcript_bus: self.transcript_bus,
378            randomness_bus: self.constraint_sumcheck_randomness_bus,
379            batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
380            stacking_module_bus: self.stacking_module_bus,
381        };
382        let eq_ns_air = EqNsAir {
383            zero_n_bus: self.zero_n_bus,
384            xi_bus: self.xi_randomness_bus,
385            r_xi_bus: self.batch_constraint_conductor_bus,
386            sel_hypercube_bus: self.sel_hypercube_bus,
387            eq_n_outer_bus: self.eq_n_outer_bus,
388            eq_n_logup_n_max_bus: self.eq_n_logup_n_max_bus,
389            l_skip,
390        };
391        let eq_3b_air = Eq3bAir {
392            eq_3b_bus: self.eq_3b_bus,
393            eq_3b_shape_bus: self.eq_3b_shape_bus,
394            batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
395            l_skip,
396        };
397        let eq_sharp_uni_air = EqSharpUniAir {
398            xi_bus: self.xi_randomness_bus,
399            eq_bus: self.eq_sharp_uni_bus,
400            batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
401            l_skip,
402            canonical_inverse_generator: F::two_adic_generator(l_skip).inverse(),
403        };
404        let eq_sharp_uni_receiver_air = EqSharpUniReceiverAir {
405            r_bus: self.batch_constraint_conductor_bus,
406            eq_bus: self.eq_sharp_uni_bus,
407            zero_n_bus: self.zero_n_bus,
408            l_skip,
409        };
410        let eq_uni_air = EqUniAir {
411            r_xi_bus: self.batch_constraint_conductor_bus,
412            zero_n_bus: self.zero_n_bus,
413            l_skip,
414        };
415        let eq_neg_air = EqNegAir {
416            result_bus: self.eq_neg_result_bus,
417            base_rand_bus: self.eq_neg_base_rand_bus,
418            internal_bus: self.eq_neg_internal_bus,
419            sel_uni_bus: self.sel_uni_bus,
420            l_skip: self.l_skip,
421        };
422        let expression_claim_air = ExpressionClaimAir {
423            expression_claim_n_max_bus: self.expression_claim_n_max_bus,
424            expr_claim_bus: self.expression_claim_bus,
425            mu_bus: self.batch_constraint_conductor_bus,
426            sumcheck_claim_bus: self.sumcheck_bus,
427            eq_n_outer_bus: self.eq_n_outer_bus,
428            pow_checker_bus: self.power_checker_bus,
429            hyperdim_bus: self.hyperdim_bus,
430        };
431        let interactions_folding_air = InteractionsFoldingAir {
432            transcript_bus: self.transcript_bus,
433            air_shape_bus: self.air_shape_bus,
434            interaction_bus: self.interactions_folding_bus,
435            interactions_folding_input_bus: self.interactions_folding_input_bus,
436            expression_claim_bus: self.expression_claim_bus,
437            eq_3b_bus: self.eq_3b_bus,
438        };
439        let constraints_folding_air = ConstraintsFoldingAir {
440            transcript_bus: self.transcript_bus,
441            constraint_bus: self.constraints_folding_bus,
442            expression_claim_bus: self.expression_claim_bus,
443            eq_n_outer_bus: self.eq_n_outer_bus,
444            n_lift_bus: self.n_lift_bus,
445            air_shape_bus: self.air_shape_bus,
446            constraints_folding_input_bus: self.constraints_folding_input_bus,
447        };
448        // WARNING: SymbolicExpressionAir MUST be the first AIR in verifier circuit
449        vec![
450            Arc::new(symbolic_expression_air) as AirRef<_>,
451            Arc::new(fraction_folder_air) as AirRef<_>,
452            Arc::new(sumcheck_uni_air) as AirRef<_>,
453            Arc::new(sumcheck_lin_air) as AirRef<_>,
454            Arc::new(eq_ns_air) as AirRef<_>,
455            Arc::new(eq_3b_air) as AirRef<_>,
456            Arc::new(eq_sharp_uni_air) as AirRef<_>,
457            Arc::new(eq_sharp_uni_receiver_air) as AirRef<_>,
458            Arc::new(eq_uni_air) as AirRef<_>,
459            Arc::new(expression_claim_air) as AirRef<_>,
460            Arc::new(interactions_folding_air) as AirRef<_>,
461            Arc::new(constraints_folding_air) as AirRef<_>,
462            Arc::new(eq_neg_air) as AirRef<_>,
463        ]
464    }
465}
466
467pub(crate) struct BatchConstraintBlob {
468    // Per proof, per air (vkey order), the evaluations. For optional AIRs without traces, the
469    // innermost vec is empty.
470    pub expr_evals: MultiVecWithBounds<EF, 2>,
471    // Per proof, per log height.
472    pub selector_counts: MultiVecWithBounds<SelectorCount, 1>,
473
474    pub eq_3b_blob: Eq3bBlob,
475    pub eq_sharp_uni_blob: EqSharpUniBlob,
476}
477
478#[derive(Clone, Copy, Debug, Default)]
479pub struct SelectorCount {
480    pub first: usize,
481    pub last: usize,
482    pub transition: usize,
483}
484
485impl BatchConstraintBlob {
486    pub fn new(
487        child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
488        proofs: &[&Proof<BabyBearPoseidon2Config>],
489        preflights: &[&Preflight],
490    ) -> Self {
491        let child_vk = &child_vk.inner;
492        let params = &child_vk.params;
493
494        let mut expr_evals_per_proof = MultiVecWithBounds::new();
495        let mut eq_r_one_counts_per_proof = MultiVecWithBounds::new();
496        for (proof, preflight) in zip(proofs, preflights) {
497            let rs = &preflight.batch_constraint.sumcheck_rnd;
498
499            let (&rs_0, rs_rest) = rs.split_first().unwrap();
500            let mut is_first_row_by_log_height = vec![];
501            let mut is_last_row_by_log_height = vec![];
502            let n_max = preflight.proof_shape.n_max;
503            let mut selector_counts =
504                vec![SelectorCount::default(); child_vk.params.l_skip + n_max + 1];
505
506            let omega = F::two_adic_generator(params.l_skip);
507            for log_height in 0..=params.l_skip {
508                is_first_row_by_log_height.push(eval_eq_uni_at_one(
509                    log_height,
510                    rs_0.exp_power_of_2(params.l_skip - log_height),
511                ));
512                is_last_row_by_log_height.push(eval_eq_uni_at_one(
513                    log_height,
514                    (rs_0 * omega).exp_power_of_2(params.l_skip - log_height),
515                ));
516            }
517            for (i, &r) in rs_rest.iter().enumerate() {
518                is_first_row_by_log_height
519                    .push(is_first_row_by_log_height[params.l_skip + i] * (EF::ONE - r));
520                is_last_row_by_log_height.push(is_last_row_by_log_height[params.l_skip + i] * r);
521            }
522
523            let mut expr_evals_per_air = vec![];
524            for (air_idx, vk) in child_vk.per_air.iter().enumerate() {
525                if proof.trace_vdata[air_idx].is_none() {
526                    expr_evals_per_air.push(vec![]);
527                    continue;
528                }
529
530                let need_rot = child_vk.per_air[air_idx].params.need_rot;
531                let openings_per_col = if need_rot { 2 } else { 1 };
532                let openings = &proof.batch_constraint_proof.column_openings;
533                let (sorted_idx, vdata) = preflight
534                    .proof_shape
535                    .sorted_trace_vdata
536                    .iter()
537                    .enumerate()
538                    .find_map(|(sorted_idx, (idx, vdata))| {
539                        if air_idx == *idx {
540                            Some((sorted_idx, vdata))
541                        } else {
542                            None
543                        }
544                    })
545                    .unwrap();
546
547                let constraints = &vk.symbolic_constraints.constraints;
548                let mut expr_evals =
549                    vec![EF::ZERO; constraints.nodes.len() + vk.unused_variables.len()];
550                let log_height = proof.trace_vdata[air_idx].as_ref().unwrap().log_height;
551
552                for (node_idx, node) in constraints.nodes.iter().enumerate() {
553                    match node {
554                        SymbolicExpressionNode::Variable(var) => match var.entry {
555                            Entry::Preprocessed { offset } => {
556                                debug_assert!(offset < openings_per_col);
557                                expr_evals[node_idx] =
558                                    openings[sorted_idx][1][var.index * openings_per_col + offset];
559                            }
560                            Entry::Main { part_index, offset } => {
561                                let part = vk.dag_main_part_index_to_commit_index(part_index);
562                                debug_assert!(offset < openings_per_col);
563                                expr_evals[node_idx] = openings[sorted_idx][part]
564                                    [var.index * openings_per_col + offset];
565                            }
566                            Entry::Public => {
567                                expr_evals[node_idx] =
568                                    EF::from(proof.public_values[air_idx][var.index]);
569                            }
570                            Entry::Challenge => unreachable!(),
571                        },
572                        SymbolicExpressionNode::IsFirstRow => {
573                            expr_evals[node_idx] = is_first_row_by_log_height[vdata.log_height];
574                            selector_counts[log_height].first += 1;
575                        }
576                        SymbolicExpressionNode::IsLastRow => {
577                            expr_evals[node_idx] = is_last_row_by_log_height[vdata.log_height];
578                            selector_counts[log_height].last += 1;
579                        }
580                        SymbolicExpressionNode::IsTransition => {
581                            expr_evals[node_idx] =
582                                EF::ONE - is_last_row_by_log_height[vdata.log_height];
583                            selector_counts[log_height].transition += 1;
584                        }
585                        SymbolicExpressionNode::Constant(val) => {
586                            expr_evals[node_idx] = EF::from(*val);
587                        }
588                        SymbolicExpressionNode::Add {
589                            left_idx,
590                            right_idx,
591                            degree_multiple: _,
592                        } => {
593                            debug_assert!(*left_idx < node_idx);
594                            debug_assert!(*right_idx < node_idx);
595                            expr_evals[node_idx] = expr_evals[*left_idx] + expr_evals[*right_idx];
596                        }
597                        SymbolicExpressionNode::Sub {
598                            left_idx,
599                            right_idx,
600                            degree_multiple: _,
601                        } => {
602                            debug_assert!(*left_idx < node_idx);
603                            debug_assert!(*right_idx < node_idx);
604                            expr_evals[node_idx] = expr_evals[*left_idx] - expr_evals[*right_idx];
605                        }
606                        SymbolicExpressionNode::Neg {
607                            idx,
608                            degree_multiple: _,
609                        } => {
610                            debug_assert!(*idx < node_idx);
611                            expr_evals[node_idx] = -expr_evals[*idx];
612                        }
613                        SymbolicExpressionNode::Mul {
614                            left_idx,
615                            right_idx,
616                            degree_multiple: _,
617                        } => {
618                            debug_assert!(*left_idx < node_idx);
619                            debug_assert!(*right_idx < node_idx);
620                            expr_evals[node_idx] = expr_evals[*left_idx] * expr_evals[*right_idx];
621                        }
622                    };
623                }
624                let mut node_idx = constraints.nodes.len();
625                for unused_var in &vk.unused_variables {
626                    match unused_var.entry {
627                        Entry::Preprocessed { offset } => {
628                            debug_assert!(offset < openings_per_col);
629                            expr_evals[node_idx] = openings[sorted_idx][1]
630                                [unused_var.index * openings_per_col + offset];
631                        }
632                        Entry::Main { part_index, offset } => {
633                            let part = vk.dag_main_part_index_to_commit_index(part_index);
634                            debug_assert!(offset < openings_per_col);
635                            expr_evals[node_idx] = openings[sorted_idx][part]
636                                [unused_var.index * openings_per_col + offset];
637                        }
638                        Entry::Public | Entry::Challenge => {
639                            unreachable!()
640                        }
641                    }
642                    node_idx += 1;
643                }
644                expr_evals_per_air.push(expr_evals);
645            }
646            for v in expr_evals_per_air {
647                expr_evals_per_proof.extend(v);
648                expr_evals_per_proof.close_level(1);
649            }
650            expr_evals_per_proof.close_level(0);
651            eq_r_one_counts_per_proof.extend(selector_counts);
652            eq_r_one_counts_per_proof.close_level(0);
653        }
654        let eq_3b_blob = eq_airs::generate_eq_3b_blob(child_vk, preflights);
655        let eq_sharp_uni_blob = generate_eq_sharp_uni_blob(child_vk, preflights);
656        Self {
657            expr_evals: expr_evals_per_proof,
658            selector_counts: eq_r_one_counts_per_proof,
659            eq_3b_blob,
660            eq_sharp_uni_blob,
661        }
662    }
663}
664
665pub(crate) struct BatchConstraintBlobCpu {
666    pub common_blob: BatchConstraintBlob,
667    pub cf_blob: Option<ConstraintsFoldingBlob>,
668    pub if_blob: Option<InteractionsFoldingBlob>,
669    pub expr_claim_blob: ExpressionClaimBlob,
670}
671
672impl BatchConstraintBlobCpu {
673    #[tracing::instrument(name = "generate_blob", skip_all)]
674    pub fn new(
675        child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
676        proofs: &[Proof<BabyBearPoseidon2Config>],
677        preflights: &[Preflight],
678    ) -> Self {
679        let proofs = proofs.iter().collect_vec();
680        let preflights = preflights.iter().collect_vec();
681        let common_blob = BatchConstraintBlob::new(child_vk, &proofs, &preflights);
682        let cf_blob =
683            ConstraintsFoldingBlob::new(&child_vk.inner, &common_blob.expr_evals, &preflights);
684        let if_blob = InteractionsFoldingBlob::new(
685            &child_vk.inner,
686            &common_blob.expr_evals,
687            &common_blob.eq_3b_blob,
688            &preflights,
689        );
690        let expr_claim_blob =
691            generate_expression_claim_blob(&cf_blob.folded_claims, &if_blob.folded_claims);
692        Self {
693            common_blob,
694            cf_blob: Some(cf_blob),
695            if_blob: Some(if_blob),
696            expr_claim_blob,
697        }
698    }
699}
700
701impl<SC: StarkProtocolConfig<F = F>> TraceGenModule<GlobalCtxCpu, CpuBackend<SC>>
702    for BatchConstraintModule
703{
704    type ModuleSpecificCtx<'a> = (
705        Option<&'a CachedTraceRecord>,
706        Arc<PowerCheckerCpuTraceGenerator<2, POW_CHECKER_HEIGHT>>,
707    );
708
709    /// **Note**: This generates all common main traces but leaves the cached trace for
710    /// `SymbolicExpressionAir` unset. The cached trace must be loaded **after** calling this
711    /// function.
712    #[tracing::instrument(skip_all)]
713    fn generate_proving_ctxs(
714        &self,
715        child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
716        proofs: &[Proof<BabyBearPoseidon2Config>],
717        preflights: &[Preflight],
718        ctx: &Self::ModuleSpecificCtx<'_>,
719        required_heights: Option<&[usize]>,
720    ) -> Option<Vec<AirProvingContext<CpuBackend<SC>>>> {
721        let blob = BatchConstraintBlobCpu::new(child_vk, proofs, preflights);
722        let pow_checker = ctx.1.clone();
723        let ctx = (
724            StandardTracegenCtx {
725                vk: child_vk,
726                proofs: &proofs.iter().collect_vec(),
727                preflights: &preflights.iter().collect_vec(),
728            },
729            blob,
730            ctx.0,
731        );
732
733        let chips = [
734            BatchConstraintModuleChip::SymbolicExpression {
735                max_num_proofs: self.max_num_proofs,
736                has_cached: self.has_cached,
737            },
738            BatchConstraintModuleChip::FractionsFolder,
739            BatchConstraintModuleChip::SumcheckUni,
740            BatchConstraintModuleChip::SumcheckLin,
741            BatchConstraintModuleChip::EqNs,
742            BatchConstraintModuleChip::Eq3b,
743            BatchConstraintModuleChip::EqSharpUni,
744            BatchConstraintModuleChip::EqSharpUniReceiver,
745            BatchConstraintModuleChip::EqUni,
746            BatchConstraintModuleChip::ExpressionClaim { pow_checker },
747            BatchConstraintModuleChip::InteractionsFolding,
748            BatchConstraintModuleChip::ConstraintsFolding,
749            BatchConstraintModuleChip::EqNeg,
750        ];
751
752        let span = tracing::Span::current();
753        chips
754            .par_iter()
755            .map(|chip| {
756                let _guard = span.enter();
757                chip.generate_proving_ctx(
758                    &ctx,
759                    required_heights.map(|heights| heights[chip.index()]),
760                )
761            })
762            .collect::<Vec<_>>()
763            .into_iter()
764            .collect()
765    }
766}
767
768// NOTE: ordering of enum must match AIR ordering
769#[derive(strum_macros::Display, EnumDiscriminants)]
770#[strum_discriminants(derive(strum_macros::EnumCount))]
771#[strum_discriminants(repr(usize))]
772enum BatchConstraintModuleChip {
773    SymbolicExpression {
774        max_num_proofs: usize,
775        has_cached: bool,
776    },
777    FractionsFolder,
778    SumcheckUni,
779    SumcheckLin,
780    EqNs,
781    Eq3b,
782    EqSharpUni,
783    EqSharpUniReceiver,
784    EqUni,
785    ExpressionClaim {
786        pow_checker: Arc<PowerCheckerCpuTraceGenerator<2, POW_CHECKER_HEIGHT>>,
787    },
788    InteractionsFolding,
789    ConstraintsFolding,
790    EqNeg,
791}
792
793impl BatchConstraintModuleChip {
794    fn index(&self) -> usize {
795        BatchConstraintModuleChipDiscriminants::from(self) as usize
796    }
797}
798
799impl RowMajorChip<F> for BatchConstraintModuleChip {
800    type Ctx<'a> = (
801        StandardTracegenCtx<'a>,
802        BatchConstraintBlobCpu,
803        Option<&'a CachedTraceRecord>,
804    );
805
806    #[tracing::instrument(
807        name = "wrapper.generate_trace",
808        level = "trace",
809        skip_all,
810        fields(air = %self)
811    )]
812    fn generate_trace(
813        &self,
814        ctx: &Self::Ctx<'_>,
815        required_height: Option<usize>,
816    ) -> Option<RowMajorMatrix<F>> {
817        use BatchConstraintModuleChip::*;
818        let child_vk = ctx.0.vk;
819        let proofs = ctx.0.proofs;
820        let preflights = ctx.0.preflights;
821        let blob = &ctx.1;
822        let cached_trace_record = ctx.2;
823        match self {
824            FractionsFolder => {
825                FractionsFolderTraceGenerator.generate_trace(&ctx.0, required_height)
826            }
827            SumcheckUni => UnivariateSumcheckTraceGenerator.generate_trace(&ctx.0, required_height),
828            SumcheckLin => {
829                MultilinearSumcheckTraceGenerator.generate_trace(&ctx.0, required_height)
830            }
831            EqNs => eq_airs::EqNsTraceGenerator.generate_trace(
832                &(child_vk, preflights, &blob.common_blob.selector_counts),
833                required_height,
834            ),
835            Eq3b => eq_airs::Eq3bTraceGenerator.generate_trace(
836                &(child_vk, &blob.common_blob.eq_3b_blob, preflights),
837                required_height,
838            ),
839            EqSharpUni => eq_airs::EqSharpUniTraceGenerator.generate_trace(
840                &(child_vk, &blob.common_blob.eq_sharp_uni_blob, preflights),
841                required_height,
842            ),
843            EqSharpUniReceiver => eq_airs::EqSharpUniReceiverTraceGenerator.generate_trace(
844                &(child_vk, &blob.common_blob.eq_sharp_uni_blob, preflights),
845                required_height,
846            ),
847            EqUni => eq_airs::EqUniTraceGenerator.generate_trace(&ctx.0, required_height),
848            SymbolicExpression {
849                max_num_proofs,
850                has_cached,
851            } => expr_eval::SymbolicExpressionTraceGenerator {
852                max_num_proofs: *max_num_proofs,
853                has_cached: *has_cached,
854            }
855            .generate_trace(
856                &expr_eval::SymbolicExpressionCtx {
857                    vk: child_vk,
858                    preflights,
859                    expr_evals: &blob.common_blob.expr_evals,
860                    cached_trace_record: &cached_trace_record,
861                },
862                required_height,
863            ),
864            ExpressionClaim { pow_checker } => ExpressionClaimTraceGenerator.generate_trace(
865                &expression_claim::ExpressionClaimCtx {
866                    blob: &blob.expr_claim_blob,
867                    proofs,
868                    preflights,
869                    pow_checker: pow_checker.as_ref(),
870                },
871                required_height,
872            ),
873            InteractionsFolding => expr_eval::InteractionsFoldingTraceGenerator
874                .generate_trace(&(child_vk, blob, preflights), required_height),
875            ConstraintsFolding => expr_eval::ConstraintsFoldingTraceGenerator.generate_trace(
876                &(blob.cf_blob.as_ref().unwrap(), preflights),
877                required_height,
878            ),
879            EqNeg => EqNegTraceGenerator.generate_trace(
880                &(child_vk, preflights, &blob.common_blob.selector_counts),
881                required_height,
882            ),
883        }
884    }
885}
886
887#[cfg(feature = "cuda")]
888pub mod cuda_tracegen {
889    use openvm_cuda_backend::{data_transporter::transport_matrix_h2d_row, GpuBackend};
890    use openvm_cuda_common::stream::GpuDeviceCtx;
891
892    use super::*;
893    use crate::{
894        batch_constraint::expr_eval::{
895            constraints_folding::cuda::ConstraintsFoldingBlobGpu,
896            interactions_folding::cuda::InteractionsFoldingBlobGpu,
897        },
898        cuda::{preflight::PreflightGpu, proof::ProofGpu, vk::VerifyingKeyGpu, GlobalCtxGpu},
899        tracegen::cuda::StandardTracegenGpuCtx,
900    };
901
902    impl ModuleChip<GpuBackend> for BatchConstraintModuleChip {
903        type Ctx<'a> = (
904            StandardTracegenGpuCtx<'a>,
905            &'a BatchConstraintBlobGpu,
906            Option<&'a CachedTraceRecord>,
907        );
908
909        fn generate_proving_ctx(
910            &self,
911            ctx: &Self::Ctx<'_>,
912            required_height: Option<usize>,
913        ) -> Option<AirProvingContext<GpuBackend>> {
914            use BatchConstraintModuleChip::*;
915            let child_vk = ctx.0.vk;
916            let proofs = ctx.0.proofs;
917            let preflights = ctx.0.preflights;
918            let blob = ctx.1;
919            let cached_trace_record = ctx.2;
920            match self {
921                Eq3b => eq_airs::Eq3bTraceGenerator.generate_proving_ctx(
922                    &(
923                        &child_vk.cpu,
924                        &blob.common_blob.eq_3b_blob,
925                        preflights,
926                        ctx.0.device_ctx,
927                    ),
928                    required_height,
929                ),
930                SymbolicExpression {
931                    max_num_proofs,
932                    has_cached,
933                } => expr_eval::SymbolicExpressionTraceGenerator {
934                    max_num_proofs: *max_num_proofs,
935                    has_cached: *has_cached,
936                }
937                .generate_proving_ctx(
938                    &expr_eval::symbolic_expression::cuda::SymbolicExpressionGpuCtx {
939                        vk: &child_vk.cpu,
940                        proofs,
941                        preflights,
942                        expr_evals: &blob.common_blob.expr_evals,
943                        cached_trace_record: &cached_trace_record,
944                        device_ctx: ctx.0.device_ctx,
945                    },
946                    required_height,
947                ),
948                InteractionsFolding => expr_eval::InteractionsFoldingTraceGenerator
949                    .generate_proving_ctx(
950                        &(child_vk, preflights, &blob.if_blob, ctx.0.device_ctx),
951                        required_height,
952                    ),
953                ConstraintsFolding => expr_eval::ConstraintsFoldingTraceGenerator
954                    .generate_proving_ctx(
955                        &(child_vk, preflights, &blob.cf_blob, ctx.0.device_ctx),
956                        required_height,
957                    ),
958                _ => unreachable!(),
959            }
960        }
961    }
962
963    pub(in crate::batch_constraint) struct BatchConstraintBlobGpu {
964        pub common_blob: BatchConstraintBlob,
965        pub cf_blob: ConstraintsFoldingBlobGpu,
966        pub if_blob: InteractionsFoldingBlobGpu,
967        pub expr_claim_blob: ExpressionClaimBlob,
968    }
969
970    impl BatchConstraintBlobGpu {
971        #[tracing::instrument(name = "generate_blob", skip_all)]
972        pub fn new(
973            child_vk: &VerifyingKeyGpu,
974            proofs: &[ProofGpu],
975            preflights: &[PreflightGpu],
976            device_ctx: &GpuDeviceCtx,
977        ) -> Self {
978            let cpu_proofs = proofs.iter().map(|p| &p.cpu).collect_vec();
979            let cpu_preflights = preflights.iter().map(|p| &p.cpu).collect_vec();
980            let common_blob = BatchConstraintBlob::new(&child_vk.cpu, &cpu_proofs, &cpu_preflights);
981            let cf_blob = ConstraintsFoldingBlobGpu::new(
982                child_vk,
983                &common_blob.expr_evals,
984                preflights,
985                device_ctx,
986            );
987            let if_blob = InteractionsFoldingBlobGpu::new(
988                child_vk,
989                &common_blob.expr_evals,
990                &common_blob.eq_3b_blob,
991                preflights,
992                device_ctx,
993            );
994            let expr_claim_blob =
995                generate_expression_claim_blob(&cf_blob.folded_claims, &if_blob.folded_claims);
996            Self {
997                common_blob,
998                cf_blob,
999                if_blob,
1000                expr_claim_blob,
1001            }
1002        }
1003    }
1004
1005    impl TraceGenModule<GlobalCtxGpu, GpuBackend> for BatchConstraintModule {
1006        type ModuleSpecificCtx<'a> = (
1007            Option<&'a CachedTraceRecord>,
1008            Arc<PowerCheckerCpuTraceGenerator<2, POW_CHECKER_HEIGHT>>,
1009            &'a openvm_cuda_common::stream::GpuDeviceCtx,
1010        );
1011
1012        #[tracing::instrument(skip_all)]
1013        fn generate_proving_ctxs(
1014            &self,
1015            child_vk: &VerifyingKeyGpu,
1016            proofs: &[ProofGpu],
1017            preflights: &[PreflightGpu],
1018            module_ctx: &Self::ModuleSpecificCtx<'_>,
1019            required_heights: Option<&[usize]>,
1020        ) -> Option<Vec<AirProvingContext<GpuBackend>>> {
1021            let cached_trace_record = module_ctx.0;
1022            let pow_checker = module_ctx.1.clone();
1023            let device_ctx = module_ctx.2;
1024            let blob = BatchConstraintBlobGpu::new(child_vk, proofs, preflights, device_ctx);
1025            let ctx = (
1026                StandardTracegenGpuCtx {
1027                    vk: child_vk,
1028                    proofs,
1029                    preflights,
1030                    device_ctx,
1031                },
1032                &blob,
1033                cached_trace_record,
1034            );
1035
1036            // Chips with cuda kernels for tracegen
1037            let gpu_chips = [
1038                BatchConstraintModuleChip::SymbolicExpression {
1039                    max_num_proofs: self.max_num_proofs,
1040                    has_cached: self.has_cached,
1041                },
1042                BatchConstraintModuleChip::Eq3b,
1043                BatchConstraintModuleChip::ConstraintsFolding,
1044                BatchConstraintModuleChip::InteractionsFolding,
1045            ];
1046            // Chips that will use fallback cpu tracegen
1047            let cpu_chips = [
1048                BatchConstraintModuleChip::FractionsFolder,
1049                BatchConstraintModuleChip::SumcheckUni,
1050                BatchConstraintModuleChip::SumcheckLin,
1051                BatchConstraintModuleChip::EqNs,
1052                BatchConstraintModuleChip::EqSharpUni,
1053                BatchConstraintModuleChip::EqSharpUniReceiver,
1054                BatchConstraintModuleChip::EqUni,
1055                BatchConstraintModuleChip::ExpressionClaim { pow_checker },
1056                BatchConstraintModuleChip::EqNeg,
1057            ];
1058            let span = tracing::Span::current();
1059            // NOTE: do NOT use par_iter since that will lead to kernels on cuda streams != default
1060            // stream, whereas previous H2D transfer was on default stream.
1061            let indexed_gpu_traces = gpu_chips
1062                .iter()
1063                .map(|chip| {
1064                    // This span is not very useful because the kernel does not synchronize on host:
1065                    let _guard = span.enter();
1066                    (
1067                        chip.index(),
1068                        chip.generate_proving_ctx(
1069                            &ctx,
1070                            required_heights.map(|heights| heights[chip.index()]),
1071                        ),
1072                    )
1073                })
1074                .collect_vec();
1075
1076            let blob = BatchConstraintBlobCpu {
1077                common_blob: blob.common_blob,
1078                cf_blob: None,
1079                if_blob: None,
1080                expr_claim_blob: blob.expr_claim_blob,
1081            };
1082            let cpu_proofs = proofs.iter().map(|p| &p.cpu).collect_vec();
1083            let cpu_preflights = preflights.iter().map(|p| &p.cpu).collect_vec();
1084            let cpu_ctx = (
1085                StandardTracegenCtx {
1086                    vk: &child_vk.cpu,
1087                    proofs: &cpu_proofs,
1088                    preflights: &cpu_preflights,
1089                },
1090                blob,
1091                cached_trace_record,
1092            );
1093
1094            // Phase 1: CPU trace generation in parallel
1095            let indexed_cpu_rm_traces = cpu_chips
1096                .par_iter()
1097                .map(|chip| {
1098                    let _guard = span.enter();
1099                    (
1100                        chip.index(),
1101                        chip.generate_trace(
1102                            &cpu_ctx,
1103                            required_heights.map(|heights| heights[chip.index()]),
1104                        ),
1105                    )
1106                })
1107                .collect::<Vec<_>>();
1108
1109            // Phase 2: H2D transfer serially on main thread
1110            let indexed_cpu_gpu_traces = indexed_cpu_rm_traces
1111                .into_iter()
1112                .map(|(idx, trace)| {
1113                    (
1114                        idx,
1115                        trace.map(|m| {
1116                            AirProvingContext::simple_no_pis(
1117                                transport_matrix_h2d_row(&m, device_ctx).unwrap(),
1118                            )
1119                        }),
1120                    )
1121                })
1122                .collect::<Vec<_>>();
1123
1124            indexed_gpu_traces
1125                .into_iter()
1126                .chain(indexed_cpu_gpu_traces)
1127                .sorted_by(|a, b| a.0.cmp(&b.0))
1128                .map(|(_index, ctx)| ctx)
1129                .collect()
1130        }
1131    }
1132}