diff --git a/extensions/native/compiler/src/ir/builder.rs b/extensions/native/compiler/src/ir/builder.rs index dd98af9705..9d5ca83251 100644 --- a/extensions/native/compiler/src/ir/builder.rs +++ b/extensions/native/compiler/src/ir/builder.rs @@ -286,32 +286,6 @@ impl Builder { } } - pub fn iter<'a, V: MemVariable>( - &'a mut self, - array: &'a Array, - ) -> IteratorBuilder<'a, C, V> { - match array { - Array::Fixed(_) => IteratorBuilder { - start: RVar::zero(), - end: array.len().into(), - step_size: 1, - builder: self, - array, - }, - Array::Dyn(ptr, len) => { - let len: RVar = len.clone().into(); - let end: Var = self.eval(ptr.address + len * RVar::from(V::size_of())); - IteratorBuilder { - start: ptr.address.into(), - end: end.into(), - step_size: V::size_of(), - builder: self, - array, - } - } - } - } - pub fn zip<'a>( &'a mut self, arrays: &'a [Box + 'a>], @@ -825,69 +799,6 @@ impl ZippedPointerIteratorBuilder<'_, C> { } } -pub struct IteratorBuilder<'a, C: Config, V: MemVariable> { - start: RVar, - end: RVar, - step_size: usize, - builder: &'a mut Builder, - array: &'a Array, -} - -impl> IteratorBuilder<'_, C, V> { - pub fn for_each(&mut self, mut f: impl FnMut(V, &mut Builder)) { - if self.start.is_const() && self.end.is_const() { - self.for_each_unrolled(|var, builder| { - f(var, builder); - }); - return; - } - self.for_each_dynamic(|var, builder| { - f(var, builder); - }); - } - - fn for_each_unrolled(&mut self, mut f: impl FnMut(V, &mut Builder)) { - let start = self.start.value(); - let end = self.end.value(); - for i in (start..end).step_by(self.step_size) { - let val = self.builder.get(self.array, i); - f(val, self.builder); - } - } - - fn for_each_dynamic(&mut self, mut f: impl FnMut(V, &mut Builder)) { - assert!( - !self.builder.flags.static_only, - "Cannot use dynamic loop in static mode" - ); - let step_size = C::N::from_canonical_usize(self.step_size); - let loop_variable: Var = self.builder.uninit(); - let mut loop_body_builder = self.builder.create_sub_builder(); - let val: V = loop_body_builder.uninit(); - loop_body_builder.load( - val.clone(), - Ptr { - address: loop_variable, - }, - MemIndex { - index: 0.into(), - offset: 0, - size: V::size_of(), - }, - ); - f(val, &mut loop_body_builder); - let loop_instructions = loop_body_builder.operations; - let op = DslIr::For( - self.start, - self.end, - step_size, - loop_variable, - loop_instructions, - ); - self.builder.operations.push(op); - } -} - /// A builder for the DSL that handles for loops. pub struct RangeBuilder<'a, C: Config> { start: RVar, diff --git a/extensions/native/compiler/src/ir/poseidon.rs b/extensions/native/compiler/src/ir/poseidon.rs index 5001985c1d..b71bb18be5 100644 --- a/extensions/native/compiler/src/ir/poseidon.rs +++ b/extensions/native/compiler/src/ir/poseidon.rs @@ -1,6 +1,7 @@ +use openvm_native_compiler_derive::compile_zip; use openvm_stark_backend::p3_field::FieldAlgebra; -use super::{Array, Builder, Config, DslIr, Ext, Felt, MemIndex, Ptr, Usize, Var}; +use super::{Array, ArrayLike, Builder, Config, DslIr, Ext, Felt, MemIndex, Ptr, Usize, Var}; pub const DIGEST_SIZE: usize = 8; pub const HASH_RATE: usize = 8; @@ -88,8 +89,10 @@ impl Builder { let address = self.eval(state.ptr().address); let start: Var<_> = self.eval(address); let end: Var<_> = self.eval(address + C::N::from_canonical_usize(HASH_RATE)); - self.iter(array).for_each(|subarray, builder| { - builder.iter(&subarray).for_each(|element, builder| { + compile_zip!(self, array).for_each(|idx_vec, builder| { + let subarray = builder.iter_ptr_get(&array, idx_vec[0]); + compile_zip!(builder, subarray).for_each(|ptr_vec, builder| { + let element = builder.iter_ptr_get(&subarray, ptr_vec[0]); builder.cycle_tracker_start("poseidon2-hash-setup"); builder.store( Ptr { address }, diff --git a/extensions/native/compiler/src/ir/utils.rs b/extensions/native/compiler/src/ir/utils.rs index 50648865d8..700f5ec27e 100644 --- a/extensions/native/compiler/src/ir/utils.rs +++ b/extensions/native/compiler/src/ir/utils.rs @@ -1,9 +1,11 @@ use std::ops::{Add, Mul}; +use openvm_native_compiler_derive::compile_zip; use openvm_stark_backend::p3_field::{FieldAlgebra, FieldExtensionAlgebra, PrimeField}; use super::{ - Array, Builder, CanSelect, Config, DslIr, Ext, Felt, MemIndex, RVar, SymbolicExt, Var, Variable, + Array, ArrayLike, Builder, CanSelect, Config, DslIr, Ext, Felt, MemIndex, RVar, SymbolicExt, + Var, Variable, }; pub const NUM_LIMBS: usize = 32; @@ -88,7 +90,8 @@ impl Builder { let one_var: V = self.eval(V::Expression::ONE); // Implements a square-and-multiply algorithm. - self.iter(power_bits).for_each(|bit, builder| { + compile_zip!(self, power_bits).for_each(|ptr_vec, builder| { + let bit = builder.iter_ptr_get(&power_bits, ptr_vec[0]); builder.assign(&result, result * result); let mul = V::select(builder, bit, power_f, one_var); builder.assign(&result, result * mul); diff --git a/extensions/native/compiler/tests/for_loops.rs b/extensions/native/compiler/tests/for_loops.rs index d3b405bd11..dba810f0aa 100644 --- a/extensions/native/compiler/tests/for_loops.rs +++ b/extensions/native/compiler/tests/for_loops.rs @@ -49,51 +49,6 @@ fn test_compiler_for_loops() { execute_program(program, vec![]); } -#[test] -fn test_compiler_iter_fixed() { - let mut builder = AsmBuilder::::default(); - let zero: Var<_> = builder.eval(F::ZERO); - let one: Var<_> = builder.eval(F::ONE); - let two: Var<_> = builder.eval(F::TWO); - let arr = builder.vec(vec![zero, one, two]); - let x: Var<_> = builder.eval(F::ZERO); - let count: Var<_> = builder.eval(F::ZERO); - builder.iter(&arr).for_each(|val: Var<_>, builder| { - builder.assign(&x, x + val); - builder.assign(&count, count + F::ONE); - }); - builder.assert_var_eq(count, F::from_canonical_usize(3)); - builder.assert_var_eq(x, F::from_canonical_usize(3)); - builder.halt(); - - let program = builder.compile_isa(); - execute_program(program, vec![]); -} - -#[test] -fn test_compiler_iter_dyn() { - let mut builder = AsmBuilder::::default(); - let zero: Var<_> = builder.eval(F::ZERO); - let one: Var<_> = builder.eval(F::ONE); - let two: Var<_> = builder.eval(F::TWO); - let arr = builder.dyn_array(3); - builder.set(&arr, 0, zero); - builder.set(&arr, 1, one); - builder.set(&arr, 2, two); - let x: Var<_> = builder.eval(F::ZERO); - let count: Var<_> = builder.eval(F::ZERO); - builder.iter(&arr).for_each(|val: Var<_>, builder| { - builder.assign(&x, x + val); - builder.assign(&count, count + F::ONE); - }); - builder.assert_var_eq(count, F::from_canonical_usize(3)); - builder.assert_var_eq(x, F::from_canonical_usize(3)); - builder.halt(); - - let program = builder.compile_isa(); - execute_program(program, vec![]); -} - #[test] fn test_compiler_zip_fixed() { let mut builder = AsmBuilder::::default(); diff --git a/extensions/native/recursion/src/challenger/duplex.rs b/extensions/native/recursion/src/challenger/duplex.rs index a7ae768f79..db23ed6c02 100644 --- a/extensions/native/recursion/src/challenger/duplex.rs +++ b/extensions/native/recursion/src/challenger/duplex.rs @@ -1,7 +1,8 @@ use openvm_native_compiler::{ ir::{RVar, DIGEST_SIZE, PERMUTATION_WIDTH}, - prelude::{Array, Builder, Config, Ext, Felt, Var}, + prelude::{Array, ArrayLike, Builder, Config, Ext, Felt, Var}, }; +use openvm_native_compiler_derive::compile_zip; use openvm_stark_backend::p3_field::{Field, FieldAlgebra}; use crate::{ @@ -146,11 +147,10 @@ impl DuplexChallengerVariable { self.observe(builder, witness); let element_bits = self.sample_bits(builder, RVar::from(nb_bits)); let element_bits_truncated = element_bits.slice(builder, 0, nb_bits); - builder - .iter(&element_bits_truncated) - .for_each(|element, builder| { - builder.assert_var_eq(element, C::N::ZERO); - }); + compile_zip!(builder, element_bits_truncated).for_each(|ptr_vec, builder| { + let element = builder.iter_ptr_get(&element_bits_truncated, ptr_vec[0]); + builder.assert_var_eq(element, C::N::ZERO); + }); } } @@ -160,7 +160,8 @@ impl CanObserveVariable> for DuplexChallengerVariable, values: Array>) { - builder.iter(&values).for_each(|element, builder| { + compile_zip!(builder, values).for_each(|ptr_vec, builder| { + let element = builder.iter_ptr_get(&values, ptr_vec[0]); self.observe(builder, element); }); } diff --git a/extensions/native/recursion/src/fri/two_adic_pcs.rs b/extensions/native/recursion/src/fri/two_adic_pcs.rs index 18b9f4ca99..13a5a05d7f 100644 --- a/extensions/native/recursion/src/fri/two_adic_pcs.rs +++ b/extensions/native/recursion/src/fri/two_adic_pcs.rs @@ -60,12 +60,12 @@ pub fn verify_two_adic_pcs( builder.iter_ptr_set(&betas, beta_ptr, sample); }); - builder - .iter(&proof.final_poly) - .for_each(|final_poly_elem, builder| { - let final_poly_elem_felts = builder.ext2felt(final_poly_elem); - challenger.observe_slice(builder, final_poly_elem_felts); - }); + compile_zip!(builder, proof.final_poly).for_each(|ptr_vec, builder| { + let final_poly_elem = builder.iter_ptr_get(&proof.final_poly, ptr_vec[0]); + let final_poly_elem_felts = builder.ext2felt(final_poly_elem); + challenger.observe_slice(builder, final_poly_elem_felts); + }); + let num_query_proofs = proof.query_proofs.len().clone(); builder .if_ne(num_query_proofs, RVar::from(config.num_queries)) @@ -78,169 +78,167 @@ pub fn verify_two_adic_pcs( let log_max_height = builder.eval_expr(proof.commit_phase_commits.len() + RVar::from(log_blowup)); - builder - .iter(&proof.query_proofs) - .for_each(|query_proof, builder| { - let index_bits = challenger.sample_bits(builder, log_max_height); - - let ro: Array> = builder.array(32); - let alpha_pow: Array> = builder.array(32); - if builder.flags.static_only { - for j in 0..32 { - // ATTENTION: don't use set_value here, Fixed will share the same variable. - builder.set(&ro, j, C::EF::ZERO.cons()); - builder.set(&alpha_pow, j, C::EF::ONE.cons()); - } - } else { - let zero_ef = builder.eval(C::EF::ZERO.cons()); - let one_ef = builder.eval(C::EF::ONE.cons()); - for j in 0..32 { - // Use set_value here to save a copy. - builder.set_value(&ro, j, zero_ef); - builder.set_value(&alpha_pow, j, one_ef); - } + compile_zip!(builder, proof.query_proofs).for_each(|ptr_vec, builder| { + let query_proof = builder.iter_ptr_get(&proof.query_proofs, ptr_vec[0]); + let index_bits = challenger.sample_bits(builder, log_max_height); + + let ro: Array> = builder.array(32); + let alpha_pow: Array> = builder.array(32); + if builder.flags.static_only { + for j in 0..32 { + // ATTENTION: don't use set_value here, Fixed will share the same variable. + builder.set(&ro, j, C::EF::ZERO.cons()); + builder.set(&alpha_pow, j, C::EF::ONE.cons()); + } + } else { + let zero_ef = builder.eval(C::EF::ZERO.cons()); + let one_ef = builder.eval(C::EF::ONE.cons()); + for j in 0..32 { + // Use set_value here to save a copy. + builder.set_value(&ro, j, zero_ef); + builder.set_value(&alpha_pow, j, one_ef); } + } - compile_zip!(builder, query_proof.input_proof, rounds).for_each(|ptr_vec, builder| { - let batch_opening = builder.iter_ptr_get(&query_proof.input_proof, ptr_vec[0]); - let round = builder.iter_ptr_get(&rounds, ptr_vec[1]); - let batch_commit = round.batch_commit; - let mats = round.mats; - let permutation = round.permutation; - let to_perm_index = |builder: &mut Builder<_>, k: RVar<_>| { - // Always no permutation in static mode - if builder.flags.static_only { - builder.eval(k) - } else { - let ret: Usize<_> = builder.uninit(); - builder.if_eq(permutation.len(), RVar::zero()).then_or_else( - |builder| { - builder.assign(&ret, k); - }, - |builder| { - let value = builder.get(&permutation, k); - builder.assign(&ret, value); - }, - ); - ret - } - }; + compile_zip!(builder, query_proof.input_proof, rounds).for_each(|ptr_vec, builder| { + let batch_opening = builder.iter_ptr_get(&query_proof.input_proof, ptr_vec[0]); + let round = builder.iter_ptr_get(&rounds, ptr_vec[1]); + let batch_commit = round.batch_commit; + let mats = round.mats; + let permutation = round.permutation; + let to_perm_index = |builder: &mut Builder<_>, k: RVar<_>| { + // Always no permutation in static mode + if builder.flags.static_only { + builder.eval(k) + } else { + let ret: Usize<_> = builder.uninit(); + builder.if_eq(permutation.len(), RVar::zero()).then_or_else( + |builder| { + builder.assign(&ret, k); + }, + |builder| { + let value = builder.get(&permutation, k); + builder.assign(&ret, value); + }, + ); + ret + } + }; - let log_batch_max_height: Usize<_> = { - let log_batch_max_index = to_perm_index(builder, RVar::zero()); - let mat = builder.get(&mats, log_batch_max_index); - let domain = mat.domain; - builder.eval(domain.log_n + RVar::from(log_blowup)) - }; + let log_batch_max_height: Usize<_> = { + let log_batch_max_index = to_perm_index(builder, RVar::zero()); + let mat = builder.get(&mats, log_batch_max_index); + let domain = mat.domain; + builder.eval(domain.log_n + RVar::from(log_blowup)) + }; - let batch_dims: Array> = builder.array(mats.len()); - // `verify_batch` requires `permed_opened_values` to be in the committed order. - let permed_opened_values = builder.array(batch_opening.opened_values.len()); - builder.range(0, mats.len()).for_each(|k, builder| { - let mat_index = to_perm_index(builder, k); - - let mat = builder.get(&mats, mat_index.clone()); - let domain = mat.domain; - let dim = DimensionsVariable:: { - height: builder.eval(domain.size() * RVar::from(blowup)), - }; - builder.set_value(&batch_dims, k, dim); - let opened_value = builder.get(&batch_opening.opened_values, mat_index); - builder.set_value(&permed_opened_values, k, opened_value); - }); - let permed_opened_values = NestedOpenedValues::Felt(permed_opened_values); - - let bits_reduced: Usize<_> = builder.eval(log_max_height - log_batch_max_height); - let index_bits_shifted_v1 = index_bits.shift(builder, bits_reduced); - - builder.cycle_tracker_start("verify-batch"); - verify_batch::( - builder, - &batch_commit, - batch_dims, - index_bits_shifted_v1, - &permed_opened_values, - &batch_opening.opening_proof, - ); - builder.cycle_tracker_end("verify-batch"); - - builder.cycle_tracker_start("compute-reduced-opening"); - // `verify_challenges` requires `opened_values` to be in the original order. - let opened_values = batch_opening.opened_values; - - compile_zip!(builder, opened_values, mats).for_each(|ptr_vec, builder| { - let mat_opening = builder.iter_ptr_get(&opened_values, ptr_vec[0]); - let mat = builder.iter_ptr_get(&mats, ptr_vec[1]); - let mat_points = mat.points; - let mat_values = mat.values; - let domain = mat.domain; - let log2_domain_size = domain.log_n; - let log_height = builder.eval_expr(log2_domain_size + RVar::from(log_blowup)); - - let cur_ro = builder.get(&ro, log_height); - let cur_alpha_pow = builder.get(&alpha_pow, log_height); - - let bits_reduced: Usize<_> = builder.eval(log_max_height - log_height); - let index_bits_shifted = index_bits.shift(builder, bits_reduced.clone()); - - let two_adic_generator = config.get_two_adic_generator(builder, log_height); - builder.cycle_tracker_start("exp-reverse-bits-len"); - - let index_bits_shifted_truncated = - index_bits_shifted.slice(builder, 0, log_height); - let two_adic_generator_exp = builder - .exp_bits_big_endian(two_adic_generator, &index_bits_shifted_truncated); - builder.cycle_tracker_end("exp-reverse-bits-len"); - let x: Felt = builder.eval(two_adic_generator_exp * g); - - compile_zip!(builder, mat_points, mat_values).for_each(|ptr_vec, builder| { - let z: Ext = builder.iter_ptr_get(&mat_points, ptr_vec[0]); - let ps_at_z = builder.iter_ptr_get(&mat_values, ptr_vec[1]); - - builder.cycle_tracker_start("single-reduced-opening-eval"); - - if builder.flags.static_only { - builder.range(0, ps_at_z.len()).for_each(|t, builder| { - let p_at_x = builder.get(&mat_opening, t); - let p_at_z = builder.get(&ps_at_z, t); - let quotient = (p_at_z - p_at_x) / (z - x); - - builder.assign(&cur_ro, cur_ro + cur_alpha_pow * quotient); - builder.assign(&cur_alpha_pow, cur_alpha_pow * alpha); - }); - } else { - let mat_ro = builder.fri_single_reduced_opening_eval( - alpha, - cur_alpha_pow, - &mat_opening, - &ps_at_z, - ); - builder.assign(&cur_ro, cur_ro + (mat_ro / (z - x))); - } - - builder.cycle_tracker_end("single-reduced-opening-eval"); - }); - - builder.set_value(&ro, log_height, cur_ro); - builder.set_value(&alpha_pow, log_height, cur_alpha_pow); - }); - builder.cycle_tracker_end("compute-reduced-opening"); + let batch_dims: Array> = builder.array(mats.len()); + // `verify_batch` requires `permed_opened_values` to be in the committed order. + let permed_opened_values = builder.array(batch_opening.opened_values.len()); + builder.range(0, mats.len()).for_each(|k, builder| { + let mat_index = to_perm_index(builder, k); + + let mat = builder.get(&mats, mat_index.clone()); + let domain = mat.domain; + let dim = DimensionsVariable:: { + height: builder.eval(domain.size() * RVar::from(blowup)), + }; + builder.set_value(&batch_dims, k, dim); + let opened_value = builder.get(&batch_opening.opened_values, mat_index); + builder.set_value(&permed_opened_values, k, opened_value); }); + let permed_opened_values = NestedOpenedValues::Felt(permed_opened_values); + + let bits_reduced: Usize<_> = builder.eval(log_max_height - log_batch_max_height); + let index_bits_shifted_v1 = index_bits.shift(builder, bits_reduced); - let folded_eval = verify_query( + builder.cycle_tracker_start("verify-batch"); + verify_batch::( builder, - config, - &proof.commit_phase_commits, - &index_bits, - &query_proof, - &betas, - &ro, - log_max_height, + &batch_commit, + batch_dims, + index_bits_shifted_v1, + &permed_opened_values, + &batch_opening.opening_proof, ); + builder.cycle_tracker_end("verify-batch"); + + builder.cycle_tracker_start("compute-reduced-opening"); + // `verify_challenges` requires `opened_values` to be in the original order. + let opened_values = batch_opening.opened_values; + + compile_zip!(builder, opened_values, mats).for_each(|ptr_vec, builder| { + let mat_opening = builder.iter_ptr_get(&opened_values, ptr_vec[0]); + let mat = builder.iter_ptr_get(&mats, ptr_vec[1]); + let mat_points = mat.points; + let mat_values = mat.values; + let domain = mat.domain; + let log2_domain_size = domain.log_n; + let log_height = builder.eval_expr(log2_domain_size + RVar::from(log_blowup)); + + let cur_ro = builder.get(&ro, log_height); + let cur_alpha_pow = builder.get(&alpha_pow, log_height); + + let bits_reduced: Usize<_> = builder.eval(log_max_height - log_height); + let index_bits_shifted = index_bits.shift(builder, bits_reduced.clone()); + + let two_adic_generator = config.get_two_adic_generator(builder, log_height); + builder.cycle_tracker_start("exp-reverse-bits-len"); + + let index_bits_shifted_truncated = index_bits_shifted.slice(builder, 0, log_height); + let two_adic_generator_exp = + builder.exp_bits_big_endian(two_adic_generator, &index_bits_shifted_truncated); + builder.cycle_tracker_end("exp-reverse-bits-len"); + let x: Felt = builder.eval(two_adic_generator_exp * g); + + compile_zip!(builder, mat_points, mat_values).for_each(|ptr_vec, builder| { + let z: Ext = builder.iter_ptr_get(&mat_points, ptr_vec[0]); + let ps_at_z = builder.iter_ptr_get(&mat_values, ptr_vec[1]); + + builder.cycle_tracker_start("single-reduced-opening-eval"); - let final_poly_elem = builder.get(&proof.final_poly, 0); - builder.assert_ext_eq(folded_eval, final_poly_elem); + if builder.flags.static_only { + builder.range(0, ps_at_z.len()).for_each(|t, builder| { + let p_at_x = builder.get(&mat_opening, t); + let p_at_z = builder.get(&ps_at_z, t); + let quotient = (p_at_z - p_at_x) / (z - x); + + builder.assign(&cur_ro, cur_ro + cur_alpha_pow * quotient); + builder.assign(&cur_alpha_pow, cur_alpha_pow * alpha); + }); + } else { + let mat_ro = builder.fri_single_reduced_opening_eval( + alpha, + cur_alpha_pow, + &mat_opening, + &ps_at_z, + ); + builder.assign(&cur_ro, cur_ro + (mat_ro / (z - x))); + } + + builder.cycle_tracker_end("single-reduced-opening-eval"); + }); + + builder.set_value(&ro, log_height, cur_ro); + builder.set_value(&alpha_pow, log_height, cur_alpha_pow); + }); + builder.cycle_tracker_end("compute-reduced-opening"); }); + + let folded_eval = verify_query( + builder, + config, + &proof.commit_phase_commits, + &index_bits, + &query_proof, + &betas, + &ro, + log_max_height, + ); + + let final_poly_elem = builder.get(&proof.final_poly, 0); + builder.assert_ext_eq(folded_eval, final_poly_elem); + }); builder.cycle_tracker_end("stage-d-verifier-verify"); } diff --git a/extensions/native/recursion/src/stark/mod.rs b/extensions/native/recursion/src/stark/mod.rs index f9c9469896..c97aff8b01 100644 --- a/extensions/native/recursion/src/stark/mod.rs +++ b/extensions/native/recursion/src/stark/mod.rs @@ -4,9 +4,10 @@ use itertools::Itertools; use openvm_circuit::arch::instructions::program::Program; use openvm_native_compiler::{ conversion::CompilerOptions, - ir::{Array, Builder, Config, Ext, ExtConst, Felt, SymbolicExt, Usize}, + ir::{Array, ArrayLike, Builder, Config, Ext, ExtConst, Felt, SymbolicExt, Usize}, prelude::RVar, }; +use openvm_native_compiler_derive::compile_zip; use openvm_stark_backend::{ air_builders::{ symbolic::symbolic_expression::SymbolicExpression, @@ -173,29 +174,28 @@ where // Count the number of main trace commitments together to save a loop. let num_cached_mains: Usize<_> = builder.eval(RVar::zero()); let num_common_main_traces: Usize<_> = builder.eval(RVar::zero()); - builder - .iter(&m_advice_var.per_air) - .for_each(|air_advice, builder| { - builder - .if_eq(air_advice.preprocessed_data.len(), RVar::one()) - .then(|builder| { - let commit = builder.get(&air_advice.preprocessed_data, RVar::zero()); - challenger.observe_digest(builder, commit); - }); + compile_zip!(builder, m_advice_var.per_air).for_each(|ptr_vec, builder| { + let air_advice = builder.iter_ptr_get(&m_advice_var.per_air, ptr_vec[0]); + builder + .if_eq(air_advice.preprocessed_data.len(), RVar::one()) + .then(|builder| { + let commit = builder.get(&air_advice.preprocessed_data, RVar::zero()); + challenger.observe_digest(builder, commit); + }); - builder.assign( - &num_cached_mains, - num_cached_mains.clone() + air_advice.width.cached_mains.len(), - ); - builder - .if_ne(air_advice.width.common_main, RVar::zero()) - .then(|builder| { - builder.assign( - &num_common_main_traces, - num_common_main_traces.clone() + RVar::one(), - ); - }); - }); + builder.assign( + &num_cached_mains, + num_cached_mains.clone() + air_advice.width.cached_mains.len(), + ); + builder + .if_ne(air_advice.width.common_main, RVar::zero()) + .then(|builder| { + builder.assign( + &num_common_main_traces, + num_common_main_traces.clone() + RVar::one(), + ); + }); + }); let CommitmentsVariable { main_trace: main_trace_commits, @@ -204,13 +204,13 @@ where } = commitments; // Observe main trace commitments - builder - .iter(main_trace_commits) - .for_each(|main_commit, builder| { - challenger.observe_digest(builder, main_commit); - }); + compile_zip!(builder, main_trace_commits).for_each(|ptr_vec, builder| { + let main_commit = builder.iter_ptr_get(&main_trace_commits, ptr_vec[0]); + challenger.observe_digest(builder, main_commit); + }); - builder.iter(air_proofs).for_each(|air_proof, builder| { + compile_zip!(builder, air_proofs).for_each(|ptr_vec, builder| { + let air_proof = builder.iter_ptr_get(&air_proofs, ptr_vec[0]); let log_degree = if builder.flags.static_only { builder.eval(C::F::from_canonical_usize(air_proof.log_degree.value())) } else { @@ -254,7 +254,8 @@ where builder.get(&air_advice.num_exposed_values_after_challenge, phase_idx); builder.assert_eq::>(values_len, values.len()); - builder.iter(&values).for_each(|value, builder| { + compile_zip!(builder, values).for_each(|ptr_vec, builder| { + let value = builder.iter_ptr_get(&values, ptr_vec[0]); let felts = builder.ext2felt(value); challenger.observe_slice(builder, felts); }); @@ -388,39 +389,38 @@ where let domain = builder.get(&domains, i); let trace_points = builder.get(&trace_points_per_domain, i); - builder - .iter(cached_main_widths) - .for_each(|cached_main_width, builder| { - let values_per_mat = builder.get(&opening.values.main, main_commit_idx.clone()); - let batch_commit = builder.get(main_trace_commits, main_commit_idx.clone()); - builder.assign(&main_commit_idx, main_commit_idx.clone() + RVar::one()); - - builder.assert_eq::>(values_per_mat.len(), RVar::one()); - let main = builder.get(&values_per_mat, RVar::zero()); - let values = builder.array::>(2); - builder.assert_eq::>(main.local.len(), cached_main_width.clone()); - builder.assert_eq::>(main.next.len(), cached_main_width); - builder.set_value(&values, 0, main.local); - builder.set_value(&values, 1, main.next); - let main_mat = TwoAdicPcsMatsVariable:: { - domain: domain.clone(), - values, - points: trace_points.clone(), - }; - let mats = builder.array(1); - builder.set_value(&mats, 0, main_mat); - - builder.set_value( - &rounds, - round_idx.clone(), - TwoAdicPcsRoundVariable { - batch_commit, - mats, - permutation: null_perm.clone(), - }, - ); - builder.assign(&round_idx, round_idx.clone() + RVar::one()); - }); + compile_zip!(builder, cached_main_widths).for_each(|ptr_vec, builder| { + let cached_main_width = builder.iter_ptr_get(&cached_main_widths, ptr_vec[0]); + let values_per_mat = builder.get(&opening.values.main, main_commit_idx.clone()); + let batch_commit = builder.get(main_trace_commits, main_commit_idx.clone()); + builder.assign(&main_commit_idx, main_commit_idx.clone() + RVar::one()); + + builder.assert_eq::>(values_per_mat.len(), RVar::one()); + let main = builder.get(&values_per_mat, RVar::zero()); + let values = builder.array::>(2); + builder.assert_eq::>(main.local.len(), cached_main_width.clone()); + builder.assert_eq::>(main.next.len(), cached_main_width); + builder.set_value(&values, 0, main.local); + builder.set_value(&values, 1, main.next); + let main_mat = TwoAdicPcsMatsVariable:: { + domain: domain.clone(), + values, + points: trace_points.clone(), + }; + let mats = builder.array(1); + builder.set_value(&mats, 0, main_mat); + + builder.set_value( + &rounds, + round_idx.clone(), + TwoAdicPcsRoundVariable { + batch_commit, + mats, + permutation: null_perm.clone(), + }, + ); + builder.assign(&round_idx, round_idx.clone() + RVar::one()); + }); let common_main_width = RVar::from(advice.width.common_main); builder @@ -522,13 +522,12 @@ where let quotient_perm = builder.array(num_quotient_mats); let perm_offset_per_air = builder.array::>(num_airs); let offset: Usize<_> = builder.eval(RVar::zero()); - builder - .iter(air_perm_by_height) - .for_each(|air_index, builder| { - builder.set(&perm_offset_per_air, air_index.clone(), offset.clone()); - let qc_domains = builder.get("ient_chunk_domains, air_index); - builder.assign(&offset, offset.clone() + qc_domains.len()); - }); + compile_zip!(builder, air_perm_by_height).for_each(|ptr_vec, builder| { + let air_index = builder.iter_ptr_get(&air_perm_by_height, ptr_vec[0]); + builder.set(&perm_offset_per_air, air_index.clone(), offset.clone()); + let qc_domains = builder.get("ient_chunk_domains, air_index); + builder.assign(&offset, offset.clone() + qc_domains.len()); + }); let quotient_mats: Array<_, TwoAdicPcsMatsVariable<_>> = builder.array(num_quotient_mats); let qc_points = builder.array::>(1);