diff --git a/ceno_recursion_v2/Cargo.lock b/ceno_recursion_v2/Cargo.lock index dab203786..1d4b0275d 100644 --- a/ceno_recursion_v2/Cargo.lock +++ b/ceno_recursion_v2/Cargo.lock @@ -1198,7 +1198,7 @@ dependencies = [ [[package]] name = "ff_ext" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "once_cell", "p3", @@ -1823,7 +1823,7 @@ dependencies = [ [[package]] name = "mpcs" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "bincode", "clap", @@ -1846,7 +1846,7 @@ dependencies = [ [[package]] name = "multilinear_extensions" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "either", "ff_ext", @@ -2469,7 +2469,7 @@ dependencies = [ [[package]] name = "p3" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "p3-air", "p3-baby-bear", @@ -2917,7 +2917,7 @@ dependencies = [ [[package]] name = "poseidon" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "ff_ext", "p3", @@ -3531,7 +3531,7 @@ dependencies = [ [[package]] name = "sp1-curves" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "cfg-if", "dashu", @@ -3657,7 +3657,7 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "sumcheck" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "either", "ff_ext", @@ -3675,7 +3675,7 @@ dependencies = [ [[package]] name = "sumcheck_macro" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "itertools 0.13.0", "p3", @@ -3919,7 +3919,7 @@ dependencies = [ [[package]] name = "transcript" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "ff_ext", "itertools 0.13.0", @@ -4122,7 +4122,7 @@ dependencies = [ [[package]] name = "whir" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "bincode", "clap", @@ -4353,7 +4353,7 @@ dependencies = [ [[package]] name = "witness" version = "0.1.0" -source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#092da2bccc8b5ff73ac0bf69927bb72163908854" +source = "git+https://github.com/scroll-tech/gkr-backend.git?branch=feat%2Fbump-p3#553f46f4a84673a9f6ee474d383c63585bb5c3ba" dependencies = [ "ff_ext", "multilinear_extensions", diff --git a/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/air.rs b/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/air.rs deleted file mode 100644 index 743d241f1..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/air.rs +++ /dev/null @@ -1,185 +0,0 @@ -use std::borrow::Borrow; - -use openvm_circuit_primitives::{SubAir, utils::assert_array_eq}; -use openvm_stark_backend::{ - BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, -}; -use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF; -use p3_air::{Air, AirBuilder, BaseAir}; -use p3_field::{PrimeCharacteristicRing, extension::BinomiallyExtendable}; -use p3_matrix::Matrix; -use stark_recursion_circuit_derive::AlignedBorrow; - -use crate::{ - batch_constraint::bus::{ - ConstraintsFoldingBus, ConstraintsFoldingMessage, EqNOuterBus, EqNOuterMessage, - ExpressionClaimBus, ExpressionClaimMessage, - }, - bus::{NLiftBus, NLiftMessage, TranscriptBus}, - subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir}, - utils::{ext_field_add, ext_field_multiply, ext_field_multiply_scalar}, -}; - -#[derive(AlignedBorrow, Copy, Clone)] -#[repr(C)] -pub struct ConstraintsFoldingCols { - pub is_valid: T, - pub is_first: T, - pub proof_idx: T, - - pub air_idx: T, - pub sort_idx: T, - pub constraint_idx: T, - pub n_lift: T, - - pub lambda_tidx: T, - pub lambda: [T; D_EF], - - pub value: [T; D_EF], - pub cur_sum: [T; D_EF], - pub eq_n: [T; D_EF], - - pub is_first_in_air: T, -} - -pub struct ConstraintsFoldingAir { - pub transcript_bus: TranscriptBus, - pub constraint_bus: ConstraintsFoldingBus, - pub expression_claim_bus: ExpressionClaimBus, - pub eq_n_outer_bus: EqNOuterBus, - pub n_lift_bus: NLiftBus, -} - -impl BaseAirWithPublicValues for ConstraintsFoldingAir {} -impl PartitionedBaseAir for ConstraintsFoldingAir {} - -impl BaseAir for ConstraintsFoldingAir { - fn width(&self) -> usize { - ConstraintsFoldingCols::::width() - } -} - -impl Air for ConstraintsFoldingAir -where - ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, -{ - fn eval(&self, builder: &mut AB) { - let main = builder.main(); - let (local, next) = ( - main.row_slice(0).expect("window should have two elements"), - main.row_slice(1).expect("window should have two elements"), - ); - - let local: &ConstraintsFoldingCols = (*local).borrow(); - let next: &ConstraintsFoldingCols = (*next).borrow(); - - type LoopSubAir = NestedForLoopSubAir<2>; - LoopSubAir {}.eval( - builder, - ( - NestedForLoopIoCols { - is_enabled: local.is_valid, - counter: [local.proof_idx, local.sort_idx], - is_first: [local.is_first, local.is_first_in_air], - } - .map_into(), - NestedForLoopIoCols { - is_enabled: next.is_valid, - counter: [next.proof_idx, next.sort_idx], - is_first: [next.is_first, next.is_first_in_air], - } - .map_into(), - ), - ); - - let is_same_proof = next.is_valid - next.is_first; - let is_same_air = next.is_valid - next.is_first_in_air; - - // =========================== indices consistency =============================== - // When we are within one air, constraint_idx increases by 0/1 - builder - .when(is_same_air.clone()) - .assert_bool(next.constraint_idx - local.constraint_idx); - // First constraint_idx within an air is zero - builder - .when(local.is_first_in_air) - .assert_zero(local.constraint_idx); - builder - .when(is_same_air.clone()) - .assert_eq(local.n_lift, next.n_lift); - - // ======================== lambda and cur sum consistency ============================ - assert_array_eq(&mut builder.when(is_same_proof), local.lambda, next.lambda); - assert_array_eq( - &mut builder.when(is_same_air.clone()), - local.cur_sum, - ext_field_add( - local.value, - ext_field_multiply::(local.lambda, next.cur_sum), - ), - ); - assert_array_eq( - &mut builder.when(is_same_air.clone()), - local.eq_n, - next.eq_n, - ); - // numerator and the last element of the message are just the corresponding values - assert_array_eq( - &mut builder.when(AB::Expr::ONE - is_same_air.clone()), - local.cur_sum, - local.value, - ); - - self.n_lift_bus.receive( - builder, - local.proof_idx, - NLiftMessage { - air_idx: local.air_idx, - n_lift: local.n_lift, - }, - local.is_first_in_air * local.is_valid, - ); - self.constraint_bus.receive( - builder, - local.proof_idx, - ConstraintsFoldingMessage { - air_idx: local.air_idx.into(), - constraint_idx: local.constraint_idx - AB::Expr::ONE, - value: local.value.map(Into::into), - }, - local.is_valid * (AB::Expr::ONE - local.is_first_in_air), - ); - let folded_sum: [AB::Expr; D_EF] = ext_field_add( - ext_field_multiply_scalar::(next.cur_sum, is_same_air.clone()), - ext_field_multiply_scalar::(local.cur_sum, AB::Expr::ONE - is_same_air), - ); - self.expression_claim_bus.send( - builder, - local.proof_idx, - ExpressionClaimMessage { - is_interaction: AB::Expr::ZERO, - idx: local.sort_idx.into(), - value: ext_field_multiply(folded_sum, local.eq_n), - }, - local.is_first_in_air * local.is_valid, - ); - self.transcript_bus.sample_ext( - builder, - local.proof_idx, - local.lambda_tidx, - local.lambda, - local.is_valid * local.is_first, - ); - - self.eq_n_outer_bus.lookup_key( - builder, - local.proof_idx, - EqNOuterMessage { - is_sharp: AB::Expr::ZERO, - n: local.n_lift.into(), - value: local.eq_n.map(Into::into), - }, - local.is_first_in_air * local.is_valid, - ); - } -} diff --git a/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/mod.rs b/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/mod.rs deleted file mode 100644 index 26ed4a40d..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -mod air; -mod trace; - -pub use air::*; -pub use trace::*; diff --git a/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/trace.rs b/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/trace.rs deleted file mode 100644 index 81b6805eb..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/trace.rs +++ /dev/null @@ -1,406 +0,0 @@ -use std::borrow::BorrowMut; - -use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, EF, F}; -use p3_field::{BasedVectorSpace, PrimeCharacteristicRing}; -use p3_matrix::dense::RowMajorMatrix; -use p3_maybe_rayon::prelude::*; - -use crate::{ - batch_constraint::expr_eval::constraints_folding::air::ConstraintsFoldingCols, - system::{Preflight, RecursionVk}, - tracegen::RowMajorChip, - utils::{MultiProofVecVec, MultiVecWithBounds}, -}; - -#[derive(Copy, Clone)] -#[repr(C)] -pub(crate) struct ConstraintsFoldingRecord { - sort_idx: usize, - air_idx: usize, - constraint_idx: usize, - node_idx: usize, - is_first_in_air: bool, - value: EF, -} - -pub(crate) struct ConstraintsFoldingBlob { - pub(crate) records: MultiProofVecVec, - // (n, value), n is before lift, can be negative - #[allow(dead_code)] - pub(crate) folded_claims: MultiProofVecVec<(isize, EF)>, -} - -impl ConstraintsFoldingBlob { - #[allow(dead_code)] - pub fn new( - child_vk: &RecursionVk, - expr_evals: &MultiVecWithBounds, - preflights: &[&Preflight], - ) -> Self { - let mut max_air_idx = 0usize; - for key in child_vk.circuit_index_to_name.keys().copied() { - max_air_idx = max_air_idx.max(key); - } - let mut constraints = vec![Vec::::new(); max_air_idx + 1]; - for (&air_idx, name) in &child_vk.circuit_index_to_name { - let expr_len = child_vk - .circuit_vks - .get(name) - .and_then(|vk| vk.cs.gkr_circuit.as_ref()) - .and_then(|circuit| circuit.layers.first()) - .map(|layer| layer.exprs.len()) - .unwrap_or_default(); - if air_idx >= constraints.len() { - constraints.resize(air_idx + 1, vec![]); - } - constraints[air_idx] = (0..expr_len).collect(); - } - - let mut records = MultiProofVecVec::new(); - let mut folded = MultiProofVecVec::new(); - for (pidx, preflight) in preflights.iter().enumerate() { - let lambda_tidx = preflight.batch_constraint.lambda_tidx; - let lambda = EF::from_basis_coefficients_slice( - &preflight.transcript.values()[lambda_tidx..lambda_tidx + D_EF], - ) - .unwrap(); - - let vdata = &preflight.proof_shape.sorted_trace_vdata; - for (sort_idx, (air_idx, v)) in vdata.iter().enumerate() { - let constrs = &constraints[*air_idx]; - records.push(ConstraintsFoldingRecord { - // dummy to avoid handling case with no constraints - sort_idx, - air_idx: *air_idx, - constraint_idx: 0, - node_idx: 0, - is_first_in_air: true, - value: EF::ZERO, - }); - let mut folded_claim = EF::ZERO; - let mut lambda_pow = EF::ONE; - for (constraint_idx, &constr) in constrs.iter().enumerate() { - let value = expr_evals[[pidx, *air_idx]][constr]; - folded_claim += lambda_pow * value; - lambda_pow *= lambda; - records.push(ConstraintsFoldingRecord { - sort_idx, - air_idx: *air_idx, - constraint_idx: constraint_idx + 1, - node_idx: constr, - is_first_in_air: false, - value, - }); - } - let l_skip = preflight.proof_shape.l_skip; - let n_lift = v.log_height.saturating_sub(l_skip); - let n = v.log_height as isize - l_skip as isize; - folded.push(( - n, - folded_claim * preflight.batch_constraint.eq_ns_frontloaded[n_lift], - )); - } - records.end_proof(); - folded.end_proof(); - } - Self { - records, - folded_claims: folded, - } - } -} - -pub struct ConstraintsFoldingTraceGenerator; - -impl RowMajorChip for ConstraintsFoldingTraceGenerator { - type Ctx<'a> = (&'a ConstraintsFoldingBlob, &'a [&'a Preflight]); - - #[tracing::instrument(level = "trace", skip_all)] - fn generate_trace( - &self, - ctx: &Self::Ctx<'_>, - required_height: Option, - ) -> Option> { - let (blob, preflights) = ctx; - let width = ConstraintsFoldingCols::::width(); - - let total_height = blob.records.len(); - debug_assert!(total_height > 0); - let padding_height = if let Some(height) = required_height { - if height < total_height { - return None; - } - height - } else { - total_height.next_power_of_two() - }; - let mut trace = vec![F::ZERO; padding_height * width]; - - let mut cur_height = 0; - for (pidx, preflight) in preflights.iter().enumerate() { - let lambda_tidx = preflight.batch_constraint.lambda_tidx; - let lambda_slice = &preflight.transcript.values()[lambda_tidx..lambda_tidx + D_EF]; - let records = &blob.records[pidx]; - - trace[cur_height * width..(cur_height + records.len()) * width] - .par_chunks_exact_mut(width) - .zip(records.par_iter()) - .for_each(|(chunk, record)| { - let cols: &mut ConstraintsFoldingCols<_> = chunk.borrow_mut(); - let n_lift = preflight.proof_shape.sorted_trace_vdata[record.sort_idx] - .1 - .log_height - .saturating_sub(preflight.proof_shape.l_skip); - - cols.is_valid = F::ONE; - cols.proof_idx = F::from_usize(pidx); - cols.air_idx = F::from_usize(record.air_idx); - cols.sort_idx = F::from_usize(record.sort_idx); - cols.constraint_idx = F::from_usize(record.constraint_idx); - cols.n_lift = F::from_usize(n_lift); - cols.lambda_tidx = F::from_usize(lambda_tidx); - cols.lambda.copy_from_slice(lambda_slice); - cols.value - .copy_from_slice(record.value.as_basis_coefficients_slice()); - cols.eq_n.copy_from_slice( - preflight.batch_constraint.eq_ns_frontloaded[n_lift] - .as_basis_coefficients_slice(), - ); - cols.is_first_in_air = F::from_bool(record.is_first_in_air); - }); - - // Setting `cur_sum` - let mut cur_sum = EF::ZERO; - let lambda = EF::from_basis_coefficients_slice(lambda_slice).unwrap(); - trace[cur_height * width..(cur_height + records.len()) * width] - .chunks_exact_mut(width) - .rev() - .for_each(|chunk| { - let cols: &mut ConstraintsFoldingCols<_> = chunk.borrow_mut(); - cur_sum = - cur_sum * lambda + EF::from_basis_coefficients_slice(&cols.value).unwrap(); - cols.cur_sum - .copy_from_slice(cur_sum.as_basis_coefficients_slice()); - if cols.is_first_in_air == F::ONE { - cur_sum = EF::ZERO; - } - }); - - { - let cols: &mut ConstraintsFoldingCols<_> = - trace[cur_height * width..(cur_height + 1) * width].borrow_mut(); - cols.is_first = F::ONE; - } - cur_height += records.len(); - } - Some(RowMajorMatrix::new(trace, width)) - } -} - -#[cfg(feature = "cuda")] -pub(in crate::batch_constraint) mod cuda { - use openvm_circuit_primitives::cuda_abi::UInt2; - use openvm_cuda_backend::{GpuBackend, base::DeviceMatrix}; - use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer}; - use openvm_stark_backend::prover::AirProvingContext; - - use super::*; - use crate::{ - batch_constraint::cuda_abi::{ - AffineFpExt, FpExtWithTidx, constraints_folding_tracegen, - constraints_folding_tracegen_temp_bytes, - }, - cuda::{preflight::PreflightGpu, vk::VerifyingKeyGpu}, - tracegen::ModuleChip, - }; - - pub struct ConstraintsFoldingBlobGpu { - // Per proof, per AIR, per constraint - pub values: Vec>>, - // Per proof - pub constraints_folding_per_proof: Vec, - // For compatibility with CPU tracegen - pub folded_claims: MultiProofVecVec<(isize, EF)>, - } - - impl ConstraintsFoldingBlobGpu { - pub fn new( - vk: &VerifyingKeyGpu, - expr_evals: &MultiVecWithBounds, - preflights: &[PreflightGpu], - ) -> Self { - let constraints = vk - .cpu - .inner - .per_air - .iter() - .map(|vk| vk.symbolic_constraints.constraints.constraint_idx.clone()) - .collect_vec(); - - let mut values = Vec::with_capacity(preflights.len()); - let mut constraints_folding_per_proof = Vec::with_capacity(preflights.len()); - let mut folded_claims = MultiProofVecVec::new(); - - for (pidx, preflight) in preflights.iter().enumerate() { - let lambda_tidx = preflight.cpu.batch_constraint.lambda_tidx; - let lambda = EF::from_basis_coefficients_slice( - &preflight.cpu.transcript.values()[lambda_tidx..lambda_tidx + D_EF], - ) - .unwrap(); - - let vdata = &preflight.cpu.proof_shape.sorted_trace_vdata; - let mut proof_values = Vec::with_capacity(vdata.len()); - - for (air_idx, v) in vdata.iter() { - let mut folded_claim = EF::ZERO; - let mut lambda_pow = EF::ONE; - - let air_values = std::iter::once(EF::ZERO) - .chain(constraints[*air_idx].iter().map(|&constr| { - let value = expr_evals[[pidx, *air_idx]][constr]; - folded_claim += lambda_pow * value; - lambda_pow *= lambda; - value - })) - .collect_vec(); - proof_values.push(air_values); - - let n_lift = v.log_height.saturating_sub(vk.system_params.l_skip); - let n = v.log_height as isize - vk.system_params.l_skip as isize; - folded_claims.push(( - n, - folded_claim * preflight.cpu.batch_constraint.eq_ns_frontloaded[n_lift], - )); - } - - values.push(proof_values); - constraints_folding_per_proof.push(FpExtWithTidx { - value: lambda, - tidx: lambda_tidx as u32, - }); - folded_claims.end_proof(); - } - - Self { - values, - constraints_folding_per_proof, - folded_claims, - } - } - } - - impl ModuleChip for ConstraintsFoldingTraceGenerator { - type Ctx<'a> = ( - &'a VerifyingKeyGpu, - &'a [PreflightGpu], - &'a ConstraintsFoldingBlobGpu, - ); - - #[tracing::instrument(level = "trace", skip_all)] - fn generate_proving_ctx( - &self, - ctx: &Self::Ctx<'_>, - required_height: Option, - ) -> Option> { - let (child_vk, preflights_gpu, blob) = ctx; - - let mut num_valid_rows = 0u32; - let mut row_bounds = Vec::with_capacity(preflights_gpu.len()); - let mut constraint_bounds = Vec::with_capacity(preflights_gpu.len()); - let mut proof_and_sort_idxs = vec![]; - - let flat_values = blob - .values - .iter() - .enumerate() - .flat_map(|(proof_idx, proof_values)| { - let mut num_constraints_in_proof = 0; - let mut proof_constraint_bounds = Vec::with_capacity(proof_values.len()); - for (sort_idx, air_values) in proof_values.iter().enumerate() { - let num_constraints = air_values.len(); - num_constraints_in_proof += num_constraints as u32; - proof_constraint_bounds.push(num_constraints_in_proof); - proof_and_sort_idxs.extend(std::iter::repeat_n( - UInt2 { - x: proof_idx as u32, - y: sort_idx as u32, - }, - num_constraints, - )); - } - num_valid_rows += num_constraints_in_proof; - row_bounds.push(num_valid_rows); - constraint_bounds.push(proof_constraint_bounds.to_device().unwrap()); - proof_values.iter().flatten().copied() - }) - .collect_vec(); - let eq_ns = preflights_gpu - .iter() - .map(|preflight| { - preflight - .cpu - .batch_constraint - .eq_ns_frontloaded - .to_device() - .unwrap() - }) - .collect_vec(); - - let height = if let Some(height) = required_height { - if height < num_valid_rows as usize { - return None; - } - height - } else { - (num_valid_rows as usize).next_power_of_two() - }; - let width = ConstraintsFoldingCols::::width(); - let d_trace = DeviceMatrix::::with_capacity(height, width); - - let d_proof_and_sort_idxs = proof_and_sort_idxs.to_device().unwrap(); - let d_values = flat_values.to_device().unwrap(); - let d_cur_sum_evals = DeviceBuffer::::with_capacity(d_values.len()); - - let d_constraint_bounds = constraint_bounds.iter().map(|b| b.as_ptr()).collect_vec(); - let d_sorted_trace_heights = preflights_gpu - .iter() - .map(|preflight| preflight.proof_shape.sorted_trace_heights.as_ptr()) - .collect_vec(); - let d_eq_ns = eq_ns.iter().map(|b| b.as_ptr()).collect_vec(); - - let d_per_proof = blob.constraints_folding_per_proof.to_device().unwrap(); - - unsafe { - let temp_bytes = constraints_folding_tracegen_temp_bytes( - &d_proof_and_sort_idxs, - &d_cur_sum_evals, - num_valid_rows, - ) - .unwrap(); - let d_temp_buffer = DeviceBuffer::::with_capacity(temp_bytes); - constraints_folding_tracegen( - d_trace.buffer(), - height, - width, - &d_proof_and_sort_idxs, - &d_cur_sum_evals, - &d_values, - &row_bounds, - d_constraint_bounds, - d_sorted_trace_heights, - d_eq_ns, - &d_per_proof, - preflights_gpu.len() as u32, - child_vk.per_air.len() as u32, - num_valid_rows, - child_vk.system_params.l_skip as u32, - &d_temp_buffer, - temp_bytes, - ) - .unwrap(); - } - - Some(AirProvingContext::simple_no_pis(d_trace)) - } - } -} diff --git a/ceno_recursion_v2/src/batch_constraint/expr_eval/mod.rs b/ceno_recursion_v2/src/batch_constraint/expr_eval/mod.rs deleted file mode 100644 index 996852018..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expr_eval/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod constraints_folding; -pub mod symbolic_expression; - -pub use constraints_folding::*; -pub use symbolic_expression::*; diff --git a/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/air.rs b/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/air.rs deleted file mode 100644 index 3a4d59054..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/air.rs +++ /dev/null @@ -1,408 +0,0 @@ -use core::array; -use std::borrow::Borrow; - -use openvm_circuit_primitives::{encoder::Encoder, utils::assert_array_eq}; -use openvm_stark_backend::{ - BaseAirWithPublicValues, PartitionedBaseAir, air_builders::PartitionedAirBuilder, - interaction::InteractionBuilder, -}; -use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF; -use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir}; -use p3_field::{Field, PrimeCharacteristicRing, extension::BinomiallyExtendable}; -use p3_matrix::Matrix; -use stark_recursion_circuit_derive::AlignedBorrow; -use strum::{EnumCount, IntoEnumIterator}; -use strum_macros::EnumIter; - -use crate::{ - batch_constraint::bus::{ - ConstraintsFoldingBus, ConstraintsFoldingMessage, InteractionsFoldingBus, - InteractionsFoldingMessage, SymbolicExpressionBus, SymbolicExpressionMessage, - }, - bus::{ - AirPresenceBus, AirPresenceBusMessage, AirShapeBus, AirShapeBusMessage, ColumnClaimsBus, - ColumnClaimsMessage, HyperdimBus, HyperdimBusMessage, PublicValuesBus, - PublicValuesBusMessage, SelHypercubeBus, SelHypercubeBusMessage, SelUniBus, - SelUniBusMessage, - }, - proof_shape::bus::AirShapeProperty, - utils::{ - base_to_ext, ext_field_add, ext_field_multiply, ext_field_multiply_scalar, - ext_field_subtract, scalar_subtract_ext_field, - }, -}; - -pub const NUM_FLAGS: usize = 5; -pub const ENCODER_MAX_DEGREE: u32 = 2; - -#[derive(Default, Debug, Clone, Copy, EnumIter, EnumCount)] -pub enum NodeKind { - #[default] - WitIn = 0, - StructuralWitIn = 1, - Fixed = 2, - Instance = 3, - SelIsFirst = 4, - SelIsLast = 5, - SelIsTransition = 6, - Constant = 7, - Add = 8, - Sub = 9, - Neg = 10, - Mul = 11, - InteractionMult = 12, - InteractionMsgComp = 13, - InteractionBusIndex = 14, -} - -#[derive(AlignedBorrow, Copy, Clone)] -#[repr(C)] -pub struct CachedSymbolicExpressionColumns { - pub flags: [T; NUM_FLAGS], - pub air_idx: T, - pub node_or_interaction_idx: T, - pub attrs: [T; 3], - pub fanout: T, - pub is_constraint: T, - pub constraint_idx: T, -} - -#[derive(AlignedBorrow, Copy, Clone)] -#[repr(C)] -pub struct SingleMainSymbolicExpressionColumns { - /// 0 = absent proof, 1 = proof present but air absent, 2 = proof+air present. - pub slot_state: T, - pub args: [T; 2 * D_EF], - pub sort_idx: T, - pub n_abs: T, - pub is_n_neg: T, -} - -pub struct SymbolicExpressionAir { - pub expr_bus: SymbolicExpressionBus, - pub hyperdim_bus: HyperdimBus, - pub air_shape_bus: AirShapeBus, - pub air_presence_bus: AirPresenceBus, - pub column_claims_bus: ColumnClaimsBus, - pub interactions_folding_bus: InteractionsFoldingBus, - pub constraints_folding_bus: ConstraintsFoldingBus, - pub public_values_bus: PublicValuesBus, - pub sel_hypercube_bus: SelHypercubeBus, - pub sel_uni_bus: SelUniBus, - - pub cnt_proofs: usize, -} - -impl BaseAirWithPublicValues for SymbolicExpressionAir {} - -impl PartitionedBaseAir for SymbolicExpressionAir { - fn cached_main_widths(&self) -> Vec { - vec![CachedSymbolicExpressionColumns::::width()] - } - - fn common_main_width(&self) -> usize { - SingleMainSymbolicExpressionColumns::::width() * self.cnt_proofs - } -} - -impl BaseAir for SymbolicExpressionAir { - fn width(&self) -> usize { - CachedSymbolicExpressionColumns::::width() - + SingleMainSymbolicExpressionColumns::::width() * self.cnt_proofs - } -} - -impl Air - for SymbolicExpressionAir -where - ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, -{ - fn eval(&self, builder: &mut AB) { - let cached_local = builder.cached_mains()[0] - .row_slice(0) - .expect("cached window should have a row") - .to_vec(); - let main_local = builder - .common_main() - .row_slice(0) - .expect("main window should have a row") - .to_vec(); - let main_next = builder - .common_main() - .row_slice(1) - .expect("main window should have two rows") - .to_vec(); - - let cached_cols: &CachedSymbolicExpressionColumns = - cached_local.as_slice().borrow(); - let main_cols: Vec<&SingleMainSymbolicExpressionColumns> = main_local - .chunks(SingleMainSymbolicExpressionColumns::::width()) - .map(|chunk| chunk.borrow()) - .collect(); - let next_main_cols: Vec<&SingleMainSymbolicExpressionColumns> = main_next - .chunks(SingleMainSymbolicExpressionColumns::::width()) - .map(|chunk| chunk.borrow()) - .collect(); - - let enc = Encoder::new(NodeKind::COUNT, ENCODER_MAX_DEGREE, true); - let flags = cached_cols.flags; - let is_valid_row = enc.is_valid::(&flags); - - let is_arg0_node_idx = enc.contains_flag::( - &flags, - &[ - NodeKind::Add, - NodeKind::Sub, - NodeKind::Mul, - NodeKind::Neg, - NodeKind::InteractionMult, - NodeKind::InteractionMsgComp, - NodeKind::WitIn, - NodeKind::StructuralWitIn, - NodeKind::Fixed, - NodeKind::Instance, - ] - .map(|x| x as usize), - ); - let is_arg1_node_idx = enc.contains_flag::( - &flags, - &[ - NodeKind::Add, - NodeKind::Sub, - NodeKind::Mul, - NodeKind::InteractionMsgComp, - ] - .map(|x| x as usize), - ); - - for (proof_idx, (&cols, &next_cols)) in main_cols.iter().zip(&next_main_cols).enumerate() { - let proof_idx = AB::F::from_usize(proof_idx); - - let slot_state: AB::Expr = cols.slot_state.into(); - let next_slot_state: AB::Expr = next_cols.slot_state.into(); - let proof_present = slot_state.clone() - * (AB::Expr::from_u8(3) - slot_state.clone()) - * AB::F::TWO.inverse(); - let next_proof_present = next_slot_state.clone() - * (AB::Expr::from_u8(3) - next_slot_state) - * AB::F::TWO.inverse(); - let air_present = - slot_state.clone() * (slot_state.clone() - AB::Expr::ONE) * AB::F::TWO.inverse(); - - let arg_ef0: [AB::Var; D_EF] = cols.args[..D_EF].try_into().unwrap(); - let arg_ef1: [AB::Var; D_EF] = cols.args[D_EF..2 * D_EF].try_into().unwrap(); - - builder.assert_tern(cols.slot_state); - builder - .when(cols.is_n_neg) - .assert_eq(cols.slot_state, AB::Expr::TWO); - builder - .when(air_present.clone()) - .assert_one(is_valid_row.clone()); - builder - .when_transition() - .assert_eq(proof_present.clone(), next_proof_present); - - let mut value = [AB::Expr::ZERO; D_EF]; - for node_kind in NodeKind::iter() { - let sel = enc.get_flag_expr::(node_kind as usize, &flags); - let expr = match node_kind { - NodeKind::Add => ext_field_add::(arg_ef0, arg_ef1), - NodeKind::Sub => ext_field_subtract::(arg_ef0, arg_ef1), - NodeKind::Neg => scalar_subtract_ext_field::(AB::Expr::ZERO, arg_ef0), - NodeKind::Mul => ext_field_multiply::(arg_ef0, arg_ef1), - NodeKind::Constant => base_to_ext(cached_cols.attrs[0]), - NodeKind::Instance => base_to_ext(cols.args[0]), - NodeKind::SelIsFirst => ext_field_multiply(arg_ef0, arg_ef1), - NodeKind::SelIsLast => ext_field_multiply(arg_ef0, arg_ef1), - NodeKind::SelIsTransition => scalar_subtract_ext_field( - AB::Expr::ONE, - ext_field_multiply(arg_ef0, arg_ef1), - ), - NodeKind::WitIn - | NodeKind::StructuralWitIn - | NodeKind::Fixed - | NodeKind::InteractionMult - | NodeKind::InteractionMsgComp => arg_ef0.map(Into::into), - NodeKind::InteractionBusIndex => { - base_to_ext(cached_cols.attrs[0] + AB::Expr::ONE) - } - }; - value = ext_field_add::( - value, - ext_field_multiply_scalar::(expr, sel), - ); - } - - self.expr_bus.add_key_with_lookups( - builder, - proof_idx, - SymbolicExpressionMessage { - air_idx: cached_cols.air_idx.into(), - node_idx: cached_cols.node_or_interaction_idx.into(), - value: value.clone(), - }, - air_present.clone() * cached_cols.fanout, - ); - self.expr_bus.lookup_key( - builder, - proof_idx, - SymbolicExpressionMessage { - air_idx: cached_cols.air_idx, - node_idx: cached_cols.attrs[0], - value: arg_ef0, - }, - air_present.clone() * is_arg0_node_idx.clone(), - ); - self.expr_bus.lookup_key( - builder, - proof_idx, - SymbolicExpressionMessage { - air_idx: cached_cols.air_idx, - node_idx: cached_cols.attrs[1], - value: arg_ef1, - }, - air_present.clone() * is_arg1_node_idx.clone(), - ); - - let is_var = enc.contains_flag::( - &flags, - &[NodeKind::WitIn, NodeKind::StructuralWitIn, NodeKind::Fixed].map(|x| x as usize), - ); - self.column_claims_bus.receive( - builder, - proof_idx, - ColumnClaimsMessage { - sort_idx: cols.sort_idx.into(), - part_idx: cached_cols.attrs[1].into(), - col_idx: cached_cols.attrs[0].into(), - claim: array::from_fn(|i| cols.args[i].into()), - is_rot: cached_cols.attrs[2].into(), - }, - is_var * air_present.clone(), - ); - self.public_values_bus.receive( - builder, - proof_idx, - PublicValuesBusMessage { - air_idx: cached_cols.air_idx, - pv_idx: cached_cols.attrs[0], - value: cols.args[0], - }, - enc.get_flag_expr::(NodeKind::Instance as usize, &flags) * air_present.clone(), - ); - self.air_shape_bus.lookup_key( - builder, - proof_idx, - AirShapeBusMessage { - sort_idx: cols.sort_idx.into(), - property_idx: AirShapeProperty::AirId.to_field(), - value: cached_cols.air_idx.into(), - }, - air_present.clone(), - ); - self.air_presence_bus.lookup_key( - builder, - proof_idx, - AirPresenceBusMessage { - air_idx: cached_cols.air_idx.into(), - is_present: air_present.clone(), - }, - proof_present * is_valid_row.clone(), - ); - self.hyperdim_bus.lookup_key( - builder, - proof_idx, - HyperdimBusMessage { - sort_idx: cols.sort_idx, - n_abs: cols.n_abs, - n_sign_bit: cols.is_n_neg, - }, - air_present.clone(), - ); - - let is_sel = enc.contains_flag::( - &flags, - &[ - NodeKind::SelIsFirst, - NodeKind::SelIsLast, - NodeKind::SelIsTransition, - ] - .map(|x| x as usize), - ); - let is_first = enc.get_flag_expr::(NodeKind::SelIsFirst as usize, &flags); - self.sel_uni_bus.lookup_key( - builder, - proof_idx, - SelUniBusMessage { - n: AB::Expr::NEG_ONE * cols.n_abs * cols.is_n_neg, - is_first: is_first.clone(), - value: arg_ef0.map(Into::into), - }, - air_present.clone() * is_sel.clone(), - ); - self.sel_hypercube_bus.lookup_key( - builder, - proof_idx, - SelHypercubeBusMessage { - n: cols.n_abs.into(), - is_first: is_first.clone(), - value: arg_ef1.map(Into::into), - }, - is_sel.clone() * (air_present.clone() - cols.is_n_neg), - ); - assert_array_eq( - &mut builder.when(is_sel.clone() * cols.is_n_neg), - arg_ef1, - [ - AB::Expr::ONE, - AB::Expr::ZERO, - AB::Expr::ZERO, - AB::Expr::ZERO, - ], - ); - - let is_mult = enc.get_flag_expr::(NodeKind::InteractionMult as usize, &flags); - let is_bus_index = - enc.get_flag_expr::(NodeKind::InteractionBusIndex as usize, &flags); - let is_interaction = enc.contains_flag::( - &flags, - &[NodeKind::InteractionMult, NodeKind::InteractionMsgComp].map(|x| x as usize), - ); - self.interactions_folding_bus.send( - builder, - proof_idx, - InteractionsFoldingMessage { - air_idx: cached_cols.air_idx.into(), - interaction_idx: cached_cols.node_or_interaction_idx.into(), - is_mult, - idx_in_message: cached_cols.attrs[1].into(), - value: value.clone(), - }, - is_interaction * air_present.clone(), - ); - self.interactions_folding_bus.send( - builder, - proof_idx, - InteractionsFoldingMessage { - air_idx: cached_cols.air_idx.into(), - interaction_idx: cached_cols.node_or_interaction_idx.into(), - is_mult: AB::Expr::ZERO, - idx_in_message: AB::Expr::NEG_ONE, - value: value.clone(), - }, - is_bus_index * air_present.clone(), - ); - self.constraints_folding_bus.send( - builder, - proof_idx, - ConstraintsFoldingMessage { - air_idx: cached_cols.air_idx.into(), - constraint_idx: cached_cols.constraint_idx.into(), - value: value.clone(), - }, - cached_cols.is_constraint * air_present, - ); - } - } -} diff --git a/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/mod.rs b/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/mod.rs deleted file mode 100644 index c68123602..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub(crate) mod air; -pub(crate) mod trace; - -pub use air::*; -pub use trace::*; diff --git a/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/trace.rs b/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/trace.rs deleted file mode 100644 index cbb141d60..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/trace.rs +++ /dev/null @@ -1,513 +0,0 @@ -use core::cmp::min; -use std::borrow::BorrowMut; - -use openvm_circuit_primitives::encoder::Encoder; -use openvm_stark_backend::{ - air_builders::symbolic::{SymbolicExpressionNode, symbolic_variable::Entry}, - poly_common::{Squarable, eval_eq_uni_at_one}, -}; -use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, EF, F}; -use p3_field::{BasedVectorSpace, PrimeCharacteristicRing, TwoAdicField}; -use p3_matrix::dense::RowMajorMatrix; -use p3_maybe_rayon::prelude::*; -use strum::EnumCount; - -use crate::{ - batch_constraint::expr_eval::symbolic_expression::air::{ - CachedSymbolicExpressionColumns, ENCODER_MAX_DEGREE, NodeKind, - SingleMainSymbolicExpressionColumns, - }, - system::{Preflight, RecursionField, RecursionVk, convert_vk_from_zkvm}, - tracegen::RowMajorChip, - utils::{MultiVecWithBounds, interaction_length}, -}; -use ceno_zkvm::structs::ComposedConstrainSystem; -use multilinear_extensions::{Expression, Fixed}; - -pub struct SymbolicExpressionTraceGenerator { - pub max_num_proofs: usize, -} - -pub struct SymbolicExpressionCtx<'a> { - pub vk: &'a RecursionVk, - pub preflights: &'a [&'a Preflight], - pub expr_evals: &'a MultiVecWithBounds, -} - -impl RowMajorChip for SymbolicExpressionTraceGenerator { - type Ctx<'a> = SymbolicExpressionCtx<'a>; - - #[tracing::instrument(level = "trace", skip_all)] - fn generate_trace( - &self, - ctx: &Self::Ctx<'_>, - required_height: Option, - ) -> Option> { - let child_vk = convert_vk_from_zkvm(ctx.vk); - let child_vk = child_vk.as_ref(); - let preflights = ctx.preflights; - let max_num_proofs = self.max_num_proofs; - let expr_evals = ctx.expr_evals; - let l_skip = child_vk.inner.params.l_skip; - - let single_main_width = SingleMainSymbolicExpressionColumns::::width(); - let main_width = single_main_width * max_num_proofs; - - struct Record { - args: [F; 2 * D_EF], - sort_idx: usize, - n_abs: usize, - is_n_neg: usize, - } - let mut records = vec![]; - - for (proof_idx, preflight) in preflights.iter().enumerate() { - let rs = &preflight.batch_constraint.sumcheck_rnd; - if rs.is_empty() { - continue; - } - let (&rs_0, rs_rest) = rs.split_first().unwrap(); - let mut is_first_uni_by_log_height = vec![]; - let mut is_last_uni_by_log_height = vec![]; - - for (log_height, &r_pow) in rs_0 - .exp_powers_of_2() - .take(l_skip + 1) - .collect::>() - .iter() - .rev() - .enumerate() - { - is_first_uni_by_log_height.push(eval_eq_uni_at_one(log_height, r_pow)); - is_last_uni_by_log_height.push(eval_eq_uni_at_one( - log_height, - r_pow * F::two_adic_generator(log_height), - )); - } - let mut is_first_mle_by_n = vec![EF::ONE]; - let mut is_last_mle_by_n = vec![EF::ONE]; - for (i, &r) in rs_rest.iter().enumerate() { - is_first_mle_by_n.push(is_first_mle_by_n[i] * (EF::ONE - r)); - is_last_mle_by_n.push(is_last_mle_by_n[i] * r); - } - - for (air_idx, vk) in child_vk.inner.per_air.iter().enumerate() { - let constraints = &vk.symbolic_constraints.constraints; - let expr_evals = &expr_evals[[proof_idx, air_idx]]; - - if expr_evals.is_empty() { - let n = constraints.nodes.len() - + vk.symbolic_constraints - .interactions - .iter() - .map(interaction_length) - .sum::() - + vk.unused_variables.len(); - records.resize_with(records.len() + n, || None); - continue; - } - - let (sort_idx, trace_vdata) = preflight - .proof_shape - .sorted_trace_vdata - .iter() - .enumerate() - .find_map(|(sort_idx, (idx, vdata))| { - (*idx == air_idx).then_some((sort_idx, vdata)) - }) - .unwrap(); - - let log_height = trace_vdata.log_height; - let (n_abs, is_n_neg) = if log_height < l_skip { - (l_skip - log_height, 1) - } else { - (log_height - l_skip, 0) - }; - - for (node_idx, node) in constraints.nodes.iter().enumerate() { - let mut record = Record { - args: [F::ZERO; 2 * D_EF], - sort_idx, - n_abs, - is_n_neg, - }; - match node { - SymbolicExpressionNode::Variable(var) => match var.entry { - Entry::Preprocessed { .. } | Entry::Main { .. } | Entry::Public => { - record.args[..D_EF].copy_from_slice( - expr_evals[node_idx].as_basis_coefficients_slice(), - ); - } - Entry::Challenge => unreachable!(), - }, - SymbolicExpressionNode::IsFirstRow => { - record.args[..D_EF].copy_from_slice( - is_first_uni_by_log_height[min(log_height, l_skip)] - .as_basis_coefficients_slice(), - ); - record.args[D_EF..2 * D_EF].copy_from_slice( - is_first_mle_by_n[log_height.saturating_sub(l_skip)] - .as_basis_coefficients_slice(), - ); - } - SymbolicExpressionNode::IsLastRow - | SymbolicExpressionNode::IsTransition => { - record.args[..D_EF].copy_from_slice( - is_last_uni_by_log_height[min(log_height, l_skip)] - .as_basis_coefficients_slice(), - ); - record.args[D_EF..2 * D_EF].copy_from_slice( - is_last_mle_by_n[log_height.saturating_sub(l_skip)] - .as_basis_coefficients_slice(), - ); - } - SymbolicExpressionNode::Constant(_) => {} - SymbolicExpressionNode::Add { - left_idx, - right_idx, - .. - } - | SymbolicExpressionNode::Sub { - left_idx, - right_idx, - .. - } - | SymbolicExpressionNode::Mul { - left_idx, - right_idx, - .. - } => { - record.args[..D_EF].copy_from_slice( - expr_evals[*left_idx].as_basis_coefficients_slice(), - ); - record.args[D_EF..2 * D_EF].copy_from_slice( - expr_evals[*right_idx].as_basis_coefficients_slice(), - ); - } - SymbolicExpressionNode::Neg { idx, .. } => { - record.args[..D_EF] - .copy_from_slice(expr_evals[*idx].as_basis_coefficients_slice()); - } - }; - records.push(Some(record)); - } - for interaction in &vk.symbolic_constraints.interactions { - let mut args = [F::ZERO; 2 * D_EF]; - args[..D_EF].copy_from_slice( - expr_evals[interaction.count].as_basis_coefficients_slice(), - ); - records.push(Some(Record { - args, - sort_idx, - n_abs, - is_n_neg, - })); - - for &node_idx in &interaction.message { - let mut args = [F::ZERO; 2 * D_EF]; - args[..D_EF] - .copy_from_slice(expr_evals[node_idx].as_basis_coefficients_slice()); - records.push(Some(Record { - args, - sort_idx, - n_abs, - is_n_neg, - })); - } - - args.fill(F::ZERO); - args[0] = F::from_u16(interaction.bus_index + 1); - records.push(Some(Record { - args, - sort_idx, - n_abs, - is_n_neg, - })); - } - - let mut node_idx = constraints.nodes.len(); - for unused_var in &vk.unused_variables { - if matches!(unused_var.entry, Entry::Public | Entry::Challenge) { - continue; - } - let mut args = [F::ZERO; 2 * D_EF]; - args[..D_EF] - .copy_from_slice(expr_evals[node_idx].as_basis_coefficients_slice()); - records.push(Some(Record { - args, - sort_idx, - n_abs, - is_n_neg, - })); - node_idx += 1; - } - } - } - - let num_valid_rows = records.len() / preflights.len(); - let height = if let Some(height) = required_height { - if height < num_valid_rows { - return None; - } - height - } else { - num_valid_rows.max(1).next_power_of_two() - }; - let mut main_trace = F::zero_vec(main_width * height); - main_trace - .par_chunks_exact_mut(main_width) - .enumerate() - .for_each(|(row_idx, row)| { - if row_idx >= num_valid_rows { - return; - } - for proof_idx in 0..max_num_proofs { - if proof_idx >= preflights.len() { - continue; - } - let record_idx = proof_idx * num_valid_rows + row_idx; - let Some(record) = records[record_idx].as_ref() else { - continue; - }; - let start = proof_idx * single_main_width; - let end = start + single_main_width; - let cols: &mut SingleMainSymbolicExpressionColumns<_> = - row[start..end].borrow_mut(); - cols.slot_state = F::from_u8(2); - cols.args = record.args; - cols.sort_idx = F::from_usize(record.sort_idx); - cols.n_abs = F::from_usize(record.n_abs); - cols.is_n_neg = F::from_usize(record.is_n_neg); - } - }); - - Some(RowMajorMatrix::new(main_trace, main_width)) - } -} - -#[derive(Debug, Clone, Copy, Default)] -pub struct CachedRecord { - pub kind: NodeKind, - pub air_idx: usize, - pub node_idx: usize, - pub attrs: [usize; 3], - pub is_constraint: bool, - pub constraint_idx: usize, - pub fanout: usize, -} - -#[derive(Debug, Clone, Default)] -pub struct CachedTraceRecord { - pub records: Vec, -} - -pub fn build_cached_trace_record(child_vk: &RecursionVk) -> CachedTraceRecord { - let mut records = Vec::new(); - for (&air_idx, circuit_name) in &child_vk.circuit_index_to_name { - let Some(circuit_vk) = child_vk.circuit_vks.get(circuit_name) else { - continue; - }; - let Some(gkr) = circuit_vk.cs.gkr_circuit.as_ref() else { - continue; - }; - let Some(layer) = gkr.layers.first() else { - continue; - }; - let counts = Counts::from_css(&circuit_vk.cs); - let offsets = Offsets::new(&counts); - let mut builder = AirBuilder::new(&mut records, air_idx); - push_base_nodes(&mut builder, &counts, &offsets); - for (constraint_idx, expr) in layer.exprs.iter().enumerate() { - let root_idx = build_expression_nodes(expr, &mut builder, &offsets); - builder.mark_constraint(root_idx, constraint_idx); - } - } - - CachedTraceRecord { records } -} - -fn push_base_nodes(builder: &mut AirBuilder<'_>, counts: &Counts, offsets: &Offsets) { - for local in 0..counts.num_witin { - let global = offsets.witin + local; - builder.push(NodeKind::WitIn, [global, local, 0]); - } - for local in 0..counts.num_structural_witin { - let global = offsets.structural + local; - builder.push(NodeKind::StructuralWitIn, [global, local, 0]); - } - for local in 0..counts.num_fixed { - let global = offsets.fixed + local; - builder.push(NodeKind::Fixed, [global, local, 0]); - } - for local in 0..counts.num_instance { - let global = offsets.instance + local; - builder.push(NodeKind::Instance, [global, local, 0]); - } -} - -struct Counts { - num_witin: usize, - num_structural_witin: usize, - num_fixed: usize, - num_instance: usize, -} - -impl Counts { - fn from_css(cs: &ComposedConstrainSystem) -> Self { - let css = &cs.zkvm_v1_css; - Self { - num_witin: css.num_witin as usize, - num_structural_witin: css.num_structural_witin as usize, - num_fixed: css.num_fixed, - num_instance: css.instance.len(), - } - } -} - -struct Offsets { - witin: usize, - structural: usize, - fixed: usize, - instance: usize, -} - -impl Offsets { - fn new(counts: &Counts) -> Self { - let witin = 0; - let structural = witin + counts.num_witin; - let fixed = structural + counts.num_structural_witin; - let instance = fixed + counts.num_fixed; - Self { - witin, - structural, - fixed, - instance, - } - } -} - -struct AirBuilder<'a> { - records: &'a mut Vec, - air_idx: usize, - air_start: usize, - next_local_idx: usize, -} - -impl<'a> AirBuilder<'a> { - fn new(records: &'a mut Vec, air_idx: usize) -> Self { - let air_start = records.len(); - Self { - records, - air_idx, - air_start, - next_local_idx: 0, - } - } - - fn push(&mut self, kind: NodeKind, attrs: [usize; 3]) -> usize { - let node_idx = self.next_local_idx; - self.next_local_idx += 1; - self.records.push(CachedRecord { - kind, - air_idx: self.air_idx, - node_idx, - attrs, - is_constraint: false, - constraint_idx: 0, - fanout: 0, - }); - node_idx - } - - fn bump_fanout(&mut self, local_idx: usize) { - let global_idx = self.air_start + local_idx; - if let Some(record) = self.records.get_mut(global_idx) { - record.fanout = record.fanout.saturating_add(1); - } - } - - fn mark_constraint(&mut self, local_idx: usize, constraint_idx: usize) { - let global_idx = self.air_start + local_idx; - if let Some(record) = self.records.get_mut(global_idx) { - record.is_constraint = true; - record.constraint_idx = constraint_idx; - } - } -} - -fn build_expression_nodes( - expr: &Expression, - builder: &mut AirBuilder<'_>, - offsets: &Offsets, -) -> usize { - match expr { - Expression::WitIn(id) => offsets.witin + (*id as usize), - Expression::StructuralWitIn(id, _) => offsets.structural + (*id as usize), - Expression::Fixed(Fixed(idx)) => offsets.fixed + *idx, - Expression::Instance(instance) | Expression::InstanceScalar(instance) => { - offsets.instance + instance.0 - } - Expression::Constant(_) => builder.push(NodeKind::Constant, [0, 0, 0]), - Expression::Challenge(ch_id, pow, _, _) => { - builder.push(NodeKind::Constant, [*ch_id as usize, *pow, 1]) - } - Expression::Sum(left, right) => { - let left_idx = build_expression_nodes(left, builder, offsets); - let right_idx = build_expression_nodes(right, builder, offsets); - builder.bump_fanout(left_idx); - builder.bump_fanout(right_idx); - builder.push(NodeKind::Add, [left_idx, right_idx, 0]) - } - Expression::Product(left, right) => { - let left_idx = build_expression_nodes(left, builder, offsets); - let right_idx = build_expression_nodes(right, builder, offsets); - builder.bump_fanout(left_idx); - builder.bump_fanout(right_idx); - builder.push(NodeKind::Mul, [left_idx, right_idx, 0]) - } - Expression::ScaledSum(x, a, b) => { - let mul_left = build_expression_nodes(a, builder, offsets); - let mul_right = build_expression_nodes(x, builder, offsets); - builder.bump_fanout(mul_left); - builder.bump_fanout(mul_right); - let mul_idx = builder.push(NodeKind::Mul, [mul_left, mul_right, 0]); - let b_idx = build_expression_nodes(b, builder, offsets); - builder.bump_fanout(mul_idx); - builder.bump_fanout(b_idx); - builder.push(NodeKind::Add, [mul_idx, b_idx, 0]) - } - } -} - -pub fn generate_symbolic_expr_cached_trace( - cached_trace_record: &CachedTraceRecord, -) -> RowMajorMatrix { - let encoder = Encoder::new(NodeKind::COUNT, ENCODER_MAX_DEGREE, true); - let cached_width = CachedSymbolicExpressionColumns::::width(); - let records = &cached_trace_record.records; - - let height = records.len().next_power_of_two(); - let mut cached_trace = F::zero_vec(cached_width * height); - cached_trace - .par_chunks_exact_mut(cached_width) - .zip(records) - .for_each(|(row, record)| { - let cols: &mut CachedSymbolicExpressionColumns<_> = row.borrow_mut(); - - for (i, x) in encoder - .get_flag_pt(record.kind as usize) - .into_iter() - .enumerate() - { - cols.flags[i] = F::from_u32(x); - } - cols.air_idx = F::from_usize(record.air_idx); - cols.node_or_interaction_idx = F::from_usize(record.node_idx); - cols.attrs = record.attrs.map(F::from_usize); - cols.is_constraint = F::from_bool(record.is_constraint); - cols.constraint_idx = F::from_usize(record.constraint_idx); - cols.fanout = F::from_usize(record.fanout); - }); - - RowMajorMatrix::new(cached_trace, cached_width) -} diff --git a/ceno_recursion_v2/src/batch_constraint/expression_claim/air.rs b/ceno_recursion_v2/src/batch_constraint/expression_claim/air.rs deleted file mode 100644 index 1c5bf2654..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expression_claim/air.rs +++ /dev/null @@ -1,222 +0,0 @@ -use std::borrow::Borrow; - -use openvm_circuit_primitives::utils::{assert_array_eq, not}; -use openvm_stark_backend::{ - BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, -}; -use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF; -use p3_air::{Air, AirBuilder, BaseAir}; -use p3_field::{PrimeCharacteristicRing, extension::BinomiallyExtendable}; -use p3_matrix::Matrix; -use stark_recursion_circuit_derive::AlignedBorrow; - -use crate::{ - batch_constraint::bus::{ - BatchConstraintConductorBus, BatchConstraintConductorMessage, - BatchConstraintInnerMessageType, EqNOuterBus, EqNOuterMessage, ExpressionClaimBus, - ExpressionClaimMessage, - }, - bus::{ - ExpressionClaimNMaxBus, ExpressionClaimNMaxMessage, HyperdimBus, HyperdimBusMessage, - MainExpressionClaimBus, MainExpressionClaimMessage, - }, - primitives::bus::{PowerCheckerBus, PowerCheckerBusMessage}, - utils::{base_to_ext, ext_field_add, ext_field_multiply, ext_field_multiply_scalar}, -}; - -/// For each proof, this AIR will receive 2t interaction claims and t constraint claims. -/// (2 interaction claims and 1 constraint claim per trace). -/// These values are folded (algebraic batching) with mu into a single value, which -/// should match the final sumcheck claim. -#[derive(AlignedBorrow, Copy, Clone, Debug)] -#[repr(C)] -pub struct ExpressionClaimCols { - pub is_valid: T, - pub is_first: T, - pub proof_idx: T, - - pub is_interaction: T, - /// Index within the proof, 0 ~ 2t-1 are interaction claims, 0~t-1 are constraint claims. - pub idx: T, - pub idx_parity: T, - pub trace_idx: T, - /// The received evaluation claim. Note that for interactions, this is without norm_factor and - /// eq_sharp_ns. These are interactions_evals (without norm_factor and eq_sharp_ns) and - /// constraint_evals in the rust verifier. - pub value: [T; D_EF], - /// Receive from eq_ns AIR - pub eq_sharp_ns: [T; D_EF], - - /// For folding with mu. - pub cur_sum: [T; D_EF], - pub mu: [T; D_EF], - pub multiplier: [T; D_EF], - - /// Need to know n as if n<0, we need to multiply some norm_factor. - pub n_abs: T, - pub n_abs_pow: T, - pub n_sign: T, - /// The round idx for final sumcheck claim. - pub num_multilinear_sumcheck_rounds: T, -} - -pub struct ExpressionClaimAir { - pub expression_claim_n_max_bus: ExpressionClaimNMaxBus, - pub expr_claim_bus: ExpressionClaimBus, - pub mu_bus: BatchConstraintConductorBus, - pub main_claim_bus: MainExpressionClaimBus, - pub eq_n_outer_bus: EqNOuterBus, - pub pow_checker_bus: PowerCheckerBus, - pub hyperdim_bus: HyperdimBus, -} - -impl BaseAirWithPublicValues for ExpressionClaimAir {} -impl PartitionedBaseAir for ExpressionClaimAir {} - -impl BaseAir for ExpressionClaimAir { - fn width(&self) -> usize { - ExpressionClaimCols::::width() - } -} - -impl Air for ExpressionClaimAir -where - ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, -{ - fn eval(&self, builder: &mut AB) { - let main = builder.main(); - let (local, next) = ( - main.row_slice(0).expect("window should have two elements"), - main.row_slice(1).expect("window should have two elements"), - ); - - let local: &ExpressionClaimCols = (*local).borrow(); - let next: &ExpressionClaimCols = (*next).borrow(); - - builder.assert_bool(local.is_valid); - builder.assert_bool(local.is_first); - builder.assert_bool(local.is_interaction); - builder.assert_bool(local.idx_parity); - builder.assert_bool(local.n_sign); - builder - .when(local.is_first) - .assert_one(local.is_interaction); - builder.when(local.is_first).assert_zero(local.idx_parity); - builder - .when(local.is_interaction) - .assert_eq(local.idx_parity + next.idx_parity, AB::Expr::ONE); - builder - .when(local.idx_parity) - .assert_one(local.is_interaction); - - // === cum sum folding === - // cur_sum = next_cur_sum * mu + value * multiplier - assert_array_eq( - &mut builder.when(local.is_valid * not(next.is_first)), - local.cur_sum, - ext_field_add::( - ext_field_multiply::(local.value, local.multiplier), - ext_field_multiply::(next.cur_sum, local.mu), - ), - ); - // multiplier = 1 if not interaction - assert_array_eq( - &mut builder.when(not(local.is_interaction)).when(local.is_valid), - local.multiplier, - base_to_ext::(AB::Expr::ONE), - ); - - // IF negative n and numerator - assert_array_eq( - &mut builder.when(local.n_sign * (local.is_interaction - local.idx_parity)), - ext_field_multiply_scalar::(local.multiplier, local.n_abs_pow), - local.eq_sharp_ns, - ); - // ELSE 1 - assert_array_eq( - &mut builder.when(local.is_interaction * (AB::Expr::ONE - local.n_sign)), - local.multiplier, - local.eq_sharp_ns, - ); - // ELSE 2 - assert_array_eq( - &mut builder.when(local.idx_parity), - local.multiplier, - local.eq_sharp_ns, - ); - - // === interactions === - self.expr_claim_bus.receive( - builder, - local.proof_idx, - ExpressionClaimMessage { - is_interaction: local.is_interaction, - idx: local.idx, - value: local.value, - }, - local.is_valid, - ); - - self.mu_bus.lookup_key( - builder, - local.proof_idx, - BatchConstraintConductorMessage { - msg_type: BatchConstraintInnerMessageType::Mu.to_field(), - idx: AB::Expr::ZERO, - value: local.mu.map(Into::into), - }, - local.is_first * local.is_valid, - ); - - // Receive n_max value from proof shape air - self.expression_claim_n_max_bus.receive( - builder, - local.proof_idx, - ExpressionClaimNMaxMessage { - n_max: local.num_multilinear_sumcheck_rounds, - }, - local.is_first * local.is_valid, - ); - - self.main_claim_bus.receive( - builder, - local.proof_idx, - MainExpressionClaimMessage { - idx: local.idx.into(), - claim: local.cur_sum.map(Into::into), - }, - local.is_first * local.is_valid, - ); - - self.hyperdim_bus.lookup_key( - builder, - local.proof_idx, - HyperdimBusMessage { - sort_idx: local.trace_idx.into(), - n_abs: local.n_abs.into(), - n_sign_bit: local.n_sign.into(), - }, - local.is_valid * (local.is_interaction - local.idx_parity), - ); - - self.eq_n_outer_bus.lookup_key( - builder, - local.proof_idx, - EqNOuterMessage { - is_sharp: AB::Expr::ONE, - n: local.n_abs * (AB::Expr::ONE - local.n_sign), - value: local.eq_sharp_ns.map(Into::into), - }, - local.is_valid * local.is_interaction, - ); - - self.pow_checker_bus.lookup_key( - builder, - PowerCheckerBusMessage { - log: local.n_abs.into(), - exp: local.n_abs_pow.into(), - }, - local.is_valid * local.is_interaction, - ); - } -} diff --git a/ceno_recursion_v2/src/batch_constraint/expression_claim/mod.rs b/ceno_recursion_v2/src/batch_constraint/expression_claim/mod.rs deleted file mode 100644 index d58dde58d..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expression_claim/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -mod air; -mod trace; - -pub use air::{ExpressionClaimAir, ExpressionClaimCols}; diff --git a/ceno_recursion_v2/src/batch_constraint/expression_claim/trace.rs b/ceno_recursion_v2/src/batch_constraint/expression_claim/trace.rs deleted file mode 100644 index f559a8848..000000000 --- a/ceno_recursion_v2/src/batch_constraint/expression_claim/trace.rs +++ /dev/null @@ -1,169 +0,0 @@ -use std::borrow::BorrowMut; - -use openvm_stark_backend::proof::Proof; -use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, D_EF, EF, F}; -use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, PrimeField32}; -use p3_matrix::dense::RowMajorMatrix; -use p3_maybe_rayon::prelude::*; - -use super::ExpressionClaimCols; -use crate::{ - primitives::pow::PowerCheckerCpuTraceGenerator, - system::{POW_CHECKER_HEIGHT, Preflight}, - tracegen::RowMajorChip, - utils::MultiProofVecVec, -}; - -#[allow(dead_code)] -pub struct ExpressionClaimBlob { - // (n, value), n is before lift, can be negative - claims: MultiProofVecVec<(isize, EF)>, -} - -#[allow(dead_code)] -pub fn generate_expression_claim_blob( - cf_folded_claims: &MultiProofVecVec<(isize, EF)>, - if_folded_claims: &MultiProofVecVec<(isize, EF)>, -) -> ExpressionClaimBlob { - let mut claims = MultiProofVecVec::new(); - for pidx in 0..cf_folded_claims.num_proofs() { - claims.extend(if_folded_claims[pidx].iter().cloned()); - claims.extend(cf_folded_claims[pidx].iter().cloned()); - claims.end_proof(); - } - ExpressionClaimBlob { claims } -} - -#[allow(dead_code)] -pub struct ExpressionClaimTraceGenerator; - -#[allow(dead_code)] -pub(crate) struct ExpressionClaimCtx<'a> { - pub blob: &'a ExpressionClaimBlob, - pub proofs: &'a [&'a Proof], - pub preflights: &'a [&'a Preflight], - pub pow_checker: &'a PowerCheckerCpuTraceGenerator<2, POW_CHECKER_HEIGHT>, -} - -impl RowMajorChip for ExpressionClaimTraceGenerator { - type Ctx<'a> = ExpressionClaimCtx<'a>; - - #[tracing::instrument(level = "trace", skip_all)] - fn generate_trace( - &self, - ctx: &Self::Ctx<'_>, - required_height: Option, - ) -> Option> { - let blob = ctx.blob; - let proofs = ctx.proofs; - let preflights = ctx.preflights; - let pow_checker = ctx.pow_checker; - let width = ExpressionClaimCols::::width(); - - let num_valid = blob.claims.len(); - let padded_height = if let Some(height) = required_height { - if height < num_valid { - return None; - } - height - } else { - num_valid.next_power_of_two() - }; - let mut trace = vec![F::ZERO; padded_height * width]; - let mut cur_height = 0; - for (pidx, preflight) in preflights.iter().enumerate() { - let claims = &blob.claims[pidx]; - - let num_rounds = proofs[pidx] - .batch_constraint_proof - .sumcheck_round_polys - .len(); - let num_present = preflight.proof_shape.sorted_trace_vdata.len(); - debug_assert_eq!(claims.len(), 3 * num_present); - let mu_tidx = preflight.batch_constraint.tidx_before_univariate - D_EF; - - trace[cur_height * width..(cur_height + claims.len()) * width] - .par_chunks_exact_mut(width) - .enumerate() - .for_each(|(i, chunk)| { - let n_lift = claims[i].0.max(0) as usize; - let n_abs = claims[i].0.unsigned_abs(); - let is_interaction = i < 2 * num_present; - if is_interaction { - pow_checker.add_pow(n_abs); - } - let cols: &mut ExpressionClaimCols<_> = chunk.borrow_mut(); - cols.is_first = F::from_bool(i == 0); - cols.is_valid = F::ONE; - cols.proof_idx = F::from_usize(pidx); - cols.is_interaction = F::from_bool(is_interaction); - cols.num_multilinear_sumcheck_rounds = F::from_usize(num_rounds); - cols.idx = F::from_usize(if i < 2 * num_present { - i - } else { - i - 2 * num_present - }); - cols.idx_parity = F::from_bool(is_interaction && i % 2 == 1); - let trace_idx = if is_interaction { - i / 2 - } else { - i - 2 * num_present - }; - cols.trace_idx = F::from_usize(trace_idx); - cols.mu - .copy_from_slice(&preflight.transcript.values()[mu_tidx..mu_tidx + D_EF]); - cols.value - .copy_from_slice(claims[i].1.as_basis_coefficients_slice()); - cols.eq_sharp_ns.copy_from_slice( - preflight.batch_constraint.eq_sharp_ns_frontloaded[n_lift] - .as_basis_coefficients_slice(), - ); - cols.multiplier - .copy_from_slice(EF::ONE.as_basis_coefficients_slice()); - cols.n_abs = F::from_usize(n_abs); - cols.n_sign = F::from_bool(claims[i].0.is_negative()); - cols.n_abs_pow = F::from_usize(1 << n_abs); - }); - - // Setting `cur_sum` - let mut cur_sum = EF::ZERO; - let mu = EF::from_basis_coefficients_slice( - &preflight.transcript.values()[mu_tidx..mu_tidx + D_EF], - ) - .unwrap(); - trace[cur_height * width..(cur_height + claims.len()) * width] - .chunks_exact_mut(width) - .rev() - .for_each(|chunk| { - let cols: &mut ExpressionClaimCols<_> = chunk.borrow_mut(); - // if it's interaction, we need to multiply by eq_sharp_ns and norm_factor - let multiplier = if cols.is_interaction == F::ONE { - let mut mult = - EF::from_basis_coefficients_slice(&cols.eq_sharp_ns).unwrap(); - if cols.n_sign == F::ONE && cols.idx.as_canonical_u32() % 2 == 0 { - mult *= F::from_u32(1 << cols.n_abs.as_canonical_u32()).inverse(); - } - mult - } else { - EF::ONE - }; - cols.multiplier - .copy_from_slice(multiplier.as_basis_coefficients_slice()); - cur_sum = cur_sum * mu - + EF::from_basis_coefficients_slice(&cols.value).unwrap() * multiplier; - cols.cur_sum - .copy_from_slice(cur_sum.as_basis_coefficients_slice()); - }); - - cur_height += claims.len(); - } - trace[cur_height * width..] - .par_chunks_mut(width) - .enumerate() - .for_each(|(i, chunk)| { - let cols: &mut ExpressionClaimCols = chunk.borrow_mut(); - cols.proof_idx = F::from_usize(preflights.len() + i); - }); - Some(RowMajorMatrix::new(trace, width)) - } -} diff --git a/ceno_recursion_v2/src/batch_constraint/mod.rs b/ceno_recursion_v2/src/batch_constraint/mod.rs index d6f2e880f..9647d33f0 100644 --- a/ceno_recursion_v2/src/batch_constraint/mod.rs +++ b/ceno_recursion_v2/src/batch_constraint/mod.rs @@ -1,260 +1,97 @@ use std::sync::Arc; -use ceno_zkvm::scheme::ZKVMChipProof; use openvm_cpu_backend::CpuBackend; use openvm_poseidon2_air::POSEIDON2_WIDTH; use openvm_stark_backend::{ - AirRef, FiatShamirTranscript, StarkEngine, StarkProtocolConfig, TranscriptHistory, - prover::{AirProvingContext, CommittedTraceData, TraceCommitter}, + AirRef, FiatShamirTranscript, StarkProtocolConfig, TranscriptHistory, prover::AirProvingContext, }; -use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, EF, F}; -use p3_field::PrimeCharacteristicRing; -use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use openvm_stark_sdk::config::baby_bear_poseidon2::F; use crate::{ - batch_constraint::{ - bus::{ - BatchConstraintConductorBus, ConstraintsFoldingBus, EqNOuterBus, ExpressionClaimBus, - InteractionsFoldingBus, SymbolicExpressionBus, - }, - expr_eval::{ - ConstraintsFoldingAir, ConstraintsFoldingCols, - symbolic_expression::{SingleMainSymbolicExpressionColumns, SymbolicExpressionAir}, - }, - expression_claim::{ExpressionClaimAir, ExpressionClaimCols}, + main::{ + MainCollectedRecords, MainModule, + eval_absorb::{MainEvalAbsorbAir, MainEvalAbsorbTraceGenerator}, + final_claim::{MainFinalClaimAir, MainFinalClaimTraceGenerator}, + frontload::{MainFrontloadTermAir, MainFrontloadTermTraceGenerator}, + global_sumcheck::{MainGlobalSumcheckAir, MainGlobalSumcheckTraceGenerator}, + tower_point::{MainTowerPointEqAir, MainTowerPointEqTraceGenerator}, }, - bus::{AirPresenceBus, ColumnClaimsBus, SelHypercubeBus, SelUniBus}, system::{ - AirModule, BatchConstraintPreflight, BusIndexManager, BusInventory, GlobalCtxCpu, - Preflight, RecursionField, RecursionProof, RecursionVk, TraceGenModule, + AirModule, BusIndexManager, BusInventory, GlobalCtxCpu, Preflight, RecursionProof, + RecursionVk, TraceGenModule, }, + tracegen::{ModuleChip, RowMajorChip}, }; -pub mod expr_eval; -pub mod expression_claim; -pub mod bus { - use p3_field::PrimeCharacteristicRing; - pub use recursion_circuit::batch_constraint::bus::*; - - #[repr(u8)] - #[derive(Debug, Copy, Clone)] - pub enum BatchConstraintInnerMessageType { - R, - Xi, - Mu, - } - - impl BatchConstraintInnerMessageType { - pub fn to_field(self) -> T { - T::from_u8(self as u8) - } - } -} - -pub use expr_eval::CachedTraceRecord; - pub struct BatchConstraintModule { transcript_bus: crate::bus::TranscriptBus, - hyperdim_bus: crate::bus::HyperdimBus, - air_shape_bus: crate::bus::AirShapeBus, - air_presence_bus: AirPresenceBus, - column_claims_bus: ColumnClaimsBus, - public_values_bus: crate::bus::PublicValuesBus, - sel_hypercube_bus: SelHypercubeBus, - sel_uni_bus: SelUniBus, - - expression_claim_n_max_bus: crate::bus::ExpressionClaimNMaxBus, - n_lift_bus: crate::bus::NLiftBus, - main_expression_claim_bus: crate::bus::MainExpressionClaimBus, - power_checker_bus: recursion_circuit::primitives::bus::PowerCheckerBus, - - batch_constraint_conductor_bus: BatchConstraintConductorBus, - eq_n_outer_bus: EqNOuterBus, - symbolic_expression_bus: SymbolicExpressionBus, - expression_claim_bus: ExpressionClaimBus, - interactions_folding_bus: InteractionsFoldingBus, - constraints_folding_bus: ConstraintsFoldingBus, - - max_num_proofs: usize, + global_claim_bus: crate::bus::MainGlobalClaimBus, + global_point_bus: crate::bus::MainGlobalPointBus, + eval_bus: crate::bus::MainEvalBus, + contribution_bus: crate::bus::MainContributionBus, + tower_point_bus: crate::bus::TowerMainPointBus, } impl BatchConstraintModule { pub fn new( - b: &mut BusIndexManager, + _b: &mut BusIndexManager, bus_inventory: BusInventory, - max_num_proofs: usize, + _max_num_proofs: usize, ) -> Self { Self { transcript_bus: bus_inventory.transcript_bus, - hyperdim_bus: bus_inventory.hyperdim_bus, - air_shape_bus: bus_inventory.air_shape_bus, - air_presence_bus: AirPresenceBus::new(b.new_bus_idx()), - column_claims_bus: ColumnClaimsBus::new(b.new_bus_idx()), - public_values_bus: bus_inventory.public_values_bus, - sel_hypercube_bus: SelHypercubeBus::new(b.new_bus_idx()), - sel_uni_bus: SelUniBus::new(b.new_bus_idx()), - - expression_claim_n_max_bus: bus_inventory.expression_claim_n_max_bus, - n_lift_bus: bus_inventory.n_lift_bus, - main_expression_claim_bus: bus_inventory.main_expression_claim_bus, - power_checker_bus: bus_inventory.power_checker_bus, - - batch_constraint_conductor_bus: BatchConstraintConductorBus::new(b.new_bus_idx()), - eq_n_outer_bus: EqNOuterBus::new(b.new_bus_idx()), - symbolic_expression_bus: SymbolicExpressionBus::new(b.new_bus_idx()), - expression_claim_bus: ExpressionClaimBus::new(b.new_bus_idx()), - interactions_folding_bus: InteractionsFoldingBus::new(b.new_bus_idx()), - constraints_folding_bus: ConstraintsFoldingBus::new(b.new_bus_idx()), - max_num_proofs, + global_claim_bus: bus_inventory.main_global_claim_bus, + global_point_bus: bus_inventory.main_global_point_bus, + eval_bus: bus_inventory.main_eval_bus, + contribution_bus: bus_inventory.main_contribution_bus, + tower_point_bus: bus_inventory.tower_main_point_bus, } } - #[tracing::instrument(level = "trace", skip_all)] pub fn run_preflight( &self, _child_vk: &RecursionVk, - proof: &RecursionProof, - preflight: &mut Preflight, - ts: &mut TS, + _proof: &RecursionProof, + _preflight: &mut Preflight, + _ts: &mut TS, ) where - TS: FiatShamirTranscript - + TranscriptHistory, + TS: FiatShamirTranscript< + openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config, + > + TranscriptHistory, { - // Constraint batching challenge. - let lambda_tidx = ts.len(); - let _lambda = FiatShamirTranscript::::sample_ext(ts); - - // Replay a lightweight subset of batch-constraint transcript observes from per-chip - // sumcheck messages, then sample mu. - for chip_proof in proof.chip_proofs.values() { - observe_main_sumcheck_msgs(ts, chip_proof); - } - let _mu = FiatShamirTranscript::::sample_ext(ts); - let tidx_before_univariate = ts.len(); - - let mut sumcheck_rnd = vec![]; - for chip_proof in proof.chip_proofs.values() { - if let Some(layer) = chip_proof - .gkr_iop_proof - .as_ref() - .and_then(|proof| proof.0.first()) - { - for msg in &layer.main.proof.proofs { - for eval in msg.evaluations.iter().take(3) { - ts.observe_ext(*eval); - } - sumcheck_rnd.push(FiatShamirTranscript::::sample(ts)); - } - } - } - if sumcheck_rnd.is_empty() { - // Keep downstream preflight consumers shape-safe when this bridge has no rounds. - sumcheck_rnd.push(F::ZERO); - } - - let n_max = preflight - .proof_shape - .sorted_trace_vdata - .iter() - .map(|(_, v)| v.log_height) - .max() - .unwrap_or(0); - let eq_ns_frontloaded = vec![EF::ONE; n_max + 1]; - let eq_sharp_ns_frontloaded = vec![EF::ONE; n_max + 1]; - - // TODO(recursion-proof-bridge): replace placeholder eq vectors with verifier-equivalent - // frontloaded eq_n / eq_sharp_n computation derived from xi and sumcheck randomness. - preflight.batch_constraint = BatchConstraintPreflight { - lambda_tidx, - tidx_before_univariate, - sumcheck_rnd, - eq_ns_frontloaded, - eq_sharp_ns_frontloaded, - }; - } - - fn placeholder_common_widths(&self) -> [usize; 3] { - [ - SingleMainSymbolicExpressionColumns::::width() * self.max_num_proofs, - ConstraintsFoldingCols::::width(), - ExpressionClaimCols::::width(), - ] - } - - pub(crate) fn generate_placeholder_proving_ctxs>( - &self, - cached_symbolic_expr_trace: CommittedTraceData>, - required_heights: Option<&[usize]>, - ) -> Option>>> { - let widths = self.placeholder_common_widths(); - let air_count = required_heights - .map(|heights| heights.len()) - .unwrap_or(self.num_airs()); - let cached_height = cached_symbolic_expr_trace.trace.height(); - - (0..air_count) - .map(|idx| { - let height = required_heights - .and_then(|heights| heights.get(idx).copied()) - .unwrap_or_else(|| if idx == 0 { cached_height } else { 2 }); - if height < 2 { - return None; - } - let width = widths.get(idx).copied().unwrap_or(1).max(1); - let matrix = RowMajorMatrix::new(vec![F::ZERO; height * width], width); - Some(if idx == 0 { - AirProvingContext { - cached_mains: vec![cached_symbolic_expr_trace.clone()], - common_main: matrix, - public_values: vec![], - } - } else { - AirProvingContext::simple_no_pis(matrix) - }) - }) - .collect::>>() + let _ = self; } } impl AirModule for BatchConstraintModule { fn num_airs(&self) -> usize { - 3 + 5 } fn airs>(&self) -> Vec> { - let symbolic_expression_air = SymbolicExpressionAir { - expr_bus: self.symbolic_expression_bus, - hyperdim_bus: self.hyperdim_bus, - air_shape_bus: self.air_shape_bus, - air_presence_bus: self.air_presence_bus, - column_claims_bus: self.column_claims_bus, - interactions_folding_bus: self.interactions_folding_bus, - constraints_folding_bus: self.constraints_folding_bus, - public_values_bus: self.public_values_bus, - sel_hypercube_bus: self.sel_hypercube_bus, - sel_uni_bus: self.sel_uni_bus, - cnt_proofs: self.max_num_proofs, - }; - let constraints_folding_air = ConstraintsFoldingAir { - transcript_bus: self.transcript_bus, - constraint_bus: self.constraints_folding_bus, - expression_claim_bus: self.expression_claim_bus, - eq_n_outer_bus: self.eq_n_outer_bus, - n_lift_bus: self.n_lift_bus, - }; - let expression_claim_air = ExpressionClaimAir { - expression_claim_n_max_bus: self.expression_claim_n_max_bus, - expr_claim_bus: self.expression_claim_bus, - mu_bus: self.batch_constraint_conductor_bus, - main_claim_bus: self.main_expression_claim_bus, - eq_n_outer_bus: self.eq_n_outer_bus, - pow_checker_bus: self.power_checker_bus, - hyperdim_bus: self.hyperdim_bus, - }; vec![ - Arc::new(symbolic_expression_air) as AirRef<_>, - Arc::new(constraints_folding_air) as AirRef<_>, - Arc::new(expression_claim_air) as AirRef<_>, + Arc::new(MainGlobalSumcheckAir { + transcript_bus: self.transcript_bus, + global_claim_bus: self.global_claim_bus, + global_point_bus: self.global_point_bus, + }) as AirRef<_>, + Arc::new(MainEvalAbsorbAir { + transcript_bus: self.transcript_bus, + eval_bus: self.eval_bus, + }) as AirRef<_>, + Arc::new(MainTowerPointEqAir { + global_point_bus: self.global_point_bus, + tower_point_bus: self.tower_point_bus, + }) as AirRef<_>, + Arc::new(MainFrontloadTermAir { + eval_bus: self.eval_bus, + global_point_bus: self.global_point_bus, + contribution_bus: self.contribution_bus, + }) as AirRef<_>, + Arc::new(MainFinalClaimAir { + global_claim_bus: self.global_claim_bus, + contribution_bus: self.contribution_bus, + }) as AirRef<_>, ] } } @@ -267,67 +104,90 @@ impl> TraceGenModule #[tracing::instrument(skip_all)] fn generate_proving_ctxs( &self, - _child_vk: &RecursionVk, - _proofs: &[RecursionProof], - _preflights: &[Preflight], + child_vk: &RecursionVk, + proofs: &[RecursionProof], + preflights: &[Preflight], _ctx: &Self::ModuleSpecificCtx<'_>, required_heights: Option<&[usize]>, - ) -> Option>>> { - let widths = self.placeholder_common_widths(); - let air_count = required_heights - .map(|heights| heights.len()) - .unwrap_or(self.num_airs()); - - (0..air_count) - .map(|idx| { - let height = required_heights - .and_then(|heights| heights.get(idx).copied()) - .unwrap_or(1); - if required_heights.is_some() && height < 2 { - return None; - } - let width = widths.get(idx).copied().unwrap_or(1); - let rows = height.max(2); - let cols = width.max(1); - let matrix = RowMajorMatrix::new(vec![F::ZERO; rows * cols], cols); - Some(openvm_stark_backend::prover::AirProvingContext::simple_no_pis(matrix)) + ) -> Option>>> { + let mut records = MainModule::collect_records(child_vk, proofs, preflights).ok()?; + records + .global_sumcheck_records + .sort_by_key(|record| record.proof_idx); + records + .eval_records + .sort_by_key(|record| (record.proof_idx, record.idx, record.eval_idx)); + records + .tower_point_eq_records + .sort_by_key(|record| (record.proof_idx, record.idx, record.round_idx)); + records + .frontload_term_records + .sort_by_key(|record| (record.proof_idx, record.idx, record.row_idx)); + records + .final_claim_records + .sort_by_key(|record| (record.proof_idx, record.idx)); + + let ctx = BatchConstraintTraceCtx { records: &records }; + let chips = [ + BatchConstraintModuleChip::GlobalSumcheck, + BatchConstraintModuleChip::EvalAbsorb, + BatchConstraintModuleChip::TowerPointEq, + BatchConstraintModuleChip::FrontloadTerm, + BatchConstraintModuleChip::FinalClaim, + ]; + chips + .into_iter() + .enumerate() + .map(|(idx, chip)| { + chip.generate_proving_ctx( + &ctx, + required_heights.and_then(|heights| heights.get(idx).copied()), + ) }) - .collect::>>() + .collect() } } -fn observe_main_sumcheck_msgs(ts: &mut TS, chip_proof: &ZKVMChipProof) -where - TS: FiatShamirTranscript, -{ - if let Some(proofs) = &chip_proof.main_sumcheck_proofs { - for msg in proofs { - for eval in msg.evaluations.iter().take(3) { - ts.observe_ext(*eval); - } - } - } +struct BatchConstraintTraceCtx<'a> { + records: &'a MainCollectedRecords, } -pub fn cached_trace_record(child_vk: &RecursionVk) -> CachedTraceRecord { - expr_eval::symbolic_expression::build_cached_trace_record(child_vk) +enum BatchConstraintModuleChip { + GlobalSumcheck, + EvalAbsorb, + TowerPointEq, + FrontloadTerm, + FinalClaim, } -pub fn commit_child_vk( - engine: &E, - child_vk: &RecursionVk, -) -> CommittedTraceData> -where - E: StarkEngine>, - SC: StarkProtocolConfig, -{ - let cached_trace = expr_eval::symbolic_expression::generate_symbolic_expr_cached_trace( - &cached_trace_record(child_vk), - ); - let (commitment, data) = engine.device().commit(&[&cached_trace]).unwrap(); - CommittedTraceData { - commitment, - data: Arc::new(data), - trace: cached_trace, +impl RowMajorChip for BatchConstraintModuleChip { + type Ctx<'a> = BatchConstraintTraceCtx<'a>; + + fn generate_trace( + &self, + ctx: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + match self { + BatchConstraintModuleChip::GlobalSumcheck => MainGlobalSumcheckTraceGenerator + .generate_trace( + &ctx.records.global_sumcheck_records.as_slice(), + required_height, + ), + BatchConstraintModuleChip::EvalAbsorb => MainEvalAbsorbTraceGenerator + .generate_trace(&ctx.records.eval_records.as_slice(), required_height), + BatchConstraintModuleChip::TowerPointEq => MainTowerPointEqTraceGenerator + .generate_trace( + &ctx.records.tower_point_eq_records.as_slice(), + required_height, + ), + BatchConstraintModuleChip::FrontloadTerm => MainFrontloadTermTraceGenerator + .generate_trace( + &ctx.records.frontload_term_records.as_slice(), + required_height, + ), + BatchConstraintModuleChip::FinalClaim => MainFinalClaimTraceGenerator + .generate_trace(&ctx.records.final_claim_records.as_slice(), required_height), + } } } diff --git a/ceno_recursion_v2/src/bus.rs b/ceno_recursion_v2/src/bus.rs index eb61e242a..1488b958a 100644 --- a/ceno_recursion_v2/src/bus.rs +++ b/ceno_recursion_v2/src/bus.rs @@ -1,14 +1,11 @@ -use openvm_poseidon2_air::POSEIDON2_WIDTH; -use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, DIGEST_SIZE}; use recursion_circuit::{ bus as upstream, define_typed_per_proof_lookup_bus, define_typed_per_proof_permutation_bus, }; pub use upstream::{ AirPresenceBus, AirPresenceBusMessage, AirShapeBus, AirShapeBusMessage, CachedCommitBus, - CachedCommitBusMessage, ColumnClaimsBus, ColumnClaimsMessage, ExpressionClaimNMaxBus, - ExpressionClaimNMaxMessage, FractionFolderInputBus, FractionFolderInputMessage, HyperdimBus, - HyperdimBusMessage, LiftedHeightsBus, LiftedHeightsBusMessage, NLiftBus, NLiftMessage, - PublicValuesBus, PublicValuesBusMessage, SelHypercubeBus, SelHypercubeBusMessage, SelUniBus, + CachedCommitBusMessage, ColumnClaimsBus, ColumnClaimsMessage, PublicValuesBus, + PublicValuesBusMessage, SelHypercubeBus, SelHypercubeBusMessage, SelUniBus, SelUniBusMessage, TranscriptBus, TranscriptBusMessage, }; @@ -110,3 +107,413 @@ pub struct MainExpressionClaimMessage { } define_typed_per_proof_permutation_bus!(MainExpressionClaimBus, MainExpressionClaimMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainGlobalClaimMessage { + pub expected: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(MainGlobalClaimBus, MainGlobalClaimMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainGlobalPointMessage { + pub round_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(MainGlobalPointBus, MainGlobalPointMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainEvalMessage { + pub idx: T, + pub eval_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(MainEvalBus, MainEvalMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainContributionMessage { + pub idx: T, + pub contribution: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(MainContributionBus, MainContributionMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainSelectorResultMessage { + pub idx: T, + pub eval_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(MainSelectorResultBus, MainSelectorResultMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainSelectorPointMessage { + pub idx: T, + pub air_idx: T, + pub selector_idx: T, + pub source_kind: T, + pub round_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(MainSelectorPointBus, MainSelectorPointMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct EccRtMessage { + pub idx: T, + pub round_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(EccRtBus, EccRtMessage); + +#[repr(u8)] +#[derive(Debug, Clone, Copy)] +pub enum MainEccRtChallengeKind { + Rt = 0, + OutRt = 1, + Alpha = 2, +} + +impl MainEccRtChallengeKind { + pub const fn as_usize(self) -> usize { + self as usize + } +} + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainEccRtChallengeMessage { + pub idx: T, + pub round_idx: T, + pub kind: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(MainEccRtChallengeBus, MainEccRtChallengeMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainEccRtSumcheckFinalMessage { + pub idx: T, + pub claim: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(MainEccRtSumcheckFinalBus, MainEccRtSumcheckFinalMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainEccRtEquationTotalsMessage { + pub idx: T, + pub add_eval: [T; D_EF], + pub bypass_eval: [T; D_EF], + pub export_eval: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(MainEccRtEquationTotalsBus, MainEccRtEquationTotalsMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainEccRtQuarkFinalMessage { + pub idx: T, + pub quark_out: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(MainEccRtQuarkFinalBus, MainEccRtQuarkFinalMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainSelectorShapeMessage { + pub air_idx: T, + pub selector_idx: T, + pub kind: T, + pub point_source: T, + pub eval_idx: T, + pub ctx_offset: T, + pub ctx_num_instances: T, + pub ctx_num_vars: T, + pub ordered_sparse_num_vars: T, + pub num_sparse_indices: T, +} + +define_typed_per_proof_permutation_bus!(MainSelectorShapeBus, MainSelectorShapeMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct MainSelectorSparseIndexShapeMessage { + pub air_idx: T, + pub selector_idx: T, + pub sparse_pos: T, + pub sparse_index: T, +} + +define_typed_per_proof_permutation_bus!( + MainSelectorSparseIndexShapeBus, + MainSelectorSparseIndexShapeMessage +); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct TowerMainPointMessage { + pub idx: T, + pub round_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(TowerMainPointBus, TowerMainPointMessage); + +#[repr(u8)] +#[derive(Debug, Clone, Copy)] +pub enum PcsBasefoldQueryStage { + ReducedOpening = 0, + FinalFolded = 1, + QueryIndex = 2, + FinalCodewordStart = 3, +} + +impl PcsBasefoldQueryStage { + pub const fn as_usize(self) -> usize { + self as usize + } +} + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsBasefoldQueryMessage { + pub query_idx: T, + pub stage: T, + pub round: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(PcsBasefoldQueryBus, PcsBasefoldQueryMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsBaseInputOpeningMessage { + pub query_idx: T, + pub opening_idx: T, + pub reduced_index: T, + pub value_idx: T, + pub elem_idx: T, + pub log2_height: T, + pub opened_value: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(PcsBaseInputOpeningBus, PcsBaseInputOpeningMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsFinalMessageMessage { + pub tidx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsFinalMessageBus, PcsFinalMessageMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsQuerySampleMessage { + pub tidx: T, + pub value: T, +} + +define_typed_per_proof_lookup_bus!(PcsQuerySampleBus, PcsQuerySampleMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsCommitmentRootMessage { + pub commit_major: T, + pub commit_minor: T, + pub root: [T; DIGEST_SIZE], +} + +define_typed_per_proof_lookup_bus!(PcsCommitmentRootBus, PcsCommitmentRootMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsCommitPhaseLeafMessage { + pub query_idx: T, + pub round: T, + pub leaf_idx: T, + pub left: [T; D_EF], + pub right: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(PcsCommitPhaseLeafBus, PcsCommitPhaseLeafMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsSumcheckClaimMessage { + pub idx: T, + pub claim: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(PcsSumcheckClaimBus, PcsSumcheckClaimMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsSumcheckInputMessage { + pub idx: T, + pub claim: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(PcsSumcheckInputBus, PcsSumcheckInputMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsSumcheckOutputMessage { + pub idx: T, + pub claim: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(PcsSumcheckOutputBus, PcsSumcheckOutputMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsBasefoldFinalExpectedMessage { + pub sumcheck_idx: T, + pub expected: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!( + PcsBasefoldFinalExpectedBus, + PcsBasefoldFinalExpectedMessage +); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsBasefoldFinalPointMessage { + pub sumcheck_idx: T, + pub point_idx: T, + pub coord_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_permutation_bus!(PcsBasefoldFinalPointBus, PcsBasefoldFinalPointMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsFoldChallengeMessage { + pub sumcheck_idx: T, + pub round: T, + pub challenge: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsFoldChallengeBus, PcsFoldChallengeMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsBatchCoeffMessage { + pub global_coeff_idx: T, + pub coeff: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsBatchCoeffBus, PcsBatchCoeffMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsBatchAlphaMessage { + pub tidx: T, + pub alpha: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsBatchAlphaBus, PcsBatchAlphaMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsBasefoldEvalMessage { + pub tidx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsBasefoldEvalBus, PcsBasefoldEvalMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsTranscriptExtMessage { + pub tidx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsTranscriptExtBus, PcsTranscriptExtMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsJaggedFEvalMessage { + pub sumcheck_idx: T, + pub tidx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsJaggedFEvalBus, PcsJaggedFEvalMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsOpeningEvalMessage { + pub opening_idx: T, + pub commit_kind: T, + pub eval_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsOpeningEvalBus, PcsOpeningEvalMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsEqProductMessage { + pub kind: T, + pub round_idx: T, + pub term_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsEqProductBus, PcsEqProductMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsSuffixProductMessage { + pub round_idx: T, + pub term_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsSuffixProductBus, PcsSuffixProductMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsJaggedAssistHMessage { + pub round_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsJaggedAssistHBus, PcsJaggedAssistHMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsJaggedAssistQMessage { + pub round_idx: T, + pub value: [T; D_EF], +} + +define_typed_per_proof_lookup_bus!(PcsJaggedAssistQBus, PcsJaggedAssistQMessage); + +#[repr(C)] +#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)] +pub struct PcsCommitHeightMessage { + pub commitment_kind: T, + pub height_idx: T, + pub value: T, +} + +define_typed_per_proof_lookup_bus!(PcsCommitHeightBus, PcsCommitHeightMessage); diff --git a/ceno_recursion_v2/src/circuit/inner/def_pvs/air.rs b/ceno_recursion_v2/src/circuit/inner/def_pvs/air.rs index 2497e02e0..aeb683c29 100644 --- a/ceno_recursion_v2/src/circuit/inner/def_pvs/air.rs +++ b/ceno_recursion_v2/src/circuit/inner/def_pvs/air.rs @@ -175,7 +175,7 @@ impl Air f deferral_flag: local.deferral_flag, has_verifier_pvs: local.has_verifier_pvs, }, - local.is_present * AB::Expr::from_bool(!crate::system::TOWER_PREFIX_ONLY), + local.is_present, ); // If this row corresponds to a direct deferral hook circuit child (i.e. diff --git a/ceno_recursion_v2/src/circuit/inner/mod.rs b/ceno_recursion_v2/src/circuit/inner/mod.rs index b2274a1de..bba440222 100644 --- a/ceno_recursion_v2/src/circuit/inner/mod.rs +++ b/ceno_recursion_v2/src/circuit/inner/mod.rs @@ -54,6 +54,7 @@ impl, S: AggregationSubCircuit> Circuit for I }; let verifier_pvs_air = Arc::new(verifier::VerifierPvsAir { + transcript_bus, public_values_bus, // cached_commit_bus, pvs_air_consistency_bus, @@ -69,8 +70,10 @@ impl, S: AggregationSubCircuit> Circuit for I deferral_enabled, instance_public_value_indices: self.instance_public_value_indices.clone(), }) as AirRef; - let vm_pvs_cumulative_heights_air = - Arc::new(vm_pvs::VmPvsCumulativeHeightsAir { transcript_bus }) as AirRef; + let vm_pvs_cumulative_heights_air = Arc::new(vm_pvs::VmPvsCumulativeHeightsAir { + transcript_bus, + commit_height_bus: bus_inventory.pcs_commit_height_bus, + }) as AirRef; let (idx2_air, post_airs): (AirRef, Vec>) = if deferral_enabled { let def_pvs_air = Arc::new(def_pvs::DeferralPvsAir { public_values_bus, diff --git a/ceno_recursion_v2/src/circuit/inner/trace.rs b/ceno_recursion_v2/src/circuit/inner/trace.rs index bc518a83f..d3d54c247 100644 --- a/ceno_recursion_v2/src/circuit/inner/trace.rs +++ b/ceno_recursion_v2/src/circuit/inner/trace.rs @@ -77,6 +77,7 @@ impl InnerTraceGen> for InnerTraceGenImpl { + TranscriptHistory, { let (verifier_ctx, poseidon2_inputs) = super::verifier::generate_proving_ctx( + child_vk, proofs, proofs_type, child_is_app, diff --git a/ceno_recursion_v2/src/circuit/inner/unset/air.rs b/ceno_recursion_v2/src/circuit/inner/unset/air.rs index c0cc6f029..c18d3fa81 100644 --- a/ceno_recursion_v2/src/circuit/inner/unset/air.rs +++ b/ceno_recursion_v2/src/circuit/inner/unset/air.rs @@ -69,7 +69,7 @@ impl Air for UnsetPvsAir { deferral_flag: AB::F::from_u32(self.def_flag), has_verifier_pvs: AB::F::ONE, }, - local.is_valid * AB::Expr::from_bool(!crate::system::TOWER_PREFIX_ONLY), + local.is_valid, ); } } diff --git a/ceno_recursion_v2/src/circuit/inner/verifier/air.rs b/ceno_recursion_v2/src/circuit/inner/verifier/air.rs index 0f5a302de..4a43950aa 100644 --- a/ceno_recursion_v2/src/circuit/inner/verifier/air.rs +++ b/ceno_recursion_v2/src/circuit/inner/verifier/air.rs @@ -1,15 +1,18 @@ use std::borrow::Borrow; +use ceno_zkvm::structs::VK_DIGEST_LEN; use openvm_circuit_primitives::utils::{and, assert_array_eq, not}; use openvm_stark_backend::{ BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, }; +use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF; use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir}; use p3_field::{Field, PrimeCharacteristicRing}; use p3_matrix::Matrix; use recursion_circuit::{ bus::{ Poseidon2CompressBus, Poseidon2CompressMessage, PublicValuesBus, PublicValuesBusMessage, + TranscriptBus, TranscriptBusMessage, }, prelude::DIGEST_SIZE, utils::assert_zeros, @@ -19,7 +22,7 @@ use verify_stark::pvs::{VERIFIER_PVS_AIR_ID, VerifierBasePvs, VerifierDefPvs, Vk use crate::{ circuit::inner::bus::{PvsAirConsistencyBus, PvsAirConsistencyMessage}, - utils::digests_to_poseidon2_input, + utils::{TranscriptLabel, digests_to_poseidon2_input}, }; /////////////////////////////////////////////////////////////////////////////// @@ -32,10 +35,12 @@ pub struct VerifierPvsCols { pub proof_idx: F, pub is_valid: F, pub has_verifier_pvs: F, + pub child_vk_digest: [[F; D_EF]; VK_DIGEST_LEN], pub child_pvs: VerifierBasePvs, } pub struct VerifierPvsAir { + pub transcript_bus: TranscriptBus, pub public_values_bus: PublicValuesBus, // pub cached_commit_bus: CachedCommitBus, pub pvs_air_consistency_bus: PvsAirConsistencyBus, @@ -107,6 +112,23 @@ impl Air f .when(local.has_verifier_pvs) .assert_one(local.is_valid); + for (digest_idx, digest) in local.child_vk_digest.iter().enumerate() { + for (limb_idx, value) in digest.iter().enumerate() { + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: AB::Expr::from_usize( + TranscriptLabel::Riscv.field_len() + digest_idx * D_EF + limb_idx, + ), + value: (*value).into(), + is_sample: AB::Expr::ZERO, + }, + local.is_valid, + ); + } + } + // We constrain the consistency of verifier-specific public values. We can determine // what layer a verifier is at using the has_verifier_pvs and internal_flag columns. // There are several cases we cover: @@ -253,9 +275,7 @@ impl Air f deferral_flag, has_verifier_pvs: local.has_verifier_pvs.into(), }, - local.is_valid - * consistency_mult - * AB::Expr::from_bool(!crate::system::TOWER_PREFIX_ONLY), + local.is_valid * consistency_mult, ); // Finally, we need to constrain that the public values this AIR produces are consistent diff --git a/ceno_recursion_v2/src/circuit/inner/verifier/mod.rs b/ceno_recursion_v2/src/circuit/inner/verifier/mod.rs index f10782180..fda9adf58 100644 --- a/ceno_recursion_v2/src/circuit/inner/verifier/mod.rs +++ b/ceno_recursion_v2/src/circuit/inner/verifier/mod.rs @@ -1,7 +1,7 @@ use openvm_stark_backend::{FiatShamirTranscript, TranscriptHistory}; use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config; -use crate::system::{Preflight, RecursionProof, RecursionVk}; +use crate::system::{Preflight, RecursionProof, RecursionVk, child_vk_digest}; mod air; mod trace; @@ -18,6 +18,8 @@ pub fn run_preflight( ) where TS: FiatShamirTranscript + TranscriptHistory, { - // Reserved verifier-owned preflight step. VmPvs currently owns transcript observations. - let _ = (child_vk, proof, ts); + for digest_elem in child_vk_digest(child_vk) { + ts.observe_ext(digest_elem); + } + let _ = proof; } diff --git a/ceno_recursion_v2/src/circuit/inner/verifier/trace.rs b/ceno_recursion_v2/src/circuit/inner/verifier/trace.rs index 0bfbecc00..3552131ce 100644 --- a/ceno_recursion_v2/src/circuit/inner/verifier/trace.rs +++ b/ceno_recursion_v2/src/circuit/inner/verifier/trace.rs @@ -4,7 +4,7 @@ use openvm_cpu_backend::CpuBackend; use openvm_poseidon2_air::POSEIDON2_WIDTH; use openvm_stark_backend::prover::AirProvingContext; use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, DIGEST_SIZE, F}; -use p3_field::PrimeCharacteristicRing; +use p3_field::{BasedVectorSpace, PrimeCharacteristicRing}; use p3_matrix::dense::RowMajorMatrix; use verify_stark::pvs::{VerifierBasePvs, VerifierDefPvs}; @@ -13,7 +13,7 @@ use crate::{ ProofsType, verifier::air::{VerifierDeferralCols, VerifierPvsCols}, }, - system::RecursionProof, + system::{RecursionProof, RecursionVk, child_vk_digest}, }; /////////////////////////////////////////////////////////////////////////////// @@ -21,6 +21,7 @@ use crate::{ /////////////////////////////////////////////////////////////////////////////// pub fn generate_proving_ctx( + child_vk: &RecursionVk, proofs: &[RecursionProof], proofs_type: ProofsType, child_is_app: bool, @@ -33,6 +34,7 @@ pub fn generate_proving_ctx( // TODO(recursion-proof-bridge): populate verifier trace/public values from RecursionProof. // Any verifier-specific values not available on RecursionProof are currently zero-mocked. let _ = (proofs, proofs_type, child_is_app, child_dag_commit); + let child_vk_digest = child_vk_digest(child_vk); let rows = proofs.len().max(1).next_power_of_two(); let width = VerifierPvsCols::::width() @@ -44,19 +46,22 @@ pub fn generate_proving_ctx( let mut trace = vec![F::ZERO; rows * width]; - if rows > 0 { - let first_row = &mut trace[..width]; + for proof_idx in 0..proofs.len().max(1) { + let row = &mut trace[proof_idx * width..(proof_idx + 1) * width]; let base_width = VerifierPvsCols::::width(); - let (base_row, def_row) = first_row.split_at_mut(base_width); + let (base_row, def_row) = row.split_at_mut(base_width); let cols: &mut VerifierPvsCols = base_row.borrow_mut(); - cols.proof_idx = F::ZERO; + cols.proof_idx = F::from_usize(proof_idx); cols.is_valid = F::ONE; cols.has_verifier_pvs = F::ZERO; + for (dst, digest_elem) in cols.child_vk_digest.iter_mut().zip(child_vk_digest) { + dst.copy_from_slice(digest_elem.as_basis_coefficients_slice()); + } if deferral_enabled { let def_cols: &mut VerifierDeferralCols = def_row.borrow_mut(); - def_cols.is_last = F::ONE; + def_cols.is_last = F::from_bool(proof_idx + 1 == proofs.len().max(1)); def_cols.child_pvs.deferral_flag = F::ZERO; } } diff --git a/ceno_recursion_v2/src/circuit/inner/vm_pvs/air.rs b/ceno_recursion_v2/src/circuit/inner/vm_pvs/air.rs index 23c362349..5a5818f4b 100644 --- a/ceno_recursion_v2/src/circuit/inner/vm_pvs/air.rs +++ b/ceno_recursion_v2/src/circuit/inner/vm_pvs/air.rs @@ -4,7 +4,7 @@ use ceno_emul::{FullTracer as Tracer, WORD_SIZE}; use ceno_zkvm::instructions::riscv::constants::{ END_CYCLE_IDX, END_PC_IDX, EXIT_CODE_IDX, EXIT_PC, HEAP_LENGTH_IDX, HEAP_START_ADDR_IDX, HINT_LENGTH_IDX, HINT_START_ADDR_IDX, INIT_CYCLE_IDX, INIT_PC_IDX, PUBIO_DIGEST_IDX, - SHARD_ID_IDX, SHARD_RW_SUM_IDX, + PUBIO_DIGEST_U16_LIMBS, SHARD_ID_IDX, SHARD_RW_SUM_IDX, }; use openvm_circuit_primitives::utils::{and, assert_array_eq, not}; use openvm_stark_backend::{ @@ -21,7 +21,10 @@ use stark_recursion_circuit_derive::AlignedBorrow; use crate::{ bus::{LookupChallengeBus, LookupChallengeKind, LookupChallengeMessage}, - circuit::inner::{bus::PvsAirConsistencyBus, vm_pvs::VmPvs}, + circuit::inner::{ + bus::{PvsAirConsistencyBus, PvsAirConsistencyMessage}, + vm_pvs::VmPvs, + }, }; #[repr(C)] @@ -160,7 +163,6 @@ impl Air f // When local and next are valid, enforce continuation consistency. let mut when_both_valid = builder.when(and(local.is_valid, not(local.is_last))); when_both_valid.assert_eq(local.child_pvs.end_pc, next.child_pvs.init_pc); - when_both_valid.assert_eq(local.child_pvs.end_cycle, next.child_pvs.init_cycle); when_both_valid.assert_eq( local.child_pvs.shard_id + AB::Expr::ONE, next.child_pvs.shard_id, @@ -187,7 +189,7 @@ impl Air f pv_idx: AB::Expr::from_usize(pv_idx), value: vm_public_value_by_index::(local, *global_pv_idx), }, - local.is_valid * AB::Expr::from_bool(!crate::system::TOWER_PREFIX_ONLY), + local.is_valid, ); } } @@ -339,15 +341,15 @@ impl Air f } // We look up proof metadata from VerifierPvsAir here to ensure consistency on each row. - // self.pvs_air_consistency_bus.lookup_key( - // builder, - // local.proof_idx, - // PvsAirConsistencyMessage { - // deferral_flag, - // has_verifier_pvs: local.has_verifier_pvs.into(), - // }, - // local.is_valid, - // ); + self.pvs_air_consistency_bus.lookup_key( + builder, + local.proof_idx, + PvsAirConsistencyMessage { + deferral_flag, + has_verifier_pvs: local.has_verifier_pvs.into(), + }, + local.is_valid, + ); // Finally, constrain that this AIR's output public values are consistent with child_pvs. let &VmPvs::<_> { @@ -389,29 +391,31 @@ impl Air f exit_code, ); - // constrain static per-proof public values + // The aggregate VM public values are first-shard init/static values plus last-shard + // exit/end values. Per-shard values are still constrained row-locally through the + // public-values and transcript buses above. builder - .when(local.is_valid) + .when_first_row() .assert_eq(local.child_pvs.shard_id, shard_id); builder - .when(local.is_valid) + .when_first_row() .assert_eq(local.child_pvs.heap_start_addr, heap_start_addr); builder - .when(local.is_valid) + .when_first_row() .assert_eq(local.child_pvs.heap_shard_len, heap_shard_len); builder - .when(local.is_valid) + .when_first_row() .assert_eq(local.child_pvs.hint_start_addr, hint_start_addr); builder - .when(local.is_valid) + .when_first_row() .assert_eq(local.child_pvs.hint_shard_len, hint_shard_len); assert_array_eq( - &mut builder.when(local.is_valid), + &mut builder.when_first_row(), local.child_pvs.public_io, public_io, ); assert_array_eq( - &mut builder.when(local.is_valid), + &mut builder.when_first_row(), local.child_pvs.shard_rw_sum, shard_rw_sum, ); @@ -423,10 +427,10 @@ impl Air f fixed_commit, ); builder - .when(local.is_valid) + .when_first_row() .assert_eq(local.child_pvs.init_pc, init_pc); builder - .when(local.is_valid) + .when_first_row() .assert_eq(local.child_pvs.init_cycle, init_cycle); assert_array_eq( &mut builder.when(local.is_valid), @@ -434,7 +438,7 @@ impl Air f fixed_no_omc_init_commit, ); assert_array_eq( - &mut builder.when(local.is_valid), + &mut builder.when_first_row(), local.child_pvs.witness_commit, witness_commit, ); @@ -463,8 +467,9 @@ where { local.child_pvs.shard_rw_sum[idx - SHARD_RW_SUM_IDX].into() } - idx if idx == PUBIO_DIGEST_IDX => local.child_pvs.public_io[0].into(), - idx if idx == PUBIO_DIGEST_IDX + 1 => local.child_pvs.public_io[1].into(), + idx if (PUBIO_DIGEST_IDX..(PUBIO_DIGEST_IDX + PUBIO_DIGEST_U16_LIMBS)).contains(&idx) => { + local.child_pvs.public_io[idx - PUBIO_DIGEST_IDX].into() + } _ => AB::Expr::ZERO, } } diff --git a/ceno_recursion_v2/src/circuit/inner/vm_pvs/metadata.rs b/ceno_recursion_v2/src/circuit/inner/vm_pvs/metadata.rs index ab9cd1cde..3f342cd93 100644 --- a/ceno_recursion_v2/src/circuit/inner/vm_pvs/metadata.rs +++ b/ceno_recursion_v2/src/circuit/inner/vm_pvs/metadata.rs @@ -1,18 +1,21 @@ use core::borrow::{Borrow, BorrowMut}; +use ceno_zkvm::structs::VK_DIGEST_LEN; use openvm_cpu_backend::CpuBackend; use openvm_stark_backend::{ BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, prover::AirProvingContext, }; -use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, DIGEST_SIZE, F}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{ + BabyBearPoseidon2Config, D_EF, DIGEST_SIZE, F, +}; use p3_air::{Air, AirBuilder, BaseAir}; use p3_field::PrimeCharacteristicRing; use p3_matrix::{Matrix, dense::RowMajorMatrix}; use stark_recursion_circuit_derive::AlignedBorrow; use crate::{ - bus::{TranscriptBus, TranscriptBusMessage}, + bus::{PcsCommitHeightBus, PcsCommitHeightMessage, TranscriptBus, TranscriptBusMessage}, circuit::inner::vm_pvs::RecursionCommitment, system::{Preflight, RecursionProof, RecursionVk}, tracegen::RowMajorChip, @@ -28,10 +31,12 @@ pub struct VmPvsCumulativeHeightsCols { pub height_idx: T, pub tidx: T, pub value: T, + pub lookup_count: T, } pub struct VmPvsCumulativeHeightsAir { pub transcript_bus: TranscriptBus, + pub commit_height_bus: PcsCommitHeightBus, } impl BaseAir for VmPvsCumulativeHeightsAir { @@ -60,6 +65,16 @@ impl Air for VmPvsCumulativeHeightsAir }, local.is_valid, ); + self.commit_height_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsCommitHeightMessage { + commitment_kind: local.commitment_kind.into(), + height_idx: local.height_idx.into(), + value: local.value.into(), + }, + local.is_valid * local.lookup_count, + ); } } @@ -86,7 +101,7 @@ impl RowMajorChip for VmPvsCumulativeHeightsTraceGenerator { }; let mut trace = vec![F::ZERO; height * width]; - for (row, (proof_idx, commitment_kind, height_idx, tidx, value)) in + for (row, (proof_idx, commitment_kind, height_idx, tidx, value, lookup_count)) in trace.chunks_exact_mut(width).zip(rows) { let cols: &mut VmPvsCumulativeHeightsCols = row.borrow_mut(); @@ -96,6 +111,7 @@ impl RowMajorChip for VmPvsCumulativeHeightsTraceGenerator { cols.height_idx = F::from_usize(height_idx); cols.tidx = F::from_usize(tidx); cols.value = value; + cols.lookup_count = F::from_usize(lookup_count); } Some(RowMajorMatrix::new(trace, width)) } @@ -118,13 +134,14 @@ fn collect_rows( child_vk: &RecursionVk, proofs: &[RecursionProof], preflights: &[Preflight], -) -> Vec<(usize, usize, usize, usize, F)> { +) -> Vec<(usize, usize, usize, usize, F, usize)> { proofs .iter() .zip(preflights) .enumerate() .flat_map(|(proof_idx, (proof, preflight))| { let mut tidx = TranscriptLabel::Riscv.field_len() + + VK_DIGEST_LEN * D_EF + child_vk .circuit_vks .values() @@ -142,6 +159,32 @@ fn collect_rows( } push_commit_cumulative_height_rows(&mut rows, proof_idx, 2, tidx, &proof.witin_commit); + let mut lookup_counts = + std::collections::BTreeMap::<(usize, usize, usize), usize>::new(); + if proof.public_values.shard_id == 0 { + if let Some(commitment) = child_vk.fixed_commit.as_ref() { + accumulate_commit_height_lookup_counts( + &mut lookup_counts, + proof_idx, + 0, + commitment, + ); + } + } else if let Some(commitment) = child_vk.fixed_no_omc_init_commit.as_ref() { + accumulate_commit_height_lookup_counts( + &mut lookup_counts, + proof_idx, + 1, + commitment, + ); + } + accumulate_commit_height_lookup_counts( + &mut lookup_counts, + proof_idx, + 2, + &proof.witin_commit, + ); + rows.into_iter() .filter_map(|(proof_idx, commitment_kind, height_idx, tidx)| { preflight @@ -149,7 +192,20 @@ fn collect_rows( .values() .get(tidx) .copied() - .map(|value| (proof_idx, commitment_kind, height_idx, tidx, value)) + .map(|value| { + let lookup_count = lookup_counts + .get(&(proof_idx, commitment_kind, height_idx)) + .copied() + .unwrap_or_default(); + ( + proof_idx, + commitment_kind, + height_idx, + tidx, + value, + lookup_count, + ) + }) }) .collect::>() }) @@ -160,6 +216,22 @@ fn commitment_transcript_len(commitment: &RecursionCommitment) -> usize { DIGEST_SIZE + 3 + commitment.cumulative_heights.len() } +fn accumulate_commit_height_lookup_counts( + lookup_counts: &mut std::collections::BTreeMap<(usize, usize, usize), usize>, + proof_idx: usize, + commitment_kind: usize, + commitment: &RecursionCommitment, +) { + for term_idx in 0..commitment.cumulative_heights.len().saturating_sub(1) { + *lookup_counts + .entry((proof_idx, commitment_kind, term_idx)) + .or_default() += 1; + *lookup_counts + .entry((proof_idx, commitment_kind, term_idx + 1)) + .or_default() += 1; + } +} + fn push_commit_cumulative_height_rows( rows: &mut Vec<(usize, usize, usize, usize)>, proof_idx: usize, diff --git a/ceno_recursion_v2/src/circuit/inner/vm_pvs/mod.rs b/ceno_recursion_v2/src/circuit/inner/vm_pvs/mod.rs index 433602fd6..393cce6c3 100644 --- a/ceno_recursion_v2/src/circuit/inner/vm_pvs/mod.rs +++ b/ceno_recursion_v2/src/circuit/inner/vm_pvs/mod.rs @@ -1,3 +1,4 @@ +use ceno_zkvm::instructions::riscv::constants::PUBIO_DIGEST_U16_LIMBS; use openvm_stark_backend::{FiatShamirTranscript, TranscriptHistory}; use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, DIGEST_SIZE, F}; use p3_field::PrimeCharacteristicRing; @@ -29,7 +30,7 @@ pub struct VmPvs { pub heap_shard_len: F, pub hint_start_addr: F, pub hint_shard_len: F, - pub public_io: [F; 2], + pub public_io: [F; PUBIO_DIGEST_U16_LIMBS], pub shard_rw_sum: [F; 2 * SEPTIC_EXTENSION_DEGREE], } diff --git a/ceno_recursion_v2/src/circuit/inner/vm_pvs/trace.rs b/ceno_recursion_v2/src/circuit/inner/vm_pvs/trace.rs index 514415e6e..ce18ab04e 100644 --- a/ceno_recursion_v2/src/circuit/inner/vm_pvs/trace.rs +++ b/ceno_recursion_v2/src/circuit/inner/vm_pvs/trace.rs @@ -1,3 +1,7 @@ +use ceno_zkvm::{ + instructions::riscv::constants::{LIMB_BITS, LIMB_MASK, PUBIO_DIGEST_U16_LIMBS, UINT_LIMBS}, + structs::VK_DIGEST_LEN, +}; use openvm_cpu_backend::CpuBackend; use openvm_stark_backend::prover::AirProvingContext; use openvm_stark_sdk::config::baby_bear_poseidon2::{ @@ -136,7 +140,7 @@ fn build_vm_pvs( heap_shard_len: F::from_u32(pv.heap_shard_len), hint_start_addr: F::from_u32(pv.hint_start_addr), hint_shard_len: F::from_u32(pv.hint_shard_len), - public_io: split_u32_lo_hi(pv.public_io_digest[0]), + public_io: split_public_io_digest(pv.public_io_digest), shard_rw_sum: pv.shard_rw_sum.map(F::from_u32), } } @@ -156,6 +160,7 @@ fn commitment_digest_tidxs( proof: &RecursionProof, ) -> (usize, usize, usize) { let mut tidx = crate::utils::TranscriptLabel::Riscv.field_len() + + VK_DIGEST_LEN * D_EF + child_vk .circuit_vks .values() @@ -195,6 +200,14 @@ fn split_u32_lo_hi(value: u32) -> [F; 2] { ] } +fn split_public_io_digest(digest: [u32; 8]) -> [F; PUBIO_DIGEST_U16_LIMBS] { + core::array::from_fn(|digest_limb_idx| { + let word_idx = digest_limb_idx / UINT_LIMBS; + let limb_idx = digest_limb_idx % UINT_LIMBS; + F::from_u32((digest[word_idx] >> (limb_idx * LIMB_BITS)) & LIMB_MASK) + }) +} + fn ef_to_limbs(value: EF) -> [F; D_EF] { let mut out = [F::ZERO; D_EF]; out.copy_from_slice(value.as_basis_coefficients_slice()); diff --git a/ceno_recursion_v2/src/continuation/prover/inner/mod.rs b/ceno_recursion_v2/src/continuation/prover/inner/mod.rs index a136a2f92..0ffa6478e 100644 --- a/ceno_recursion_v2/src/continuation/prover/inner/mod.rs +++ b/ceno_recursion_v2/src/continuation/prover/inner/mod.rs @@ -2,6 +2,8 @@ use std::{sync::Arc, time::Instant}; use continuations_v2::SC; use eyre::Result; +#[cfg(test)] +use openvm_cpu_backend::CpuBackend; use openvm_poseidon2_air::POSEIDON2_WIDTH; use openvm_stark_backend::{ StarkEngine, @@ -17,6 +19,8 @@ use openvm_stark_sdk::config::baby_bear_poseidon2::{ }; use verify_stark::pvs::DeferralPvs; +#[cfg(test)] +use crate::system::VerifierSubCircuit; use crate::{ circuit::{ Circuit, @@ -48,7 +52,7 @@ pub struct InnerAggregationProver< agg_node_tracegen: T, child_vk: Arc, - child_vk_pcs_data: CommittedTraceData, + child_vk_pcs_data: Option>, circuit: Arc>, self_vk_pcs_data: Option>, @@ -178,6 +182,7 @@ impl< > InnerAggregationProver where PB::Matrix: Clone, + PB::Commitment: Default, { pub fn agg_prove_no_def>( &self, @@ -235,6 +240,11 @@ where let mut transcript = default_duplex_sponge_recorder(); transcript_observe_label(&mut transcript, TranscriptLabel::Riscv.as_bytes()); + let child_dag_commit = child_vk_pcs_data + .as_ref() + .map(|data| data.commitment.clone()) + .unwrap_or_default(); + let (pre_ctxs, poseidon2_compress_inputs, subcircuit_initial_transcripts) = self .agg_node_tracegen .generate_pre_verifier_subcircuit_ctxs( @@ -243,7 +253,7 @@ where absent_trace_pvs, child_is_app, child_vk, - child_vk_pcs_data.commitment, + child_dag_commit, transcript, ); @@ -305,6 +315,101 @@ where } } +#[cfg(test)] +impl + InnerAggregationProver, VerifierSubCircuit, T> +where + T: InnerTraceGen>, +{ + pub(crate) fn debug_with_preflight_mutation( + &self, + proofs: &[RecursionProof], + child_vk_kind: ChildVkKind, + mutate: M, + ) -> bool + where + Eg: StarkEngine>, + M: FnOnce(&mut [crate::system::Preflight]) -> bool, + { + assert!(proofs.len() <= self.circuit.verifier_circuit.max_num_proofs()); + + let (child_vk, child_vk_pcs_data) = match child_vk_kind { + ChildVkKind::RecursiveSelf => { + unimplemented!("RecursiveSelf proving is not wired for RecursionVk yet") + } + _ => (&self.child_vk, self.child_vk_pcs_data.clone()), + }; + let child_is_app = matches!(child_vk_kind, ChildVkKind::App); + + let mut transcript = default_duplex_sponge_recorder(); + transcript_observe_label(&mut transcript, TranscriptLabel::Riscv.as_bytes()); + let child_dag_commit = child_vk_pcs_data + .as_ref() + .map(|data| data.commitment.clone()) + .unwrap_or_default(); + + let (pre_ctxs, poseidon2_compress_inputs, subcircuit_initial_transcripts) = self + .agg_node_tracegen + .generate_pre_verifier_subcircuit_ctxs( + proofs, + ProofsType::Vm, + None, + child_is_app, + child_vk, + child_dag_commit, + transcript, + ); + let poseidon2_permute_inputs: Vec<[F; POSEIDON2_WIDTH]> = vec![]; + let range_check_inputs = vec![]; + let power_check_inputs = vec![]; + let external_data = VerifierExternalData { + poseidon2_compress_inputs: &poseidon2_compress_inputs, + poseidon2_permute_inputs: &poseidon2_permute_inputs, + range_check_inputs: &range_check_inputs, + power_check_inputs: &power_check_inputs, + required_heights: None, + final_transcript_state: None, + }; + + let mut preflights = proofs + .iter() + .zip(subcircuit_initial_transcripts) + .map(|(proof, sponge)| { + self.circuit + .verifier_circuit + .test_run_preflight(sponge, child_vk, proof) + }) + .collect::>(); + if !mutate(&mut preflights) { + return false; + } + + let subcircuit_ctxs = self + .circuit + .verifier_circuit + .test_generate_proving_ctxs_from_preflights::( + child_vk, + proofs, + &preflights, + &external_data, + ); + let post_ctxs = self + .agg_node_tracegen + .generate_post_verifier_subcircuit_ctxs(proofs, ProofsType::Vm, child_is_app); + let ctx = ProvingContext { + per_trace: pre_ctxs + .into_iter() + .chain(subcircuit_ctxs) + .chain(post_ctxs) + .enumerate() + .collect(), + }; + let engine = Eg::new(self.pk.params.clone()); + debug_constraints(&self.circuit, &ctx, &engine); + true + } +} + fn trace_heights_tracing_info( ctxs: &[(usize, AirProvingContext)], airs: &[openvm_stark_backend::AirRef], diff --git a/ceno_recursion_v2/src/continuation/tests/mod.rs b/ceno_recursion_v2/src/continuation/tests/mod.rs index 1ef8823ba..015e28202 100644 --- a/ceno_recursion_v2/src/continuation/tests/mod.rs +++ b/ceno_recursion_v2/src/continuation/tests/mod.rs @@ -1,21 +1,22 @@ #[cfg(test)] mod prover_integration { use crate::{ - continuation::prover::{AggProver, AggregationOptions}, + continuation::prover::{AggProver, AggregationOptions, ChildVkKind, InnerCpuProver}, system::{ AggregationSubCircuit, RecursionField, RecursionProof, RecursionVk, VerifierSubCircuit, - utils::test_system_params_zero_pow, + VerifierTraceGen, utils::test_system_params_zero_pow, }, }; use bincode; - use ceno_zkvm::{ - scheme::{MainConstraintProof, PublicValues, ZKVMChipProof, ZKVMProof}, - structs::ZKVMVerifyingKey, - }; + use ceno_zkvm::scheme::{MainConstraintProof, PublicValues, ZKVMChipProof, ZKVMProof}; use eyre::{Result, eyre}; use ff_ext::GoldilocksExt2; use mpcs::{Basefold, BasefoldRSParams, Jagged, Whir, WhirDefaultSpec}; - use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2CpuEngine, DuplexSponge}; + use openvm_stark_sdk::config::baby_bear_poseidon2::{ + BabyBearPoseidon2CpuEngine, DuplexSponge, DuplexSpongeRecorder, F, + default_duplex_sponge_recorder, + }; + use p3_field::PrimeCharacteristicRing; use std::{ collections::BTreeMap, io::Cursor, @@ -222,6 +223,148 @@ mod prover_integration { Ok(()) } + fn first_fixture_proof_and_vk() -> Result> { + let Some((proofs, child_vk)) = load_fixtures()? else { + return Ok(None); + }; + Ok(Some((select_proofs(&proofs, 1)?.remove(0), child_vk))) + } + + fn assert_recursion_pcs_rejects(child_vk: &ZkvmVk, proof: ZkvmProof) { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + const MAX_NUM_PROOFS: usize = 2; + let subcircuit = VerifierSubCircuit::::new(Arc::new(child_vk.clone())); + let mut transcript = default_duplex_sponge_recorder(); + crate::utils::transcript_observe_label( + &mut transcript, + crate::utils::TranscriptLabel::Riscv.as_bytes(), + ); + as VerifierTraceGen< + openvm_cpu_backend::CpuBackend, + continuations_v2::SC, + >>::generate_proving_ctxs_base::( + &subcircuit, + child_vk, + None, + &[proof], + transcript, + ) + })); + assert!( + result.is_err(), + "mutated PCS proof unexpectedly generated recursion-v2 traces" + ); + } + + fn assert_recursion_pcs_air_rejects( + child_vk: &ZkvmVk, + proof: ZkvmProof, + case_name: &str, + mutate: impl FnOnce(&mut crate::system::Preflight), + ) -> bool { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let prover = InnerCpuProver::<2>::new::( + Arc::new(child_vk.clone()), + test_system_params_zero_pow(5, 16, 3), + false, + None, + ); + prover.debug_with_preflight_mutation::( + &[proof], + ChildVkKind::App, + |preflights| { + mutate(&mut preflights[0]); + true + }, + ) + })); + match result { + Ok(false) => false, + Ok(true) => { + panic!("{case_name}: mutated PCS preflight unexpectedly satisfied AIR/LogUp") + } + Err(_) => true, + } + } + + fn mutate_first_jagged_col_eval(proof: &mut ZkvmProof) { + proof.opening_proof.rounds[0].col_evals[0] += E::ONE; + } + + fn mutate_first_jagged_assist_round_eval(proof: &mut ZkvmProof) { + proof.opening_proof.rounds[0].assist_proof.proofs[0].evaluations[0] += E::ONE; + } + + fn mutate_last_jagged_assist_round_eval(proof: &mut ZkvmProof) { + let round = proof.opening_proof.rounds[0] + .assist_proof + .proofs + .last_mut() + .expect("fixture has assist sumcheck rounds"); + round.evaluations[0] += E::ONE; + } + + fn mutate_second_final_message_value(proof: &mut ZkvmProof) { + let row = proof + .opening_proof + .inner_proof + .final_message + .first_mut() + .expect("fixture has basefold final message rows"); + let idx = usize::from(row.len() > 1); + row[idx] += E::ONE; + } + + fn mutate_first_final_expected( + preflight: &mut crate::system::Preflight, + mutate: impl FnOnce(&mut crate::system::PcsBasefoldFinalExpectedRecord), + ) { + let record = preflight + .pcs + .basefold_final_expected + .first_mut() + .expect("fixture has basefold final expected rows"); + mutate(record); + } + + fn mutate_first_final_codeword( + preflight: &mut crate::system::Preflight, + mutate: impl FnOnce(&mut crate::system::PcsBasefoldFinalCodewordRecord), + ) { + let record = preflight + .pcs + .basefold_final_codeword + .first_mut() + .expect("fixture has basefold final codeword rows"); + mutate(record); + } + + fn mutate_first_final_message_transcript( + preflight: &mut crate::system::Preflight, + mutate: impl FnOnce(&mut crate::system::PcsTranscriptValueRecord), + ) { + let record = preflight + .pcs + .transcript_values + .iter_mut() + .find(|record| record.is_final_message) + .expect("fixture has PCS final-message transcript rows"); + mutate(record); + } + + fn mutate_first_basefold_sumcheck_round( + preflight: &mut crate::system::Preflight, + mutate: impl FnOnce(&mut crate::system::PcsSumcheckRoundRecord), + ) { + let record = preflight + .pcs + .sumcheck_rounds + .iter_mut() + .find(|record| record.idx == 8_000_000) + .expect("fixture has basefold sumcheck rows"); + mutate(record); + } + #[test] fn agg_prover_single_shard() -> Result<()> { agg_prove_with_count(1) @@ -232,6 +375,366 @@ mod prover_integration { agg_prove_with_count(2) } + #[test] + fn pcs_rejects_jagged_round_count_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + proof.opening_proof.rounds.clear(); + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_jagged_claimed_sum_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + proof.opening_proof.rounds[0].sumcheck_proof.proofs[0].evaluations[0] += E::ONE; + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_jagged_subclaim_multiplication_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + proof.opening_proof.rounds[0].f_at_rho += E::ONE; + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_jagged_q_eval_col_eval_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + mutate_first_jagged_col_eval(&mut proof); + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_jagged_assist_final_check_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + mutate_first_jagged_assist_round_eval(&mut proof); + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_jagged_assist_operand_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + mutate_last_jagged_assist_round_eval(&mut proof); + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_basefold_final_claim_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + proof.opening_proof.inner_proof.final_message[0][0] += E::ONE; + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_basefold_final_codeword_folding_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + mutate_second_final_message_value(&mut proof); + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_basefold_merkle_query_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + proof.opening_proof.inner_proof.query_opening_proof[0].input_proofs[0].opened_values[0] + [0] += F::ONE; + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_basefold_base_sibling_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + proof.opening_proof.inner_proof.query_opening_proof[0].input_proofs[0].opening_proof[0] + [0] += F::ONE; + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_basefold_base_root_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + let mut root: [F; 8] = proof.witin_commit.inner.commit.into(); + root[0] += F::ONE; + proof.witin_commit.inner.commit = root.into(); + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_basefold_commit_phase_sibling_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + proof.opening_proof.inner_proof.query_opening_proof[0].commit_phase_openings[0] + .opening_proof[0][0] += F::ONE; + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_rejects_basefold_commit_phase_root_mismatch() -> Result<()> { + let Some((mut proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + let mut root: [F; 8] = proof.opening_proof.inner_proof.commits[0].into(); + root[0] += F::ONE; + proof.opening_proof.inner_proof.commits[0] = root.into(); + assert_recursion_pcs_rejects(&child_vk, proof); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_final_expected_values() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "final_expected.point_value", + |p| { + mutate_first_final_expected(p, |record| record.point_value += E::ONE); + }, + ); + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "final_expected.acc_in", |p| { + mutate_first_final_expected(p, |record| record.acc_in += E::ONE); + }); + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "final_expected.acc_out", |p| { + mutate_first_final_expected(p, |record| record.acc_out += E::ONE); + }); + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "final_expected.coeff_in", + |p| { + mutate_first_final_expected(p, |record| record.coeff_in += E::ONE); + }, + ); + assert_recursion_pcs_air_rejects(&child_vk, proof, "final_expected.coeff_out", |p| { + mutate_first_final_expected(p, |record| record.coeff_out += E::ONE); + }); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_final_codeword_shape_values() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "final_codeword.coeff", |p| { + mutate_first_final_codeword(p, |record| record.coeff += E::ONE); + }); + assert_recursion_pcs_air_rejects(&child_vk, proof, "final_codeword.elem_idx", |p| { + mutate_first_final_codeword(p, |record| record.elem_idx += 1); + }); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_final_point_multiplicity_and_keys() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "drop final point row", |p| { + p.pcs.basefold_final_points.remove(0); + }); + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "duplicate final point row", + |p| { + let row = p.pcs.basefold_final_points[0].clone(); + p.pcs.basefold_final_points.push(row); + }, + ); + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "final point sumcheck_idx", + |p| { + p.pcs.basefold_final_points[0].sumcheck_idx += 1; + }, + ); + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "final point point_idx", |p| { + p.pcs.basefold_final_points[0].point_idx += 1; + }); + assert_recursion_pcs_air_rejects(&child_vk, proof, "final point coord_idx", |p| { + p.pcs.basefold_final_points[0].coord_idx += 1; + }); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_final_expected_multiplicity() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "drop final expected row", + |p| { + p.pcs.basefold_final_expected.remove(0); + }, + ); + assert_recursion_pcs_air_rejects(&child_vk, proof, "duplicate final expected row", |p| { + let row = p.pcs.basefold_final_expected[0].clone(); + p.pcs.basefold_final_expected.push(row); + }); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_final_message_transcript_binding() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "final message tidx", |p| { + mutate_first_final_message_transcript(p, |record| record.tidx += 1); + }); + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "final message value", |p| { + mutate_first_final_message_transcript(p, |record| record.value += E::ONE); + }); + assert_recursion_pcs_air_rejects(&child_vk, proof, "drop final message row", |p| { + let idx = p + .pcs + .transcript_values + .iter() + .position(|record| record.is_final_message) + .expect("fixture has PCS final-message transcript rows"); + p.pcs.transcript_values.remove(idx); + }); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_basefold_sumcheck_count_and_keys() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "basefold fold challenge multiplicity", + |p| { + mutate_first_basefold_sumcheck_round(p, |record| { + record.fold_challenge_lookup_count += 1; + }); + }, + ); + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "basefold sumcheck round", + |p| { + mutate_first_basefold_sumcheck_round(p, |record| record.round += 1); + }, + ); + assert_recursion_pcs_air_rejects(&child_vk, proof, "drop basefold sumcheck round", |p| { + let idx = p + .pcs + .sumcheck_rounds + .iter() + .position(|record| record.idx == 8_000_000) + .expect("fixture has basefold sumcheck rows"); + p.pcs.sumcheck_rounds.remove(idx); + }); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_basefold_query_and_opening_counts() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects(&child_vk, proof.clone(), "drop query index row", |p| { + p.pcs.basefold_query_indices.remove(0); + }); + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "duplicate query index row", + |p| { + let row = p.pcs.basefold_query_indices[0].clone(); + p.pcs.basefold_query_indices.push(row); + }, + ); + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "drop base input leaf row", + |p| { + p.pcs.base_input_leaf_hashes.remove(0); + }, + ); + assert_recursion_pcs_air_rejects( + &child_vk, + proof, + "drop basefold query opening row", + |p| { + p.pcs.basefold_query_opens.remove(0); + }, + ); + Ok(()) + } + + #[test] + fn pcs_air_rejects_mutated_commit_phase_counts() -> Result<()> { + let Some((proof, child_vk)) = first_fixture_proof_and_vk()? else { + return Ok(()); + }; + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "drop commit phase leaf row", + |p| { + p.pcs.commit_phase_leaf_hashes.remove(0); + }, + ); + assert_recursion_pcs_air_rejects( + &child_vk, + proof.clone(), + "duplicate commit phase leaf row", + |p| { + let row = p.pcs.commit_phase_leaf_hashes[0].clone(); + p.pcs.commit_phase_leaf_hashes.push(row); + }, + ); + assert_recursion_pcs_air_rejects(&child_vk, proof, "drop commit phase merkle row", |p| { + p.pcs.commit_phase_merkle_rows.remove(0); + }); + Ok(()) + } + // ---- AIR registration test ---- #[test] @@ -248,24 +751,40 @@ mod prover_integration { .map(|air| air.name().to_string()) .collect::>(); air_names.sort_unstable(); - assert!( - air_names - .iter() - .any(|name| name.contains("SymbolicExpressionAir")), - "SymbolicExpressionAir missing from recursion v2 verifier" - ); - assert!( - air_names - .iter() - .any(|name| name.contains("ConstraintsFoldingAir")), - "ConstraintsFoldingAir missing from recursion v2 verifier" - ); - assert!( - air_names - .iter() - .any(|name| name.contains("ExpressionClaimAir")), - "ExpressionClaimAir missing from recursion v2 verifier" - ); + let batch_main_airs = [ + "MainGlobalSumcheckAir", + "MainEvalAbsorbAir", + "MainTowerPointEqAir", + "MainFrontloadTermAir", + "MainFinalClaimAir", + ]; + for air in batch_main_airs { + assert!( + air_names.iter().any(|name| name.contains(air)), + "{air} missing from recursion v2 verifier" + ); + } + for air in [ + "PcsEqProductAir", + "PcsSuffixProductAir", + "PcsJaggedAssistHAir", + "PcsJaggedAssistQAir", + ] { + assert!( + air_names.iter().any(|name| name.contains(air)), + "{air} missing from recursion v2 verifier" + ); + } + for removed_air in [ + "SymbolicExpressionAir", + "ConstraintsFoldingAir", + "ExpressionClaimAir", + ] { + assert!( + !air_names.iter().any(|name| name.contains(removed_air)), + "{removed_air} should not be registered in recursion v2 verifier" + ); + } println!("registered recursion v2 airs={}", air_names.len()); Ok(()) diff --git a/ceno_recursion_v2/src/cuda/types.rs b/ceno_recursion_v2/src/cuda/types.rs index d6c3d1436..d57990245 100644 --- a/ceno_recursion_v2/src/cuda/types.rs +++ b/ceno_recursion_v2/src/cuda/types.rs @@ -12,7 +12,6 @@ pub struct TraceHeight { pub struct TraceMetadata { pub cached_idx: usize, pub total_interactions: usize, - pub num_air_id_lookups: usize, } #[repr(C)] @@ -32,5 +31,4 @@ pub struct AirData { pub num_interactions_per_row: usize, pub total_width: usize, pub has_preprocessed: bool, - pub need_rot: bool, } diff --git a/ceno_recursion_v2/src/lib.rs b/ceno_recursion_v2/src/lib.rs index 623d77720..1672911d1 100644 --- a/ceno_recursion_v2/src/lib.rs +++ b/ceno_recursion_v2/src/lib.rs @@ -6,6 +6,7 @@ pub mod bn254; pub mod circuit; pub mod continuation; pub mod main; +pub mod pcs; pub mod proof_shape; pub mod system; pub mod tower; diff --git a/ceno_recursion_v2/src/main/air.rs b/ceno_recursion_v2/src/main/air.rs index d94927630..81cfc0f83 100644 --- a/ceno_recursion_v2/src/main/air.rs +++ b/ceno_recursion_v2/src/main/air.rs @@ -6,20 +6,18 @@ use openvm_stark_backend::{ }; use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF; use p3_air::{Air, AirBuilder, BaseAir}; -use p3_field::Field; +use p3_field::{Field, PrimeCharacteristicRing}; use p3_matrix::Matrix; use recursion_circuit::subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir}; use stark_recursion_circuit_derive::AlignedBorrow; -use crate::bus::{ - MainBus, MainExpressionClaimBus, MainExpressionClaimMessage, MainMessage, MainSumcheckInputBus, - MainSumcheckInputMessage, MainSumcheckOutputBus, MainSumcheckOutputMessage, -}; +use crate::bus::{MainBus, MainExpressionClaimBus, MainExpressionClaimMessage, MainMessage}; #[repr(C)] #[derive(AlignedBorrow, Debug)] pub struct MainCols { pub is_enabled: T, + pub is_present: T, pub proof_idx: T, pub idx: T, pub is_first_idx: T, @@ -31,9 +29,8 @@ pub struct MainCols { pub struct MainAir { pub main_bus: MainBus, - pub sumcheck_input_bus: MainSumcheckInputBus, - pub sumcheck_output_bus: MainSumcheckOutputBus, pub expression_claim_bus: MainExpressionClaimBus, + pub send_expression_claim: bool, } impl BaseAir for MainAir { @@ -55,8 +52,15 @@ impl Air for MainAir { let local: &MainCols = (*local_row).borrow(); let next: &MainCols = (*next_row).borrow(); - type LoopSubAir = NestedForLoopSubAir<2>; - LoopSubAir {}.eval( + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_present); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_first_idx); + builder + .when(local.is_present) + .assert_one(local.is_enabled); + + NestedForLoopSubAir::<2> {}.eval( builder, ( NestedForLoopIoCols { @@ -74,7 +78,6 @@ impl Air for MainAir { ), ); - let receive_mask = local.is_enabled * local.is_first; self.main_bus.receive( builder, local.proof_idx, @@ -83,32 +86,11 @@ impl Air for MainAir { tidx: local.tidx.into(), claim: local.claim_in.map(Into::into), }, - receive_mask, - ); - - self.sumcheck_input_bus.send( - builder, - local.proof_idx, - MainSumcheckInputMessage { - idx: local.idx.into(), - tidx: local.tidx.into(), - claim: local.claim_in.map(Into::into), - }, - local.is_enabled, - ); - - self.sumcheck_output_bus.receive( - builder, - local.proof_idx, - MainSumcheckOutputMessage { - idx: local.idx.into(), - claim: local.claim_out.map(Into::into), - }, - local.is_enabled, + local.is_present * local.is_first, ); assert_array_eq( - &mut builder.when(local.is_enabled), + &mut builder.when(local.is_present), local.claim_in, local.claim_out, ); @@ -120,7 +102,7 @@ impl Air for MainAir { idx: local.idx.into(), claim: local.claim_out.map(Into::into), }, - local.is_enabled * local.is_first, + local.is_present * local.is_first * AB::Expr::from_bool(self.send_expression_claim), ); } } diff --git a/ceno_recursion_v2/src/main/ecc_rt.rs b/ceno_recursion_v2/src/main/ecc_rt.rs new file mode 100644 index 000000000..7133369e9 --- /dev/null +++ b/ceno_recursion_v2/src/main/ecc_rt.rs @@ -0,0 +1,1903 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_circuit_primitives::utils::assert_array_eq; +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, extension::BinomiallyExtendable}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use recursion_circuit::utils::{ + ext_field_add, ext_field_multiply, ext_field_multiply_scalar, ext_field_one_minus, + ext_field_subtract, +}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{ + EccRtBus, EccRtMessage, ForkedTranscriptBus, ForkedTranscriptBusMessage, + MainEccRtChallengeBus, MainEccRtChallengeKind, MainEccRtChallengeMessage, + MainEccRtEquationTotalsBus, MainEccRtEquationTotalsMessage, MainEccRtQuarkFinalBus, + MainEccRtQuarkFinalMessage, MainEccRtSumcheckFinalBus, MainEccRtSumcheckFinalMessage, + }, + system::MainEccRtRecord, + tracegen::RowMajorChip, +}; + +const SEPTIC_DEGREE: usize = 7; +const ECC_EQUATION_FAMILIES: usize = 7; +const MAX_ECC_VARS: usize = 32; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainEccRtChallengeCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub fork_id: T, + pub round_idx: T, + pub num_rounds: T, + pub is_first: T, + pub tidx: T, + pub out_tidx: T, + pub alpha_tidx: T, + pub lookup_count: T, + pub rt: [T; D_EF], + pub out_rt: [T; D_EF], + pub alpha: [T; D_EF], +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainEccRtCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub round_idx: T, + pub is_first: T, + pub is_last: T, + pub rt: [T; D_EF], + pub out_rt: [T; D_EF], + pub sumcheck_final_claim: [T; D_EF], + pub sel_add: [T; D_EF], + pub sel_bypass: [T; D_EF], + pub sel_export: [T; D_EF], + pub eq_in: [T; D_EF], + pub eq_out: [T; D_EF], + pub last_in: [T; D_EF], + pub last_out: [T; D_EF], + pub export_out_in: [T; D_EF], + pub export_out_out: [T; D_EF], + pub export_rt_in: [T; D_EF], + pub export_rt_out: [T; D_EF], + pub quark_out: [T; D_EF], + pub add_eval: [T; D_EF], + pub bypass_eval: [T; D_EF], + pub export_eval: [T; D_EF], +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainEccRtEquationCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub family_idx: T, + pub septic_idx: T, + pub is_first: T, + pub is_last: T, + pub family_flags: [T; ECC_EQUATION_FAMILIES], + pub septic_flags: [T; SEPTIC_DEGREE], + pub alpha: [T; D_EF], + pub alpha_pow: [T; D_EF], + pub s0: [T; D_EF], + pub s0_all: [[T; D_EF]; SEPTIC_DEGREE], + pub x0: [T; D_EF], + pub x0_all: [[T; D_EF]; SEPTIC_DEGREE], + pub y0: [T; D_EF], + pub x1: [T; D_EF], + pub x1_all: [[T; D_EF]; SEPTIC_DEGREE], + pub y1: [T; D_EF], + pub x3: [T; D_EF], + pub x3_all: [[T; D_EF]; SEPTIC_DEGREE], + pub y3: [T; D_EF], + pub sum_x: [T; D_EF], + pub sum_y: [T; D_EF], + pub s0_x0_x1: [T; D_EF], + pub s0_squared: [T; D_EF], + pub s0_x0_x3: [T; D_EF], + pub add_in: [T; D_EF], + pub add_out: [T; D_EF], + pub bypass_in: [T; D_EF], + pub bypass_out: [T; D_EF], + pub export_in: [T; D_EF], + pub export_out: [T; D_EF], +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainEccRtSumcheckCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub round_idx: T, + pub is_first: T, + pub is_last: T, + pub rt: [T; D_EF], + pub ev1: [T; D_EF], + pub ev2: [T; D_EF], + pub ev3: [T; D_EF], + pub claim_in: [T; D_EF], + pub claim_out: [T; D_EF], +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainEccRtQuarkCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub round_idx: T, + pub bit_idx: T, + pub bit_weight: T, + pub is_first_round: T, + pub is_last_round: T, + pub is_first_bit: T, + pub is_last_bit: T, + pub is_active: T, + pub bit: T, + pub current_rt: [T; D_EF], + pub current_out_rt: [T; D_EF], + pub point_rt: [T; D_EF], + pub point_out_rt: [T; D_EF], + pub same_one: [T; D_EF], + pub same_zero: [T; D_EF], + pub equal_choice: [T; D_EF], + pub active_prefix: [T; D_EF], + pub prefix_in: [T; D_EF], + pub prefix_out: [T; D_EF], + pub less_in: [T; D_EF], + pub less_from_prior: [T; D_EF], + pub less_from_equal: [T; D_EF], + pub active_less: [T; D_EF], + pub less_out: [T; D_EF], + pub active_count_in: T, + pub active_count_out: T, + pub bit_value_in: T, + pub bit_value_out: T, + pub quark_prefix_count: T, + pub quark_prefix_is_zero: T, + pub quark_prefix_inv: T, + pub quark_layer_n: T, + pub quark_parity: T, + pub quark_in: [T; D_EF], + pub quark_factor: [T; D_EF], + pub quark_zero: [T; D_EF], + pub quark_one: [T; D_EF], + pub quark_lhs: [T; D_EF], + pub quark_rhs: [T; D_EF], + pub quark_out: [T; D_EF], +} + +pub struct MainEccRtChallengeAir { + pub forked_transcript_bus: ForkedTranscriptBus, + pub ecc_rt_bus: EccRtBus, + pub challenge_bus: MainEccRtChallengeBus, +} + +impl BaseAir for MainEccRtChallengeAir { + fn width(&self) -> usize { + MainEccRtChallengeCols::::width() + } +} + +impl BaseAirWithPublicValues for MainEccRtChallengeAir {} +impl PartitionedBaseAir for MainEccRtChallengeAir {} + +impl Air for MainEccRtChallengeAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &MainEccRtChallengeCols = (*local_row).borrow(); + let next: &MainEccRtChallengeCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + + for i in 0..D_EF { + self.forked_transcript_bus.receive( + builder, + local.proof_idx, + ForkedTranscriptBusMessage { + fork_id: local.fork_id.into(), + tidx: local.tidx + AB::Expr::from_usize(i), + value: local.rt[i].into(), + is_sample: AB::Expr::ONE, + }, + local.is_enabled, + ); + self.forked_transcript_bus.receive( + builder, + local.proof_idx, + ForkedTranscriptBusMessage { + fork_id: local.fork_id.into(), + tidx: local.out_tidx + AB::Expr::from_usize(i), + value: local.out_rt[i].into(), + is_sample: AB::Expr::ONE, + }, + local.is_enabled, + ); + self.forked_transcript_bus.receive( + builder, + local.proof_idx, + ForkedTranscriptBusMessage { + fork_id: local.fork_id.into(), + tidx: local.alpha_tidx + AB::Expr::from_usize(i), + value: local.alpha[i].into(), + is_sample: AB::Expr::ONE, + }, + local.is_enabled * local.is_first, + ); + } + + self.ecc_rt_bus.add_key_with_lookups( + builder, + local.proof_idx, + EccRtMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + value: local.rt.map(Into::into), + }, + local.is_enabled * local.lookup_count, + ); + + self.challenge_bus.add_key_with_lookups( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::Rt.as_usize()), + value: local.rt.map(Into::into), + }, + local.is_enabled * (local.num_rounds - local.round_idx + AB::Expr::from_usize(2)), + ); + self.challenge_bus.add_key_with_lookups( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::OutRt.as_usize()), + value: local.out_rt.map(Into::into), + }, + local.is_enabled * (local.num_rounds - local.round_idx + AB::Expr::ONE), + ); + self.challenge_bus.add_key_with_lookups( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::Alpha.as_usize()), + value: local.alpha.map(Into::into), + }, + local.is_enabled + * local.is_first + * AB::Expr::from_usize(ECC_EQUATION_FAMILIES * SEPTIC_DEGREE), + ); + + let same_ecc = local.is_enabled * next.is_enabled * (AB::Expr::ONE - next.is_first); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.idx, local.idx); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.fork_id, local.fork_id); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.round_idx, local.round_idx + AB::Expr::ONE); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.num_rounds, local.num_rounds); + assert_array_eq( + &mut builder.when_transition().when(same_ecc), + local.alpha, + next.alpha, + ); + } +} + +pub struct MainEccRtAir { + pub challenge_bus: MainEccRtChallengeBus, + pub sumcheck_final_bus: MainEccRtSumcheckFinalBus, + pub equation_totals_bus: MainEccRtEquationTotalsBus, + pub quark_final_bus: MainEccRtQuarkFinalBus, +} + +pub struct MainEccRtEquationAir { + pub challenge_bus: MainEccRtChallengeBus, + pub equation_totals_bus: MainEccRtEquationTotalsBus, +} + +impl BaseAir for MainEccRtEquationAir { + fn width(&self) -> usize { + MainEccRtEquationCols::::width() + } +} + +impl BaseAirWithPublicValues for MainEccRtEquationAir {} +impl PartitionedBaseAir for MainEccRtEquationAir {} + +impl Air for MainEccRtEquationAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &MainEccRtEquationCols = (*local_row).borrow(); + let next: &MainEccRtEquationCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + + let mut family_sum = AB::Expr::ZERO; + let mut family_idx = AB::Expr::ZERO; + for (i, flag) in local.family_flags.iter().enumerate() { + builder.assert_bool(*flag); + family_sum += *flag; + family_idx += *flag * AB::Expr::from_usize(i); + } + builder + .when(local.is_enabled) + .assert_eq(family_sum, AB::Expr::ONE); + builder + .when(local.is_enabled) + .assert_eq(local.family_idx, family_idx); + + let mut septic_sum = AB::Expr::ZERO; + let mut septic_idx = AB::Expr::ZERO; + for (i, flag) in local.septic_flags.iter().enumerate() { + builder.assert_bool(*flag); + septic_sum += *flag; + septic_idx += *flag * AB::Expr::from_usize(i); + } + builder + .when(local.is_enabled) + .assert_eq(septic_sum, AB::Expr::ONE); + builder + .when(local.is_enabled) + .assert_eq(local.septic_idx, septic_idx); + + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: AB::Expr::ZERO, + kind: AB::Expr::from_usize(MainEccRtChallengeKind::Alpha.as_usize()), + value: local.alpha.map(Into::into), + }, + local.is_enabled, + ); + + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.alpha_pow, + ext_one::(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.add_in, + ext_zero::(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.bypass_in, + ext_zero::(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.export_in, + ext_zero::(), + ); + + let x0_minus_x1 = core::array::from_fn(|i| { + ext_field_subtract::(local.x0_all[i], local.x1_all[i]) + }); + let x0_minus_x3 = core::array::from_fn(|i| { + ext_field_subtract::(local.x0_all[i], local.x3_all[i]) + }); + let s0_expr = core::array::from_fn(|i| local.s0_all[i].map(Into::into)); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.s0_x0_x1, + septic_mul_selected::(&local.s0_all, &x0_minus_x1, local.septic_flags.clone()), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.s0_squared, + septic_mul_selected::(&local.s0_all, &s0_expr, local.septic_flags.clone()), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.s0_x0_x3, + septic_mul_selected::(&local.s0_all, &x0_minus_x3, local.septic_flags), + ); + let v1 = ext_field_subtract::( + local.s0_x0_x1, + ext_field_subtract::(local.y0, local.y1), + ); + let v2 = ext_field_subtract::( + ext_field_subtract::( + ext_field_subtract::(local.s0_squared, local.x0), + local.x1, + ), + local.x3, + ); + let v3 = ext_field_subtract::( + local.s0_x0_x3, + ext_field_add::(local.y0, local.y3), + ); + let v4 = ext_field_subtract::(local.x3, local.x0); + let v5 = ext_field_subtract::(local.y3, local.y0); + let v6 = ext_field_subtract::(local.x3, local.sum_x); + let v7 = ext_field_subtract::(local.y3, local.sum_y); + let add_term = [v1.clone(), v2.clone(), v3.clone()] + .into_iter() + .enumerate() + .fold(ext_zero::(), |acc, (i, term)| { + ext_field_add::( + acc, + ext_field_multiply_scalar::(term, local.family_flags[i]), + ) + }); + let bypass_term = [v4.clone(), v5.clone()].into_iter().enumerate().fold( + ext_zero::(), + |acc, (i, term)| { + ext_field_add::( + acc, + ext_field_multiply_scalar::(term, local.family_flags[3 + i]), + ) + }, + ); + let export_term = + [v6, v7] + .into_iter() + .enumerate() + .fold(ext_zero::(), |acc, (i, term)| { + ext_field_add::( + acc, + ext_field_multiply_scalar::(term, local.family_flags[5 + i]), + ) + }); + let add_weighted = ext_field_multiply::(add_term, local.alpha_pow); + let bypass_weighted = ext_field_multiply::(bypass_term, local.alpha_pow); + let export_weighted = ext_field_multiply::(export_term, local.alpha_pow); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.add_out, + ext_field_add::(local.add_in.map(Into::into), add_weighted), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.bypass_out, + ext_field_add::(local.bypass_in.map(Into::into), bypass_weighted), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.export_out, + ext_field_add::(local.export_in.map(Into::into), export_weighted), + ); + + let same_ecc = local.is_enabled * (AB::Expr::ONE - local.is_last); + let is_septic_last = local.septic_flags[SEPTIC_DEGREE - 1]; + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.idx, local.idx); + builder.when_transition().when(same_ecc.clone()).assert_eq( + next.septic_idx, + local.septic_idx + AB::Expr::ONE - is_septic_last * AB::Expr::from_usize(SEPTIC_DEGREE), + ); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.family_idx, local.family_idx + is_septic_last); + assert_array_eq( + &mut builder.when_transition().when(same_ecc.clone()), + next.alpha_pow, + ext_field_multiply::(local.alpha_pow, local.alpha), + ); + assert_array_eq( + &mut builder.when_transition().when(same_ecc.clone()), + local.add_out, + next.add_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_ecc.clone()), + local.bypass_out, + next.bypass_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_ecc), + local.export_out, + next.export_in, + ); + + self.equation_totals_bus.send( + builder, + local.proof_idx, + MainEccRtEquationTotalsMessage { + idx: local.idx.into(), + add_eval: local.add_out.map(Into::into), + bypass_eval: local.bypass_out.map(Into::into), + export_eval: local.export_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +pub struct MainEccRtQuarkAir { + pub challenge_bus: MainEccRtChallengeBus, + pub quark_final_bus: MainEccRtQuarkFinalBus, +} + +impl BaseAir for MainEccRtQuarkAir { + fn width(&self) -> usize { + MainEccRtQuarkCols::::width() + } +} + +impl BaseAirWithPublicValues for MainEccRtQuarkAir {} +impl PartitionedBaseAir for MainEccRtQuarkAir {} + +impl Air for MainEccRtQuarkAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &MainEccRtQuarkCols = (*local_row).borrow(); + let next: &MainEccRtQuarkCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first_round); + builder.assert_bool(local.is_last_round); + builder.assert_bool(local.is_first_bit); + builder.assert_bool(local.is_last_bit); + builder.assert_bool(local.is_active); + builder.assert_bool(local.bit); + builder.assert_bool(local.quark_parity); + builder.assert_bool(local.quark_prefix_is_zero); + builder + .when(local.is_enabled * local.is_first_bit) + .assert_eq( + local.bit_weight, + AB::Expr::from_usize(1usize << (MAX_ECC_VARS - 1)), + ); + builder + .when(local.is_enabled) + .assert_zero(local.bit * (AB::Expr::ONE - local.is_active)); + builder + .when(local.is_enabled) + .assert_zero(local.quark_prefix_is_zero * local.quark_prefix_count); + builder.when(local.is_enabled).assert_eq( + local.quark_prefix_count * local.quark_prefix_inv, + AB::Expr::ONE - local.quark_prefix_is_zero, + ); + builder.when(local.is_enabled).assert_eq( + local.quark_layer_n, + local.quark_prefix_count * AB::Expr::from_usize(2) + local.quark_parity, + ); + + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::Rt.as_usize()), + value: local.current_rt.map(Into::into), + }, + local.is_enabled * local.is_first_bit, + ); + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::OutRt.as_usize()), + value: local.current_out_rt.map(Into::into), + }, + local.is_enabled * local.is_first_bit, + ); + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.bit_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::Rt.as_usize()), + value: local.point_rt.map(Into::into), + }, + local.is_enabled * local.is_active, + ); + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.bit_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::OutRt.as_usize()), + value: local.point_out_rt.map(Into::into), + }, + local.is_enabled * local.is_active, + ); + + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first_bit), + local.prefix_in, + ext_one::(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first_bit), + local.less_in, + ext_zero::(), + ); + builder + .when(local.is_enabled * local.is_first_bit) + .assert_zero(local.active_count_in); + builder + .when(local.is_enabled * local.is_first_bit) + .assert_zero(local.bit_value_in); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first_round), + local.quark_in, + ext_zero::(), + ); + + assert_array_eq( + &mut builder.when(local.is_enabled), + local.same_one, + ext_field_multiply::(local.point_out_rt, local.point_rt), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.same_zero, + ext_field_multiply::( + one_minus_ext::(local.point_out_rt), + one_minus_ext::(local.point_rt), + ), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.equal_choice, + choose_ext::( + local.bit, + local.same_one.map(Into::into), + local.same_zero.map(Into::into), + ), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.active_prefix, + ext_field_multiply::(local.prefix_in, local.equal_choice), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.prefix_out, + choose_ext::( + local.is_active, + local.active_prefix.map(Into::into), + local.prefix_in.map(Into::into), + ), + ); + let same_any = ext_field_add::( + local.same_one.map(Into::into), + local.same_zero.map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.less_from_prior, + ext_field_multiply::(local.less_in, same_any), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.less_from_equal, + ext_field_multiply::(local.prefix_in, local.same_zero.map(Into::into)), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.active_less, + ext_field_add::( + local.less_from_prior.map(Into::into), + ext_field_multiply_scalar::( + local.less_from_equal.map(Into::into), + local.bit, + ), + ), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.less_out, + choose_ext::( + local.is_active, + local.active_less.map(Into::into), + local.less_in.map(Into::into), + ), + ); + builder.when(local.is_enabled).assert_eq( + local.active_count_out, + local.active_count_in + local.is_active, + ); + builder.when(local.is_enabled).assert_eq( + local.bit_value_out, + local.bit_value_in + local.bit * local.bit_weight, + ); + builder + .when(local.is_enabled * local.is_last_bit) + .assert_eq(local.active_count_out, local.round_idx); + builder + .when(local.is_enabled * local.is_last_bit) + .assert_zero( + (AB::Expr::ONE - local.quark_prefix_is_zero) + * (local.bit_value_out + AB::Expr::ONE - local.quark_prefix_count), + ); + + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first_round), + local.quark_factor, + ext_one::(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.quark_prefix_is_zero), + local.quark_factor, + ext_zero::(), + ); + let lte_value = ext_field_add::(local.prefix_out, local.less_out); + assert_array_eq( + &mut builder.when( + local.is_enabled + * local.is_last_bit + * (AB::Expr::ONE - local.is_first_round) + * (AB::Expr::ONE - local.quark_prefix_is_zero), + ), + local.quark_factor, + lte_value, + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last_bit), + local.quark_zero, + ext_field_multiply::( + one_minus_ext::(local.current_out_rt), + one_minus_ext::(local.current_rt), + ), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last_bit), + local.quark_one, + ext_field_multiply::(local.current_out_rt, local.current_rt), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last_bit), + local.quark_lhs, + ext_field_multiply::(local.quark_zero, local.quark_factor), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last_bit), + local.quark_rhs, + ext_field_multiply::(local.quark_one, local.quark_in), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last_bit), + local.quark_out, + ext_field_add::( + local.quark_lhs.map(Into::into), + local.quark_rhs.map(Into::into), + ), + ); + + let same_round = local.is_enabled * (AB::Expr::ONE - local.is_last_bit); + builder + .when_transition() + .when(same_round.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.idx, local.idx); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.round_idx, local.round_idx); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.bit_idx + AB::Expr::ONE, local.bit_idx); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(local.bit_weight, next.bit_weight * AB::Expr::from_usize(2)); + builder + .when_transition() + .when(same_round.clone()) + .assert_zero(local.is_active * (AB::Expr::ONE - next.is_active)); + assert_array_eq( + &mut builder.when_transition().when(same_round.clone()), + local.prefix_out, + next.prefix_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_round.clone()), + local.less_out, + next.less_in, + ); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.active_count_in, local.active_count_out); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.bit_value_in, local.bit_value_out); + assert_array_eq( + &mut builder.when_transition().when(same_round.clone()), + local.current_rt, + next.current_rt, + ); + assert_array_eq( + &mut builder.when_transition().when(same_round.clone()), + local.current_out_rt, + next.current_out_rt, + ); + assert_array_eq( + &mut builder.when_transition().when(same_round.clone()), + local.quark_in, + next.quark_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_round.clone()), + local.quark_factor, + next.quark_factor, + ); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.quark_prefix_count, local.quark_prefix_count); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.quark_prefix_is_zero, local.quark_prefix_is_zero); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.quark_prefix_inv, local.quark_prefix_inv); + builder + .when_transition() + .when(same_round.clone()) + .assert_eq(next.quark_layer_n, local.quark_layer_n); + builder + .when_transition() + .when(same_round) + .assert_eq(next.quark_parity, local.quark_parity); + + let next_round = + local.is_enabled * local.is_last_bit * (AB::Expr::ONE - local.is_last_round); + builder + .when_transition() + .when(next_round.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(next_round.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(next_round.clone()) + .assert_eq(next.idx, local.idx); + builder + .when_transition() + .when(next_round.clone()) + .assert_eq(next.round_idx, local.round_idx + AB::Expr::ONE); + builder + .when_transition() + .when(next_round.clone()) + .assert_eq(next.bit_idx, AB::Expr::from_usize(MAX_ECC_VARS - 1)); + assert_array_eq( + &mut builder.when_transition().when(next_round.clone()), + local.quark_out, + next.quark_in, + ); + builder.when_transition().when(next_round).assert_eq( + next.quark_layer_n, + local.quark_layer_n * AB::Expr::from_usize(2) - next.quark_parity, + ); + + self.quark_final_bus.send( + builder, + local.proof_idx, + MainEccRtQuarkFinalMessage { + idx: local.idx.into(), + quark_out: local.quark_out.map(Into::into), + }, + local.is_enabled * local.is_last_round * local.is_last_bit, + ); + } +} + +pub struct MainEccRtSumcheckAir { + pub challenge_bus: MainEccRtChallengeBus, + pub sumcheck_final_bus: MainEccRtSumcheckFinalBus, +} + +impl BaseAir for MainEccRtSumcheckAir { + fn width(&self) -> usize { + MainEccRtSumcheckCols::::width() + } +} + +impl BaseAirWithPublicValues for MainEccRtSumcheckAir {} +impl PartitionedBaseAir for MainEccRtSumcheckAir {} + +impl Air for MainEccRtSumcheckAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &MainEccRtSumcheckCols = (*local_row).borrow(); + let next: &MainEccRtSumcheckCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::Rt.as_usize()), + value: local.rt.map(Into::into), + }, + local.is_enabled, + ); + + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.claim_in, + ext_zero::(), + ); + let ev0 = ext_field_subtract(local.claim_in, local.ev1); + let claim_out = interpolate_cubic_at_0123(ev0, local.ev1, local.ev2, local.ev3, local.rt); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.claim_out, + claim_out, + ); + + let same_ecc = local.is_enabled * next.is_enabled * (AB::Expr::ONE - next.is_first); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.idx, local.idx); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.round_idx, local.round_idx + AB::Expr::ONE); + assert_array_eq( + &mut builder.when_transition().when(same_ecc), + local.claim_out, + next.claim_in, + ); + + self.sumcheck_final_bus.send( + builder, + local.proof_idx, + MainEccRtSumcheckFinalMessage { + idx: local.idx.into(), + claim: local.claim_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +impl BaseAir for MainEccRtAir { + fn width(&self) -> usize { + MainEccRtCols::::width() + } +} + +impl BaseAirWithPublicValues for MainEccRtAir {} +impl PartitionedBaseAir for MainEccRtAir {} + +impl Air for MainEccRtAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &MainEccRtCols = (*local_row).borrow(); + let next: &MainEccRtCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::Rt.as_usize()), + value: local.rt.map(Into::into), + }, + local.is_enabled, + ); + self.challenge_bus.lookup_key( + builder, + local.proof_idx, + MainEccRtChallengeMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + kind: AB::Expr::from_usize(MainEccRtChallengeKind::OutRt.as_usize()), + value: local.out_rt.map(Into::into), + }, + local.is_enabled, + ); + + self.sumcheck_final_bus.receive( + builder, + local.proof_idx, + MainEccRtSumcheckFinalMessage { + idx: local.idx.into(), + claim: local.sumcheck_final_claim.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + self.equation_totals_bus.receive( + builder, + local.proof_idx, + MainEccRtEquationTotalsMessage { + idx: local.idx.into(), + add_eval: local.add_eval.map(Into::into), + bypass_eval: local.bypass_eval.map(Into::into), + export_eval: local.export_eval.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + self.quark_final_bus.receive( + builder, + local.proof_idx, + MainEccRtQuarkFinalMessage { + idx: local.idx.into(), + quark_out: local.quark_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + + let eq_factor = eq_factor::(local.out_rt, local.rt); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.eq_out, + ext_field_multiply::(local.eq_in, eq_factor), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.eq_in, + ext_one::(), + ); + + let last_factor = ext_field_multiply::(local.out_rt, local.rt); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.last_out, + ext_field_multiply::(local.last_in, last_factor), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.last_in, + ext_one::(), + ); + + let export_out_factor = choose_ext::( + local.is_first, + ext_field_one_minus(local.out_rt.map(Into::into)), + local.out_rt.map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.export_out_out, + ext_field_multiply::(local.export_out_in, export_out_factor), + ); + let export_rt_factor = choose_ext::( + local.is_first, + ext_field_one_minus(local.rt.map(Into::into)), + local.rt.map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.export_rt_out, + ext_field_multiply::(local.export_rt_in, export_rt_factor), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.export_out_in, + ext_one::(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.export_rt_in, + ext_one::(), + ); + + let same_ecc = local.is_enabled * next.is_enabled * (AB::Expr::ONE - next.is_first); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.idx, local.idx); + builder + .when_transition() + .when(same_ecc.clone()) + .assert_eq(next.round_idx, local.round_idx + AB::Expr::ONE); + assert_array_eq( + &mut builder.when_transition().when(same_ecc.clone()), + local.eq_out, + next.eq_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_ecc.clone()), + local.last_out, + next.last_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_ecc.clone()), + local.export_out_out, + next.export_out_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_ecc.clone()), + local.export_rt_out, + next.export_rt_in, + ); + + let sel_bypass = ext_field_subtract::( + ext_field_subtract::(local.eq_out, local.sel_add), + local.last_out, + ); + let sel_export = ext_field_multiply::(local.export_out_out, local.export_rt_out); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.sel_add, + local.quark_out, + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.sel_bypass, + sel_bypass, + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.sel_export, + sel_export, + ); + let expected = ext_field_add::( + ext_field_add::( + ext_field_multiply::(local.add_eval, local.sel_add), + ext_field_multiply::(local.bypass_eval, local.sel_bypass), + ), + ext_field_multiply::(local.export_eval, local.sel_export), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.sumcheck_final_claim, + expected, + ); + } +} + +pub struct MainEccRtChallengeTraceGenerator; + +impl RowMajorChip for MainEccRtChallengeTraceGenerator { + type Ctx<'a> = &'a [MainEccRtRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainEccRtChallengeCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + if records.is_empty() { + return Some(RowMajorMatrix::new(trace, width)); + } + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainEccRtChallengeCols = row.borrow_mut(); + fill_challenge_cols(record, cols); + } + + Some(RowMajorMatrix::new(trace, width)) + } +} + +pub struct MainEccRtTraceGenerator; + +pub struct MainEccRtEquationTraceGenerator; + +pub struct MainEccRtQuarkTraceGenerator; + +impl RowMajorChip for MainEccRtQuarkTraceGenerator { + type Ctx<'a> = &'a [MainEccRtRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainEccRtQuarkCols::::width(); + let quark_rows = build_quark_rows(records); + let num_valid_rows = quark_rows.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + if quark_rows.is_empty() { + return Some(RowMajorMatrix::new(trace, width)); + } + for (row_idx, row_record) in quark_rows.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainEccRtQuarkCols = row.borrow_mut(); + fill_quark_cols(row_record, cols); + } + + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for MainEccRtEquationTraceGenerator { + type Ctx<'a> = &'a [MainEccRtRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainEccRtEquationCols::::width(); + let eq_rows = build_equation_rows(records); + let num_valid_rows = eq_rows.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + if eq_rows.is_empty() { + return Some(RowMajorMatrix::new(trace, width)); + } + for (row_idx, row_record) in eq_rows.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainEccRtEquationCols = row.borrow_mut(); + fill_equation_cols(row_record, cols); + } + + Some(RowMajorMatrix::new(trace, width)) + } +} + +pub struct MainEccRtSumcheckTraceGenerator; + +impl RowMajorChip for MainEccRtSumcheckTraceGenerator { + type Ctx<'a> = &'a [MainEccRtRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainEccRtSumcheckCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + if records.is_empty() { + return Some(RowMajorMatrix::new(trace, width)); + } + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainEccRtSumcheckCols = row.borrow_mut(); + fill_sumcheck_cols(record, cols); + } + + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for MainEccRtTraceGenerator { + type Ctx<'a> = &'a [MainEccRtRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainEccRtCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + if records.is_empty() { + return Some(RowMajorMatrix::new(trace, width)); + } + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainEccRtCols = row.borrow_mut(); + fill_cols(record, cols); + } + + Some(RowMajorMatrix::new(trace, width)) + } +} + +fn fill_challenge_cols(record: &MainEccRtRecord, cols: &mut MainEccRtChallengeCols) { + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.fork_id = F::from_usize(record.fork_id); + cols.round_idx = F::from_usize(record.round_idx); + cols.num_rounds = F::from_usize(record.num_rounds); + cols.is_first = F::from_bool(record.is_first); + cols.tidx = F::from_usize(record.tidx); + cols.out_tidx = F::from_usize(record.out_tidx); + cols.alpha_tidx = F::from_usize(record.alpha_tidx); + cols.lookup_count = F::from_usize(record.lookup_count); + cols.rt = ext_to_basis(record.value); + cols.out_rt = ext_to_basis(record.out_value); + cols.alpha = ext_to_basis(record.alpha); +} + +struct MainEccRtEquationRow<'a> { + record: &'a MainEccRtRecord, + family_idx: usize, + septic_idx: usize, + alpha_pow: openvm_stark_sdk::config::baby_bear_poseidon2::EF, + add_in: openvm_stark_sdk::config::baby_bear_poseidon2::EF, + add_out: openvm_stark_sdk::config::baby_bear_poseidon2::EF, + bypass_in: openvm_stark_sdk::config::baby_bear_poseidon2::EF, + bypass_out: openvm_stark_sdk::config::baby_bear_poseidon2::EF, + export_in: openvm_stark_sdk::config::baby_bear_poseidon2::EF, + export_out: openvm_stark_sdk::config::baby_bear_poseidon2::EF, + is_first: bool, + is_last: bool, +} + +fn build_equation_rows(records: &[MainEccRtRecord]) -> Vec> { + let mut rows = Vec::new(); + for record in records.iter().filter(|record| record.is_first) { + let mut add = openvm_stark_sdk::config::baby_bear_poseidon2::EF::ZERO; + let mut bypass = openvm_stark_sdk::config::baby_bear_poseidon2::EF::ZERO; + let mut export = openvm_stark_sdk::config::baby_bear_poseidon2::EF::ZERO; + for family_idx in 0..ECC_EQUATION_FAMILIES { + for septic_idx in 0..SEPTIC_DEGREE { + let alpha_pow = record.alpha_pows[family_idx * SEPTIC_DEGREE + septic_idx]; + let weighted = equation_term(record, family_idx, septic_idx) * alpha_pow; + let add_in = add; + let bypass_in = bypass; + let export_in = export; + match family_idx { + 0..=2 => add += weighted, + 3..=4 => bypass += weighted, + 5..=6 => export += weighted, + _ => unreachable!("invalid ECC equation family"), + } + rows.push(MainEccRtEquationRow { + record, + family_idx, + septic_idx, + alpha_pow, + add_in, + add_out: add, + bypass_in, + bypass_out: bypass, + export_in, + export_out: export, + is_first: family_idx == 0 && septic_idx == 0, + is_last: family_idx + 1 == ECC_EQUATION_FAMILIES + && septic_idx + 1 == SEPTIC_DEGREE, + }); + } + } + } + rows +} + +fn equation_term( + record: &MainEccRtRecord, + family_idx: usize, + septic_idx: usize, +) -> openvm_stark_sdk::config::baby_bear_poseidon2::EF { + match family_idx { + 0 => { + septic_mul_coeff_native( + &record.s0, + &core::array::from_fn(|i| record.x0[i] - record.x1[i]), + septic_idx, + ) - (record.y0[septic_idx] - record.y1[septic_idx]) + } + 1 => { + septic_mul_coeff_native(&record.s0, &record.s0, septic_idx) + - record.x0[septic_idx] + - record.x1[septic_idx] + - record.x3[septic_idx] + } + 2 => { + septic_mul_coeff_native( + &record.s0, + &core::array::from_fn(|i| record.x0[i] - record.x3[i]), + septic_idx, + ) - (record.y0[septic_idx] + record.y3[septic_idx]) + } + 3 => record.x3[septic_idx] - record.x0[septic_idx], + 4 => record.y3[septic_idx] - record.y0[septic_idx], + 5 => record.x3[septic_idx] - record.sum_x[septic_idx], + 6 => record.y3[septic_idx] - record.sum_y[septic_idx], + _ => unreachable!("invalid ECC equation family"), + } +} + +fn septic_mul_coeff_native( + a: &[openvm_stark_sdk::config::baby_bear_poseidon2::EF; SEPTIC_DEGREE], + b: &[openvm_stark_sdk::config::baby_bear_poseidon2::EF; SEPTIC_DEGREE], + coeff_idx: usize, +) -> openvm_stark_sdk::config::baby_bear_poseidon2::EF { + let mut out = openvm_stark_sdk::config::baby_bear_poseidon2::EF::ZERO; + let two = openvm_stark_sdk::config::baby_bear_poseidon2::EF::from_usize(2); + let five = openvm_stark_sdk::config::baby_bear_poseidon2::EF::from_usize(5); + for i in 0..SEPTIC_DEGREE { + for j in 0..SEPTIC_DEGREE { + let term = a[i] * b[j]; + let mut index = i + j; + if index < SEPTIC_DEGREE { + if index == coeff_idx { + out += term; + } + } else { + index -= SEPTIC_DEGREE; + if index == coeff_idx { + out += five * term; + } + if index + 1 == coeff_idx { + out += two * term; + } + } + } + } + out +} + +fn fill_equation_cols(row: &MainEccRtEquationRow<'_>, cols: &mut MainEccRtEquationCols) { + let record = row.record; + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.family_idx = F::from_usize(row.family_idx); + cols.septic_idx = F::from_usize(row.septic_idx); + cols.is_first = F::from_bool(row.is_first); + cols.is_last = F::from_bool(row.is_last); + cols.family_flags = core::array::from_fn(|i| F::from_bool(i == row.family_idx)); + cols.septic_flags = core::array::from_fn(|i| F::from_bool(i == row.septic_idx)); + cols.alpha = ext_to_basis(record.alpha); + cols.alpha_pow = ext_to_basis(row.alpha_pow); + cols.s0 = ext_to_basis(record.s0[row.septic_idx]); + cols.s0_all = record.s0.map(ext_to_basis); + cols.x0 = ext_to_basis(record.x0[row.septic_idx]); + cols.x0_all = record.x0.map(ext_to_basis); + cols.y0 = ext_to_basis(record.y0[row.septic_idx]); + cols.x1 = ext_to_basis(record.x1[row.septic_idx]); + cols.x1_all = record.x1.map(ext_to_basis); + cols.y1 = ext_to_basis(record.y1[row.septic_idx]); + cols.x3 = ext_to_basis(record.x3[row.septic_idx]); + cols.x3_all = record.x3.map(ext_to_basis); + cols.y3 = ext_to_basis(record.y3[row.septic_idx]); + cols.sum_x = ext_to_basis(record.sum_x[row.septic_idx]); + cols.sum_y = ext_to_basis(record.sum_y[row.septic_idx]); + cols.s0_x0_x1 = ext_to_basis(septic_mul_coeff_native( + &record.s0, + &core::array::from_fn(|i| record.x0[i] - record.x1[i]), + row.septic_idx, + )); + cols.s0_squared = ext_to_basis(septic_mul_coeff_native( + &record.s0, + &record.s0, + row.septic_idx, + )); + cols.s0_x0_x3 = ext_to_basis(septic_mul_coeff_native( + &record.s0, + &core::array::from_fn(|i| record.x0[i] - record.x3[i]), + row.septic_idx, + )); + cols.add_in = ext_to_basis(row.add_in); + cols.add_out = ext_to_basis(row.add_out); + cols.bypass_in = ext_to_basis(row.bypass_in); + cols.bypass_out = ext_to_basis(row.bypass_out); + cols.export_in = ext_to_basis(row.export_in); + cols.export_out = ext_to_basis(row.export_out); +} + +struct MainEccRtQuarkRow<'a> { + record: &'a MainEccRtRecord, + bit_idx: usize, + active_count_in: usize, + active_count_out: usize, + bit_value_in: usize, + bit_value_out: usize, +} + +fn build_quark_rows(records: &[MainEccRtRecord]) -> Vec> { + let mut rows = Vec::new(); + for record in records { + let mut active_count = 0usize; + let mut bit_value = 0usize; + for bit_idx in (0..MAX_ECC_VARS).rev() { + let active_count_in = active_count; + let bit_value_in = bit_value; + if record.lte_active[bit_idx] { + active_count += 1; + } + if record.lte_bits[bit_idx] { + bit_value += 1usize << bit_idx; + } + rows.push(MainEccRtQuarkRow { + record, + bit_idx, + active_count_in, + active_count_out: active_count, + bit_value_in, + bit_value_out: bit_value, + }); + } + } + rows +} + +fn fill_quark_cols(row: &MainEccRtQuarkRow<'_>, cols: &mut MainEccRtQuarkCols) { + let record = row.record; + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.bit_idx = F::from_usize(row.bit_idx); + cols.bit_weight = F::from_usize(1usize << row.bit_idx); + cols.is_first_round = F::from_bool(record.is_first); + cols.is_last_round = F::from_bool(record.is_last); + cols.is_first_bit = F::from_bool(row.bit_idx + 1 == MAX_ECC_VARS); + cols.is_last_bit = F::from_bool(row.bit_idx == 0); + cols.is_active = F::from_bool(record.lte_active[row.bit_idx]); + cols.bit = F::from_bool(record.lte_bits[row.bit_idx]); + cols.current_rt = ext_to_basis(record.value); + cols.current_out_rt = ext_to_basis(record.out_value); + cols.point_rt = ext_to_basis(record.lte_rt_point[row.bit_idx]); + cols.point_out_rt = ext_to_basis(record.lte_out_point[row.bit_idx]); + let same_one = record.lte_out_point[row.bit_idx] * record.lte_rt_point[row.bit_idx]; + let same_zero = (EF::ONE - record.lte_out_point[row.bit_idx]) + * (EF::ONE - record.lte_rt_point[row.bit_idx]); + let equal_choice = if record.lte_bits[row.bit_idx] { + same_one + } else { + same_zero + }; + cols.same_one = ext_to_basis(same_one); + cols.same_zero = ext_to_basis(same_zero); + cols.equal_choice = ext_to_basis(equal_choice); + cols.active_prefix = ext_to_basis(record.lte_prefix_acc[row.bit_idx + 1] * equal_choice); + cols.prefix_in = ext_to_basis(record.lte_prefix_acc[row.bit_idx + 1]); + cols.prefix_out = ext_to_basis(record.lte_prefix_acc[row.bit_idx]); + cols.less_in = ext_to_basis(record.lte_less_acc[row.bit_idx + 1]); + let less_from_prior = record.lte_less_acc[row.bit_idx + 1] * (same_one + same_zero); + let less_from_equal = record.lte_prefix_acc[row.bit_idx + 1] * same_zero; + let active_less = less_from_prior + + if record.lte_bits[row.bit_idx] { + less_from_equal + } else { + EF::ZERO + }; + cols.less_from_prior = ext_to_basis(less_from_prior); + cols.less_from_equal = ext_to_basis(less_from_equal); + cols.active_less = ext_to_basis(active_less); + cols.less_out = ext_to_basis(record.lte_less_acc[row.bit_idx]); + cols.active_count_in = F::from_usize(row.active_count_in); + cols.active_count_out = F::from_usize(row.active_count_out); + cols.bit_value_in = F::from_usize(row.bit_value_in); + cols.bit_value_out = F::from_usize(row.bit_value_out); + cols.quark_prefix_count = F::from_usize(record.quark_prefix_count); + cols.quark_prefix_is_zero = F::from_bool(record.quark_prefix_count == 0); + cols.quark_prefix_inv = if record.quark_prefix_count == 0 { + F::ZERO + } else { + F::from_usize(record.quark_prefix_count).inverse() + }; + cols.quark_layer_n = F::from_usize(record.quark_layer_n); + cols.quark_parity = F::from_bool(record.quark_parity); + cols.quark_in = ext_to_basis(record.quark_in); + cols.quark_factor = ext_to_basis(record.quark_factor); + let quark_zero = (EF::ONE - record.out_value) * (EF::ONE - record.value); + let quark_one = record.out_value * record.value; + cols.quark_zero = ext_to_basis(quark_zero); + cols.quark_one = ext_to_basis(quark_one); + cols.quark_lhs = ext_to_basis(quark_zero * record.quark_factor); + cols.quark_rhs = ext_to_basis(quark_one * record.quark_in); + cols.quark_out = ext_to_basis(record.quark_out); +} + +fn fill_sumcheck_cols(record: &MainEccRtRecord, cols: &mut MainEccRtSumcheckCols) { + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.rt = ext_to_basis(record.value); + cols.ev1 = ext_to_basis(record.ev1); + cols.ev2 = ext_to_basis(record.ev2); + cols.ev3 = ext_to_basis(record.ev3); + cols.claim_in = ext_to_basis(record.claim_in); + cols.claim_out = ext_to_basis(record.claim_out); +} + +fn fill_cols(record: &MainEccRtRecord, cols: &mut MainEccRtCols) { + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.rt = ext_to_basis(record.value); + cols.out_rt = ext_to_basis(record.out_value); + cols.sumcheck_final_claim = ext_to_basis(record.claim_out); + cols.sel_add = ext_to_basis(record.sel_add); + cols.sel_bypass = ext_to_basis(record.sel_bypass); + cols.sel_export = ext_to_basis(record.sel_export); + cols.eq_in = ext_to_basis(record.eq_in); + cols.eq_out = ext_to_basis(record.eq_out); + cols.last_in = ext_to_basis(record.last_in); + cols.last_out = ext_to_basis(record.last_out); + cols.export_out_in = ext_to_basis(record.export_out_in); + cols.export_out_out = ext_to_basis(record.export_out_out); + cols.export_rt_in = ext_to_basis(record.export_rt_in); + cols.export_rt_out = ext_to_basis(record.export_rt_out); + cols.quark_out = ext_to_basis(record.quark_out); + cols.add_eval = ext_to_basis(record.add_eval); + cols.bypass_eval = ext_to_basis(record.bypass_eval); + cols.export_eval = ext_to_basis(record.export_eval); +} + +fn ext_to_basis(value: openvm_stark_sdk::config::baby_bear_poseidon2::EF) -> [F; D_EF] { + value.as_basis_coefficients_slice().try_into().unwrap() +} + +fn eq_factor(a: [AB::Var; D_EF], b: [AB::Var; D_EF]) -> [AB::Expr; D_EF] +where + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + ext_field_add::( + ext_field_multiply::(a.clone(), b.clone()), + ext_field_multiply::(one_minus_ext::(a), one_minus_ext::(b)), + ) +} + +fn one_minus_ext(value: [AB::Var; D_EF]) -> [AB::Expr; D_EF] { + ext_field_one_minus(value.map(Into::into)) +} + +fn choose_ext( + flag: AB::Var, + when_one: [AB::Expr; D_EF], + when_zero: [AB::Expr; D_EF], +) -> [AB::Expr; D_EF] { + core::array::from_fn(|i| { + flag.clone() * when_one[i].clone() + (AB::Expr::ONE - flag.clone()) * when_zero[i].clone() + }) +} + +fn septic_mul_selected( + a: &[[AB::Var; D_EF]; SEPTIC_DEGREE], + b: &[[AB::Expr; D_EF]; SEPTIC_DEGREE], + septic_flags: [AB::Var; SEPTIC_DEGREE], +) -> [AB::Expr; D_EF] +where + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + (0..SEPTIC_DEGREE).fold(ext_zero::(), |acc, coeff_idx| { + let coeff = septic_mul_coeff::(a, b, coeff_idx); + ext_field_add::( + acc, + ext_field_multiply_scalar::(coeff, septic_flags[coeff_idx].clone()), + ) + }) +} + +fn septic_mul_coeff( + a: &[[AB::Var; D_EF]; SEPTIC_DEGREE], + b: &[[AB::Expr; D_EF]; SEPTIC_DEGREE], + coeff_idx: usize, +) -> [AB::Expr; D_EF] +where + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + let mut out = ext_zero::(); + for i in 0..SEPTIC_DEGREE { + for j in 0..SEPTIC_DEGREE { + let term = ext_field_multiply::(a[i].clone().map(Into::into), b[j].clone()); + let mut index = i + j; + if index < SEPTIC_DEGREE { + if index == coeff_idx { + out = ext_field_add::(out, term); + } + } else { + index -= SEPTIC_DEGREE; + if index == coeff_idx { + out = ext_field_add::( + out, + ext_field_multiply_scalar::( + term.clone(), + AB::Expr::from_usize(5), + ), + ); + } + if index + 1 == coeff_idx { + out = ext_field_add::( + out, + ext_field_multiply_scalar::(term, AB::Expr::from_usize(2)), + ); + } + } + } + } + out +} + +fn ext_zero() -> [FA; D_EF] { + core::array::from_fn(|_| FA::ZERO) +} + +fn ext_one() -> [FA; D_EF] { + let mut out = ext_zero::(); + out[0] = FA::ONE; + out +} + +fn interpolate_cubic_at_0123( + ev0: [FA; D_EF], + ev1: [F; D_EF], + ev2: [F; D_EF], + ev3: [F; D_EF], + x: [F; D_EF], +) -> [FA; D_EF] +where + F: Into + Copy, + FA: PrimeCharacteristicRing, + FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + let three: FA = FA::from_usize(3); + let inv2: FA = FA::from_prime_subfield(FA::PrimeSubfield::from_usize(2).inverse()); + let inv6: FA = FA::from_prime_subfield(FA::PrimeSubfield::from_usize(6).inverse()); + + let s1: [FA; D_EF] = ext_field_subtract(ev1, ev0.clone()); + let s2: [FA; D_EF] = ext_field_subtract(ev2, ev0.clone()); + let s3: [FA; D_EF] = ext_field_subtract(ev3, ev0.clone()); + + let d3: [FA; D_EF] = ext_field_subtract::( + s3, + ext_field_multiply_scalar::(ext_field_subtract::(s2.clone(), s1.clone()), three), + ); + let p: [FA; D_EF] = ext_field_multiply_scalar(d3.clone(), inv6); + let q: [FA; D_EF] = ext_field_subtract::( + ext_field_multiply_scalar::(ext_field_subtract::(s2, d3), inv2), + s1.clone(), + ); + let r: [FA; D_EF] = ext_field_subtract::(s1, ext_field_add::(p.clone(), q.clone())); + + ext_field_add::( + ext_field_multiply::( + ext_field_add::( + ext_field_multiply::(ext_field_add::(ext_field_multiply::(p, x), q), x), + r, + ), + x, + ), + ev0, + ) +} diff --git a/ceno_recursion_v2/src/main/eval_absorb.rs b/ceno_recursion_v2/src/main/eval_absorb.rs new file mode 100644 index 000000000..5e18fa7a2 --- /dev/null +++ b/ceno_recursion_v2/src/main/eval_absorb.rs @@ -0,0 +1,114 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{MainEvalBus, MainEvalMessage, TranscriptBus, TranscriptBusMessage}, + system::MainEvalRecord, + tracegen::RowMajorChip, +}; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainEvalAbsorbCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub eval_idx: T, + pub tidx: T, + pub lookup_count: T, + pub value: [T; D_EF], +} + +pub struct MainEvalAbsorbAir { + pub transcript_bus: TranscriptBus, + pub eval_bus: MainEvalBus, +} + +impl BaseAir for MainEvalAbsorbAir { + fn width(&self) -> usize { + MainEvalAbsorbCols::::width() + } +} + +impl BaseAirWithPublicValues for MainEvalAbsorbAir {} +impl PartitionedBaseAir for MainEvalAbsorbAir {} + +impl Air for MainEvalAbsorbAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &MainEvalAbsorbCols = (*local_row).borrow(); + + builder.assert_bool(local.is_enabled); + self.eval_bus.add_key_with_lookups( + builder, + local.proof_idx, + MainEvalMessage { + idx: local.idx.into(), + eval_idx: local.eval_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.lookup_count, + ); + for i in 0..D_EF { + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.tidx + AB::Expr::from_usize(i), + value: local.value[i].into(), + is_sample: AB::Expr::ZERO, + }, + local.is_enabled, + ); + } + } +} + +pub struct MainEvalAbsorbTraceGenerator; + +impl RowMajorChip for MainEvalAbsorbTraceGenerator { + type Ctx<'a> = &'a [MainEvalRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainEvalAbsorbCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainEvalAbsorbCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.eval_idx = F::from_usize(record.eval_idx); + cols.tidx = F::from_usize(record.tidx); + cols.lookup_count = F::from_usize(record.lookup_count); + cols.value = record + .value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + } + Some(RowMajorMatrix::new(trace, width)) + } +} diff --git a/ceno_recursion_v2/src/main/final_claim.rs b/ceno_recursion_v2/src/main/final_claim.rs new file mode 100644 index 000000000..7bc35362f --- /dev/null +++ b/ceno_recursion_v2/src/main/final_claim.rs @@ -0,0 +1,177 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_circuit_primitives::utils::assert_array_eq; +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use recursion_circuit::utils::ext_field_add; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{ + MainContributionBus, MainContributionMessage, MainGlobalClaimBus, MainGlobalClaimMessage, + }, + system::MainFinalClaimRecord, + tracegen::RowMajorChip, +}; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainFinalClaimCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub is_first_idx: T, + pub is_first: T, + pub is_last: T, + pub contribution: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], + pub expected: [T; D_EF], +} + +pub struct MainFinalClaimAir { + pub global_claim_bus: MainGlobalClaimBus, + pub contribution_bus: MainContributionBus, +} + +impl BaseAir for MainFinalClaimAir { + fn width(&self) -> usize { + MainFinalClaimCols::::width() + } +} + +impl BaseAirWithPublicValues for MainFinalClaimAir {} +impl PartitionedBaseAir for MainFinalClaimAir {} + +impl Air for MainFinalClaimAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("window should have two elements"), + main.row_slice(1).expect("window should have two elements"), + ); + let local: &MainFinalClaimCols = (*local_row).borrow(); + let next: &MainFinalClaimCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder + .when_transition() + .when_ne(local.is_enabled, AB::Expr::ONE) + .assert_zero(next.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_first_idx); + builder.assert_bool(local.is_last); + + let computed_is_last = local.is_enabled - next.is_enabled + next.is_first; + builder + .when(local.is_enabled) + .assert_eq(local.is_last, computed_is_last); + + let acc_out = ext_field_add(local.acc_in, local.contribution); + assert_array_eq(&mut builder.when(local.is_enabled), local.acc_out, acc_out); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + [AB::Expr::ZERO; D_EF], + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.acc_out, + local.expected, + ); + self.global_claim_bus.receive( + builder, + local.proof_idx, + MainGlobalClaimMessage { + expected: local.expected.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + self.contribution_bus.receive( + builder, + local.proof_idx, + MainContributionMessage { + idx: local.idx.into(), + contribution: local.contribution.map(Into::into), + }, + local.is_enabled, + ); + + let is_transition = next.is_enabled - next.is_first; + assert_array_eq(&mut builder.when(is_transition), local.acc_out, next.acc_in); + } +} + +pub struct MainFinalClaimTraceGenerator; + +impl RowMajorChip for MainFinalClaimTraceGenerator { + type Ctx<'a> = &'a [MainFinalClaimRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainFinalClaimCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + if records.is_empty() { + return Some(RowMajorMatrix::new(trace, width)); + } + + let mut prev_proof_idx = usize::MAX; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainFinalClaimCols = row.borrow_mut(); + let is_first_idx = prev_proof_idx != record.proof_idx; + let is_last = records + .get(row_idx + 1) + .is_none_or(|next| next.proof_idx != record.proof_idx); + + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.is_first_idx = F::from_bool(is_first_idx); + cols.is_first = F::from_bool(is_first_idx); + cols.is_last = F::from_bool(is_last); + cols.contribution = record + .contribution + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.acc_in = record + .acc_in + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.acc_out = record + .acc_out + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.expected = record + .expected + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + + prev_proof_idx = record.proof_idx; + } + + Some(RowMajorMatrix::new(trace, width)) + } +} diff --git a/ceno_recursion_v2/src/main/frontload.rs b/ceno_recursion_v2/src/main/frontload.rs new file mode 100644 index 000000000..71308e435 --- /dev/null +++ b/ceno_recursion_v2/src/main/frontload.rs @@ -0,0 +1,278 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_circuit_primitives::utils::assert_array_eq; +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, extension::BinomiallyExtendable}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use recursion_circuit::utils::{ + ext_field_add, ext_field_multiply, ext_field_multiply_scalar, ext_field_subtract, + scalar_subtract_ext_field, +}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{ + MainContributionBus, MainContributionMessage, MainEvalBus, MainEvalMessage, + MainGlobalPointBus, MainGlobalPointMessage, + }, + system::MainFrontloadTermRecord, + tracegen::RowMajorChip, +}; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainFrontloadTermCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub row_idx: T, + pub node_idx: T, + pub eval_idx: T, + pub has_eval_factor: T, + pub instance_idx: T, + pub challenge_idx: T, + pub global_round_idx: T, + pub has_global_factor: T, + pub is_wit: T, + pub is_const: T, + pub is_instance: T, + pub is_challenge: T, + pub is_add: T, + pub is_sub: T, + pub is_neg: T, + pub is_mul: T, + pub is_fold: T, + pub is_tail: T, + pub constraint_idx: T, + pub alpha: [T; D_EF], + pub arg0: [T; D_EF], + pub arg1: [T; D_EF], + pub value: [T; D_EF], + pub chip_acc_in: [T; D_EF], + pub chip_acc_out: [T; D_EF], + pub is_last_chip_step: T, +} + +pub struct MainFrontloadTermAir { + pub eval_bus: MainEvalBus, + pub global_point_bus: MainGlobalPointBus, + pub contribution_bus: MainContributionBus, +} + +impl BaseAir for MainFrontloadTermAir { + fn width(&self) -> usize { + MainFrontloadTermCols::::width() + } +} + +impl BaseAirWithPublicValues for MainFrontloadTermAir {} +impl PartitionedBaseAir for MainFrontloadTermAir {} + +impl Air for MainFrontloadTermAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &MainFrontloadTermCols = (*local_row).borrow(); + let next: &MainFrontloadTermCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.has_eval_factor); + builder.assert_bool(local.has_global_factor); + builder.assert_bool(local.is_wit); + builder.assert_bool(local.is_const); + builder.assert_bool(local.is_instance); + builder.assert_bool(local.is_challenge); + builder.assert_bool(local.is_add); + builder.assert_bool(local.is_sub); + builder.assert_bool(local.is_neg); + builder.assert_bool(local.is_mul); + builder.assert_bool(local.is_fold); + builder.assert_bool(local.is_tail); + builder.assert_bool(local.is_last_chip_step); + builder + .when_transition() + .when_ne(local.is_enabled, AB::Expr::ONE) + .assert_zero(next.is_enabled); + + let selector_sum = local.is_wit + + local.is_const + + local.is_instance + + local.is_challenge + + local.is_add + + local.is_sub + + local.is_neg + + local.is_mul + + local.is_fold + + local.is_tail; + builder + .when(local.is_enabled) + .assert_one(selector_sum.clone()); + + let add_value = ext_field_add(local.arg0, local.arg1); + let sub_value = ext_field_subtract(local.arg0, local.arg1); + let neg_value = scalar_subtract_ext_field(AB::Expr::ZERO, local.arg0); + let mul_value = ext_field_multiply(local.arg0, local.arg1); + let fold_value = ext_field_add( + local.chip_acc_in, + ext_field_multiply(local.alpha, local.arg0), + ); + let tail_value = ext_field_multiply(local.arg1, local.arg0); + + let mut expected = [AB::Expr::ZERO; D_EF]; + for (sel, value) in [ + (local.is_wit, local.arg0.map(Into::into)), + (local.is_const, local.arg0.map(Into::into)), + (local.is_instance, local.arg0.map(Into::into)), + (local.is_challenge, local.arg0.map(Into::into)), + (local.is_add, add_value), + (local.is_sub, sub_value), + (local.is_neg, neg_value), + (local.is_mul, mul_value), + (local.is_fold, fold_value), + (local.is_tail, tail_value), + ] { + expected = ext_field_add(expected, ext_field_multiply_scalar(value, sel)); + } + assert_array_eq(&mut builder.when(local.is_enabled), local.value, expected); + + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_fold), + local.chip_acc_out, + local.value.map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_tail), + local.chip_acc_out, + local.value.map(Into::into), + ); + let is_expr_node = selector_sum - local.is_fold - local.is_tail; + assert_array_eq( + &mut builder.when(local.is_enabled * is_expr_node), + local.chip_acc_out, + local.chip_acc_in.map(Into::into), + ); + + self.eval_bus.lookup_key( + builder, + local.proof_idx, + MainEvalMessage { + idx: local.idx.into(), + eval_idx: local.eval_idx.into(), + value: local.arg0.map(Into::into), + }, + local.is_enabled * local.has_eval_factor, + ); + self.global_point_bus.lookup_key( + builder, + local.proof_idx, + MainGlobalPointMessage { + round_idx: local.global_round_idx.into(), + value: local.arg0.map(Into::into), + }, + local.is_enabled * local.has_global_factor, + ); + self.contribution_bus.send( + builder, + local.proof_idx, + MainContributionMessage { + idx: local.idx.into(), + contribution: local.chip_acc_out.map(Into::into), + }, + local.is_enabled * local.is_last_chip_step, + ); + } +} + +pub struct MainFrontloadTermTraceGenerator; + +impl RowMajorChip for MainFrontloadTermTraceGenerator { + type Ctx<'a> = &'a [MainFrontloadTermRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainFrontloadTermCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainFrontloadTermCols = row.borrow_mut(); + + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.row_idx = F::from_usize(record.row_idx); + cols.node_idx = F::from_usize(record.node_idx); + cols.eval_idx = F::from_usize(record.eval_idx); + cols.has_eval_factor = F::from_bool(record.has_eval_factor); + cols.instance_idx = F::from_usize(record.instance_idx); + cols.challenge_idx = F::from_usize(record.challenge_idx); + cols.global_round_idx = F::from_usize(record.global_round_idx); + cols.has_global_factor = F::from_bool(record.has_global_factor); + cols.is_wit = F::from_bool(record.is_wit); + cols.is_const = F::from_bool(record.is_const); + cols.is_instance = F::from_bool(record.is_instance); + cols.is_challenge = F::from_bool(record.is_challenge); + cols.is_add = F::from_bool(record.is_add); + cols.is_sub = F::from_bool(record.is_sub); + cols.is_neg = F::from_bool(record.is_neg); + cols.is_mul = F::from_bool(record.is_mul); + cols.is_fold = F::from_bool(record.is_fold); + cols.is_tail = F::from_bool(record.is_tail); + cols.constraint_idx = F::from_usize(record.constraint_idx); + cols.alpha = record + .alpha + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.arg0 = record + .arg0 + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.arg1 = record + .arg1 + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.value = record + .value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.chip_acc_in = record + .chip_acc_in + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.chip_acc_out = record + .chip_acc_out + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.is_last_chip_step = F::from_bool(record.is_last_chip_step); + } + Some(RowMajorMatrix::new(trace, width)) + } +} diff --git a/ceno_recursion_v2/src/main/global_sumcheck.rs b/ceno_recursion_v2/src/main/global_sumcheck.rs new file mode 100644 index 000000000..d9863536d --- /dev/null +++ b/ceno_recursion_v2/src/main/global_sumcheck.rs @@ -0,0 +1,297 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_circuit_primitives::{SubAir, utils::assert_array_eq}; +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, extension::BinomiallyExtendable}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use recursion_circuit::{ + subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir}, + utils::{ext_field_add, ext_field_multiply, ext_field_multiply_scalar, ext_field_subtract}, +}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{ + MainGlobalClaimBus, MainGlobalClaimMessage, MainGlobalPointBus, MainGlobalPointMessage, + TranscriptBus, TranscriptBusMessage, + }, + system::MainGlobalSumcheckRecord, + tracegen::RowMajorChip, +}; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainGlobalSumcheckCols { + pub is_enabled: T, + pub proof_idx: T, + pub is_first: T, + pub is_last: T, + pub is_dummy: T, + pub round: T, + pub point_lookup_count: T, + pub challenge_tidx: T, + pub ev1: [T; D_EF], + pub ev2: [T; D_EF], + pub ev3: [T; D_EF], + pub ev4: [T; D_EF], + pub claim_in: [T; D_EF], + pub claim_out: [T; D_EF], + pub challenge: [T; D_EF], + pub expected: [T; D_EF], +} + +pub struct MainGlobalSumcheckAir { + pub transcript_bus: TranscriptBus, + pub global_claim_bus: MainGlobalClaimBus, + pub global_point_bus: MainGlobalPointBus, +} + +impl BaseAir for MainGlobalSumcheckAir { + fn width(&self) -> usize { + MainGlobalSumcheckCols::::width() + } +} + +impl BaseAirWithPublicValues for MainGlobalSumcheckAir {} +impl PartitionedBaseAir for MainGlobalSumcheckAir {} + +impl Air for MainGlobalSumcheckAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &MainGlobalSumcheckCols = (*local_row).borrow(); + let next: &MainGlobalSumcheckCols = (*next_row).borrow(); + + builder.assert_bool(local.is_dummy); + builder.assert_bool(local.is_last); + + type LoopSubAir = NestedForLoopSubAir<1>; + LoopSubAir {}.eval( + builder, + ( + NestedForLoopIoCols { + is_enabled: local.is_enabled, + counter: [local.proof_idx], + is_first: [local.is_first], + } + .map_into(), + NestedForLoopIoCols { + is_enabled: next.is_enabled, + counter: [next.proof_idx], + is_first: [next.is_first], + } + .map_into(), + ), + ); + + let is_transition = LoopSubAir::local_is_transition(next.is_enabled, next.is_first); + builder.when(local.is_first).assert_zero(local.round); + builder + .when(is_transition.clone()) + .assert_eq(next.round, local.round + AB::Expr::ONE); + builder + .when(is_transition.clone()) + .assert_zero(local.is_last); + + let ev0 = ext_field_subtract(local.claim_in, local.ev1); + let claim_out = interpolate_quartic_at_01234( + ev0, + local.ev1, + local.ev2, + local.ev3, + local.ev4, + local.challenge, + ); + assert_array_eq(builder, local.claim_out, claim_out); + assert_array_eq( + &mut builder.when(is_transition.clone()), + local.claim_out, + next.claim_in, + ); + assert_array_eq( + &mut builder.when(is_transition), + local.expected, + next.expected, + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.claim_out, + local.expected, + ); + + for i in 0..D_EF { + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.challenge_tidx + AB::Expr::from_usize(i), + value: local.challenge[i].into(), + is_sample: AB::Expr::ONE, + }, + local.is_enabled * (AB::Expr::ONE - local.is_dummy), + ); + } + + self.global_claim_bus.send( + builder, + local.proof_idx, + MainGlobalClaimMessage { + expected: local.expected.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + self.global_point_bus.add_key_with_lookups( + builder, + local.proof_idx, + MainGlobalPointMessage { + round_idx: local.round.into(), + value: local.challenge.map(Into::into), + }, + local.is_enabled * (AB::Expr::ONE - local.is_dummy) * local.point_lookup_count, + ); + } +} + +fn interpolate_quartic_at_01234( + ev0: [FA; D_EF], + ev1: [F; D_EF], + ev2: [F; D_EF], + ev3: [F; D_EF], + ev4: [F; D_EF], + x: [F; D_EF], +) -> [FA; D_EF] +where + F: Into + Copy, + FA: PrimeCharacteristicRing, + FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + let x: [FA; D_EF] = x.map(Into::into); + let scalar_ext = |value: usize| { + let mut scalar = core::array::from_fn(|_| FA::ZERO); + scalar[0] = FA::from_usize(value); + scalar + }; + let xm0 = x.clone(); + let xm1 = ext_field_subtract::(x.clone(), scalar_ext(1)); + let xm2 = ext_field_subtract::(x.clone(), scalar_ext(2)); + let xm3 = ext_field_subtract::(x.clone(), scalar_ext(3)); + let xm4 = ext_field_subtract::(x, scalar_ext(4)); + + let inv4: FA = FA::from_prime_subfield(FA::PrimeSubfield::from_usize(4).inverse()); + let inv6: FA = FA::from_prime_subfield(FA::PrimeSubfield::from_usize(6).inverse()); + let inv24: FA = FA::from_prime_subfield(FA::PrimeSubfield::from_usize(24).inverse()); + + let prod4 = |a: &[FA; D_EF], b: &[FA; D_EF], c: &[FA; D_EF], d: &[FA; D_EF]| { + ext_field_multiply::( + ext_field_multiply::(ext_field_multiply::(a.clone(), b.clone()), c.clone()), + d.clone(), + ) + }; + + let basis0 = ext_field_multiply_scalar::(prod4(&xm1, &xm2, &xm3, &xm4), inv24.clone()); + let basis1 = ext_field_multiply_scalar::(prod4(&xm0, &xm2, &xm3, &xm4), -inv6.clone()); + let basis2 = ext_field_multiply_scalar::(prod4(&xm0, &xm1, &xm3, &xm4), inv4); + let basis3 = ext_field_multiply_scalar::(prod4(&xm0, &xm1, &xm2, &xm4), -inv6); + let basis4 = ext_field_multiply_scalar::(prod4(&xm0, &xm1, &xm2, &xm3), inv24); + + ext_field_add::( + ext_field_add::( + ext_field_add::( + ext_field_multiply::(ev0, basis0), + ext_field_multiply::(ev1, basis1), + ), + ext_field_add::( + ext_field_multiply::(ev2, basis2), + ext_field_multiply::(ev3, basis3), + ), + ), + ext_field_multiply::(ev4, basis4), + ) +} + +pub struct MainGlobalSumcheckTraceGenerator; + +impl RowMajorChip for MainGlobalSumcheckTraceGenerator { + type Ctx<'a> = &'a [MainGlobalSumcheckRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainGlobalSumcheckCols::::width(); + let num_valid_rows: usize = records + .iter() + .map(MainGlobalSumcheckRecord::total_rows) + .sum(); + let num_valid_rows = num_valid_rows.max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + let mut row_offset = 0; + for record in records.iter() { + let rows = record.total_rows(); + let is_dummy = record.rounds.is_empty(); + for round_idx in 0..rows { + let row = &mut trace[row_offset * width..(row_offset + 1) * width]; + let cols: &mut MainGlobalSumcheckCols = row.borrow_mut(); + let round = record.rounds.get(round_idx); + + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.is_first = F::from_bool(round_idx == 0); + cols.is_last = F::from_bool(round_idx + 1 == rows); + cols.is_dummy = F::from_bool(is_dummy); + cols.round = F::from_usize(round_idx); + cols.point_lookup_count = + F::from_usize(round.map(|round| round.point_lookup_count).unwrap_or(0)); + cols.challenge_tidx = + F::from_usize(round.map(|round| round.challenge_tidx).unwrap_or(0)); + + let evs = round + .map(|round| round.evaluations) + .unwrap_or([EF::ZERO; 4]); + cols.ev1 = evs[0].as_basis_coefficients_slice().try_into().unwrap(); + cols.ev2 = evs[1].as_basis_coefficients_slice().try_into().unwrap(); + cols.ev3 = evs[2].as_basis_coefficients_slice().try_into().unwrap(); + cols.ev4 = evs[3].as_basis_coefficients_slice().try_into().unwrap(); + + let claim_in = round.map(|round| round.claim_in).unwrap_or(record.expected); + let claim_out = round + .map(|round| round.claim_out) + .unwrap_or(record.expected); + let challenge = round.map(|round| round.challenge).unwrap_or(EF::ZERO); + cols.claim_in = claim_in.as_basis_coefficients_slice().try_into().unwrap(); + cols.claim_out = claim_out.as_basis_coefficients_slice().try_into().unwrap(); + cols.challenge = challenge.as_basis_coefficients_slice().try_into().unwrap(); + cols.expected = record + .expected + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + + row_offset += 1; + } + } + + Some(RowMajorMatrix::new(trace, width)) + } +} diff --git a/ceno_recursion_v2/src/main/mod.rs b/ceno_recursion_v2/src/main/mod.rs index 8e61fbffd..9870ccb3b 100644 --- a/ceno_recursion_v2/src/main/mod.rs +++ b/ceno_recursion_v2/src/main/mod.rs @@ -1,69 +1,151 @@ mod air; -mod sumcheck; +pub(crate) mod ecc_rt; +pub(crate) mod eval_absorb; +pub(crate) mod final_claim; +pub(crate) mod frontload; +pub(crate) mod global_sumcheck; +pub(crate) mod selector; +pub(crate) mod tower_point; mod trace; +mod transcript_bind; -use std::sync::Arc; +use std::{iter, sync::Arc}; -use ceno_zkvm::scheme::ZKVMChipProof; use eyre::{Result, bail, eyre}; +use ff_ext::ExtensionField; +use itertools::{Either, Itertools, chain}; +use multilinear_extensions::{ + Expression, StructuralWitIn, + StructuralWitInType::{ + Empty, EqualDistanceDynamicSequence, EqualDistanceSequence, + InnerRepeatingIncrementalSequence, OuterRepeatingIncrementalSequence, + StackedConstantSequence, StackedIncrementalSequence, + }, + mle::PointAndEval, + util::ceil_log2, + utils::{eval_by_expr, eval_by_expr_with_instance}, +}; use openvm_cpu_backend::CpuBackend; use openvm_poseidon2_air::POSEIDON2_WIDTH; use openvm_stark_backend::{ - AirRef, FiatShamirTranscript, ReadOnlyTranscript, StarkProtocolConfig, TranscriptHistory, - prover::AirProvingContext, + AirRef, FiatShamirTranscript, StarkProtocolConfig, TranscriptHistory, prover::AirProvingContext, }; -use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, EF, F}; -use p3_field::PrimeCharacteristicRing; +use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, D_EF, F}; +use p3_field::{BasedVectorSpace, PrimeCharacteristicRing}; use p3_matrix::dense::RowMajorMatrix; +use sumcheck::util::extrapolate_uni_poly; +use witness::next_pow2_instance_padding; use self::{ air::MainAir, - sumcheck::{ - MainSumcheckAir, MainSumcheckRecord, MainSumcheckRoundRecord, MainSumcheckTraceGenerator, + ecc_rt::{ + MainEccRtAir, MainEccRtChallengeAir, MainEccRtChallengeTraceGenerator, + MainEccRtEquationAir, MainEccRtEquationTraceGenerator, MainEccRtQuarkAir, + MainEccRtQuarkTraceGenerator, MainEccRtSumcheckAir, MainEccRtSumcheckTraceGenerator, + MainEccRtTraceGenerator, + }, + selector::{ + MAX_SELECTOR_POINT_VARS, MAX_SELECTOR_SPARSE_INDICES, MainSelectorEvalAir, + MainSelectorEvalTraceGenerator, MainSelectorFormulaAir, MainSelectorFormulaTraceGenerator, + MainSelectorPointAir, MainSelectorPointTraceGenerator, + selector_formula_global_point_lookups, selector_formula_point_lookup_counts, }, trace::{MainRecord, MainTraceGenerator}, + transcript_bind::{MainTranscriptBindAir, MainTranscriptBindTraceGenerator}, }; use crate::{ - bus::{MainBus, MainExpressionClaimBus, MainSumcheckInputBus, MainSumcheckOutputBus}, + bus::{ + AirPresenceBus, EccRtBus, ForkedTranscriptBus, MainBus, MainEccRtChallengeBus, + MainEccRtEquationTotalsBus, MainEccRtQuarkFinalBus, MainEccRtSumcheckFinalBus, MainEvalBus, + MainExpressionClaimBus, MainGlobalPointBus, MainSelectorPointBus, MainSelectorResultBus, + MainSelectorShapeBus, MainSelectorSparseIndexShapeBus, TowerMainPointBus, TranscriptBus, + }, system::{ - AirModule, BusIndexManager, BusInventory, ChipTranscriptRange, GlobalCtxCpu, Preflight, - RecursionField, RecursionProof, RecursionVk, TraceGenModule, + AirModule, BusIndexManager, BusInventory, EccReplayClaims, GlobalCtxCpu, MainEccRtRecord, + MainEvalRecord, MainFinalClaimRecord, MainFrontloadTermRecord, MainSelectorEvalRecord, + MainSelectorKind, MainSelectorPointDeriveKind, MainSelectorPointRecord, + MainSelectorPointSourceKind, MainTowerPointEqRecord, PcsOpeningClaimRecord, + PcsOpeningCommitKind, PcsOpeningEvalRecord, PcsOpeningPointRecord, Preflight, + RecursionField, RecursionProof, RecursionVk, RotationReplayClaims, TraceGenModule, + }, + tower::{ + TowerInputRecord, TowerReplayResult, build_tower_input_records, + build_tower_main_point_records, tower_transcript_len, }, tracegen::{ModuleChip, RowMajorChip}, + utils::transcript_observe_label, }; pub use air::MainCols; -pub use sumcheck::MainSumcheckCols; #[derive(Clone)] pub struct MainModule { main_bus: MainBus, - sumcheck_input_bus: MainSumcheckInputBus, - sumcheck_output_bus: MainSumcheckOutputBus, expression_claim_bus: MainExpressionClaimBus, + transcript_bus: TranscriptBus, + forked_transcript_bus: ForkedTranscriptBus, + air_presence_bus: AirPresenceBus, + main_global_point_bus: MainGlobalPointBus, + main_eval_bus: MainEvalBus, + main_selector_point_bus: MainSelectorPointBus, + main_selector_result_bus: MainSelectorResultBus, + main_selector_shape_bus: MainSelectorShapeBus, + main_selector_sparse_index_shape_bus: MainSelectorSparseIndexShapeBus, + main_ecc_rt_challenge_bus: MainEccRtChallengeBus, + main_ecc_rt_sumcheck_final_bus: MainEccRtSumcheckFinalBus, + main_ecc_rt_equation_totals_bus: MainEccRtEquationTotalsBus, + main_ecc_rt_quark_final_bus: MainEccRtQuarkFinalBus, + ecc_rt_bus: EccRtBus, + tower_main_point_bus: TowerMainPointBus, } impl MainModule { pub fn new(b: &mut BusIndexManager, bus_inventory: BusInventory) -> Self { let _ = b; let main_bus = bus_inventory.main_bus; - let sumcheck_input_bus = bus_inventory.main_sumcheck_input_bus; - let sumcheck_output_bus = bus_inventory.main_sumcheck_output_bus; let expression_claim_bus = bus_inventory.main_expression_claim_bus; + let transcript_bus = bus_inventory.transcript_bus; + let forked_transcript_bus = bus_inventory.forked_transcript_bus; + let air_presence_bus = bus_inventory.air_presence_bus; + let main_global_point_bus = bus_inventory.main_global_point_bus; + let main_eval_bus = bus_inventory.main_eval_bus; + let main_selector_point_bus = bus_inventory.main_selector_point_bus; + let main_selector_result_bus = bus_inventory.main_selector_result_bus; + let main_selector_shape_bus = bus_inventory.main_selector_shape_bus; + let main_selector_sparse_index_shape_bus = + bus_inventory.main_selector_sparse_index_shape_bus; + let main_ecc_rt_challenge_bus = bus_inventory.main_ecc_rt_challenge_bus; + let main_ecc_rt_sumcheck_final_bus = bus_inventory.main_ecc_rt_sumcheck_final_bus; + let main_ecc_rt_equation_totals_bus = bus_inventory.main_ecc_rt_equation_totals_bus; + let main_ecc_rt_quark_final_bus = bus_inventory.main_ecc_rt_quark_final_bus; + let ecc_rt_bus = bus_inventory.ecc_rt_bus; + let tower_main_point_bus = bus_inventory.tower_main_point_bus; Self { main_bus, - sumcheck_input_bus, - sumcheck_output_bus, expression_claim_bus, + transcript_bus, + forked_transcript_bus, + air_presence_bus, + main_global_point_bus, + main_eval_bus, + main_selector_point_bus, + main_selector_result_bus, + main_selector_shape_bus, + main_selector_sparse_index_shape_bus, + main_ecc_rt_challenge_bus, + main_ecc_rt_sumcheck_final_bus, + main_ecc_rt_equation_totals_bus, + main_ecc_rt_quark_final_bus, + ecc_rt_bus, + tower_main_point_bus, } } - fn collect_records( - &self, - _child_vk: &RecursionVk, + pub(crate) fn collect_records( + child_vk: &RecursionVk, proofs: &[RecursionProof], preflights: &[Preflight], - ) -> Result> { + ) -> Result { if proofs.len() != preflights.len() { bail!( "proof/preflight length mismatch ({} proofs vs {} preflights)", @@ -72,80 +154,409 @@ impl MainModule { ); } - let mut paired = Vec::new(); - for (proof_idx, (proof, preflight)) in proofs.iter().zip(preflights).enumerate() { - let mut chip_pf_iter = preflight.main.chips.iter(); - let mut saw_chip = false; - for (&chip_idx, chip_proof) in &proof.chip_proofs { - saw_chip = true; - let pf_entry = chip_pf_iter - .next() - .ok_or_else(|| eyre!("missing main preflight entry for chip {chip_idx}"))?; - if pf_entry.chip_idx != chip_idx { - bail!( - "main preflight chip mismatch: expected {}, got {}", - chip_idx, - pf_entry.chip_idx - ); - } - let claim = proof.main_constraint_proof.claimed_sum; - // Access the fork log directly using fork_idx and fork-local tidx. - let fork_log = preflight.fork_log(pf_entry.fork_idx); - let mut ts = ReadOnlyTranscript::new(fork_log, pf_entry.tidx); - record_main_transcript(&mut ts, chip_idx, claim); + let tower_input_records = build_tower_input_records(child_vk, proofs, preflights) + .map_err(|err| eyre!("failed to build tower input records for main prefix: {err}"))?; + let tower_main_point_records = build_tower_main_point_records(child_vk, proofs, preflights) + .map_err(|err| eyre!("failed to build tower point records for main prefix: {err}"))?; - let main_record = MainRecord { - proof_idx, - idx: chip_idx, - tidx: pf_entry.tidx, - claim, - }; - let sumcheck_record = build_sumcheck_record_from_chip( + let mut main_records = Vec::new(); + for input in &tower_input_records { + let tidx = main_tidx_from_tower_input(input); + main_records.push(MainRecord { + is_present: input.n_logup != 0, + proof_idx: input.proof_idx, + idx: input.idx, + tidx, + claim: input.input_layer_claim, + }); + } + + if main_records.is_empty() { + main_records.push(MainRecord::default()); + } + + let mut global_sumcheck_records = Vec::new(); + let mut eval_records = Vec::new(); + let mut selector_eval_records = Vec::new(); + let mut tower_point_eq_records = Vec::new(); + let mut frontload_term_records = Vec::new(); + let mut final_claim_records = Vec::new(); + for (proof_idx, preflight) in preflights.iter().enumerate() { + if let Some(mut record) = preflight.main.global_sumchecks.last().cloned() { + record.proof_idx = proof_idx; + global_sumcheck_records.push(record); + } + eval_records.extend(preflight.main.evals.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + })); + selector_eval_records.extend(preflight.main.selector_evals.iter().cloned().map( + move |mut record| { + record.proof_idx = proof_idx; + record + }, + )); + frontload_term_records.extend(preflight.main.frontload_terms.iter().cloned().map( + move |mut record| { + record.proof_idx = proof_idx; + record + }, + )); + final_claim_records.extend(preflight.main.final_claims.iter().cloned().map( + move |mut record| { + record.proof_idx = proof_idx; + record + }, + )); + } + if final_claim_records.is_empty() { + final_claim_records.push(MainFinalClaimRecord::default()); + } + if global_sumcheck_records.is_empty() { + global_sumcheck_records.push(crate::system::MainGlobalSumcheckRecord::default()); + } + let global_by_proof = global_sumcheck_records + .iter() + .map(|record| { + ( + record.proof_idx, + record + .rounds + .iter() + .map(|round| round.challenge) + .collect_vec(), + ) + }) + .collect::>(); + let mut eq_acc_by_chip = + std::collections::BTreeMap::<(usize, usize), RecursionField>::new(); + for tower_point in &tower_main_point_records { + let Some(global_point) = global_by_proof + .get(&tower_point.proof_idx) + .and_then(|points| points.get(tower_point.round_idx)) + .copied() + else { + continue; + }; + let acc = eq_acc_by_chip + .entry((tower_point.proof_idx, tower_point.idx)) + .or_insert(RecursionField::ONE); + let eq_in = *acc; + *acc *= global_point * tower_point.value + + (RecursionField::ONE - global_point) * (RecursionField::ONE - tower_point.value); + tower_point_eq_records.push(MainTowerPointEqRecord { + proof_idx: tower_point.proof_idx, + idx: tower_point.idx, + round_idx: tower_point.round_idx, + global_value: global_point, + tower_value: tower_point.value, + eq_in, + eq_out: *acc, + }); + } + let mut global_lookup_counts = std::collections::BTreeMap::<(usize, usize), usize>::new(); + for record in &tower_point_eq_records { + *global_lookup_counts + .entry((record.proof_idx, record.round_idx)) + .or_default() += 1; + } + for record in &frontload_term_records { + if record.has_global_factor { + *global_lookup_counts + .entry((record.proof_idx, record.global_round_idx)) + .or_default() += 1; + } + } + for (proof_idx, round_idx) in selector_formula_global_point_lookups(&selector_eval_records) + { + *global_lookup_counts + .entry((proof_idx, round_idx)) + .or_default() += 1; + } + for (proof_idx, preflight) in preflights.iter().enumerate() { + for record in &preflight.pcs.opening_points { + *global_lookup_counts + .entry((proof_idx, record.global_round_idx)) + .or_default() += 1; + } + for record in &preflight.pcs.suffix_products { + if record.has_factor { + *global_lookup_counts + .entry((proof_idx, record.coord_idx)) + .or_default() += 1; + } + } + for record in &preflight.pcs.jagged_assist_h { + if record.has_z_row { + *global_lookup_counts + .entry((proof_idx, record.robp_idx)) + .or_default() += 1; + } + } + } + let mut eval_lookup_counts = + std::collections::BTreeMap::<(usize, usize, usize), usize>::new(); + for record in &frontload_term_records { + if record.has_eval_factor { + *eval_lookup_counts + .entry((record.proof_idx, record.idx, record.eval_idx)) + .or_default() += 1; + } + } + for record in &selector_eval_records { + if !record.has_eval { + continue; + } + *eval_lookup_counts + .entry((record.proof_idx, record.idx, record.eval_idx)) + .or_default() += 1; + } + for (proof_idx, preflight) in preflights.iter().enumerate() { + for record in &preflight.pcs.opening_evals { + *eval_lookup_counts + .entry((proof_idx, record.main_idx, record.main_eval_idx)) + .or_default() += 1; + } + } + for record in &mut eval_records { + record.lookup_count = eval_lookup_counts + .get(&(record.proof_idx, record.idx, record.eval_idx)) + .copied() + .unwrap_or(0); + } + let tower_idx_by_air = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, preflight)| { + let sorted_idx_by_chip: std::collections::BTreeMap = preflight + .proof_shape + .sorted_trace_vdata + .iter() + .enumerate() + .map(|(sorted_idx, (chip_idx, _))| (*chip_idx, sorted_idx)) + .collect(); + let mut entries = preflight.gkr.chips.iter().collect_vec(); + entries.sort_by_key(|entry| { + ( + sorted_idx_by_chip + .get(&entry.chip_idx) + .copied() + .unwrap_or(usize::MAX), + entry.chip_idx, + ) + }); + entries + .into_iter() + .enumerate() + .map(move |(tower_idx, entry)| ((proof_idx, entry.chip_idx), tower_idx)) + }) + .collect::>(); + let selector_point_records = + build_main_selector_point_records(&selector_eval_records, &tower_idx_by_air); + let ecc_rt_records = build_main_ecc_rt_records( + proofs, + preflights, + &tower_idx_by_air, + &selector_point_records, + ); + for global in &mut global_sumcheck_records { + for (round_idx, round) in global.rounds.iter_mut().enumerate() { + round.point_lookup_count = global_lookup_counts + .get(&(global.proof_idx, round_idx)) + .copied() + .unwrap_or(0); + } + } + let mut transcript_records = Vec::new(); + for (proof_idx, preflight) in preflights.iter().enumerate() { + let eval_tidxs = eval_records + .iter() + .filter(|record| record.proof_idx == proof_idx) + .flat_map(|record| record.tidx..record.tidx + D_EF) + .collect::>(); + let global_challenge_tidxs = global_sumcheck_records + .iter() + .filter(|record| record.proof_idx == proof_idx) + .flat_map(|record| { + record + .rounds + .iter() + .flat_map(|round| round.challenge_tidx..round.challenge_tidx + D_EF) + }) + .collect::>(); + let semantic_trunk_tidxs = selector_point_records + .iter() + .filter(|record| { + record.proof_idx == proof_idx + && record.has_transcript + && record.source_kind == MainSelectorPointSourceKind::EccXY + }) + .flat_map(|record| record.transcript_tidx..record.transcript_tidx + D_EF) + .collect::>(); + let values = preflight.transcript.values(); + let samples = preflight.transcript.samples(); + for tidx in preflight.main.transcript_start..preflight.main.transcript_end { + if eval_tidxs.contains(&tidx) + || global_challenge_tidxs.contains(&tidx) + || semantic_trunk_tidxs.contains(&tidx) + { + continue; + } + transcript_records.push(crate::system::MainTranscriptRecord { proof_idx, - chip_idx, - claim, - chip_proof, - pf_entry.tidx, - ); - paired.push((main_record, sumcheck_record)); + fork_id: 0, + is_fork: false, + tidx, + value: values[tidx], + is_sample: samples[tidx], + }); } - - if !saw_chip { - paired.push(( - MainRecord { - proof_idx, - ..MainRecord::default() - }, - MainSumcheckRecord::default(), - )); + } + let semantic_fork_tidxs = selector_point_records + .iter() + .filter(|record| { + record.has_transcript + && record.source_kind == MainSelectorPointSourceKind::RotationOrigin + }) + .flat_map(|record| { + (0..D_EF).map(move |offset| { + ( + record.proof_idx, + record.fork_id, + record.transcript_tidx + offset, + ) + }) + }) + .chain(ecc_rt_records.iter().flat_map(|record| { + (0..D_EF).flat_map(move |offset| { + [ + (record.proof_idx, record.fork_id, record.tidx + offset), + (record.proof_idx, record.fork_id, record.out_tidx + offset), + (record.proof_idx, record.fork_id, record.alpha_tidx + offset), + ] + }) + })) + .collect::>(); + for input in tower_input_records + .iter() + .filter(|record| record.n_logup != 0) + { + let Some(fork) = preflights + .get(input.proof_idx) + .and_then(|preflight| preflight.fork_transcripts.get(input.fork_id)) + else { + continue; + }; + let fork_values = fork.log.values(); + let fork_samples = fork.log.samples(); + let tail_start = main_tidx_from_tower_input(input); + let tail_end = fork.log.len().saturating_sub(D_EF); + for tidx in tail_start..tail_end { + if semantic_fork_tidxs.contains(&(input.proof_idx, input.fork_id, tidx)) { + continue; + } + transcript_records.push(crate::system::MainTranscriptRecord { + proof_idx: input.proof_idx, + fork_id: input.fork_id, + is_fork: true, + tidx, + value: fork_values[tidx], + is_sample: fork_samples[tidx], + }); } } - - if paired.is_empty() { - paired.push((MainRecord::default(), MainSumcheckRecord::default())); + if transcript_records.is_empty() { + transcript_records.push(crate::system::MainTranscriptRecord::default()); } - Ok(paired) + Ok(MainCollectedRecords { + main_records, + global_sumcheck_records, + eval_records, + selector_eval_records, + selector_point_records, + ecc_rt_records, + tower_point_eq_records, + frontload_term_records, + final_claim_records, + transcript_records, + }) } } +pub(crate) struct MainCollectedRecords { + pub(crate) main_records: Vec, + pub(crate) global_sumcheck_records: Vec, + pub(crate) eval_records: Vec, + pub(crate) selector_eval_records: Vec, + pub(crate) selector_point_records: Vec, + pub(crate) ecc_rt_records: Vec, + pub(crate) tower_point_eq_records: Vec, + pub(crate) frontload_term_records: Vec, + pub(crate) final_claim_records: Vec, + pub(crate) transcript_records: Vec, +} + impl AirModule for MainModule { fn num_airs(&self) -> usize { - 2 + 10 } fn airs>(&self) -> Vec> { let main_air = MainAir { main_bus: self.main_bus, - sumcheck_input_bus: self.sumcheck_input_bus, - sumcheck_output_bus: self.sumcheck_output_bus, expression_claim_bus: self.expression_claim_bus, + send_expression_claim: false, }; - let main_sumcheck_air = MainSumcheckAir { - sumcheck_input_bus: self.sumcheck_input_bus, - sumcheck_output_bus: self.sumcheck_output_bus, - }; - vec![Arc::new(main_air) as AirRef<_>, Arc::new(main_sumcheck_air)] + vec![ + Arc::new(main_air) as AirRef<_>, + Arc::new(MainTranscriptBindAir { + transcript_bus: self.transcript_bus, + forked_transcript_bus: self.forked_transcript_bus, + }) as AirRef<_>, + Arc::new(MainEccRtChallengeAir { + forked_transcript_bus: self.forked_transcript_bus, + ecc_rt_bus: self.ecc_rt_bus, + challenge_bus: self.main_ecc_rt_challenge_bus, + }) as AirRef<_>, + Arc::new(MainEccRtEquationAir { + challenge_bus: self.main_ecc_rt_challenge_bus, + equation_totals_bus: self.main_ecc_rt_equation_totals_bus, + }) as AirRef<_>, + Arc::new(MainEccRtSumcheckAir { + challenge_bus: self.main_ecc_rt_challenge_bus, + sumcheck_final_bus: self.main_ecc_rt_sumcheck_final_bus, + }) as AirRef<_>, + Arc::new(MainEccRtQuarkAir { + challenge_bus: self.main_ecc_rt_challenge_bus, + quark_final_bus: self.main_ecc_rt_quark_final_bus, + }) as AirRef<_>, + Arc::new(MainEccRtAir { + challenge_bus: self.main_ecc_rt_challenge_bus, + sumcheck_final_bus: self.main_ecc_rt_sumcheck_final_bus, + equation_totals_bus: self.main_ecc_rt_equation_totals_bus, + quark_final_bus: self.main_ecc_rt_quark_final_bus, + }) as AirRef<_>, + Arc::new(MainSelectorPointAir { + selector_point_bus: self.main_selector_point_bus, + tower_point_bus: self.tower_main_point_bus, + transcript_bus: self.transcript_bus, + forked_transcript_bus: self.forked_transcript_bus, + ecc_rt_bus: self.ecc_rt_bus, + }) as AirRef<_>, + Arc::new(MainSelectorFormulaAir { + global_point_bus: self.main_global_point_bus, + selector_point_bus: self.main_selector_point_bus, + air_presence_bus: self.air_presence_bus, + selector_result_bus: self.main_selector_result_bus, + selector_shape_bus: self.main_selector_shape_bus, + selector_sparse_index_shape_bus: self.main_selector_sparse_index_shape_bus, + }) as AirRef<_>, + Arc::new(MainSelectorEvalAir { + eval_bus: self.main_eval_bus, + selector_result_bus: self.main_selector_result_bus, + }) as AirRef<_>, + ] } } @@ -160,16 +571,15 @@ impl MainModule { TS: FiatShamirTranscript + TranscriptHistory, { - let _ = (self, child_vk); - for &chip_idx in proof.chip_proofs.keys() { - let tidx = ts.len(); - record_main_transcript(ts, chip_idx, proof.main_constraint_proof.claimed_sum); - preflight.main.chips.push(ChipTranscriptRange { - chip_idx, - tidx, - fork_idx: 0, // unused in forked flow - }); + let _ = self; + preflight.main.transcript_start = ts.len(); + match replay_batched_main_preflight(child_vk, proof, ts, preflight) { + Ok(()) => {} + Err(err) => { + panic!("main preflight replay failed: {err}"); + } } + preflight.main.transcript_end = ts.len(); } } @@ -184,30 +594,59 @@ impl> TraceGenModule _ctx: &Self::ModuleSpecificCtx<'_>, required_heights: Option<&[usize]>, ) -> Option>>> { - let mut paired = self.collect_records(child_vk, proofs, preflights).ok()?; - paired.sort_by_key(|(record, _)| (record.proof_idx, record.idx)); - // Replace chip_idx with sequential index per proof_idx group to satisfy - // NestedForLoopSubAir counter constraints (must increment by 0 or 1). - // This matches the tower module which also uses entry_idx. - { - let mut prev_proof_idx = usize::MAX; - let mut seq_idx = 0usize; - for (record, sumcheck_record) in paired.iter_mut() { - if record.proof_idx != prev_proof_idx { - seq_idx = 0; - prev_proof_idx = record.proof_idx; - } - record.idx = seq_idx; - sumcheck_record.idx = seq_idx; - seq_idx += 1; - } - } - let (main_records, sumcheck_records): (Vec<_>, Vec<_>) = paired.into_iter().unzip(); + let mut records = Self::collect_records(child_vk, proofs, preflights).ok()?; + let MainCollectedRecords { + ref mut main_records, + global_sumcheck_records: _, + eval_records: _, + ref mut selector_eval_records, + ref mut selector_point_records, + ref mut ecc_rt_records, + tower_point_eq_records: _, + frontload_term_records: _, + final_claim_records: _, + ref mut transcript_records, + } = records; + main_records.sort_by_key(|record| (record.proof_idx, record.idx)); + selector_eval_records.sort_by_key(|record| { + ( + record.proof_idx, + record.idx, + record.air_idx, + record.selector_idx, + ) + }); + selector_point_records.sort_by_key(|record| { + ( + record.proof_idx, + record.idx, + record.air_idx, + record.selector_idx, + record.round_idx, + ) + }); + ecc_rt_records + .sort_by_key(|record| (record.proof_idx, record.idx, record.round_idx, record.tidx)); + transcript_records.sort_by_key(|record| (record.proof_idx, record.tidx)); let ctx = MainTraceCtx { main_records: &main_records, - sumcheck_records: &sumcheck_records, + selector_eval_records: &selector_eval_records, + selector_point_records: &selector_point_records, + ecc_rt_records: &ecc_rt_records, + transcript_records: &transcript_records, }; - let chips = [MainModuleChip::Main, MainModuleChip::Sumcheck]; + let chips = [ + MainModuleChip::Main, + MainModuleChip::TranscriptBind, + MainModuleChip::EccRtChallenge, + MainModuleChip::EccRtEquation, + MainModuleChip::EccRtSumcheck, + MainModuleChip::EccRtQuark, + MainModuleChip::EccRt, + MainModuleChip::SelectorPoint, + MainModuleChip::SelectorFormula, + MainModuleChip::SelectorEval, + ]; let span = tracing::Span::current(); let contexts = chips .into_iter() @@ -227,12 +666,23 @@ impl> TraceGenModule struct MainTraceCtx<'a> { main_records: &'a [MainRecord], - sumcheck_records: &'a [MainSumcheckRecord], + selector_eval_records: &'a [MainSelectorEvalRecord], + selector_point_records: &'a [MainSelectorPointRecord], + ecc_rt_records: &'a [MainEccRtRecord], + transcript_records: &'a [crate::system::MainTranscriptRecord], } enum MainModuleChip { Main, - Sumcheck, + TranscriptBind, + EccRtChallenge, + EccRtEquation, + EccRtSumcheck, + EccRtQuark, + EccRt, + SelectorPoint, + SelectorFormula, + SelectorEval, } impl RowMajorChip for MainModuleChip { @@ -247,53 +697,2546 @@ impl RowMajorChip for MainModuleChip { MainModuleChip::Main => { MainTraceGenerator.generate_trace(&ctx.main_records, required_height) } - MainModuleChip::Sumcheck => { - MainSumcheckTraceGenerator.generate_trace(&ctx.sumcheck_records, required_height) + MainModuleChip::TranscriptBind => MainTranscriptBindTraceGenerator + .generate_trace(&ctx.transcript_records, required_height), + MainModuleChip::EccRtChallenge => MainEccRtChallengeTraceGenerator + .generate_trace(&ctx.ecc_rt_records, required_height), + MainModuleChip::EccRtEquation => { + MainEccRtEquationTraceGenerator.generate_trace(&ctx.ecc_rt_records, required_height) + } + MainModuleChip::EccRtSumcheck => { + MainEccRtSumcheckTraceGenerator.generate_trace(&ctx.ecc_rt_records, required_height) + } + MainModuleChip::EccRtQuark => { + MainEccRtQuarkTraceGenerator.generate_trace(&ctx.ecc_rt_records, required_height) + } + MainModuleChip::EccRt => { + MainEccRtTraceGenerator.generate_trace(&ctx.ecc_rt_records, required_height) } + MainModuleChip::SelectorPoint => MainSelectorPointTraceGenerator + .generate_trace(&ctx.selector_point_records, required_height), + MainModuleChip::SelectorFormula => MainSelectorFormulaTraceGenerator + .generate_trace(&ctx.selector_eval_records, required_height), + MainModuleChip::SelectorEval => MainSelectorEvalTraceGenerator + .generate_trace(&ctx.selector_eval_records, required_height), } } } -fn build_sumcheck_record_from_chip( - proof_idx: usize, +fn main_tidx_from_tower_input(record: &TowerInputRecord) -> usize { + let tidx_after_alpha_beta = record.tidx + tower_transcript_len::ALPHA_BETA_LEN; + let read_active_layers = record.read_tower_vars - usize::from(record.has_read_out); + let write_active_layers = record.write_tower_vars - usize::from(record.has_write_out); + let logup_active_layers = record.logup_tower_vars - usize::from(record.has_logup_out); + let claim_span = + (read_active_layers + write_active_layers) * 2 * D_EF + logup_active_layers * 4 * D_EF; + let num_layers = record.n_logup; + let fixed_span = + num_layers * (tower_transcript_len::SUMCHECK_INIT_LEN + tower_transcript_len::MERGE_LEN); + let round_span = num_layers * (num_layers + 1) * (tower_transcript_len::ROUND_LEN / 2); + let alpha_span = num_layers * tower_transcript_len::ALPHA_LEN; + tidx_after_alpha_beta + fixed_span + round_span + alpha_span + claim_span +} + +pub(crate) fn replay_chip_pre_main_tail_transcript( + ts: &mut TS, + child_vk: &RecursionVk, chip_idx: usize, - claim: EF, - chip_proof: &ZKVMChipProof, - tidx: usize, -) -> MainSumcheckRecord { - let rounds = chip_proof - .gkr_iop_proof + chip_proof: &ceno_zkvm::scheme::ZKVMChipProof, + tower_replay: &TowerReplayResult, + challenges: [RecursionField; 2], +) -> Result<(Option, Option)> +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + let name = child_vk + .circuit_index_to_name + .get(&chip_idx) + .ok_or_else(|| eyre!("missing circuit name for chip index {chip_idx}"))?; + let circuit_vk = child_vk + .circuit_vks + .get(name) + .ok_or_else(|| eyre!("missing circuit vk for {name}"))?; + let composed_cs = circuit_vk.get_cs(); + + let mut ecc_replay = None; + if composed_cs.has_ecc_ops() { + let ecc_proof = chip_proof + .ecc_proof + .as_ref() + .ok_or_else(|| eyre!("{name} missing ecc proof"))?; + let num_vars = ceil_log2(next_pow2_instance_padding(ecc_proof.num_instances)); + let (_, out_rt_tidxs) = sample_vec_with_tidxs(ts, b"ecc", num_vars); + let (_, alpha_tidx) = sample_challenge_pows_with_tidx( + ts, + 7 * ceno_zkvm::scheme::constants::SEPTIC_EXTENSION_DEGREE, + b"ecc_alpha", + ); + let (_, _, rt_tidxs) = replay_main_sumcheck( + ts, + RecursionField::ZERO, + &ecc_proof.zerocheck_proof, + num_vars, + 3, + )?; + ecc_replay = Some(EccReplayClaims { + out_rt_tidxs, + alpha_tidx, + rt_tidxs, + }); + } + + let gkr_circuit = composed_cs + .gkr_circuit .as_ref() - .and_then(|proof| proof.0.first()) - .map(|layer| { - layer - .main - .proof - .proofs - .iter() - .map(|msg| { - let mut evals = [EF::ZERO; 3]; - for (dst, src) in evals.iter_mut().zip(msg.evaluations.iter().take(3)) { - *dst = *src; + .ok_or_else(|| eyre!("{name} missing gkr circuit in vk"))?; + let first_layer = gkr_circuit + .layers + .first() + .ok_or_else(|| eyre!("{name} empty gkr circuit layers"))?; + if first_layer.rotation_exprs.1.is_empty() { + return Ok((None, ecc_replay)); + } + + let rotation_proof = chip_proof + .rotation_proof + .as_ref() + .ok_or_else(|| eyre!("{name} missing rotation proof"))?; + let num_rotations = first_layer.rotation_exprs.1.len(); + if rotation_proof.evals.len() != num_rotations * 3 { + bail!( + "{name} rotation eval length mismatch: {} != {}", + rotation_proof.evals.len(), + num_rotations * 3 + ); + } + + let num_instances: usize = chip_proof.num_instances.iter().sum(); + let mut log2_num_instances = ceil_log2(next_pow2_instance_padding(num_instances)); + if composed_cs.has_ecc_ops() { + log2_num_instances += 1; + } + let num_var_with_rotation = log2_num_instances + composed_cs.rotation_vars().unwrap_or(0); + let rt_main = rt_main_from_tower_replay(tower_replay, num_var_with_rotation) + .ok_or_else(|| eyre!("{name} missing tower main point for rotation replay"))?; + + let (rotation_power_challenges, _) = + sample_challenge_pows_with_tidx(ts, num_rotations, b"combine subset evals"); + let rotation_challenges = chain!( + challenges.iter().copied(), + rotation_power_challenges.iter().copied() + ) + .collect_vec(); + let (origin_point, expected_evaluation, origin_tidxs) = replay_main_sumcheck( + ts, + RecursionField::ZERO, + &rotation_proof.proof, + num_var_with_rotation, + 2, + )?; + for eval in &rotation_proof.evals { + ts.observe_ext(*eval); + } + + let rotation_sumcheck_expression = first_layer + .rotation_sumcheck_expression + .as_ref() + .ok_or_else(|| eyre!("{name} missing rotation sumcheck expression"))?; + let bh = gkr_iop::gkr::booleanhypercube::BooleanHypercube::new( + first_layer.rotation_cyclic_group_log2, + ); + let selector_eval = gkr_iop::utils::rotation_selector_eval( + &bh, + &rt_main, + &origin_point, + first_layer.rotation_cyclic_subgroup_size, + first_layer.rotation_cyclic_group_log2, + ); + let mut left_evals = Vec::with_capacity(num_rotations); + let mut right_evals = Vec::with_capacity(num_rotations); + let mut target_evals = Vec::with_capacity(num_rotations); + let got_claim = eval_by_expr( + &rotation_proof + .evals + .chunks_exact(3) + .flat_map(|evals| { + let [left_eval, right_eval, target_eval] = evals else { + unreachable!() + }; + left_evals.push(*left_eval); + right_evals.push(*right_eval); + target_evals.push(*target_eval); + [ + (RecursionField::ONE + - origin_point[first_layer.rotation_cyclic_group_log2 - 1]) + * *left_eval + + origin_point[first_layer.rotation_cyclic_group_log2 - 1] * *right_eval, + *target_eval, + ] + }) + .chain(std::iter::once(selector_eval)) + .collect_vec(), + &[], + &rotation_challenges, + rotation_sumcheck_expression, + ); + if got_claim != expected_evaluation { + bail!("{name} rotation verify failed: {expected_evaluation} != {got_claim}"); + } + let (left_point, right_point) = bh.get_rotation_points(&origin_point); + + Ok(( + Some(RotationReplayClaims { + left_point, + right_point, + origin_point, + origin_tidxs, + left_evals, + right_evals, + target_evals, + }), + ecc_replay, + )) +} + +fn sample_vec_with_tidxs( + ts: &mut TS, + label: &[u8], + len: usize, +) -> (Vec, Vec) +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + transcript_observe_label(ts, label); + let mut values = Vec::with_capacity(len); + let mut tidxs = Vec::with_capacity(len); + for _ in 0..len { + tidxs.push(ts.len()); + values.push(FiatShamirTranscript::::sample_ext( + ts, + )); + } + (values, tidxs) +} + +struct MainReplayLayer<'a> { + air_idx: usize, + fork_id: usize, + layer: &'a gkr_iop::gkr::layer::Layer, + eval_and_dedup_points: Vec<(Vec, Option>)>, + selector_point_sources: Vec, + selector_ctxs: Vec, + eval_start: usize, + eval_len: usize, + alpha_start: usize, + num_var_with_rotation: usize, + pi: Vec, + rotation_cyclic_group_log2: usize, + rotation_origin_selector_idx: Option, + rotation_origin_tidxs: Vec, + ecc_sample_tidx: Option, + ecc_rt_tidxs: Vec, + ecc_xy_selector_idx: Option, + ecc_x3y3_selector_idx: Option, +} + +fn replay_batched_main_preflight( + child_vk: &RecursionVk, + proof: &RecursionProof, + ts: &mut TS, + preflight: &mut Preflight, +) -> Result<()> +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + let mut layers = Vec::new(); + let mut total_exprs = 0usize; + let mut total_evals = 0usize; + let mut max_num_variables = 0usize; + let mut max_degree = 0usize; + let tower_record_by_chip = preflight + .gkr + .chips + .iter() + .map(|record| (record.chip_idx, record)) + .collect::>(); + + for (chip_idx, chip_proof) in &proof.chip_proofs { + let name = child_vk + .circuit_index_to_name + .get(chip_idx) + .ok_or_else(|| eyre!("missing circuit name for chip index {chip_idx}"))?; + if proof.public_values.shard_id > 0 { + let Some(circuit_vk) = child_vk.circuit_vks.get(name) else { + continue; + }; + if circuit_vk.get_cs().with_omc_init_only() { + continue; + } + } + + let circuit_vk = child_vk + .circuit_vks + .get(name) + .ok_or_else(|| eyre!("missing circuit vk for {name}"))?; + let composed_cs = circuit_vk.get_cs(); + let gkr_circuit = composed_cs + .gkr_circuit + .as_ref() + .ok_or_else(|| eyre!("{name} missing gkr circuit in vk"))?; + let layer = gkr_circuit + .layers + .first() + .ok_or_else(|| eyre!("{name} empty gkr circuit layers"))?; + let num_instances: usize = chip_proof.num_instances.iter().sum(); + let mut log2_num_instances = ceil_log2(next_pow2_instance_padding(num_instances)); + if composed_cs.has_ecc_ops() { + log2_num_instances += 1; + } + let num_var_with_rotation = log2_num_instances + composed_cs.rotation_vars().unwrap_or(0); + let pi = composed_cs + .zkvm_v1_css + .instance + .iter() + .map(|instance| { + RecursionField::from( + proof + .public_values + .query_by_index::(instance.0), + ) + }) + .collect_vec(); + + let tower_record = tower_record_by_chip + .get(chip_idx) + .copied() + .ok_or_else(|| eyre!("{name} missing tower replay for main replay"))?; + let rt_main = rt_main_from_tower_replay(&tower_record.tower_replay, num_var_with_rotation); + if chip_proof.ecc_proof.is_some() && rt_main.is_none() { + bail!("{name} missing tower main point for ecc bridge main replay"); + } + let has_rt_main = rt_main.is_some(); + let rt_main = rt_main.unwrap_or_else(|| vec![RecursionField::ZERO; num_var_with_rotation]); + let selector_ctxs = first_layer_selector_contexts( + composed_cs, + gkr_circuit, + chip_proof.num_instances, + num_var_with_rotation, + )?; + let mut rotation_origin_selector_idx = None; + let mut rotation_origin_tidxs = Vec::new(); + let mut ecc_sample_tidx = None; + let mut ecc_rt_tidxs = Vec::new(); + let mut ecc_xy_selector_idx = None; + let mut ecc_x3y3_selector_idx = None; + let mut selector_point_sources = + vec![MainSelectorPointSourceKind::TowerMain; layer.out_sel_and_eval_exprs.len()]; + let mut out_evals = + vec![PointAndEval::new(rt_main, RecursionField::ZERO); gkr_circuit.n_evaluations]; + if chip_proof.main_out_evals.len() > gkr_circuit.n_evaluations { + bail!( + "{name} main output eval length {} exceeds gkr output length {}", + chip_proof.main_out_evals.len(), + gkr_circuit.n_evaluations + ); + } + for (out_eval, eval) in out_evals.iter_mut().zip(chip_proof.main_out_evals.iter()) { + out_eval.eval = *eval; + } + + if !layer.rotation_exprs.1.is_empty() { + let rotation_claims = tower_record + .rotation_replay + .as_ref() + .ok_or_else(|| eyre!("{name} missing rotation replay claims"))?; + let Some([left_group_idx, right_group_idx, point_group_idx]) = + layer.rotation_selector_group_indices() + else { + bail!("rotation claims expected but selectors are missing"); + }; + rotation_origin_selector_idx = Some(point_group_idx); + rotation_origin_tidxs = rotation_claims.origin_tidxs.clone(); + selector_point_sources[left_group_idx] = MainSelectorPointSourceKind::RotationLeft; + selector_point_sources[right_group_idx] = MainSelectorPointSourceKind::RotationRight; + selector_point_sources[point_group_idx] = MainSelectorPointSourceKind::RotationOrigin; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[left_group_idx].1, + &rotation_claims.left_evals, + &rotation_claims.left_point, + )?; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[right_group_idx].1, + &rotation_claims.right_evals, + &rotation_claims.right_point, + )?; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[point_group_idx].1, + &rotation_claims.target_evals, + &rotation_claims.origin_point, + )?; + } + + if let Some(ecc_proof) = chip_proof.ecc_proof.as_ref() { + let Some( + [ + x_group_idx, + y_group_idx, + slope_group_idx, + x3_group_idx, + y3_group_idx, + ], + ) = layer.ecc_bridge_group_indices() + else { + bail!("ecc bridge claims expected but selectors are missing"); + }; + selector_point_sources[x_group_idx] = MainSelectorPointSourceKind::EccXY; + selector_point_sources[y_group_idx] = MainSelectorPointSourceKind::EccXY; + selector_point_sources[slope_group_idx] = MainSelectorPointSourceKind::EccSlope; + selector_point_sources[x3_group_idx] = MainSelectorPointSourceKind::EccX3Y3; + selector_point_sources[y3_group_idx] = MainSelectorPointSourceKind::EccX3Y3; + ecc_xy_selector_idx = Some(x_group_idx); + ecc_x3y3_selector_idx = Some(x3_group_idx); + ecc_rt_tidxs = tower_record + .ecc_replay + .as_ref() + .ok_or_else(|| eyre!("{name} missing ecc replay claims"))? + .rt_tidxs + .clone(); + + let (sample_r_values, sample_r_tidxs) = + sample_vec_with_tidxs(ts, b"ecc_gkr_bridge_r", 1); + let sample_r = sample_r_values[0]; + ecc_sample_tidx = Some(sample_r_tidxs[0]); + let claims = derive_ecc_bridge_claims(ecc_proof, sample_r, num_var_with_rotation)?; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[x_group_idx].1, + &claims.x_evals, + &claims.xy_point, + )?; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[y_group_idx].1, + &claims.y_evals, + &claims.xy_point, + )?; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[slope_group_idx].1, + &claims.s_evals, + &claims.s_point, + )?; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[x3_group_idx].1, + &claims.x3_evals, + &claims.x3y3_point, + )?; + assign_group_evals( + &mut out_evals, + &layer.out_sel_and_eval_exprs[y3_group_idx].1, + &claims.y3_evals, + &claims.x3y3_point, + )?; + } + out_evals.resize(gkr_circuit.n_evaluations, PointAndEval::default()); + let eval_and_dedup_points = layer + .out_sel_and_eval_exprs + .iter() + .map(|(_, out_eval_exprs)| { + let evals = out_eval_exprs + .iter() + .map(|out_eval| { + out_eval + .evaluate(&out_evals, &pcs_challenges_from_preflight(preflight)) + .eval + }) + .collect_vec(); + let point = out_eval_exprs.first().and_then(|out_eval| { + if !has_rt_main { + return None; } - MainSumcheckRoundRecord { evaluations: evals } - }) - .collect::>() + Some( + out_eval + .evaluate(&out_evals, &pcs_challenges_from_preflight(preflight)) + .point, + ) + }); + (evals, point) + }) + .collect_vec(); + + let eval_len = layer.n_witin + layer.n_fixed + layer.n_structural_witin; + max_num_variables = max_num_variables.max(num_var_with_rotation); + max_degree = max_degree.max(layer.max_expr_degree + 1); + layers.push(MainReplayLayer { + air_idx: *chip_idx, + fork_id: tower_record.fork_idx, + layer, + eval_and_dedup_points, + selector_point_sources, + selector_ctxs, + eval_start: total_evals, + eval_len, + alpha_start: total_exprs, + num_var_with_rotation, + pi, + rotation_cyclic_group_log2: layer.rotation_cyclic_group_log2, + rotation_origin_selector_idx, + rotation_origin_tidxs, + ecc_sample_tidx, + ecc_rt_tidxs, + ecc_xy_selector_idx, + ecc_x3y3_selector_idx, + }); + total_evals += eval_len; + total_exprs += layer.exprs.len(); + } + + let main_proof = &proof.main_constraint_proof; + if layers.is_empty() { + if !main_proof.proof.proof.proofs.is_empty() + || !main_proof.proof.evals.is_empty() + || main_proof.claimed_sum != RecursionField::ZERO + { + bail!("empty main constraints with non-empty proof"); + } + preflight + .main + .global_sumchecks + .push(crate::system::MainGlobalSumcheckRecord { + proof_idx: 0, + expected: RecursionField::ZERO, + rounds: Vec::new(), + }); + preflight + .main + .final_claims + .push(MainFinalClaimRecord::default()); + return Ok(()); + } + if main_proof.proof.evals.len() != total_evals { + bail!( + "main constraint eval length mismatch: {} != {}", + main_proof.proof.evals.len(), + total_evals + ); + } + + let alpha_pows = sample_challenge_pows(ts, total_exprs, b"combine subset evals"); + let (global_in_point, expected_evaluation, challenge_tidxs) = replay_main_sumcheck( + ts, + main_proof.claimed_sum, + &main_proof.proof.proof, + max_num_variables, + max_degree, + )?; + + let eval_tidx_start = ts.len(); + for eval in &main_proof.proof.evals { + ts.observe_ext(*eval); + } + + let pcs_challenges = [ + preflight.vm_pvs.lookup_challenge_alpha, + preflight.vm_pvs.lookup_challenge_beta, + ]; + let mut eval_records = + build_main_eval_records(&layers, &main_proof.proof.evals, eval_tidx_start); + let tower_point_eqs = build_main_tower_point_eq_records(&layers, &global_in_point); + let selector_evals = + build_main_selector_eval_records(&layers, &main_proof.proof.evals, &global_in_point)?; + let (acc, frontload_terms, final_claims) = build_final_claim_records( + &layers, + &main_proof.proof.evals, + &pcs_challenges, + &alpha_pows, + &global_in_point, + expected_evaluation, + )?; + if acc != expected_evaluation { + bail!("main constraint claim mismatch: {expected_evaluation} != {acc}"); + } + preflight + .pcs + .opening_claims + .extend(layers.iter().map(|layer| { + let in_point = global_in_point[..layer.num_var_with_rotation].to_vec(); + let layer_evals = + &main_proof.proof.evals[layer.eval_start..layer.eval_start + layer.eval_len]; + PcsOpeningClaimRecord { + input_opening_point: in_point, + wits_in_evals: layer_evals[..layer.layer.n_witin].to_vec(), + fixed_in_evals: layer_evals + [layer.layer.n_witin..layer.layer.n_witin + layer.layer.n_fixed] + .to_vec(), + } + })); + let mut global_point_lookup_counts = vec![0usize; global_in_point.len()]; + for (opening_idx, layer) in layers.iter().enumerate() { + for (coord_idx, value) in global_in_point[..layer.num_var_with_rotation] + .iter() + .copied() + .enumerate() + { + global_point_lookup_counts[coord_idx] += 1; + preflight.pcs.opening_points.push(PcsOpeningPointRecord { + proof_idx: 0, + opening_idx, + coord_idx, + global_round_idx: coord_idx, + value, + }); + } + } + for record in &tower_point_eqs { + global_point_lookup_counts[record.round_idx] += 1; + } + for record in &frontload_terms { + if record.has_global_factor { + global_point_lookup_counts[record.global_round_idx] += 1; + } + } + for (_, round_idx) in selector_formula_global_point_lookups(&selector_evals) { + global_point_lookup_counts[round_idx] += 1; + } + let mut eval_lookup_counts = vec![0usize; main_proof.proof.evals.len()]; + for (opening_idx, layer) in layers.iter().enumerate() { + for eval_idx in 0..layer.layer.n_witin { + let global_eval_idx = layer.eval_start + eval_idx; + eval_lookup_counts[global_eval_idx] += 1; + preflight.pcs.opening_evals.push(PcsOpeningEvalRecord { + proof_idx: 0, + opening_idx, + commit_kind: PcsOpeningCommitKind::Witin, + eval_idx, + main_idx: opening_idx, + main_eval_idx: eval_idx, + value: main_proof.proof.evals[global_eval_idx], + raw_value: main_proof.proof.evals[global_eval_idx], + }); + } + for eval_idx in 0..layer.layer.n_fixed { + let main_eval_idx = layer.layer.n_witin + eval_idx; + let global_eval_idx = layer.eval_start + main_eval_idx; + eval_lookup_counts[global_eval_idx] += 1; + preflight.pcs.opening_evals.push(PcsOpeningEvalRecord { + proof_idx: 0, + opening_idx, + commit_kind: PcsOpeningCommitKind::Fixed, + eval_idx, + main_idx: opening_idx, + main_eval_idx, + value: main_proof.proof.evals[global_eval_idx], + raw_value: main_proof.proof.evals[global_eval_idx], + }); + } + } + for record in &frontload_terms { + if record.has_eval_factor { + let global_eval_idx = layers[record.idx].eval_start + record.eval_idx; + eval_lookup_counts[global_eval_idx] += 1; + } + } + for record in &selector_evals { + if !record.has_eval { + continue; + } + let global_eval_idx = layers[record.idx].eval_start + record.eval_idx; + eval_lookup_counts[global_eval_idx] += 1; + } + for record in &mut eval_records { + record.lookup_count = eval_lookup_counts[layers[record.idx].eval_start + record.eval_idx]; + } + preflight + .main + .global_sumchecks + .push(build_global_sumcheck_record( + 0, + main_proof.claimed_sum, + &main_proof.proof.proof, + &global_in_point, + &challenge_tidxs, + expected_evaluation, + &global_point_lookup_counts, + )?); + preflight.main.evals.extend(eval_records); + preflight.main.selector_evals.extend(selector_evals); + preflight.main.tower_point_eqs.extend(tower_point_eqs); + preflight.main.frontload_terms.extend(frontload_terms); + preflight.main.final_claims.extend(final_claims); + Ok(()) +} + +fn pcs_challenges_from_preflight(preflight: &Preflight) -> [RecursionField; 2] { + [ + preflight.vm_pvs.lookup_challenge_alpha, + preflight.vm_pvs.lookup_challenge_beta, + ] +} + +fn rt_main_from_tower_replay( + replay: &TowerReplayResult, + num_var_with_rotation: usize, +) -> Option> { + if num_var_with_rotation == 0 { + return Some(Vec::new()); + } + let final_layer = replay.layers.last()?; + let mut rt_tower = final_layer.challenges.clone(); + rt_tower.push(final_layer.mu); + if rt_tower.len() < num_var_with_rotation { + return None; + } + Some(rt_tower[rt_tower.len() - num_var_with_rotation..].to_vec()) +} + +struct EccBridgeClaims { + xy_point: Vec, + s_point: Vec, + x3y3_point: Vec, + x_evals: Vec, + y_evals: Vec, + s_evals: Vec, + x3_evals: Vec, + y3_evals: Vec, +} + +fn derive_ecc_bridge_claims( + ecc_proof: &ceno_zkvm::structs::EccQuarkProof, + sample_r: RecursionField, + num_var_with_rotation: usize, +) -> Result { + let degree = ceno_zkvm::scheme::constants::SEPTIC_EXTENSION_DEGREE; + if ecc_proof.evals.len() < 3 { + bail!("ecc proof evals shorter than selector prefix"); + } + let evals = &ecc_proof.evals[3..]; + if evals.len() != degree * 7 { + bail!( + "invalid ecc proof eval length: expected {}, got {}", + degree * 7, + evals.len() + ); + } + + let s1 = &evals[0..degree]; + let x0 = &evals[degree..2 * degree]; + let y0 = &evals[2 * degree..3 * degree]; + let x1 = &evals[3 * degree..4 * degree]; + let y1 = &evals[4 * degree..5 * degree]; + let x3 = &evals[5 * degree..6 * degree]; + let y3 = &evals[6 * degree..7 * degree]; + + let one_minus_r = RecursionField::ONE - sample_r; + let x_evals = x0 + .iter() + .zip_eq(x1.iter()) + .map(|(a, b)| *a * one_minus_r + *b * sample_r) + .collect_vec(); + let y_evals = y0 + .iter() + .zip_eq(y1.iter()) + .map(|(a, b)| *a * one_minus_r + *b * sample_r) + .collect_vec(); + let s_evals = s1.iter().map(|v| *v * sample_r).collect_vec(); + let x3_evals = x3.to_vec(); + let y3_evals = y3.to_vec(); + + let mut xy_point = vec![sample_r]; + xy_point.extend(ecc_proof.rt.iter().copied()); + if xy_point.len() != num_var_with_rotation { + bail!( + "invalid ecc xy point length: expected {}, got {}", + num_var_with_rotation, + xy_point.len() + ); + } + + let mut s_point = ecc_proof.rt.clone(); + s_point.push(sample_r); + if s_point.len() != num_var_with_rotation { + bail!( + "invalid ecc slope point length: expected {}, got {}", + num_var_with_rotation, + s_point.len() + ); + } + + let mut x3y3_point = ecc_proof.rt.clone(); + x3y3_point.push(RecursionField::ONE); + if x3y3_point.len() != num_var_with_rotation { + bail!( + "invalid ecc x3/y3 point length: expected {}, got {}", + num_var_with_rotation, + x3y3_point.len() + ); + } + + Ok(EccBridgeClaims { + xy_point, + s_point, + x3y3_point, + x_evals, + y_evals, + s_evals, + x3_evals, + y3_evals, + }) +} + +fn assign_group_evals( + out_evals: &mut [PointAndEval], + eval_exprs: &[gkr_iop::evaluation::EvalExpression], + evals: &[RecursionField], + point: &[RecursionField], +) -> Result<()> { + if eval_exprs.len() != evals.len() { + bail!( + "ecc bridge group eval length mismatch: {} != {}", + eval_exprs.len(), + evals.len() + ); + } + for (eval_expr, eval) in eval_exprs.iter().zip_eq(evals.iter()) { + let gkr_iop::evaluation::EvalExpression::Single(index) = eval_expr else { + bail!("ecc bridge group must use EvalExpression::Single"); + }; + let Some(out_eval) = out_evals.get_mut(*index) else { + bail!("ecc bridge output eval index {index} out of range"); + }; + *out_eval = PointAndEval::new(point.to_vec(), *eval); + } + Ok(()) +} + +#[derive(Clone, Copy, Default)] +struct GkrOutputStageMask(u8); + +impl GkrOutputStageMask { + const TOWER: Self = Self(1 << 0); + const ECC: Self = Self(1 << 1); + const ROTATION: Self = Self(1 << 2); + const ZERO: Self = Self(1 << 3); + + const fn union(self, other: Self) -> Self { + Self(self.0 | other.0) + } + + const fn contains(self, other: Self) -> bool { + (self.0 & other.0) == other.0 + } +} + +fn first_layer_output_group_stage_masks( + composed_cs: &ceno_zkvm::structs::ComposedConstrainSystem, + circuit: &gkr_iop::gkr::GKRCircuit, +) -> Result> { + let first_layer = circuit + .layers + .first() + .ok_or_else(|| eyre!("empty gkr circuit layer"))?; + let mut group_masks = vec![GkrOutputStageMask::ZERO; first_layer.out_sel_and_eval_exprs.len()]; + + if let Some(rotation_groups) = first_layer.rotation_selector_group_indices() { + for group_idx in rotation_groups { + let Some(mask) = group_masks.get_mut(group_idx) else { + bail!("rotation selector group index {group_idx} out of range"); + }; + *mask = GkrOutputStageMask::ROTATION; + } + } + if let Some(ecc_groups) = first_layer.ecc_bridge_group_indices() { + for group_idx in ecc_groups { + let Some(mask) = group_masks.get_mut(group_idx) else { + bail!("ecc selector group index {group_idx} out of range"); + }; + *mask = GkrOutputStageMask::ECC; + } + } + + let tower_outputs = tower_output_count(composed_cs); + let mut seen_tower_outputs = 0usize; + for (group_mask, (_, outputs)) in group_masks + .iter_mut() + .zip(first_layer.out_sel_and_eval_exprs.iter()) + { + if seen_tower_outputs >= tower_outputs { + break; + } + *group_mask = group_mask.union(GkrOutputStageMask::TOWER); + seen_tower_outputs += outputs.len(); + } + if seen_tower_outputs < tower_outputs { + bail!( + "failed to cover all tower outputs: layer={}, seen_tower_outputs={}, tower_outputs={}", + first_layer.name, + seen_tower_outputs, + tower_outputs + ); + } + + Ok(group_masks) +} + +fn first_layer_selector_contexts( + composed_cs: &ceno_zkvm::structs::ComposedConstrainSystem, + circuit: &gkr_iop::gkr::GKRCircuit, + num_instances: [usize; 2], + num_vars: usize, +) -> Result> { + let cs = &composed_cs.zkvm_v1_css; + let total_num_instances = num_instances.iter().sum(); + let first_layer = circuit + .layers + .first() + .ok_or_else(|| eyre!("empty gkr circuit layer"))?; + let group_stage_masks = first_layer_output_group_stage_masks(composed_cs, circuit)?; + let distinct_rw_selectors = + cs.r_selector.is_some() && cs.w_selector.is_some() && cs.r_selector != cs.w_selector; + + Ok(first_layer + .out_sel_and_eval_exprs + .iter() + .zip_eq(group_stage_masks.iter()) + .map(|((selector, _), stage_mask)| { + if stage_mask.contains(GkrOutputStageMask::TOWER) + && distinct_rw_selectors + && matches!(selector, gkr_iop::selector::SelectorType::Prefix(_)) + { + if cs.r_selector.as_ref() == Some(selector) { + return gkr_iop::selector::SelectorContext::new(0, num_instances[0], num_vars); + } + if cs.w_selector.as_ref() == Some(selector) { + return gkr_iop::selector::SelectorContext::new( + num_instances[0], + num_instances[1], + num_vars, + ); + } + } + + gkr_iop::selector::SelectorContext::new(0, total_num_instances, num_vars) }) - .unwrap_or_default(); + .collect_vec()) +} - MainSumcheckRecord { +fn tower_output_count( + composed_cs: &ceno_zkvm::structs::ComposedConstrainSystem, +) -> usize { + let cs = &composed_cs.zkvm_v1_css; + let num_reads = cs.r_expressions.len() + cs.r_table_expressions.len(); + let num_writes = cs.w_expressions.len() + cs.w_table_expressions.len(); + let num_lk_num = cs.lk_table_expressions.len(); + let num_lk_den = if !cs.lk_table_expressions.is_empty() { + cs.lk_table_expressions.len() + } else { + cs.lk_expressions.len() + }; + num_reads + num_writes + num_lk_num + num_lk_den +} + +fn build_global_sumcheck_record( + proof_idx: usize, + claimed_sum: RecursionField, + proof: &sumcheck::structs::IOPProof, + global_in_point: &[RecursionField], + challenge_tidxs: &[usize], + expected_evaluation: RecursionField, + point_lookup_counts: &[usize], +) -> Result { + if proof.proofs.len() != global_in_point.len() { + bail!( + "main global sumcheck proof/point length mismatch: {} != {}", + proof.proofs.len(), + global_in_point.len() + ); + } + if proof.proofs.len() != challenge_tidxs.len() { + bail!( + "main global sumcheck proof/challenge-tidx length mismatch: {} != {}", + proof.proofs.len(), + challenge_tidxs.len() + ); + } + let mut claim = claimed_sum; + let mut rounds = Vec::with_capacity(proof.proofs.len()); + for (round_idx, ((prover_msg, challenge), challenge_tidx)) in proof + .proofs + .iter() + .zip_eq(global_in_point.iter().copied()) + .zip_eq(challenge_tidxs.iter().copied()) + .enumerate() + { + if prover_msg.evaluations.len() > 4 { + bail!( + "main global sumcheck round {round_idx} has degree-width {}, max supported 4", + prover_msg.evaluations.len() + ); + } + if prover_msg.evaluations.is_empty() { + bail!("main global sumcheck round {round_idx} has no evaluations"); + } + let mut evaluations = [RecursionField::ZERO; 4]; + for (dst, src) in evaluations + .iter_mut() + .zip(prover_msg.evaluations.iter().copied()) + { + *dst = src; + } + if prover_msg.evaluations.len() < 4 { + evaluations[3] = extrapolate_uni_poly( + claim - evaluations[0], + &prover_msg.evaluations, + RecursionField::from_u64(4), + ); + } + let claim_out = extrapolate_uni_poly(claim - evaluations[0], &evaluations, challenge); + rounds.push(crate::system::MainGlobalSumcheckRoundRecord { + evaluations, + challenge, + challenge_tidx, + claim_in: claim, + claim_out, + point_lookup_count: point_lookup_counts.get(round_idx).copied().unwrap_or(0), + }); + claim = claim_out; + } + if claim != expected_evaluation { + bail!("main global sumcheck fold mismatch: {expected_evaluation} != {claim}"); + } + Ok(crate::system::MainGlobalSumcheckRecord { proof_idx, - idx: chip_idx, - tidx, - claim, + expected: expected_evaluation, rounds, + }) +} + +fn build_main_eval_records( + layers: &[MainReplayLayer<'_>], + main_evals: &[RecursionField], + eval_tidx_start: usize, +) -> Vec { + let mut records = Vec::new(); + for (idx, layer) in layers.iter().enumerate() { + for local_eval_idx in 0..layer.eval_len { + let global_eval_idx = layer.eval_start + local_eval_idx; + records.push(MainEvalRecord { + proof_idx: 0, + idx, + eval_idx: local_eval_idx, + tidx: eval_tidx_start + global_eval_idx * D_EF, + value: main_evals[global_eval_idx], + lookup_count: 0, + }); + } + } + records +} + +fn build_main_selector_point_records( + selector_evals: &[MainSelectorEvalRecord], + tower_idx_by_air: &std::collections::BTreeMap<(usize, usize), usize>, +) -> Vec { + let lookup_counts = selector_formula_point_lookup_counts(selector_evals); + let mut records = Vec::new(); + for record in selector_evals { + if !record.has_eval { + continue; + } + for (round_idx, value) in record.out_point.iter().copied().enumerate() { + let lookup_count = lookup_counts + .get(&( + record.proof_idx, + record.idx, + record.air_idx, + record.selector_idx, + round_idx, + )) + .copied() + .unwrap_or(0); + records.push(MainSelectorPointRecord { + proof_idx: record.proof_idx, + idx: record.idx, + tower_idx: tower_idx_by_air + .get(&(record.proof_idx, record.air_idx)) + .copied() + .unwrap_or(record.tower_idx), + air_idx: record.air_idx, + selector_idx: record.selector_idx, + round_idx, + value, + source_kind: record.point_source, + lookup_count, + fork_id: record.fork_id, + has_transcript: false, + transcript_tidx: 0, + has_ecc_rt: false, + has_source: false, + source_selector_idx: 0, + source_source_kind: MainSelectorPointSourceKind::TowerMain, + source_round_idx: 0, + source_value: RecursionField::ZERO, + derive_kind: MainSelectorPointDeriveKind::Identity, + }); + } + } + let values_by_key = records + .iter() + .map(|record| { + ( + ( + record.proof_idx, + record.idx, + record.air_idx, + record.selector_idx, + record.round_idx, + ), + record.value, + ) + }) + .collect::>(); + let evals_by_key = selector_evals + .iter() + .map(|record| { + ( + ( + record.proof_idx, + record.idx, + record.air_idx, + record.selector_idx, + ), + record, + ) + }) + .collect::>(); + let mut source_lookups = Vec::new(); + for point in &mut records { + let Some(eval) = evals_by_key.get(&( + point.proof_idx, + point.idx, + point.air_idx, + point.selector_idx, + )) else { + continue; + }; + match point.source_kind { + MainSelectorPointSourceKind::TowerMain => {} + MainSelectorPointSourceKind::RotationOrigin => { + if let Some(tidx) = eval.rotation_origin_tidxs.get(point.round_idx).copied() { + point.has_transcript = true; + point.transcript_tidx = tidx; + } + } + MainSelectorPointSourceKind::RotationLeft + | MainSelectorPointSourceKind::RotationRight => { + let Some(origin_selector_idx) = eval.rotation_origin_selector_idx else { + continue; + }; + let g = eval.rotation_cyclic_group_log2; + let (derive_kind, source_round_idx) = + rotation_point_source(point.source_kind, g, point.round_idx); + point.derive_kind = derive_kind; + if matches!( + derive_kind, + MainSelectorPointDeriveKind::Identity | MainSelectorPointDeriveKind::OneMinus + ) { + point.has_source = true; + point.source_selector_idx = origin_selector_idx; + point.source_source_kind = MainSelectorPointSourceKind::RotationOrigin; + point.source_round_idx = source_round_idx; + if let Some(value) = values_by_key.get(&( + point.proof_idx, + point.idx, + point.air_idx, + origin_selector_idx, + source_round_idx, + )) { + point.source_value = *value; + source_lookups.push(( + point.proof_idx, + point.idx, + point.air_idx, + origin_selector_idx, + source_round_idx, + )); + } + } + } + MainSelectorPointSourceKind::EccXY => { + if point.round_idx == 0 { + if Some(point.selector_idx) == eval.ecc_xy_selector_idx { + if let Some(tidx) = eval.ecc_sample_tidx { + point.has_transcript = true; + point.transcript_tidx = tidx; + } + } else if let Some(source_selector_idx) = eval.ecc_xy_selector_idx { + point.has_source = true; + point.source_selector_idx = source_selector_idx; + point.source_source_kind = MainSelectorPointSourceKind::EccXY; + point.source_round_idx = 0; + if let Some(value) = values_by_key.get(&( + point.proof_idx, + point.idx, + point.air_idx, + source_selector_idx, + point.source_round_idx, + )) { + point.source_value = *value; + source_lookups.push(( + point.proof_idx, + point.idx, + point.air_idx, + source_selector_idx, + point.source_round_idx, + )); + } + } else if let Some(tidx) = eval.ecc_sample_tidx { + point.has_transcript = true; + point.transcript_tidx = tidx; + } + } else if let Some(source_selector_idx) = eval.ecc_x3y3_selector_idx { + point.has_source = true; + point.source_selector_idx = source_selector_idx; + point.source_source_kind = MainSelectorPointSourceKind::EccX3Y3; + point.source_round_idx = point.round_idx - 1; + if let Some(value) = values_by_key.get(&( + point.proof_idx, + point.idx, + point.air_idx, + source_selector_idx, + point.source_round_idx, + )) { + point.source_value = *value; + source_lookups.push(( + point.proof_idx, + point.idx, + point.air_idx, + source_selector_idx, + point.source_round_idx, + )); + } + } + } + MainSelectorPointSourceKind::EccSlope => { + if point.round_idx + 1 == eval.out_point.len() { + if let Some(source_selector_idx) = eval.ecc_xy_selector_idx { + point.has_source = true; + point.source_selector_idx = source_selector_idx; + point.source_source_kind = MainSelectorPointSourceKind::EccXY; + point.source_round_idx = 0; + } + } else if let Some(source_selector_idx) = eval.ecc_x3y3_selector_idx { + point.has_source = true; + point.source_selector_idx = source_selector_idx; + point.source_source_kind = MainSelectorPointSourceKind::EccX3Y3; + point.source_round_idx = point.round_idx; + } + if point.has_source { + if let Some(value) = values_by_key.get(&( + point.proof_idx, + point.idx, + point.air_idx, + point.source_selector_idx, + point.source_round_idx, + )) { + point.source_value = *value; + source_lookups.push(( + point.proof_idx, + point.idx, + point.air_idx, + point.source_selector_idx, + point.source_round_idx, + )); + } + } + } + MainSelectorPointSourceKind::EccX3Y3 => { + if point.round_idx + 1 == eval.out_point.len() { + point.derive_kind = MainSelectorPointDeriveKind::One; + } else if Some(point.selector_idx) == eval.ecc_x3y3_selector_idx { + let Some(tidx) = eval.ecc_rt_tidxs.get(point.round_idx).copied() else { + continue; + }; + point.has_ecc_rt = true; + point.transcript_tidx = tidx; + } else if let Some(source_selector_idx) = eval.ecc_x3y3_selector_idx { + point.has_source = true; + point.source_selector_idx = source_selector_idx; + point.source_source_kind = MainSelectorPointSourceKind::EccX3Y3; + point.source_round_idx = point.round_idx; + if let Some(value) = values_by_key.get(&( + point.proof_idx, + point.idx, + point.air_idx, + source_selector_idx, + point.source_round_idx, + )) { + point.source_value = *value; + source_lookups.push(( + point.proof_idx, + point.idx, + point.air_idx, + source_selector_idx, + point.source_round_idx, + )); + } + } + } + } + } + let mut record_idx_by_key = records + .iter() + .enumerate() + .map(|(idx, record)| { + ( + ( + record.proof_idx, + record.idx, + record.air_idx, + record.selector_idx, + record.round_idx, + ), + idx, + ) + }) + .collect::>(); + for key in source_lookups { + if let Some(idx) = record_idx_by_key.remove(&key) { + records[idx].lookup_count += 1; + record_idx_by_key.insert(key, idx); + } + } + records.retain(|record| record.lookup_count != 0); + records +} + +fn build_main_ecc_rt_records( + proofs: &[RecursionProof], + preflights: &[Preflight], + tower_idx_by_air: &std::collections::BTreeMap<(usize, usize), usize>, + selector_point_records: &[MainSelectorPointRecord], +) -> Vec { + let lookup_counts = selector_point_records + .iter() + .filter(|record| record.has_ecc_rt) + .fold( + std::collections::BTreeMap::<(usize, usize, usize), usize>::new(), + |mut counts, record| { + *counts + .entry((record.proof_idx, record.tower_idx, record.round_idx)) + .or_default() += 1; + counts + }, + ); + let mut records = Vec::new(); + for (proof_idx, preflight) in preflights.iter().enumerate() { + let Some(proof) = proofs.get(proof_idx) else { + continue; + }; + for chip in &preflight.gkr.chips { + let Some(ecc_replay) = chip.ecc_replay.as_ref() else { + continue; + }; + let Some(ecc_proof) = proof + .chip_proofs + .get(&chip.chip_idx) + .and_then(|chip_proof| chip_proof.ecc_proof.as_ref()) + else { + continue; + }; + let Some(idx) = tower_idx_by_air.get(&(proof_idx, chip.chip_idx)).copied() else { + continue; + }; + let Some(fork) = preflight.fork_transcripts.get(chip.fork_idx) else { + continue; + }; + let fork_values = fork.log.values(); + let num_vars = ecc_replay.rt_tidxs.len(); + if num_vars == 0 { + continue; + } + let out_rt = ecc_replay + .out_rt_tidxs + .iter() + .map(|tidx| { + fork_values + .get(*tidx..*tidx + D_EF) + .and_then(RecursionField::from_basis_coefficients_slice) + .unwrap_or(RecursionField::ZERO) + }) + .collect_vec(); + let rt = ecc_replay + .rt_tidxs + .iter() + .map(|tidx| { + fork_values + .get(*tidx..*tidx + D_EF) + .and_then(RecursionField::from_basis_coefficients_slice) + .unwrap_or(RecursionField::ZERO) + }) + .collect_vec(); + let alpha = fork_values + .get(ecc_replay.alpha_tidx..ecc_replay.alpha_tidx + D_EF) + .and_then(RecursionField::from_basis_coefficients_slice) + .unwrap_or(RecursionField::ZERO); + let mut alpha_pows = [RecursionField::ZERO; 49]; + let mut pow = RecursionField::ONE; + for dst in &mut alpha_pows { + *dst = pow; + pow *= alpha; + } + if ecc_proof.evals.len() < 52 { + continue; + } + let evals = &ecc_proof.evals[3..]; + let to_septic = |start: usize| -> [RecursionField; 7] { + core::array::from_fn(|i| evals[start + i]) + }; + let s0 = to_septic(0); + let x0 = to_septic(7); + let y0 = to_septic(14); + let x1 = to_septic(21); + let y1 = to_septic(28); + let x3 = to_septic(35); + let y3 = to_septic(42); + let sum_x = core::array::from_fn(|i| RecursionField::from(ecc_proof.sum.x.0[i])); + let sum_y = core::array::from_fn(|i| RecursionField::from(ecc_proof.sum.y.0[i])); + let (add_eval, bypass_eval, export_eval) = native_ecc_equation_evals( + &s0, + &x0, + &y0, + &x1, + &y1, + &x3, + &y3, + &sum_x, + &sum_y, + &alpha_pows, + ); + + let mut layer_ns = (0..num_vars) + .scan(ecc_proof.num_instances, |n_instance, _| { + let current = *n_instance; + *n_instance = (*n_instance).div_ceil(2); + Some(current) + }) + .collect_vec(); + layer_ns.reverse(); + let mut claim = RecursionField::ZERO; + let mut eq_acc = RecursionField::ONE; + let mut last_acc = RecursionField::ONE; + let mut export_out_acc = RecursionField::ONE; + let mut export_rt_acc = RecursionField::ONE; + let mut quark_acc = RecursionField::ZERO; + for (round_idx, tidx) in ecc_replay.rt_tidxs.iter().copied().enumerate() { + let prover_msg = ecc_proof.zerocheck_proof.proofs.get(round_idx); + let mut round_evals = [RecursionField::ZERO; 3]; + if let Some(prover_msg) = prover_msg { + for (dst, src) in round_evals + .iter_mut() + .zip(prover_msg.evaluations.iter().copied()) + { + *dst = src; + } + } + let claim_in = claim; + let claim_out = + extrapolate_uni_poly(claim - round_evals[0], &round_evals, rt[round_idx]); + claim = claim_out; + + let eq_in = eq_acc; + eq_acc *= eq_binary_factor(out_rt[round_idx], rt[round_idx]); + let last_in = last_acc; + last_acc *= out_rt[round_idx] * rt[round_idx]; + let export_out_in = export_out_acc; + export_out_acc *= if round_idx == 0 { + RecursionField::ONE - out_rt[round_idx] + } else { + out_rt[round_idx] + }; + let export_rt_in = export_rt_acc; + export_rt_acc *= if round_idx == 0 { + RecursionField::ONE - rt[round_idx] + } else { + rt[round_idx] + }; + + let layer_n = layer_ns[round_idx]; + let prefix_count = layer_n / 2; + let quark_factor = if prefix_count == 0 { + RecursionField::ZERO + } else if round_idx == 0 { + RecursionField::ONE + } else { + native_eq_lte(prefix_count - 1, &out_rt[..round_idx], &rt[..round_idx]) + }; + let lte_witness = build_lte_witness(prefix_count, round_idx, &out_rt, &rt); + let quark_in = quark_acc; + quark_acc = (RecursionField::ONE - out_rt[round_idx]) + * (RecursionField::ONE - rt[round_idx]) + * quark_factor + + out_rt[round_idx] * rt[round_idx] * quark_acc; + + let value = fork_values + .get(tidx..tidx + D_EF) + .and_then(RecursionField::from_basis_coefficients_slice) + .unwrap_or(RecursionField::ZERO); + records.push(MainEccRtRecord { + proof_idx, + idx, + fork_id: chip.fork_idx, + round_idx, + num_rounds: ecc_replay.rt_tidxs.len(), + is_first: round_idx == 0, + is_last: round_idx + 1 == ecc_replay.rt_tidxs.len(), + tidx, + out_tidx: ecc_replay.out_rt_tidxs.get(round_idx).copied().unwrap_or(0), + alpha_tidx: ecc_replay.alpha_tidx, + value, + out_value: out_rt[round_idx], + alpha, + alpha_pows, + ev1: round_evals[0], + ev2: round_evals[1], + ev3: round_evals[2], + claim_in, + claim_out, + sel_add: ecc_proof.evals[0], + sel_bypass: ecc_proof.evals[1], + sel_export: ecc_proof.evals[2], + s0, + x0, + y0, + x1, + y1, + x3, + y3, + sum_x, + sum_y, + eq_in, + eq_out: eq_acc, + last_in, + last_out: last_acc, + export_out_in, + export_out_out: export_out_acc, + export_rt_in, + export_rt_out: export_rt_acc, + quark_in, + quark_factor, + quark_out: quark_acc, + add_eval, + bypass_eval, + export_eval, + lte_out_point: lte_witness.out_point, + lte_rt_point: lte_witness.rt_point, + lte_prefix_acc: lte_witness.prefix_acc, + lte_less_acc: lte_witness.less_acc, + lte_bits: lte_witness.bits, + lte_active: lte_witness.active, + quark_prefix_count: prefix_count, + quark_layer_n: layer_n, + quark_parity: layer_n % 2 == 1, + lookup_count: lookup_counts + .get(&(proof_idx, idx, round_idx)) + .copied() + .unwrap_or(0), + }); + } + } + } + records +} + +struct EccLteWitness { + out_point: [RecursionField; 32], + rt_point: [RecursionField; 32], + prefix_acc: [RecursionField; 33], + less_acc: [RecursionField; 33], + bits: [bool; 32], + active: [bool; 32], +} + +fn build_lte_witness( + prefix_count: usize, + round_idx: usize, + out_rt: &[RecursionField], + rt: &[RecursionField], +) -> EccLteWitness { + let max_idx = prefix_count.saturating_sub(1); + let out_point = core::array::from_fn(|i| out_rt.get(i).copied().unwrap_or_default()); + let rt_point = core::array::from_fn(|i| rt.get(i).copied().unwrap_or_default()); + let bits = core::array::from_fn(|i| i < round_idx && ((max_idx >> i) & 1) == 1); + let active = core::array::from_fn(|i| i < round_idx); + let mut prefix_acc = [RecursionField::ZERO; 33]; + let mut less_acc = [RecursionField::ZERO; 33]; + prefix_acc[32] = RecursionField::ONE; + for i in (0..32).rev() { + if active[i] { + let same_one = out_point[i] * rt_point[i]; + let same_zero = + (RecursionField::ONE - out_point[i]) * (RecursionField::ONE - rt_point[i]); + let same_any = same_one + same_zero; + let equal_choice = if bits[i] { same_one } else { same_zero }; + prefix_acc[i] = prefix_acc[i + 1] * equal_choice; + less_acc[i] = less_acc[i + 1] * same_any + + if bits[i] { + prefix_acc[i + 1] * same_zero + } else { + RecursionField::ZERO + }; + } else { + prefix_acc[i] = prefix_acc[i + 1]; + less_acc[i] = less_acc[i + 1]; + } + } + EccLteWitness { + out_point, + rt_point, + prefix_acc, + less_acc, + bits, + active, + } +} + +fn eq_binary_factor(a: RecursionField, b: RecursionField) -> RecursionField { + a * b + (RecursionField::ONE - a) * (RecursionField::ONE - b) +} + +fn native_eq_lte(max_idx: usize, a: &[RecursionField], b: &[RecursionField]) -> RecursionField { + let mut running_product = Vec::with_capacity(b.len() + 1); + running_product.push(RecursionField::ONE); + for i in 0..b.len() { + running_product.push(running_product[i] * eq_binary_factor(a[i], b[i])); + } + let mut running_product2 = vec![RecursionField::ZERO; b.len() + 1]; + running_product2[b.len()] = RecursionField::ONE; + for i in (0..b.len()).rev() { + let bit = RecursionField::from_usize((max_idx >> i) & 1); + running_product2[i] = running_product2[i + 1] + * (a[i] * b[i] * bit + + (RecursionField::ONE - a[i]) + * (RecursionField::ONE - b[i]) + * (RecursionField::ONE - bit)); + } + let mut ans = running_product[b.len()]; + for i in 0..b.len() { + if ((max_idx >> i) & 1) == 0 { + ans -= running_product[i] * running_product2[i + 1] * a[i] * b[i]; + } + } + ans +} + +#[allow(clippy::too_many_arguments)] +fn native_ecc_equation_evals( + s0: &[RecursionField; 7], + x0: &[RecursionField; 7], + y0: &[RecursionField; 7], + x1: &[RecursionField; 7], + y1: &[RecursionField; 7], + x3: &[RecursionField; 7], + y3: &[RecursionField; 7], + sum_x: &[RecursionField; 7], + sum_y: &[RecursionField; 7], + alpha_pows: &[RecursionField; 49], +) -> (RecursionField, RecursionField, RecursionField) { + let mut add_eval = RecursionField::ZERO; + let mut bypass_eval = RecursionField::ZERO; + let mut export_eval = RecursionField::ZERO; + let s0_x0_x1 = native_septic_mul(s0, &core::array::from_fn(|i| x0[i] - x1[i])); + let s0_squared = native_septic_mul(s0, s0); + let s0_x0_x3 = native_septic_mul(s0, &core::array::from_fn(|i| x0[i] - x3[i])); + for i in 0..7 { + let v1 = s0_x0_x1[i] - (y0[i] - y1[i]); + let v2 = s0_squared[i] - x0[i] - x1[i] - x3[i]; + let v3 = s0_x0_x3[i] - (y0[i] + y3[i]); + let v4 = x3[i] - x0[i]; + let v5 = y3[i] - y0[i]; + add_eval += v1 * alpha_pows[i] + v2 * alpha_pows[7 + i] + v3 * alpha_pows[14 + i]; + bypass_eval += v4 * alpha_pows[21 + i] + v5 * alpha_pows[28 + i]; + export_eval += + (x3[i] - sum_x[i]) * alpha_pows[35 + i] + (y3[i] - sum_y[i]) * alpha_pows[42 + i]; + } + (add_eval, bypass_eval, export_eval) +} + +fn native_septic_mul(a: &[RecursionField; 7], b: &[RecursionField; 7]) -> [RecursionField; 7] { + let mut out = [RecursionField::ZERO; 7]; + let two = RecursionField::from_usize(2); + let five = RecursionField::from_usize(5); + for i in 0..7 { + for j in 0..7 { + let term = a[i] * b[j]; + let mut index = i + j; + if index < 7 { + out[index] += term; + } else { + index -= 7; + out[index] += five * term; + out[index + 1] += two * term; + } + } + } + out +} + +fn rotation_point_source( + source_kind: MainSelectorPointSourceKind, + rotation_cyclic_group_log2: usize, + round_idx: usize, +) -> (MainSelectorPointDeriveKind, usize) { + match (source_kind, rotation_cyclic_group_log2, round_idx) { + (MainSelectorPointSourceKind::RotationLeft, _, 0) => (MainSelectorPointDeriveKind::Zero, 0), + (MainSelectorPointSourceKind::RotationLeft, g, round) if round < g => { + (MainSelectorPointDeriveKind::Identity, round - 1) + } + (MainSelectorPointSourceKind::RotationLeft, _, round) => { + (MainSelectorPointDeriveKind::Identity, round) + } + (MainSelectorPointSourceKind::RotationRight, _, 0) => (MainSelectorPointDeriveKind::One, 0), + (MainSelectorPointSourceKind::RotationRight, 5, 2) + | (MainSelectorPointSourceKind::RotationRight, 6, 1) => { + (MainSelectorPointDeriveKind::OneMinus, round_idx - 1) + } + (MainSelectorPointSourceKind::RotationRight, g, round) if round < g => { + (MainSelectorPointDeriveKind::Identity, round - 1) + } + (MainSelectorPointSourceKind::RotationRight, _, round) => { + (MainSelectorPointDeriveKind::Identity, round) + } + _ => (MainSelectorPointDeriveKind::Identity, round_idx), + } +} + +fn build_main_tower_point_eq_records( + layers: &[MainReplayLayer<'_>], + global_in_point: &[RecursionField], +) -> Vec { + let mut records = Vec::new(); + for (idx, layer) in layers.iter().enumerate() { + for round_idx in 0..layer.num_var_with_rotation { + records.push(MainTowerPointEqRecord { + proof_idx: 0, + idx, + round_idx, + global_value: global_in_point[round_idx], + tower_value: RecursionField::ZERO, + eq_in: RecursionField::ZERO, + eq_out: RecursionField::ZERO, + }); + } } + records } -pub(crate) fn record_main_transcript(ts: &mut TS, _chip_idx: usize, claim: EF) +fn build_main_selector_eval_records( + layers: &[MainReplayLayer<'_>], + main_evals: &[RecursionField], + global_in_point: &[RecursionField], +) -> Result> { + let mut records = Vec::new(); + for (idx, layer) in layers.iter().enumerate() { + let structural_witin_offset = layer.layer.n_witin + layer.layer.n_fixed; + let in_point = global_in_point[..layer.num_var_with_rotation].to_vec(); + if in_point.len() > MAX_SELECTOR_POINT_VARS { + bail!( + "{} selector point width {} exceeds AIR cap {}", + layer.layer.name, + in_point.len(), + MAX_SELECTOR_POINT_VARS + ); + } + for (selector_idx, (((sel_type, _), (_, out_point)), selector_ctx)) in layer + .layer + .out_sel_and_eval_exprs + .iter() + .zip(layer.eval_and_dedup_points.iter()) + .zip(layer.selector_ctxs.iter()) + .enumerate() + { + if matches!(sel_type, gkr_iop::selector::SelectorType::None) { + bail!( + "{} selector group {selector_idx} uses SelectorType::None; migrate this selector before proving recursion-v2 main", + layer.layer.name + ); + } + let (kind, ordered_sparse_num_vars, sparse_indices, wit_id) = selector_shape(sel_type)?; + if sparse_indices.len() > MAX_SELECTOR_SPARSE_INDICES { + bail!( + "{} selector group {selector_idx} sparse index count {} exceeds AIR cap {}", + layer.layer.name, + sparse_indices.len(), + MAX_SELECTOR_SPARSE_INDICES + ); + } + let eval_idx = wit_id as usize + structural_witin_offset; + let Some(out_point) = out_point.as_ref() else { + records.push(MainSelectorEvalRecord { + proof_idx: 0, + idx, + tower_idx: 0, + air_idx: layer.air_idx, + selector_idx, + has_eval: false, + eval_idx, + kind, + ctx_offset: selector_ctx.offset, + ctx_num_instances: selector_ctx.num_instances, + ctx_num_vars: selector_ctx.num_vars, + ordered_sparse_num_vars, + sparse_indices, + in_point: Vec::new(), + out_point: Vec::new(), + point_source: layer.selector_point_sources[selector_idx], + fork_id: layer.fork_id, + rotation_cyclic_group_log2: layer.rotation_cyclic_group_log2, + rotation_origin_selector_idx: layer.rotation_origin_selector_idx, + rotation_origin_tidxs: layer.rotation_origin_tidxs.clone(), + ecc_sample_tidx: layer.ecc_sample_tidx, + ecc_rt_tidxs: layer.ecc_rt_tidxs.clone(), + ecc_xy_selector_idx: layer.ecc_xy_selector_idx, + ecc_x3y3_selector_idx: layer.ecc_x3y3_selector_idx, + value: RecursionField::ZERO, + }); + continue; + }; + if out_point.len() > MAX_SELECTOR_POINT_VARS { + bail!( + "{} selector output point width {} exceeds AIR cap {}", + layer.layer.name, + out_point.len(), + MAX_SELECTOR_POINT_VARS + ); + } + let Some((expected_eval, evaluated_wit_id)) = + sel_type.evaluate(out_point, &in_point, selector_ctx) + else { + records.push(MainSelectorEvalRecord { + proof_idx: 0, + idx, + tower_idx: 0, + air_idx: layer.air_idx, + selector_idx, + has_eval: false, + eval_idx, + kind, + ctx_offset: selector_ctx.offset, + ctx_num_instances: selector_ctx.num_instances, + ctx_num_vars: selector_ctx.num_vars, + ordered_sparse_num_vars, + sparse_indices, + in_point: Vec::new(), + out_point: Vec::new(), + point_source: layer.selector_point_sources[selector_idx], + fork_id: layer.fork_id, + rotation_cyclic_group_log2: layer.rotation_cyclic_group_log2, + rotation_origin_selector_idx: layer.rotation_origin_selector_idx, + rotation_origin_tidxs: layer.rotation_origin_tidxs.clone(), + ecc_sample_tidx: layer.ecc_sample_tidx, + ecc_rt_tidxs: layer.ecc_rt_tidxs.clone(), + ecc_xy_selector_idx: layer.ecc_xy_selector_idx, + ecc_x3y3_selector_idx: layer.ecc_x3y3_selector_idx, + value: RecursionField::ZERO, + }); + continue; + }; + if evaluated_wit_id != wit_id { + bail!( + "{} selector group {selector_idx} witness id mismatch: shape {wit_id} evaluate {evaluated_wit_id}", + layer.layer.name + ); + } + let Some(_actual_eval) = main_evals.get(layer.eval_start + eval_idx).copied() else { + bail!("main selector structural witin index {eval_idx} out of range"); + }; + records.push(MainSelectorEvalRecord { + proof_idx: 0, + idx, + tower_idx: 0, + air_idx: layer.air_idx, + selector_idx, + has_eval: true, + eval_idx, + kind, + ctx_offset: selector_ctx.offset, + ctx_num_instances: selector_ctx.num_instances, + ctx_num_vars: selector_ctx.num_vars, + ordered_sparse_num_vars, + sparse_indices, + in_point: in_point.clone(), + out_point: out_point.clone(), + point_source: layer.selector_point_sources[selector_idx], + fork_id: layer.fork_id, + rotation_cyclic_group_log2: layer.rotation_cyclic_group_log2, + rotation_origin_selector_idx: layer.rotation_origin_selector_idx, + rotation_origin_tidxs: layer.rotation_origin_tidxs.clone(), + ecc_sample_tidx: layer.ecc_sample_tidx, + ecc_rt_tidxs: layer.ecc_rt_tidxs.clone(), + ecc_xy_selector_idx: layer.ecc_xy_selector_idx, + ecc_x3y3_selector_idx: layer.ecc_x3y3_selector_idx, + value: expected_eval, + }); + } + } + Ok(records) +} + +fn selector_shape( + sel_type: &gkr_iop::selector::SelectorType, +) -> Result<( + MainSelectorKind, + usize, + Vec, + multilinear_extensions::WitnessId, +)> { + use gkr_iop::selector::SelectorType; + let (kind, ordered_sparse_num_vars, sparse_indices, expr) = match sel_type { + SelectorType::None => bail!("SelectorType::None is not supported in recursion-v2 main"), + SelectorType::Whole(expr) => (MainSelectorKind::Whole, 0, Vec::new(), expr), + SelectorType::Prefix(expr) => (MainSelectorKind::Prefix, 0, Vec::new(), expr), + SelectorType::OrderedSparse { + num_vars, + indices, + expression, + } => ( + MainSelectorKind::OrderedSparse, + *num_vars, + indices.clone(), + expression, + ), + SelectorType::QuarkBinaryTreeLessThan(expr) => ( + MainSelectorKind::QuarkBinaryTreeLessThan, + 0, + Vec::new(), + expr, + ), + }; + let Expression::StructuralWitIn(wit_id, _) = expr else { + bail!("selector expression must be StructuralWitIn"); + }; + Ok((kind, ordered_sparse_num_vars, sparse_indices, *wit_id)) +} + +fn build_final_claim_records( + layers: &[MainReplayLayer<'_>], + main_evals: &[RecursionField], + pcs_challenges: &[RecursionField; 2], + alpha_pows: &[RecursionField], + global_in_point: &[RecursionField], + expected_evaluation: RecursionField, +) -> Result<( + RecursionField, + Vec, + Vec, +)> { + let mut acc = RecursionField::ZERO; + let mut frontload_records = Vec::new(); + let mut records = Vec::with_capacity(layers.len()); + for (idx, layer) in layers.iter().enumerate() { + let layer_evals = &main_evals[layer.eval_start..layer.eval_start + layer.eval_len]; + validate_direct_structural_evals(layer, layer_evals, global_in_point)?; + let main_sumcheck_challenges = chain!( + pcs_challenges.iter().copied(), + alpha_pows[layer.alpha_start..layer.alpha_start + layer.layer.exprs.len()] + .iter() + .copied() + ) + .collect_vec(); + let terms = layer + .layer + .main_sumcheck_expression_monomial_terms + .as_ref() + .ok_or_else(|| eyre!("missing main sumcheck expression monomial terms"))?; + let contribution = build_frontload_term_records_for_layer( + idx, + layer, + layer_evals, + &layer.pi, + &main_sumcheck_challenges, + global_in_point, + terms, + &mut frontload_records, + )?; + let monomial_contribution = eval_batched_main_frontload_terms_oracle( + layer_evals, + &layer.pi, + &main_sumcheck_challenges, + global_in_point, + layer.num_var_with_rotation, + terms, + ); + if contribution != monomial_contribution { + bail!( + "layer monomial contribution mismatch at chip {idx}: {monomial_contribution} != {contribution}" + ); + } + let acc_in = acc; + acc += contribution; + records.push(MainFinalClaimRecord { + proof_idx: 0, + idx, + contribution, + acc_in, + acc_out: acc, + expected: expected_evaluation, + }); + } + Ok((acc, frontload_records, records)) +} + +fn validate_direct_structural_evals( + layer: &MainReplayLayer<'_>, + layer_evals: &[RecursionField], + global_in_point: &[RecursionField], +) -> Result<()> { + let structural_witin_offset = layer.layer.n_witin + layer.layer.n_fixed; + let in_point = global_in_point[..layer.num_var_with_rotation].to_vec(); + for (((sel_type, _), (_, out_point)), selector_ctx) in layer + .layer + .out_sel_and_eval_exprs + .iter() + .zip(layer.eval_and_dedup_points.iter()) + .zip(layer.selector_ctxs.iter()) + { + let Some(out_point) = out_point.as_ref() else { + continue; + }; + if let Some((_expected_eval, wit_id)) = + sel_type.evaluate(out_point, &in_point, selector_ctx) + { + let wit_id = wit_id as usize + structural_witin_offset; + let Some(_actual_eval) = layer_evals.get(wit_id).copied() else { + bail!("main selector structural witin index {wit_id} out of range"); + }; + } + } + for StructuralWitIn { id, witin_type } in &layer.layer.structural_witins { + let wit_id = *id as usize + structural_witin_offset; + let Some(actual_eval) = layer_evals.get(wit_id).copied() else { + bail!("main structural witin index {wit_id} out of range"); + }; + let expected_eval = match witin_type { + EqualDistanceSequence { + offset, + multi_factor, + descending, + .. + } => gkr_iop::utils::eval_wellform_address_vec( + *offset as u64, + *multi_factor as u64, + &in_point, + *descending, + ), + EqualDistanceDynamicSequence { + offset_instance_id, + multi_factor, + descending, + .. + } => { + let offset = layer.pi[*offset_instance_id as usize].to_canonical_u64(); + gkr_iop::utils::eval_wellform_address_vec( + offset, + *multi_factor as u64, + &in_point, + *descending, + ) + } + StackedIncrementalSequence { .. } => { + gkr_iop::utils::eval_stacked_wellform_address_vec(&in_point) + } + StackedConstantSequence { .. } => gkr_iop::utils::eval_stacked_constant_vec(&in_point), + InnerRepeatingIncrementalSequence { k, .. } => { + gkr_iop::utils::eval_inner_repeated_incremental_vec(*k as u64, &in_point) + } + OuterRepeatingIncrementalSequence { k, .. } => { + gkr_iop::utils::eval_outer_repeated_incremental_vec(*k as u64, &in_point) + } + Empty => continue, + }; + if actual_eval != expected_eval { + bail!( + "{} structural witin mismatch: {expected_eval} != {actual_eval}", + layer.layer.name + ); + } + } + Ok(()) +} + +fn build_frontload_term_records_for_layer( + idx: usize, + layer: &MainReplayLayer<'_>, + layer_evals: &[RecursionField], + pi: &[RecursionField], + challenges: &[RecursionField], + global_in_point: &[RecursionField], + terms: &[multilinear_extensions::monomial::Term< + Expression, + Expression, + >], + records: &mut Vec, +) -> Result { + let tail_start = layer.num_var_with_rotation; + let tail_point = &global_in_point[tail_start..]; + let mut row_idx = 0usize; + let mut node_idx = 0usize; + let mut layer_acc = RecursionField::ZERO; + + for (term_idx, term) in terms.iter().enumerate() { + let mut term_value = emit_scalar_expr_row( + idx, + term_idx, + &term.scalar, + pi, + challenges, + records, + &mut row_idx, + &mut node_idx, + )?; + + for expr in &term.product { + let Expression::WitIn(wit_id) = expr else { + bail!("main monomial product must be converted to WitIn"); + }; + let eval_idx = *wit_id as usize; + let raw_eval = *layer_evals + .get(eval_idx) + .ok_or_else(|| eyre!("main monomial wit index {eval_idx} out of range"))?; + let weighted_eval = emit_monomial_wit_factor_rows( + idx, + term_idx, + eval_idx, + raw_eval, + tail_start, + tail_point, + records, + &mut row_idx, + &mut node_idx, + ); + let next = term_value * weighted_eval; + records.push(MainFrontloadTermRecord { + proof_idx: 0, + idx, + row_idx, + node_idx, + is_mul: true, + constraint_idx: term_idx, + arg0: term_value, + arg1: weighted_eval, + value: next, + ..MainFrontloadTermRecord::default() + }); + row_idx += 1; + node_idx += 1; + term_value = next; + } + + let next_acc = layer_acc + term_value; + records.push(MainFrontloadTermRecord { + proof_idx: 0, + idx, + row_idx, + node_idx, + is_fold: true, + constraint_idx: term_idx, + alpha: RecursionField::ONE, + arg0: term_value, + value: next_acc, + chip_acc_in: layer_acc, + chip_acc_out: next_acc, + is_last_chip_step: term_idx + 1 == terms.len(), + ..MainFrontloadTermRecord::default() + }); + row_idx += 1; + node_idx += 1; + layer_acc = next_acc; + } + + if terms.is_empty() { + records.push(MainFrontloadTermRecord { + proof_idx: 0, + idx, + row_idx, + node_idx, + is_fold: true, + alpha: RecursionField::ONE, + arg0: RecursionField::ZERO, + value: RecursionField::ZERO, + chip_acc_in: RecursionField::ZERO, + chip_acc_out: RecursionField::ZERO, + is_last_chip_step: true, + ..MainFrontloadTermRecord::default() + }); + } + + Ok(layer_acc) +} + +fn emit_scalar_expr_row( + idx: usize, + term_idx: usize, + expr: &Expression, + pi: &[RecursionField], + challenges: &[RecursionField], + records: &mut Vec, + row_idx: &mut usize, + node_idx: &mut usize, +) -> Result { + let this_node = *node_idx; + *node_idx += 1; + let mut record = MainFrontloadTermRecord { + proof_idx: 0, + idx, + row_idx: *row_idx, + node_idx: this_node, + constraint_idx: term_idx, + ..MainFrontloadTermRecord::default() + }; + + let value = match expr { + Expression::WitIn(_) | Expression::StructuralWitIn(_, _) | Expression::Fixed(_) => { + bail!("main monomial scalar must not contain witness-backed expressions") + } + Expression::Instance(instance) | Expression::InstanceScalar(instance) => { + let value = *pi + .get(instance.0 as usize) + .ok_or_else(|| eyre!("main scalar instance index {} out of range", instance.0))?; + record.is_instance = true; + record.instance_idx = instance.0 as usize; + record.arg0 = value; + value + } + Expression::Constant(value) => { + let value = either_to_ext(*value); + record.is_const = true; + record.arg0 = value; + value + } + Expression::Challenge(ch_id, pow, scalar, offset) => { + let challenge_idx = *ch_id as usize; + let challenge = *challenges + .get(challenge_idx) + .ok_or_else(|| eyre!("main scalar challenge index {challenge_idx} out of range"))?; + let value = challenge.exp_u64(*pow as u64) * *scalar + *offset; + record.is_challenge = true; + record.challenge_idx = challenge_idx; + record.arg0 = value; + value + } + Expression::Sum(left, right) => { + let left = emit_scalar_expr_row( + idx, term_idx, left, pi, challenges, records, row_idx, node_idx, + )?; + let right = emit_scalar_expr_row( + idx, term_idx, right, pi, challenges, records, row_idx, node_idx, + )?; + record.is_add = true; + record.arg0 = left; + record.arg1 = right; + left + right + } + Expression::Product(left, right) => { + let left = emit_scalar_expr_row( + idx, term_idx, left, pi, challenges, records, row_idx, node_idx, + )?; + let right = emit_scalar_expr_row( + idx, term_idx, right, pi, challenges, records, row_idx, node_idx, + )?; + record.is_mul = true; + record.arg0 = left; + record.arg1 = right; + left * right + } + Expression::ScaledSum(x, a, b) => { + let x = + emit_scalar_expr_row(idx, term_idx, x, pi, challenges, records, row_idx, node_idx)?; + let a = + emit_scalar_expr_row(idx, term_idx, a, pi, challenges, records, row_idx, node_idx)?; + let b = + emit_scalar_expr_row(idx, term_idx, b, pi, challenges, records, row_idx, node_idx)?; + let mul_node = *node_idx; + *node_idx += 1; + records.push(MainFrontloadTermRecord { + proof_idx: 0, + idx, + row_idx: *row_idx, + node_idx: mul_node, + is_mul: true, + constraint_idx: term_idx, + arg0: a, + arg1: x, + value: a * x, + ..MainFrontloadTermRecord::default() + }); + *row_idx += 1; + record.is_add = true; + record.arg0 = a * x; + record.arg1 = b; + a * x + b + } + }; + + record.value = value; + record.chip_acc_out = record.chip_acc_in; + records.push(record); + *row_idx += 1; + Ok(value) +} + +#[allow(clippy::too_many_arguments)] +fn emit_monomial_wit_factor_rows( + idx: usize, + term_idx: usize, + eval_idx: usize, + mut acc: RecursionField, + tail_start: usize, + tail_point: &[RecursionField], + records: &mut Vec, + row_idx: &mut usize, + node_idx: &mut usize, +) -> RecursionField { + records.push(MainFrontloadTermRecord { + proof_idx: 0, + idx, + row_idx: *row_idx, + node_idx: *node_idx, + eval_idx, + has_eval_factor: true, + is_wit: true, + constraint_idx: term_idx, + arg0: acc, + value: acc, + ..MainFrontloadTermRecord::default() + }); + *row_idx += 1; + *node_idx += 1; + + for (tail_offset, tail_value) in tail_point.iter().copied().enumerate() { + let next = acc * tail_value; + records.push(MainFrontloadTermRecord { + proof_idx: 0, + idx, + row_idx: *row_idx, + node_idx: *node_idx, + global_round_idx: tail_start + tail_offset, + has_global_factor: true, + is_tail: true, + constraint_idx: term_idx, + arg0: tail_value, + arg1: acc, + value: next, + chip_acc_out: next, + ..MainFrontloadTermRecord::default() + }); + *row_idx += 1; + *node_idx += 1; + acc = next; + } + acc +} + +fn eval_batched_main_frontload_terms_oracle( + layer_evals: &[RecursionField], + pi: &[RecursionField], + challenges: &[RecursionField], + global_in_point: &[RecursionField], + num_var_with_rotation: usize, + terms: &[multilinear_extensions::monomial::Term< + Expression, + Expression, + >], +) -> RecursionField { + let tail_point = &global_in_point[num_var_with_rotation..]; + let mut acc = RecursionField::ZERO; + for term in terms { + let mut value = either_to_ext(eval_by_expr_with_instance( + &[], + &[], + &[], + pi, + challenges, + &term.scalar, + )); + for expr in &term.product { + let Expression::WitIn(wit_id) = expr else { + panic!("main monomial product must be converted to WitIn") + }; + value *= layer_evals[*wit_id as usize]; + for tail in tail_point { + value *= *tail; + } + } + acc += value; + } + acc +} + +fn either_to_ext( + value: Either<::BaseField, RecursionField>, +) -> RecursionField { + match value { + Either::Left(base) => RecursionField::from(base), + Either::Right(ext) => ext, + } +} + +fn sample_challenge_pows(ts: &mut TS, size: usize, label: &[u8]) -> Vec where TS: FiatShamirTranscript, { - ts.observe_ext(claim); + transcript_observe_label(ts, label); + let alpha = FiatShamirTranscript::::sample_ext(ts); + iter::successors(Some(RecursionField::ONE), move |prev| Some(*prev * alpha)) + .take(size) + .collect() +} + +fn sample_challenge_pows_with_tidx( + ts: &mut TS, + size: usize, + label: &[u8], +) -> (Vec, usize) +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + transcript_observe_label(ts, label); + let tidx = ts.len(); + let alpha = FiatShamirTranscript::::sample_ext(ts); + let pows = iter::successors(Some(RecursionField::ONE), move |prev| Some(*prev * alpha)) + .take(size) + .collect(); + (pows, tidx) +} + +fn replay_main_sumcheck( + ts: &mut TS, + claimed_sum: RecursionField, + proof: &sumcheck::structs::IOPProof, + max_num_variables: usize, + max_degree: usize, +) -> Result<(Vec, RecursionField, Vec)> +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + if max_num_variables == 0 { + return Ok((vec![], claimed_sum, vec![])); + } + + transcript_observe_label(ts, &max_num_variables.to_le_bytes()); + transcript_observe_label(ts, &max_degree.to_le_bytes()); + let mut challenges = Vec::with_capacity(max_num_variables); + let mut challenge_tidxs = Vec::with_capacity(max_num_variables); + let mut expected = claimed_sum; + for round in 0..max_num_variables { + let prover_msg = proof + .proofs + .get(round) + .ok_or_else(|| eyre!("main sumcheck proof is incomplete"))?; + if prover_msg.evaluations.len() != max_degree { + bail!( + "main sumcheck round {round} eval count {} != {max_degree}", + prover_msg.evaluations.len() + ); + } + for eval in &prover_msg.evaluations { + ts.observe_ext(*eval); + } + transcript_observe_label(ts, b"Internal round"); + let challenge_tidx = ts.len(); + let challenge = FiatShamirTranscript::::sample_ext(ts); + let eval_0 = expected - prover_msg.evaluations[0]; + expected = extrapolate_uni_poly(eval_0, &prover_msg.evaluations, challenge); + challenges.push(challenge); + challenge_tidxs.push(challenge_tidx); + } + Ok((challenges, expected, challenge_tidxs)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ecc_x3y3_selector_eval_record(selector_idx: usize) -> MainSelectorEvalRecord { + MainSelectorEvalRecord { + proof_idx: 0, + idx: 0, + air_idx: 7, + selector_idx, + has_eval: true, + kind: MainSelectorKind::Whole, + ctx_num_vars: 3, + out_point: vec![ + RecursionField::from_usize(11), + RecursionField::from_usize(13), + RecursionField::ONE, + ], + point_source: MainSelectorPointSourceKind::EccX3Y3, + ecc_rt_tidxs: vec![100, 104], + ecc_x3y3_selector_idx: Some(10), + ..Default::default() + } + } + + #[test] + fn ecc_x3y3_rt_coordinates_are_bound_to_ecc_sumcheck_challenges_once() { + let records = build_main_selector_point_records( + &[ + ecc_x3y3_selector_eval_record(10), + ecc_x3y3_selector_eval_record(11), + ], + &std::collections::BTreeMap::new(), + ); + + let find = |selector_idx, round_idx| { + records + .iter() + .find(|record| record.selector_idx == selector_idx && record.round_idx == round_idx) + .expect("selector point record exists") + }; + + let canonical_r0 = find(10, 0); + assert!(!canonical_r0.has_transcript); + assert!(canonical_r0.has_ecc_rt); + assert_eq!(canonical_r0.transcript_tidx, 100); + assert!(!canonical_r0.has_source); + + let canonical_r1 = find(10, 1); + assert!(!canonical_r1.has_transcript); + assert!(canonical_r1.has_ecc_rt); + assert_eq!(canonical_r1.transcript_tidx, 104); + assert!(!canonical_r1.has_source); + + let sibling_r0 = find(11, 0); + assert!(!sibling_r0.has_transcript); + assert!(sibling_r0.has_source); + assert_eq!(sibling_r0.source_selector_idx, 10); + assert_eq!( + sibling_r0.source_source_kind, + MainSelectorPointSourceKind::EccX3Y3 + ); + assert_eq!(sibling_r0.source_round_idx, 0); + + let sibling_r1 = find(11, 1); + assert!(!sibling_r1.has_transcript); + assert!(sibling_r1.has_source); + assert_eq!(sibling_r1.source_selector_idx, 10); + assert_eq!( + sibling_r1.source_source_kind, + MainSelectorPointSourceKind::EccX3Y3 + ); + assert_eq!(sibling_r1.source_round_idx, 1); + + assert_eq!(find(10, 2).derive_kind, MainSelectorPointDeriveKind::One); + assert_eq!(find(11, 2).derive_kind, MainSelectorPointDeriveKind::One); + } } diff --git a/ceno_recursion_v2/src/main/selector.rs b/ceno_recursion_v2/src/main/selector.rs new file mode 100644 index 000000000..8d54673df --- /dev/null +++ b/ceno_recursion_v2/src/main/selector.rs @@ -0,0 +1,2136 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_circuit_primitives::utils::assert_array_eq; +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, extension::BinomiallyExtendable}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{ + AirPresenceBus, AirPresenceBusMessage, EccRtBus, EccRtMessage, ForkedTranscriptBus, + ForkedTranscriptBusMessage, MainEvalBus, MainEvalMessage, MainGlobalPointBus, + MainGlobalPointMessage, MainSelectorPointBus, MainSelectorPointMessage, + MainSelectorResultBus, MainSelectorResultMessage, MainSelectorShapeBus, + MainSelectorShapeMessage, MainSelectorSparseIndexShapeBus, + MainSelectorSparseIndexShapeMessage, TowerMainPointBus, TowerMainPointMessage, + TranscriptBus, TranscriptBusMessage, + }, + system::{ + MainSelectorEvalRecord, MainSelectorKind, MainSelectorPointDeriveKind, + MainSelectorPointRecord, MainSelectorPointSourceKind, RecursionField, + }, + tracegen::RowMajorChip, +}; + +pub const MAX_SELECTOR_POINT_VARS: usize = 32; +pub const MAX_SELECTOR_SPARSE_INDICES: usize = 256; + +const STEP_SHAPE: usize = 0; +const STEP_EQ_PRODUCT: usize = 1; +const STEP_SPARSE_INDEX: usize = 2; +const STEP_ACCUMULATE: usize = 3; +const STEP_FINAL: usize = 4; +const STEP_MULTIPLY: usize = 5; +const STEP_QUARK: usize = 6; +const STEP_EQ_LTE: usize = 7; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainSelectorFormulaCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub tower_idx: T, + pub air_idx: T, + pub selector_idx: T, + pub eval_idx: T, + pub kind: T, + pub source_kind: T, + pub is_whole: T, + pub is_prefix: T, + pub is_ordered_sparse: T, + pub is_quark_binary_tree_less_than: T, + pub ctx_offset: T, + pub ctx_num_instances: T, + pub ctx_num_vars: T, + pub ordered_sparse_num_vars: T, + pub num_sparse_indices: T, + pub step_kind: T, + pub step_idx: T, + pub is_shape_step: T, + pub is_eq_product_step: T, + pub is_sparse_index_step: T, + pub is_accumulate_step: T, + pub is_final_step: T, + pub is_multiply_step: T, + pub is_quark_step: T, + pub is_eq_lte_step: T, + pub is_first_eq_lte_step: T, + pub is_last_eq_lte_step: T, + pub eq_lte_output_to_value: T, + pub eq_lte_output_to_acc_in: T, + pub eq_lte_output_to_factor: T, + pub eq_lte_output_to_neg_factor: T, + pub eq_lte_value_is_zero: T, + pub is_first_quark_step: T, + pub is_last_quark_step: T, + pub carry_accumulator: T, + pub round_idx: T, + pub sparse_pos: T, + pub sparse_index: T, + pub sparse_index_bits_value: T, + pub point_active: T, + pub lhs_point: [T; D_EF], + pub rhs_point: [T; D_EF], + pub factor: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], + pub aux: [T; D_EF], + pub aux2: [T; D_EF], + pub aux3: [T; D_EF], + pub value: [T; D_EF], +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainSelectorEvalCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub eval_idx: T, + pub value: [T; D_EF], +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainSelectorPointCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub tower_idx: T, + pub air_idx: T, + pub selector_idx: T, + pub round_idx: T, + pub source_kind: T, + pub is_tower_main: T, + pub is_rotation_left: T, + pub is_rotation_right: T, + pub is_rotation_origin: T, + pub is_ecc_xy: T, + pub is_ecc_slope: T, + pub is_ecc_x3y3: T, + pub lookup_count: T, + pub fork_id: T, + pub has_transcript: T, + pub transcript_tidx: T, + pub has_ecc_rt: T, + pub has_source: T, + pub source_selector_idx: T, + pub source_source_kind: T, + pub source_round_idx: T, + pub source_value: [T; D_EF], + pub derive_identity: T, + pub derive_one_minus: T, + pub derive_zero: T, + pub derive_one: T, + pub value: [T; D_EF], +} + +pub struct MainSelectorFormulaAir { + pub global_point_bus: MainGlobalPointBus, + pub selector_point_bus: MainSelectorPointBus, + pub air_presence_bus: AirPresenceBus, + pub selector_result_bus: MainSelectorResultBus, + pub selector_shape_bus: MainSelectorShapeBus, + pub selector_sparse_index_shape_bus: MainSelectorSparseIndexShapeBus, +} + +impl BaseAir for MainSelectorFormulaAir { + fn width(&self) -> usize { + MainSelectorFormulaCols::::width() + } +} + +impl BaseAirWithPublicValues for MainSelectorFormulaAir {} +impl PartitionedBaseAir for MainSelectorFormulaAir {} + +impl Air for MainSelectorFormulaAir +where + AB: AirBuilder + InteractionBuilder, + AB::Expr: Clone, + AB::Var: Clone, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("main row exists"), + ); + let local: &MainSelectorFormulaCols = (*local_row).borrow(); + let next: &MainSelectorFormulaCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_whole); + builder.assert_bool(local.is_prefix); + builder.assert_bool(local.is_ordered_sparse); + builder.assert_bool(local.is_quark_binary_tree_less_than); + builder.assert_bool(local.is_shape_step); + builder.assert_bool(local.is_eq_product_step); + builder.assert_bool(local.is_sparse_index_step); + builder.assert_bool(local.is_accumulate_step); + builder.assert_bool(local.is_final_step); + builder.assert_bool(local.is_multiply_step); + builder.assert_bool(local.is_quark_step); + builder.assert_bool(local.is_eq_lte_step); + builder.assert_bool(local.is_first_eq_lte_step); + builder.assert_bool(local.is_last_eq_lte_step); + builder.assert_bool(local.eq_lte_output_to_value); + builder.assert_bool(local.eq_lte_output_to_acc_in); + builder.assert_bool(local.eq_lte_output_to_factor); + builder.assert_bool(local.eq_lte_output_to_neg_factor); + builder.assert_bool(local.eq_lte_value_is_zero); + builder.assert_bool(local.is_first_quark_step); + builder.assert_bool(local.is_last_quark_step); + builder.assert_bool(local.carry_accumulator); + builder.assert_bool(local.point_active); + builder + .when_transition() + .when_ne(local.is_enabled, AB::Expr::ONE) + .assert_zero(next.is_enabled); + builder.when(local.is_enabled).assert_one( + local.is_whole + + local.is_prefix + + local.is_ordered_sparse + + local.is_quark_binary_tree_less_than, + ); + builder.when(local.is_enabled).assert_one( + local.is_shape_step + + local.is_eq_product_step + + local.is_sparse_index_step + + local.is_accumulate_step + + local.is_final_step + + local.is_multiply_step + + local.is_quark_step + + local.is_eq_lte_step, + ); + builder.when(local.is_enabled).assert_eq( + local.kind, + local.is_prefix + + local.is_ordered_sparse * AB::Expr::from_usize(2) + + local.is_quark_binary_tree_less_than * AB::Expr::from_usize(3), + ); + builder.when(local.is_enabled).assert_eq( + local.step_kind, + local.is_eq_product_step * AB::Expr::from_usize(STEP_EQ_PRODUCT) + + local.is_sparse_index_step * AB::Expr::from_usize(STEP_SPARSE_INDEX) + + local.is_accumulate_step * AB::Expr::from_usize(STEP_ACCUMULATE) + + local.is_final_step * AB::Expr::from_usize(STEP_FINAL) + + local.is_multiply_step * AB::Expr::from_usize(STEP_MULTIPLY) + + local.is_quark_step * AB::Expr::from_usize(STEP_QUARK) + + local.is_eq_lte_step * AB::Expr::from_usize(STEP_EQ_LTE), + ); + builder + .when(local.carry_accumulator) + .assert_one(local.is_enabled); + builder + .when(local.is_eq_lte_step) + .assert_one(local.is_enabled); + builder.when(local.carry_accumulator).assert_one( + local.is_eq_product_step + + local.is_accumulate_step + + local.is_multiply_step + + local.is_quark_step, + ); + builder + .when(local.is_first_eq_lte_step) + .assert_one(local.is_eq_lte_step); + builder + .when(local.is_last_eq_lte_step) + .assert_one(local.is_eq_lte_step); + builder.when(local.is_eq_lte_step).assert_one( + local.eq_lte_output_to_value + + local.eq_lte_output_to_acc_in + + local.eq_lte_output_to_factor + + local.eq_lte_output_to_neg_factor + + (AB::Expr::ONE - local.is_last_eq_lte_step), + ); + builder + .when(local.is_first_quark_step) + .assert_one(local.is_quark_step); + builder + .when(local.is_last_quark_step) + .assert_one(local.is_quark_step); + + self.selector_shape_bus.receive( + builder, + local.proof_idx, + MainSelectorShapeMessage { + air_idx: local.air_idx.into(), + selector_idx: local.selector_idx.into(), + kind: local.kind.into(), + point_source: local.source_kind.into(), + eval_idx: local.eval_idx.into(), + ctx_offset: local.ctx_offset.into(), + ctx_num_instances: local.ctx_num_instances.into(), + ctx_num_vars: local.ctx_num_vars.into(), + ordered_sparse_num_vars: local.ordered_sparse_num_vars.into(), + num_sparse_indices: local.num_sparse_indices.into(), + }, + local.is_enabled * local.is_shape_step, + ); + self.air_presence_bus.lookup_key( + builder, + local.proof_idx, + AirPresenceBusMessage { + air_idx: local.air_idx.into(), + is_present: AB::Expr::ONE, + }, + local.is_enabled * local.is_shape_step, + ); + self.selector_sparse_index_shape_bus.receive( + builder, + local.proof_idx, + MainSelectorSparseIndexShapeMessage { + air_idx: local.air_idx.into(), + selector_idx: local.selector_idx.into(), + sparse_pos: local.sparse_pos.into(), + sparse_index: local.sparse_index.into(), + }, + local.is_enabled * local.is_sparse_index_step, + ); + builder + .when(local.is_enabled * local.is_sparse_index_step) + .assert_eq(local.sparse_index_bits_value, local.sparse_index); + + self.global_point_bus.lookup_key( + builder, + local.proof_idx, + MainGlobalPointMessage { + round_idx: local.round_idx.into(), + value: local.lhs_point.map(Into::into), + }, + local.is_enabled + * (local.is_eq_product_step * local.point_active + + local.is_quark_step + + local.is_eq_lte_step), + ); + self.selector_point_bus.lookup_key( + builder, + local.proof_idx, + MainSelectorPointMessage { + idx: local.idx.into(), + air_idx: local.air_idx.into(), + selector_idx: local.selector_idx.into(), + source_kind: local.source_kind.into(), + round_idx: local.round_idx.into(), + value: local.rhs_point.map(Into::into), + }, + local.is_enabled + * (local.is_eq_product_step * local.point_active + + local.is_quark_step + + local.is_eq_lte_step), + ); + let expected_factor = eq_factor::(local.rhs_point.clone(), local.lhs_point.clone()); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_eq_product_step * local.point_active), + local.factor, + expected_factor, + ); + assert_array_eq( + &mut builder.when( + local.is_enabled * local.is_eq_product_step * (AB::Expr::ONE - local.point_active), + ), + local.factor, + ext_one::(), + ); + let product = ext_mul::( + local.acc_in.clone().map(Into::into), + local.factor.clone().map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_eq_product_step), + local.acc_out, + product, + ); + let sum = ext_add::( + local.acc_in.clone().map(Into::into), + local.factor.clone().map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_accumulate_step), + local.acc_out, + sum, + ); + let multiply_product = ext_mul::( + local.acc_in.clone().map(Into::into), + local.factor.clone().map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_multiply_step), + local.acc_out, + multiply_product, + ); + let ordered_sparse_tail_len = local.ctx_num_vars - local.ordered_sparse_num_vars; + let ordered_sparse_tail_is_zero = local.aux[0]; + let ordered_sparse_tail_inverse = local.aux2[0]; + builder + .when(local.is_enabled * local.is_multiply_step * local.is_ordered_sparse) + .assert_zero( + ordered_sparse_tail_is_zero * (ordered_sparse_tail_is_zero - AB::Expr::ONE), + ); + builder + .when(local.is_enabled * local.is_multiply_step * local.is_ordered_sparse) + .assert_zero(ordered_sparse_tail_is_zero * ordered_sparse_tail_len.clone()); + builder + .when(local.is_enabled * local.is_multiply_step * local.is_ordered_sparse) + .assert_eq( + ordered_sparse_tail_len * ordered_sparse_tail_inverse, + AB::Expr::ONE - ordered_sparse_tail_is_zero, + ); + assert_array_eq( + &mut builder.when( + local.is_enabled + * local.is_multiply_step + * local.is_ordered_sparse + * ordered_sparse_tail_is_zero, + ), + local.factor, + ext_one::(), + ); + builder + .when(local.is_enabled * local.is_quark_step) + .assert_eq( + local.sparse_index, + local.sparse_index_bits_value * AB::Expr::from_usize(2) + local.point_active, + ); + builder + .when(local.is_enabled * local.is_first_quark_step) + .assert_zero(local.step_idx); + builder + .when(local.is_enabled * local.is_first_quark_step) + .assert_zero(local.round_idx); + builder + .when(local.is_enabled * local.is_first_quark_step) + .assert_zero( + local.sparse_index_bits_value * (local.sparse_index_bits_value - AB::Expr::ONE), + ); + let quark_prefix_is_zero = local.aux[0]; + let quark_prefix_inverse = local.aux2[0]; + builder + .when(local.is_enabled * local.is_quark_step) + .assert_zero(quark_prefix_is_zero * (quark_prefix_is_zero - AB::Expr::ONE)); + builder + .when(local.is_enabled * local.is_quark_step) + .assert_zero(quark_prefix_is_zero * local.sparse_index_bits_value); + builder + .when(local.is_enabled * local.is_quark_step) + .assert_eq( + local.sparse_index_bits_value * quark_prefix_inverse, + AB::Expr::ONE - quark_prefix_is_zero, + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_quark_step * quark_prefix_is_zero), + local.factor, + [AB::Expr::ZERO; D_EF], + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first_quark_step), + local.acc_in, + [AB::Expr::ZERO; D_EF], + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first_quark_step), + local.factor, + [ + local.sparse_index_bits_value.into(), + AB::Expr::ZERO, + AB::Expr::ZERO, + AB::Expr::ZERO, + ], + ); + builder + .when(local.is_enabled * local.is_last_quark_step) + .assert_eq(local.sparse_index, local.ctx_num_instances); + let quark_zero_factor = ext_mul::( + ext_one_minus(local.rhs_point.clone().map(Into::into)), + ext_one_minus(local.lhs_point.clone().map(Into::into)), + ); + let quark_one_factor = ext_mul::( + local.rhs_point.clone().map(Into::into), + local.lhs_point.clone().map(Into::into), + ); + let quark_lhs = + ext_mul::(quark_zero_factor, local.factor.clone().map(Into::into)); + let quark_rhs = ext_mul::(quark_one_factor, local.acc_in.clone().map(Into::into)); + let quark_out = ext_add::(quark_lhs, quark_rhs); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_quark_step), + local.acc_out, + quark_out, + ); + builder + .when(local.is_eq_lte_step) + .assert_zero(local.point_active * (local.point_active - AB::Expr::ONE)); + builder.when(local.is_eq_lte_step).assert_eq( + local.sparse_index, + local.sparse_index_bits_value * AB::Expr::from_usize(2) + local.point_active, + ); + builder + .when(local.is_first_eq_lte_step) + .assert_zero(local.step_idx); + builder + .when(local.is_first_eq_lte_step) + .assert_zero(local.sparse_index_bits_value); + assert_array_eq( + &mut builder.when(local.is_first_eq_lte_step), + local.acc_in, + ext_one::(), + ); + assert_array_eq( + &mut builder.when(local.is_first_eq_lte_step), + local.aux, + [AB::Expr::ZERO; D_EF], + ); + assert_array_eq( + &mut builder.when(local.eq_lte_value_is_zero), + local.value, + [AB::Expr::ZERO; D_EF], + ); + builder + .when(local.is_last_eq_lte_step) + .assert_eq(local.sparse_index, local.sparse_pos); + let eq_lte_same_one = ext_mul::( + local.rhs_point.clone().map(Into::into), + local.lhs_point.clone().map(Into::into), + ); + let eq_lte_same_zero = ext_mul::( + ext_one_minus(local.rhs_point.clone().map(Into::into)), + ext_one_minus(local.lhs_point.clone().map(Into::into)), + ); + let eq_lte_any = ext_add::(eq_lte_same_one.clone(), eq_lte_same_zero.clone()); + let eq_lte_equal_choice = ext_add::( + ext_mul_scalar::(eq_lte_same_one, local.point_active.into()), + ext_mul_scalar::( + eq_lte_same_zero.clone(), + AB::Expr::ONE - local.point_active, + ), + ); + assert_array_eq( + &mut builder.when(local.is_eq_lte_step), + local.aux3, + eq_lte_equal_choice, + ); + let eq_lte_prefix_out = ext_mul::( + local.acc_in.clone().map(Into::into), + local.aux3.clone().map(Into::into), + ); + let eq_lte_less_from_prior = + ext_mul::(local.aux.clone().map(Into::into), eq_lte_any); + let eq_lte_less_from_equal_prefix_unmasked = + ext_mul::(local.acc_in.clone().map(Into::into), eq_lte_same_zero); + assert_array_eq( + &mut builder.when(local.is_eq_lte_step), + local.aux2, + eq_lte_less_from_equal_prefix_unmasked, + ); + let eq_lte_less_from_equal_prefix = + ext_mul_scalar::(local.aux2.clone().map(Into::into), local.point_active); + let eq_lte_less_out = + ext_add::(eq_lte_less_from_prior, eq_lte_less_from_equal_prefix); + assert_array_eq( + &mut builder.when(local.is_eq_lte_step), + local.factor, + eq_lte_less_out.clone(), + ); + assert_array_eq( + &mut builder.when(local.is_eq_lte_step * (AB::Expr::ONE - local.is_last_eq_lte_step)), + local.acc_out, + eq_lte_prefix_out.clone(), + ); + assert_array_eq( + &mut builder.when(local.is_last_eq_lte_step), + local.acc_out, + ext_add::(eq_lte_prefix_out, eq_lte_less_out), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_final_step), + local.acc_out, + local.value.clone().map(Into::into), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_final_step), + local.acc_in, + local.value.clone().map(Into::into), + ); + let same_selector_continuation = + local.is_enabled * next.is_enabled * (AB::Expr::ONE - next.is_shape_step); + builder + .when_transition() + .when(same_selector_continuation.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(same_selector_continuation.clone()) + .assert_eq(next.idx, local.idx); + builder + .when_transition() + .when(same_selector_continuation.clone()) + .assert_eq(next.air_idx, local.air_idx); + builder + .when_transition() + .when(same_selector_continuation.clone()) + .assert_eq(next.selector_idx, local.selector_idx); + builder + .when_transition() + .when(same_selector_continuation.clone()) + .assert_eq(next.eval_idx, local.eval_idx); + builder + .when_transition() + .when(same_selector_continuation.clone()) + .assert_eq(next.kind, local.kind); + assert_array_eq( + &mut builder.when_transition().when(local.carry_accumulator), + local.acc_out, + next.acc_in, + ); + builder + .when_transition() + .when(local.is_eq_product_step * next.is_eq_product_step) + .assert_one(local.carry_accumulator); + builder + .when_transition() + .when(local.is_eq_product_step * next.is_final_step * local.is_whole) + .assert_one(local.carry_accumulator); + builder + .when_transition() + .when(local.is_accumulate_step * next.is_accumulate_step) + .assert_one(local.carry_accumulator); + builder + .when_transition() + .when(local.is_accumulate_step * next.is_multiply_step) + .assert_one(local.carry_accumulator); + builder + .when_transition() + .when(local.is_accumulate_step * next.is_final_step) + .assert_one(local.carry_accumulator); + builder + .when_transition() + .when(local.is_multiply_step * next.is_final_step) + .assert_one(local.carry_accumulator); + builder + .when_transition() + .when(local.is_quark_step * next.is_quark_step) + .assert_one(local.carry_accumulator); + builder + .when_transition() + .when(local.is_quark_step * next.is_quark_step) + .assert_eq(next.step_idx, local.step_idx + AB::Expr::ONE); + builder + .when_transition() + .when(local.is_quark_step * next.is_quark_step) + .assert_eq(next.round_idx, local.round_idx + AB::Expr::ONE); + builder + .when_transition() + .when(local.is_quark_step * next.is_quark_step) + .assert_eq( + next.sparse_index, + local.sparse_index * AB::Expr::from_usize(2) - next.point_active, + ); + builder + .when_transition() + .when(local.is_quark_step * next.is_final_step) + .assert_one(local.carry_accumulator); + let same_eq_lte_block = local.is_eq_lte_step + * next.is_eq_lte_step + * (AB::Expr::ONE - local.is_last_eq_lte_step); + builder + .when_transition() + .when(same_eq_lte_block.clone()) + .assert_eq(next.step_idx, local.step_idx + AB::Expr::ONE); + builder + .when_transition() + .when(same_eq_lte_block.clone()) + .assert_eq(next.sparse_pos, local.sparse_pos); + builder + .when_transition() + .when(same_eq_lte_block.clone()) + .assert_eq(next.sparse_index_bits_value, local.sparse_index); + builder + .when_transition() + .when(same_eq_lte_block.clone()) + .assert_eq(next.round_idx + AB::Expr::ONE, local.round_idx); + assert_array_eq( + &mut builder.when_transition().when(same_eq_lte_block.clone()), + local.acc_out, + next.acc_in, + ); + assert_array_eq( + &mut builder.when_transition().when(same_eq_lte_block.clone()), + local.factor, + next.aux, + ); + assert_array_eq( + &mut builder.when_transition().when(same_eq_lte_block), + local.value, + next.value, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_eq_lte_step * local.eq_lte_output_to_value), + local.acc_out, + next.value, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_eq_lte_step * local.eq_lte_output_to_acc_in), + local.acc_out, + next.acc_in, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_eq_lte_step * local.eq_lte_output_to_factor), + local.value, + next.acc_in, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_eq_lte_step * local.eq_lte_output_to_factor), + local.acc_out, + next.factor, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_eq_lte_step * local.eq_lte_output_to_neg_factor), + local.value, + next.acc_in, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_eq_lte_step * local.eq_lte_output_to_neg_factor), + local.acc_out.map(|value| -value.into()), + next.factor, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_accumulate_step * next.is_eq_lte_step * local.is_ordered_sparse), + local.acc_out, + next.value, + ); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_quark_step * next.is_eq_lte_step), + local.acc_out, + next.value, + ); + builder + .when_transition() + .when(local.is_enabled * local.is_final_step) + .assert_zero(next.is_enabled * (AB::Expr::ONE - next.is_shape_step)); + + self.selector_result_bus.send( + builder, + local.proof_idx, + MainSelectorResultMessage { + idx: local.idx.into(), + eval_idx: local.eval_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.is_final_step, + ); + } +} + +pub struct MainSelectorEvalAir { + pub eval_bus: MainEvalBus, + pub selector_result_bus: MainSelectorResultBus, +} + +pub struct MainSelectorPointAir { + pub selector_point_bus: MainSelectorPointBus, + pub tower_point_bus: TowerMainPointBus, + pub transcript_bus: TranscriptBus, + pub forked_transcript_bus: ForkedTranscriptBus, + pub ecc_rt_bus: EccRtBus, +} + +impl BaseAir for MainSelectorPointAir { + fn width(&self) -> usize { + MainSelectorPointCols::::width() + } +} + +impl BaseAirWithPublicValues for MainSelectorPointAir {} +impl PartitionedBaseAir for MainSelectorPointAir {} + +impl Air for MainSelectorPointAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &MainSelectorPointCols = (*local_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_tower_main); + builder.assert_bool(local.is_rotation_left); + builder.assert_bool(local.is_rotation_right); + builder.assert_bool(local.is_rotation_origin); + builder.assert_bool(local.is_ecc_xy); + builder.assert_bool(local.is_ecc_slope); + builder.assert_bool(local.is_ecc_x3y3); + builder.assert_bool(local.has_transcript); + builder.assert_bool(local.has_ecc_rt); + builder.assert_bool(local.has_source); + builder.assert_bool(local.derive_identity); + builder.assert_bool(local.derive_one_minus); + builder.assert_bool(local.derive_zero); + builder.assert_bool(local.derive_one); + builder.when(local.is_enabled).assert_one( + local.is_tower_main + + local.is_rotation_left + + local.is_rotation_right + + local.is_rotation_origin + + local.is_ecc_xy + + local.is_ecc_slope + + local.is_ecc_x3y3, + ); + builder.when(local.is_enabled).assert_eq( + local.source_kind, + local.is_rotation_left * AB::Expr::from_usize(1) + + local.is_rotation_right * AB::Expr::from_usize(2) + + local.is_rotation_origin * AB::Expr::from_usize(3) + + local.is_ecc_xy * AB::Expr::from_usize(4) + + local.is_ecc_slope * AB::Expr::from_usize(5) + + local.is_ecc_x3y3 * AB::Expr::from_usize(6), + ); + builder + .when(local.has_transcript) + .assert_one(local.is_rotation_origin + local.is_ecc_xy); + builder.when(local.has_ecc_rt).assert_one(local.is_ecc_x3y3); + builder.when(local.has_source).assert_one( + local.is_rotation_left + + local.is_rotation_right + + local.is_ecc_xy + + local.is_ecc_slope + + local.is_ecc_x3y3, + ); + builder + .when(local.is_tower_main + local.is_rotation_origin) + .assert_zero(local.has_source); + builder.when(local.has_ecc_rt).assert_zero(local.has_source); + builder + .when( + local.is_tower_main + + local.is_rotation_left + + local.is_rotation_right + + local.is_ecc_slope + + local.is_ecc_x3y3, + ) + .assert_zero(local.has_transcript); + builder + .when( + local.is_tower_main + + local.is_rotation_left + + local.is_rotation_right + + local.is_rotation_origin + + local.is_ecc_xy + + local.is_ecc_slope, + ) + .assert_zero(local.has_ecc_rt); + builder.when(local.is_enabled).assert_one( + local.derive_identity + local.derive_one_minus + local.derive_zero + local.derive_one, + ); + self.tower_point_bus.lookup_key( + builder, + local.proof_idx, + TowerMainPointMessage { + idx: local.tower_idx.into(), + round_idx: local.round_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.is_tower_main, + ); + self.selector_point_bus.lookup_key( + builder, + local.proof_idx, + MainSelectorPointMessage { + idx: local.idx.into(), + air_idx: local.air_idx.into(), + selector_idx: local.source_selector_idx.into(), + source_kind: local.source_source_kind.into(), + round_idx: local.source_round_idx.into(), + value: local.source_value.map(Into::into), + }, + local.is_enabled * local.has_source, + ); + for i in 0..D_EF { + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.transcript_tidx + AB::Expr::from_usize(i), + value: local.value[i].into(), + is_sample: AB::Expr::ONE, + }, + local.is_enabled * local.has_transcript * local.is_ecc_xy, + ); + self.forked_transcript_bus.receive( + builder, + local.proof_idx, + ForkedTranscriptBusMessage { + fork_id: local.fork_id.into(), + tidx: local.transcript_tidx + AB::Expr::from_usize(i), + value: local.value[i].into(), + is_sample: AB::Expr::ONE, + }, + local.is_enabled * local.has_transcript * local.is_rotation_origin, + ); + } + self.ecc_rt_bus.lookup_key( + builder, + local.proof_idx, + EccRtMessage { + idx: local.tower_idx.into(), + round_idx: local.round_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.has_ecc_rt, + ); + for i in 0..D_EF { + let one_coeff = if i == 0 { + AB::Expr::ONE + } else { + AB::Expr::ZERO + }; + let expected = local.derive_identity * local.source_value[i] + + local.derive_one_minus * (one_coeff.clone() - local.source_value[i]) + + local.derive_one * one_coeff; + builder + .when(local.is_enabled * (local.has_source + local.derive_zero + local.derive_one)) + .assert_eq(local.value[i], expected); + } + self.selector_point_bus.add_key_with_lookups( + builder, + local.proof_idx, + MainSelectorPointMessage { + idx: local.idx.into(), + air_idx: local.air_idx.into(), + selector_idx: local.selector_idx.into(), + source_kind: local.source_kind.into(), + round_idx: local.round_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.lookup_count, + ); + } +} + +impl BaseAir for MainSelectorEvalAir { + fn width(&self) -> usize { + MainSelectorEvalCols::::width() + } +} + +impl BaseAirWithPublicValues for MainSelectorEvalAir {} +impl PartitionedBaseAir for MainSelectorEvalAir {} + +impl Air for MainSelectorEvalAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &MainSelectorEvalCols = (*local_row).borrow(); + + builder.assert_bool(local.is_enabled); + self.selector_result_bus.receive( + builder, + local.proof_idx, + MainSelectorResultMessage { + idx: local.idx.into(), + eval_idx: local.eval_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled, + ); + self.eval_bus.lookup_key( + builder, + local.proof_idx, + MainEvalMessage { + idx: local.idx.into(), + eval_idx: local.eval_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled, + ); + } +} + +fn ext_one() -> [FA; D_EF] { + [FA::ONE, FA::ZERO, FA::ZERO, FA::ZERO] +} + +fn ext_add(x: [FA; D_EF], y: [FA; D_EF]) -> [FA; D_EF] { + let [x0, x1, x2, x3] = x; + let [y0, y1, y2, y3] = y; + [x0 + y0, x1 + y1, x2 + y2, x3 + y3] +} + +fn ext_one_minus(x: [FA; D_EF]) -> [FA; D_EF] { + let [x0, x1, x2, x3] = x; + [FA::ONE - x0, -x1, -x2, -x3] +} + +fn ext_mul(x: [FA; D_EF], y: [FA; D_EF]) -> [FA; D_EF] +where + FA: PrimeCharacteristicRing, + FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + let [x0, x1, x2, x3] = x; + let [y0, y1, y2, y3] = y; + let w = FA::from_prime_subfield(FA::PrimeSubfield::W); + let z0_beta_terms = x1.clone() * y3.clone() + x2.clone() * y2.clone() + x3.clone() * y1.clone(); + let z1_beta_terms = x2.clone() * y3.clone() + x3.clone() * y2.clone(); + let z2_beta_terms = x3.clone() * y3.clone(); + [ + x0.clone() * y0.clone() + z0_beta_terms * w.clone(), + x0.clone() * y1.clone() + x1.clone() * y0.clone() + z1_beta_terms * w.clone(), + x0.clone() * y2.clone() + + x1.clone() * y1.clone() + + x2.clone() * y0.clone() + + z2_beta_terms * w, + x0 * y3 + x1 * y2 + x2 * y1 + x3 * y0, + ] +} + +fn ext_mul_scalar(x: [FA; D_EF], y: impl Into) -> [FA; D_EF] +where + FA: PrimeCharacteristicRing, +{ + let [x0, x1, x2, x3] = x; + let y = y.into(); + [x0 * y.clone(), x1 * y.clone(), x2 * y.clone(), x3 * y] +} + +fn eq_factor(out: [AB::Var; D_EF], input: [AB::Var; D_EF]) -> [AB::Expr; D_EF] +where + AB: AirBuilder, + AB::Expr: Clone, + AB::Var: Clone, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + let out_expr = out.map(Into::into); + let input_expr = input.map(Into::into); + let same_one: [AB::Expr; D_EF] = ext_mul(out_expr.clone(), input_expr.clone()); + let same_zero: [AB::Expr; D_EF] = ext_mul(ext_one_minus(out_expr), ext_one_minus(input_expr)); + ext_add::(same_one, same_zero) +} + +pub struct MainSelectorFormulaTraceGenerator; +pub struct MainSelectorEvalTraceGenerator; +pub struct MainSelectorPointTraceGenerator; + +impl RowMajorChip for MainSelectorFormulaTraceGenerator { + type Ctx<'a> = &'a [MainSelectorEvalRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let rows = build_formula_rows(records); + generate_selector_trace(&rows, required_height, fill_formula_cols) + } +} + +impl RowMajorChip for MainSelectorEvalTraceGenerator { + type Ctx<'a> = &'a [MainSelectorEvalRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + generate_selector_trace(records, required_height, fill_eval_cols) + } +} + +impl RowMajorChip for MainSelectorPointTraceGenerator { + type Ctx<'a> = &'a [MainSelectorPointRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + generate_selector_trace(records, required_height, fill_point_cols) + } +} + +fn generate_selector_trace( + records: &[R], + required_height: Option, + mut fill: Fill, +) -> Option> +where + C: SelectorColumnAccess, + Fill: FnMut(&R, &mut C), +{ + let width = C::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + fill(record, C::from_bytes(row)); + } + Some(RowMajorMatrix::new(trace, width)) +} + +trait SelectorColumnAccess: Sized { + fn width() -> usize; + fn from_bytes(slice: &mut [F]) -> &mut Self; +} + +impl SelectorColumnAccess for MainSelectorFormulaCols { + fn width() -> usize { + MainSelectorFormulaCols::::width() + } + + fn from_bytes(slice: &mut [F]) -> &mut Self { + slice.borrow_mut() + } +} + +impl SelectorColumnAccess for MainSelectorEvalCols { + fn width() -> usize { + MainSelectorEvalCols::::width() + } + + fn from_bytes(slice: &mut [F]) -> &mut Self { + slice.borrow_mut() + } +} + +impl SelectorColumnAccess for MainSelectorPointCols { + fn width() -> usize { + MainSelectorPointCols::::width() + } + + fn from_bytes(slice: &mut [F]) -> &mut Self { + slice.borrow_mut() + } +} + +fn fill_point_cols(record: &MainSelectorPointRecord, cols: &mut MainSelectorPointCols) { + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.tower_idx = F::from_usize(record.tower_idx); + cols.air_idx = F::from_usize(record.air_idx); + cols.selector_idx = F::from_usize(record.selector_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.source_kind = F::from_usize(selector_point_source_code(record.source_kind)); + cols.is_tower_main = F::from_bool(record.source_kind == MainSelectorPointSourceKind::TowerMain); + cols.is_rotation_left = + F::from_bool(record.source_kind == MainSelectorPointSourceKind::RotationLeft); + cols.is_rotation_right = + F::from_bool(record.source_kind == MainSelectorPointSourceKind::RotationRight); + cols.is_rotation_origin = + F::from_bool(record.source_kind == MainSelectorPointSourceKind::RotationOrigin); + cols.is_ecc_xy = F::from_bool(record.source_kind == MainSelectorPointSourceKind::EccXY); + cols.is_ecc_slope = F::from_bool(record.source_kind == MainSelectorPointSourceKind::EccSlope); + cols.is_ecc_x3y3 = F::from_bool(record.source_kind == MainSelectorPointSourceKind::EccX3Y3); + cols.lookup_count = F::from_usize(record.lookup_count); + cols.fork_id = F::from_usize(record.fork_id); + cols.has_transcript = F::from_bool(record.has_transcript); + cols.transcript_tidx = F::from_usize(record.transcript_tidx); + cols.has_ecc_rt = F::from_bool(record.has_ecc_rt); + cols.has_source = F::from_bool(record.has_source); + cols.source_selector_idx = F::from_usize(record.source_selector_idx); + cols.source_source_kind = F::from_usize(selector_point_source_code(record.source_source_kind)); + cols.source_round_idx = F::from_usize(record.source_round_idx); + cols.source_value = record + .source_value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.derive_identity = + F::from_bool(record.derive_kind == MainSelectorPointDeriveKind::Identity); + cols.derive_one_minus = + F::from_bool(record.derive_kind == MainSelectorPointDeriveKind::OneMinus); + cols.derive_zero = F::from_bool(record.derive_kind == MainSelectorPointDeriveKind::Zero); + cols.derive_one = F::from_bool(record.derive_kind == MainSelectorPointDeriveKind::One); + cols.value = record + .value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); +} + +fn fill_eval_cols(record: &MainSelectorEvalRecord, cols: &mut MainSelectorEvalCols) { + cols.is_enabled = F::from_bool(record.has_eval); + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.eval_idx = F::from_usize(record.eval_idx); + cols.value = record + .value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); +} + +#[derive(Clone)] +struct FormulaRow<'a> { + record: &'a MainSelectorEvalRecord, + step_kind: usize, + step_idx: usize, + round_idx: usize, + sparse_pos: usize, + sparse_index: usize, + sparse_index_bits_value: usize, + eq_lte_is_first: bool, + eq_lte_is_last: bool, + eq_lte_output_to_value: bool, + eq_lte_output_to_acc_in: bool, + eq_lte_output_to_factor: bool, + eq_lte_output_to_neg_factor: bool, + eq_lte_value_is_zero: bool, + quark_is_first: bool, + quark_is_last: bool, + point_active: bool, + lhs_point: RecursionField, + rhs_point: RecursionField, + factor: RecursionField, + acc_in: RecursionField, + acc_out: RecursionField, + aux: RecursionField, + aux2: RecursionField, + aux3: RecursionField, + value: RecursionField, +} + +#[derive(Clone, Copy)] +enum EqLteOutput { + ToValue, + ToAccIn, + ToFactor, + ToNegFactor, +} + +fn build_formula_rows(records: &[MainSelectorEvalRecord]) -> Vec> { + if records.is_empty() { + return Vec::new(); + } + let mut rows = Vec::new(); + for record in records { + let one = RecursionField::ONE; + let zero = RecursionField::ZERO; + rows.push(FormulaRow { + record, + step_kind: STEP_SHAPE, + step_idx: 0, + round_idx: 0, + sparse_pos: 0, + sparse_index: 0, + sparse_index_bits_value: 0, + eq_lte_is_first: false, + eq_lte_is_last: false, + eq_lte_output_to_value: false, + eq_lte_output_to_acc_in: false, + eq_lte_output_to_factor: false, + eq_lte_output_to_neg_factor: false, + eq_lte_value_is_zero: false, + quark_is_first: false, + quark_is_last: false, + point_active: false, + lhs_point: zero, + rhs_point: zero, + factor: one, + acc_in: zero, + acc_out: zero, + aux: zero, + aux2: zero, + aux3: zero, + value: record.value, + }); + + for (sparse_pos, sparse_index) in record.sparse_indices.iter().copied().enumerate() { + rows.push(FormulaRow { + record, + step_kind: STEP_SPARSE_INDEX, + step_idx: sparse_pos, + round_idx: 0, + sparse_pos, + sparse_index, + sparse_index_bits_value: sparse_index, + eq_lte_is_first: false, + eq_lte_is_last: false, + eq_lte_output_to_value: false, + eq_lte_output_to_acc_in: false, + eq_lte_output_to_factor: false, + eq_lte_output_to_neg_factor: false, + eq_lte_value_is_zero: false, + quark_is_first: false, + quark_is_last: false, + point_active: false, + lhs_point: zero, + rhs_point: zero, + factor: one, + acc_in: zero, + acc_out: zero, + aux: zero, + aux2: zero, + aux3: zero, + value: record.value, + }); + } + + if !record.has_eval { + continue; + } + + let computed_value = match record.kind { + MainSelectorKind::Whole => { + let acc = + push_eq_product_rows(&mut rows, record, 0, record.ctx_num_vars, one, true); + acc + } + MainSelectorKind::Prefix => { + let _ = push_eq_product_rows(&mut rows, record, 0, record.ctx_num_vars, one, true); + let end = record.ctx_offset + record.ctx_num_instances; + let end_eval = if end == 0 { + zero + } else { + native_eq_lte(end - 1, &record.out_point, &record.in_point) + }; + let start_eval = if record.ctx_offset == 0 { + zero + } else { + native_eq_lte(record.ctx_offset - 1, &record.out_point, &record.in_point) + }; + let value = end_eval - start_eval; + if record.ctx_offset == 0 { + if end != 0 { + let _ = push_eq_lte_rows( + &mut rows, + record, + 0, + record.ctx_num_vars, + end - 1, + zero, + true, + EqLteOutput::ToAccIn, + ); + } + } else { + let _ = push_eq_lte_rows( + &mut rows, + record, + 0, + record.ctx_num_vars, + end - 1, + zero, + true, + EqLteOutput::ToValue, + ); + let _ = push_eq_lte_rows( + &mut rows, + record, + 0, + record.ctx_num_vars, + record.ctx_offset - 1, + end_eval, + false, + EqLteOutput::ToNegFactor, + ); + } + rows.push(accumulate_row(record, 0, end_eval, -start_eval, value)); + value + } + MainSelectorKind::OrderedSparse => { + let _ = push_eq_product_rows(&mut rows, record, 0, record.ctx_num_vars, one, true); + let subgroup_vars = record.ordered_sparse_num_vars; + let out_subgroup_eq = native_eq_vec(&record.out_point[..subgroup_vars]); + let in_subgroup_eq = native_eq_vec(&record.in_point[..subgroup_vars]); + let mut subgroup_acc = zero; + for (sparse_pos, sparse_index) in record.sparse_indices.iter().copied().enumerate() + { + let term = out_subgroup_eq[sparse_index] * in_subgroup_eq[sparse_index]; + let acc_in = subgroup_acc; + subgroup_acc += term; + rows.push(accumulate_row( + record, + sparse_pos, + acc_in, + term, + subgroup_acc, + )); + } + let tail_eval = native_eq_lte( + record.ctx_num_instances - 1, + &record.out_point[subgroup_vars..], + &record.in_point[subgroup_vars..], + ); + let value = subgroup_acc * tail_eval; + let _ = push_eq_lte_rows( + &mut rows, + record, + subgroup_vars, + record.ctx_num_vars - subgroup_vars, + record.ctx_num_instances - 1, + subgroup_acc, + record.sparse_indices.is_empty(), + EqLteOutput::ToFactor, + ); + rows.push(multiply_row( + record, + record.sparse_indices.len(), + subgroup_acc, + tail_eval, + value, + )); + value + } + MainSelectorKind::QuarkBinaryTreeLessThan => push_quark_rows(&mut rows, record), + }; + rows.push(FormulaRow { + record, + step_kind: STEP_FINAL, + step_idx: 0, + round_idx: 0, + sparse_pos: 0, + sparse_index: 0, + sparse_index_bits_value: 0, + eq_lte_is_first: false, + eq_lte_is_last: false, + eq_lte_output_to_value: false, + eq_lte_output_to_acc_in: false, + eq_lte_output_to_factor: false, + eq_lte_output_to_neg_factor: false, + eq_lte_value_is_zero: false, + quark_is_first: false, + quark_is_last: false, + point_active: false, + lhs_point: zero, + rhs_point: zero, + factor: one, + acc_in: computed_value, + acc_out: computed_value, + aux: zero, + aux2: zero, + aux3: zero, + value: record.value, + }); + } + rows +} + +pub(crate) fn selector_formula_global_point_lookups( + records: &[MainSelectorEvalRecord], +) -> Vec<(usize, usize)> { + build_formula_rows(records) + .into_iter() + .filter(|row| { + row.record.has_eval + && ((row.step_kind == STEP_EQ_PRODUCT && row.point_active) + || row.step_kind == STEP_QUARK + || row.step_kind == STEP_EQ_LTE) + }) + .map(|row| (row.record.proof_idx, row.round_idx)) + .collect() +} + +pub(crate) fn selector_formula_point_lookup_counts( + records: &[MainSelectorEvalRecord], +) -> std::collections::BTreeMap<(usize, usize, usize, usize, usize), usize> { + let mut counts = std::collections::BTreeMap::new(); + for row in build_formula_rows(records).into_iter().filter(|row| { + row.record.has_eval + && ((row.step_kind == STEP_EQ_PRODUCT && row.point_active) + || row.step_kind == STEP_QUARK + || row.step_kind == STEP_EQ_LTE) + }) { + *counts + .entry(( + row.record.proof_idx, + row.record.idx, + row.record.air_idx, + row.record.selector_idx, + row.round_idx, + )) + .or_default() += 1; + } + counts +} + +fn accumulate_row( + record: &MainSelectorEvalRecord, + step_idx: usize, + acc_in: RecursionField, + factor: RecursionField, + acc_out: RecursionField, +) -> FormulaRow<'_> { + FormulaRow { + record, + step_kind: STEP_ACCUMULATE, + step_idx, + round_idx: 0, + sparse_pos: step_idx, + sparse_index: 0, + sparse_index_bits_value: 0, + eq_lte_is_first: false, + eq_lte_is_last: false, + eq_lte_output_to_value: false, + eq_lte_output_to_acc_in: false, + eq_lte_output_to_factor: false, + eq_lte_output_to_neg_factor: false, + eq_lte_value_is_zero: false, + quark_is_first: false, + quark_is_last: false, + point_active: false, + lhs_point: RecursionField::ZERO, + rhs_point: RecursionField::ZERO, + factor, + acc_in, + acc_out, + aux: RecursionField::ZERO, + aux2: RecursionField::ZERO, + aux3: RecursionField::ZERO, + value: record.value, + } +} + +fn multiply_row( + record: &MainSelectorEvalRecord, + step_idx: usize, + acc_in: RecursionField, + factor: RecursionField, + acc_out: RecursionField, +) -> FormulaRow<'_> { + let ordered_sparse_tail_len = record + .ctx_num_vars + .saturating_sub(record.ordered_sparse_num_vars); + FormulaRow { + record, + step_kind: STEP_MULTIPLY, + step_idx, + round_idx: 0, + sparse_pos: step_idx, + sparse_index: 0, + sparse_index_bits_value: 0, + eq_lte_is_first: false, + eq_lte_is_last: false, + eq_lte_output_to_value: false, + eq_lte_output_to_acc_in: false, + eq_lte_output_to_factor: false, + eq_lte_output_to_neg_factor: false, + eq_lte_value_is_zero: false, + quark_is_first: false, + quark_is_last: false, + point_active: false, + lhs_point: RecursionField::ZERO, + rhs_point: RecursionField::ZERO, + factor, + acc_in, + acc_out, + aux: RecursionField::from_usize(usize::from( + record.kind == MainSelectorKind::OrderedSparse && ordered_sparse_tail_len == 0, + )), + aux2: if record.kind == MainSelectorKind::OrderedSparse && ordered_sparse_tail_len != 0 { + RecursionField::from_usize(ordered_sparse_tail_len).inverse() + } else { + RecursionField::ZERO + }, + aux3: RecursionField::ZERO, + value: record.value, + } +} + +fn push_eq_lte_rows<'a>( + rows: &mut Vec>, + record: &'a MainSelectorEvalRecord, + point_offset: usize, + len: usize, + max_idx: usize, + carried_value: RecursionField, + carried_value_is_zero: bool, + output: EqLteOutput, +) -> RecursionField { + if len == 0 { + return RecursionField::ONE; + } + + let zero = RecursionField::ZERO; + let one = RecursionField::ONE; + let mut eq_prefix = one; + let mut less_sum = zero; + let mut bound_prefix = 0usize; + let mut output_value = one; + for (step_idx, bit_idx) in (0..len).rev().enumerate() { + let bit = (max_idx >> bit_idx) & 1; + let lhs = record.in_point[point_offset + bit_idx]; + let rhs = record.out_point[point_offset + bit_idx]; + let same_one = rhs * lhs; + let same_zero = (one - rhs) * (one - lhs); + let any = same_one + same_zero; + let equal_choice = if bit == 1 { same_one } else { same_zero }; + let less_from_equal_prefix = if bit == 1 { + eq_prefix * same_zero + } else { + zero + }; + let less_from_equal_prefix_unmasked = eq_prefix * same_zero; + let next_less_sum = less_sum * any + less_from_equal_prefix; + let next_eq_prefix = eq_prefix * equal_choice; + let is_last = step_idx + 1 == len; + output_value = next_eq_prefix + next_less_sum; + bound_prefix = bound_prefix * 2 + bit; + rows.push(FormulaRow { + record, + step_kind: STEP_EQ_LTE, + step_idx, + round_idx: point_offset + bit_idx, + sparse_pos: max_idx, + sparse_index: bound_prefix, + sparse_index_bits_value: (bound_prefix - bit) / 2, + eq_lte_is_first: step_idx == 0, + eq_lte_is_last: is_last, + eq_lte_output_to_value: is_last && matches!(output, EqLteOutput::ToValue), + eq_lte_output_to_acc_in: is_last && matches!(output, EqLteOutput::ToAccIn), + eq_lte_output_to_factor: is_last && matches!(output, EqLteOutput::ToFactor), + eq_lte_output_to_neg_factor: is_last && matches!(output, EqLteOutput::ToNegFactor), + eq_lte_value_is_zero: step_idx == 0 && carried_value_is_zero, + quark_is_first: false, + quark_is_last: false, + point_active: bit == 1, + lhs_point: lhs, + rhs_point: rhs, + factor: next_less_sum, + acc_in: eq_prefix, + acc_out: if is_last { + output_value + } else { + next_eq_prefix + }, + aux: less_sum, + aux2: less_from_equal_prefix_unmasked, + aux3: equal_choice, + value: carried_value, + }); + eq_prefix = next_eq_prefix; + less_sum = next_less_sum; + } + output_value +} + +fn push_quark_rows<'a>( + rows: &mut Vec>, + record: &'a MainSelectorEvalRecord, +) -> RecursionField { + assert!(record.ctx_num_instances > 0); + assert!(!record.out_point.is_empty()); + assert_eq!(record.out_point.len(), record.in_point.len()); + + let mut layer_ns = (0..record.out_point.len()) + .scan(record.ctx_num_instances, |n_instance, _| { + let current = *n_instance; + *n_instance = (*n_instance).div_ceil(2); + Some(current) + }) + .collect::>(); + layer_ns.reverse(); + + let zero = RecursionField::ZERO; + let one = RecursionField::ONE; + let mut acc = zero; + for (i, layer_n) in layer_ns.iter().copied().enumerate() { + let prefix_count = layer_n / 2; + let parity = layer_n % 2; + let lhs = record.in_point[i]; + let rhs = record.out_point[i]; + let mut factor = if prefix_count == 0 { + zero + } else if i == 0 { + one + } else { + native_eq_lte( + prefix_count - 1, + &record.out_point[..i], + &record.in_point[..i], + ) + }; + if prefix_count > 0 && i > 0 { + factor = push_eq_lte_rows( + rows, + record, + 0, + i, + prefix_count - 1, + acc, + false, + EqLteOutput::ToFactor, + ); + } + let zero_factor = (one - rhs) * (one - lhs); + let one_factor = rhs * lhs; + let acc_in = acc; + acc = zero_factor * factor + one_factor * acc; + rows.push(FormulaRow { + record, + step_kind: STEP_QUARK, + step_idx: i, + round_idx: i, + sparse_pos: 0, + sparse_index: layer_n, + sparse_index_bits_value: prefix_count, + eq_lte_is_first: false, + eq_lte_is_last: false, + eq_lte_output_to_value: false, + eq_lte_output_to_acc_in: false, + eq_lte_output_to_factor: false, + eq_lte_output_to_neg_factor: false, + eq_lte_value_is_zero: false, + quark_is_first: i == 0, + quark_is_last: i + 1 == layer_ns.len(), + point_active: parity == 1, + lhs_point: lhs, + rhs_point: rhs, + factor, + acc_in, + acc_out: acc, + aux: RecursionField::from_usize(usize::from(prefix_count == 0)), + aux2: if prefix_count == 0 { + zero + } else { + RecursionField::from_usize(prefix_count).inverse() + }, + aux3: zero, + value: record.value, + }); + } + acc +} + +fn push_eq_product_rows<'a>( + rows: &mut Vec>, + record: &'a MainSelectorEvalRecord, + start_round_idx: usize, + len: usize, + mut acc: RecursionField, + point_active: bool, +) -> RecursionField { + let zero = RecursionField::ZERO; + for i in 0..len.min(MAX_SELECTOR_POINT_VARS) { + let point_idx = start_round_idx + i; + let lhs = record.in_point.get(point_idx).copied().unwrap_or(zero); + let rhs = record.out_point.get(point_idx).copied().unwrap_or(zero); + let factor = native_eq_factor(rhs, lhs); + let acc_in = acc; + acc *= factor; + rows.push(FormulaRow { + record, + step_kind: STEP_EQ_PRODUCT, + step_idx: point_idx, + round_idx: point_idx, + sparse_pos: 0, + sparse_index: 0, + sparse_index_bits_value: 0, + eq_lte_is_first: false, + eq_lte_is_last: false, + eq_lte_output_to_value: false, + eq_lte_output_to_acc_in: false, + eq_lte_output_to_factor: false, + eq_lte_output_to_neg_factor: false, + eq_lte_value_is_zero: false, + quark_is_first: false, + quark_is_last: false, + point_active, + lhs_point: lhs, + rhs_point: rhs, + factor, + acc_in, + acc_out: acc, + aux: zero, + aux2: zero, + aux3: zero, + value: record.value, + }); + } + acc +} + +fn fill_formula_cols(row: &FormulaRow<'_>, cols: &mut MainSelectorFormulaCols) { + let record = row.record; + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.air_idx = F::from_usize(record.air_idx); + cols.selector_idx = F::from_usize(record.selector_idx); + cols.eval_idx = F::from_usize(record.eval_idx); + cols.kind = F::from_usize(selector_kind_code(record.kind)); + cols.source_kind = F::from_usize(selector_point_source_code(record.point_source)); + cols.is_whole = F::from_bool(record.kind == MainSelectorKind::Whole); + cols.is_prefix = F::from_bool(record.kind == MainSelectorKind::Prefix); + cols.is_ordered_sparse = F::from_bool(record.kind == MainSelectorKind::OrderedSparse); + cols.is_quark_binary_tree_less_than = + F::from_bool(record.kind == MainSelectorKind::QuarkBinaryTreeLessThan); + cols.ctx_offset = F::from_usize(record.ctx_offset); + cols.ctx_num_instances = F::from_usize(record.ctx_num_instances); + cols.ctx_num_vars = F::from_usize(record.ctx_num_vars); + cols.ordered_sparse_num_vars = F::from_usize(record.ordered_sparse_num_vars); + cols.num_sparse_indices = F::from_usize(record.sparse_indices.len()); + cols.step_kind = F::from_usize(row.step_kind); + cols.step_idx = F::from_usize(row.step_idx); + cols.is_shape_step = F::from_bool(row.step_kind == STEP_SHAPE); + cols.is_eq_product_step = F::from_bool(row.step_kind == STEP_EQ_PRODUCT); + cols.is_sparse_index_step = F::from_bool(row.step_kind == STEP_SPARSE_INDEX); + cols.is_accumulate_step = F::from_bool(row.step_kind == STEP_ACCUMULATE); + cols.is_final_step = F::from_bool(row.step_kind == STEP_FINAL); + cols.is_multiply_step = F::from_bool(row.step_kind == STEP_MULTIPLY); + cols.is_quark_step = F::from_bool(row.step_kind == STEP_QUARK); + cols.is_eq_lte_step = F::from_bool(row.step_kind == STEP_EQ_LTE); + cols.is_first_eq_lte_step = F::from_bool(row.eq_lte_is_first); + cols.is_last_eq_lte_step = F::from_bool(row.eq_lte_is_last); + cols.eq_lte_output_to_value = F::from_bool(row.eq_lte_output_to_value); + cols.eq_lte_output_to_acc_in = F::from_bool(row.eq_lte_output_to_acc_in); + cols.eq_lte_output_to_factor = F::from_bool(row.eq_lte_output_to_factor); + cols.eq_lte_output_to_neg_factor = F::from_bool(row.eq_lte_output_to_neg_factor); + cols.eq_lte_value_is_zero = F::from_bool(row.eq_lte_value_is_zero); + cols.is_first_quark_step = F::from_bool(row.quark_is_first); + cols.is_last_quark_step = F::from_bool(row.quark_is_last); + let ordered_sparse_tail_len = row + .record + .ctx_num_vars + .saturating_sub(row.record.ordered_sparse_num_vars); + let is_last_ordered_sparse_accumulate = row.record.kind == MainSelectorKind::OrderedSparse + && row.step_kind == STEP_ACCUMULATE + && row.step_idx + 1 == row.record.sparse_indices.len(); + cols.carry_accumulator = F::from_bool(match row.step_kind { + STEP_EQ_PRODUCT => { + row.record.kind == MainSelectorKind::Whole || row.step_idx + 1 < row.record.ctx_num_vars + } + STEP_ACCUMULATE => !(is_last_ordered_sparse_accumulate && ordered_sparse_tail_len > 0), + STEP_MULTIPLY | STEP_QUARK => true, + _ => false, + }); + cols.round_idx = F::from_usize(row.round_idx); + cols.sparse_pos = F::from_usize(row.sparse_pos); + cols.sparse_index = F::from_usize(row.sparse_index); + cols.sparse_index_bits_value = F::from_usize(row.sparse_index_bits_value); + cols.point_active = F::from_bool(row.point_active); + cols.lhs_point = row + .lhs_point + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.rhs_point = row + .rhs_point + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.factor = row.factor.as_basis_coefficients_slice().try_into().unwrap(); + cols.acc_in = row.acc_in.as_basis_coefficients_slice().try_into().unwrap(); + cols.acc_out = row + .acc_out + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.aux = row.aux.as_basis_coefficients_slice().try_into().unwrap(); + cols.aux2 = row.aux2.as_basis_coefficients_slice().try_into().unwrap(); + cols.aux3 = row.aux3.as_basis_coefficients_slice().try_into().unwrap(); + cols.value = row.value.as_basis_coefficients_slice().try_into().unwrap(); +} + +fn selector_kind_code(kind: MainSelectorKind) -> usize { + match kind { + MainSelectorKind::Whole => 0, + MainSelectorKind::Prefix => 1, + MainSelectorKind::OrderedSparse => 2, + MainSelectorKind::QuarkBinaryTreeLessThan => 3, + } +} + +pub(crate) fn selector_point_source_code(kind: MainSelectorPointSourceKind) -> usize { + match kind { + MainSelectorPointSourceKind::TowerMain => 0, + MainSelectorPointSourceKind::RotationLeft => 1, + MainSelectorPointSourceKind::RotationRight => 2, + MainSelectorPointSourceKind::RotationOrigin => 3, + MainSelectorPointSourceKind::EccXY => 4, + MainSelectorPointSourceKind::EccSlope => 5, + MainSelectorPointSourceKind::EccX3Y3 => 6, + } +} + +fn native_eq_factor(out: RecursionField, input: RecursionField) -> RecursionField { + out * input + (RecursionField::ONE - out) * (RecursionField::ONE - input) +} + +fn native_eq_vec(point: &[RecursionField]) -> Vec { + let mut eq = vec![RecursionField::ONE]; + for &value in point.iter().rev() { + let mut next = Vec::with_capacity(eq.len() * 2); + for prev in eq { + next.push(prev * (RecursionField::ONE - value)); + next.push(prev * value); + } + eq = next; + } + eq +} + +fn native_eq_lte( + max_idx: usize, + out_point: &[RecursionField], + in_point: &[RecursionField], +) -> RecursionField { + assert!(out_point.len() >= in_point.len()); + let mut running_product = Vec::with_capacity(in_point.len() + 1); + running_product.push(RecursionField::ONE); + for i in 0..in_point.len() { + running_product.push(running_product[i] * native_eq_factor(out_point[i], in_point[i])); + } + + let mut running_product2 = vec![RecursionField::ZERO; in_point.len() + 1]; + running_product2[in_point.len()] = RecursionField::ONE; + for i in (0..in_point.len()).rev() { + let bit = if ((max_idx >> i) & 1) == 1 { + RecursionField::ONE + } else { + RecursionField::ZERO + }; + running_product2[i] = running_product2[i + 1] + * (out_point[i] * in_point[i] * bit + + (RecursionField::ONE - out_point[i]) + * (RecursionField::ONE - in_point[i]) + * (RecursionField::ONE - bit)); + } + + let mut ans = running_product[in_point.len()]; + for i in 0..in_point.len() { + if ((max_idx >> i) & 1) == 0 { + ans -= running_product[i] * running_product2[i + 1] * out_point[i] * in_point[i]; + } + } + for value in out_point.iter().skip(in_point.len()) { + ans *= RecursionField::ONE - *value; + } + ans +} + +#[cfg(test)] +mod tests { + use super::*; + use gkr_iop::selector::{SelectorContext, SelectorType}; + use multilinear_extensions::{ + Expression, StructuralWitIn, StructuralWitInType, ToExpr, WitnessId, + }; + + fn structural_expr(wit_id: WitnessId) -> Expression { + StructuralWitIn { + id: wit_id, + witin_type: StructuralWitInType::Empty, + } + .expr() + } + + fn point(values: &[usize]) -> Vec { + values + .iter() + .map(|&value| RecursionField::from_usize(value)) + .collect() + } + + fn record( + kind: MainSelectorKind, + ctx: &SelectorContext, + ordered_sparse_num_vars: usize, + sparse_indices: Vec, + out_point: Vec, + in_point: Vec, + value: RecursionField, + ) -> MainSelectorEvalRecord { + MainSelectorEvalRecord { + proof_idx: 0, + idx: 7, + tower_idx: 7, + air_idx: 11, + selector_idx: 13, + has_eval: true, + eval_idx: 17, + kind, + ctx_offset: ctx.offset, + ctx_num_instances: ctx.num_instances, + ctx_num_vars: ctx.num_vars, + ordered_sparse_num_vars, + sparse_indices, + in_point, + out_point, + point_source: MainSelectorPointSourceKind::TowerMain, + value, + ..Default::default() + } + } + + fn formula_final_value(record: &MainSelectorEvalRecord) -> RecursionField { + build_formula_rows(core::slice::from_ref(record)) + .into_iter() + .find(|row| row.step_kind == STEP_FINAL) + .map(|row| row.acc_out) + .expect("active selector has final row") + } + + fn assert_matches_native( + selector: SelectorType, + kind: MainSelectorKind, + ctx: SelectorContext, + ordered_sparse_num_vars: usize, + sparse_indices: Vec, + out_point: Vec, + in_point: Vec, + ) { + let (expected, _) = selector.evaluate(&out_point, &in_point, &ctx).unwrap(); + let record = record( + kind, + &ctx, + ordered_sparse_num_vars, + sparse_indices, + out_point, + in_point, + expected, + ); + assert_eq!(formula_final_value(&record), expected); + } + + #[test] + fn selector_formula_rows_match_native_whole() { + let ctx = SelectorContext::new(0, 8, 3); + assert_matches_native( + SelectorType::Whole(structural_expr(0)), + MainSelectorKind::Whole, + ctx, + 0, + Vec::new(), + point(&[2, 3, 4]), + point(&[5, 6, 7]), + ); + } + + #[test] + fn selector_formula_rows_match_native_prefix() { + let ctx = SelectorContext::new(2, 5, 3); + assert_matches_native( + SelectorType::Prefix(structural_expr(0)), + MainSelectorKind::Prefix, + ctx, + 0, + Vec::new(), + point(&[2, 3, 4]), + point(&[5, 6, 7]), + ); + } + + #[test] + fn selector_formula_rows_match_native_empty_prefix() { + let ctx = SelectorContext::new(0, 0, 3); + assert_matches_native( + SelectorType::Prefix(structural_expr(0)), + MainSelectorKind::Prefix, + ctx, + 0, + Vec::new(), + point(&[2, 3, 4]), + point(&[5, 6, 7]), + ); + } + + #[test] + fn selector_formula_rows_match_native_ordered_sparse() { + let ctx = SelectorContext::new(0, 3, 4); + let sparse_indices = vec![0, 2, 3]; + assert_matches_native( + SelectorType::OrderedSparse { + num_vars: 2, + indices: sparse_indices.clone(), + expression: structural_expr(0), + }, + MainSelectorKind::OrderedSparse, + ctx, + 2, + sparse_indices, + point(&[2, 3, 4, 5]), + point(&[6, 7, 8, 9]), + ); + } + + #[test] + fn selector_formula_rows_match_native_quark_binary_tree_less_than() { + let ctx = SelectorContext::new(0, 5, 3); + assert_matches_native( + SelectorType::QuarkBinaryTreeLessThan(structural_expr(0)), + MainSelectorKind::QuarkBinaryTreeLessThan, + ctx, + 0, + Vec::new(), + point(&[2, 3, 4]), + point(&[5, 6, 7]), + ); + } + + #[test] + fn inactive_selector_rows_do_not_emit_final_result() { + let mut record = MainSelectorEvalRecord { + proof_idx: 0, + idx: 7, + tower_idx: 7, + air_idx: 11, + selector_idx: 13, + has_eval: false, + eval_idx: 17, + kind: MainSelectorKind::OrderedSparse, + ctx_offset: 0, + ctx_num_instances: 3, + ctx_num_vars: 4, + ordered_sparse_num_vars: 2, + sparse_indices: vec![0, 2], + in_point: Vec::new(), + out_point: Vec::new(), + point_source: MainSelectorPointSourceKind::TowerMain, + value: RecursionField::ZERO, + ..Default::default() + }; + let rows = build_formula_rows(core::slice::from_ref(&record)); + assert!(rows.iter().any(|row| row.step_kind == STEP_SHAPE)); + assert_eq!( + rows.iter() + .filter(|row| row.step_kind == STEP_SPARSE_INDEX) + .count(), + 2 + ); + assert!(rows.iter().all(|row| row.step_kind != STEP_FINAL)); + + record.has_eval = true; + record.in_point = point(&[1, 2, 3, 4]); + record.out_point = point(&[5, 6, 7, 8]); + let rows = build_formula_rows(core::slice::from_ref(&record)); + assert!(rows.iter().any(|row| row.step_kind == STEP_FINAL)); + } +} diff --git a/ceno_recursion_v2/src/main/tower_point.rs b/ceno_recursion_v2/src/main/tower_point.rs new file mode 100644 index 000000000..28af52cd7 --- /dev/null +++ b/ceno_recursion_v2/src/main/tower_point.rs @@ -0,0 +1,140 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_circuit_primitives::utils::assert_array_eq; +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, extension::BinomiallyExtendable}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{MainGlobalPointBus, MainGlobalPointMessage, TowerMainPointBus, TowerMainPointMessage}, + system::MainTowerPointEqRecord, + tracegen::RowMajorChip, +}; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainTowerPointEqCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub round_idx: T, + pub global_value: [T; D_EF], + pub tower_value: [T; D_EF], + pub eq_in: [T; D_EF], + pub eq_out: [T; D_EF], +} + +pub struct MainTowerPointEqAir { + pub global_point_bus: MainGlobalPointBus, + pub tower_point_bus: TowerMainPointBus, +} + +impl BaseAir for MainTowerPointEqAir { + fn width(&self) -> usize { + MainTowerPointEqCols::::width() + } +} + +impl BaseAirWithPublicValues for MainTowerPointEqAir {} +impl PartitionedBaseAir for MainTowerPointEqAir {} + +impl Air for MainTowerPointEqAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &MainTowerPointEqCols = (*local_row).borrow(); + + builder.assert_bool(local.is_enabled); + self.global_point_bus.lookup_key( + builder, + local.proof_idx, + MainGlobalPointMessage { + round_idx: local.round_idx.into(), + value: local.global_value.map(Into::into), + }, + local.is_enabled, + ); + self.tower_point_bus.lookup_key( + builder, + local.proof_idx, + TowerMainPointMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + value: local.tower_value.map(Into::into), + }, + local.is_enabled, + ); + let same_bit: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_multiply(local.global_value, local.tower_value); + let opposite_zero: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply( + recursion_circuit::utils::ext_field_one_minus(local.global_value), + recursion_circuit::utils::ext_field_one_minus(local.tower_value), + ); + let eq_factor: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_add(same_bit, opposite_zero); + let eq_out = recursion_circuit::utils::ext_field_multiply(local.eq_in, eq_factor); + assert_array_eq(&mut builder.when(local.is_enabled), local.eq_out, eq_out); + } +} + +pub struct MainTowerPointEqTraceGenerator; + +impl RowMajorChip for MainTowerPointEqTraceGenerator { + type Ctx<'a> = &'a [MainTowerPointEqRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainTowerPointEqCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainTowerPointEqCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.global_value = record + .global_value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.tower_value = record + .tower_value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.eq_in = record + .eq_in + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + cols.eq_out = record + .eq_out + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + } + Some(RowMajorMatrix::new(trace, width)) + } +} diff --git a/ceno_recursion_v2/src/main/trace.rs b/ceno_recursion_v2/src/main/trace.rs index 5d7e36e53..9df763ade 100644 --- a/ceno_recursion_v2/src/main/trace.rs +++ b/ceno_recursion_v2/src/main/trace.rs @@ -9,6 +9,7 @@ use crate::tracegen::RowMajorChip; #[derive(Clone, Debug, Default)] pub struct MainRecord { + pub is_present: bool, pub proof_idx: usize, pub idx: usize, pub tidx: usize, @@ -39,7 +40,8 @@ where Fill: FnMut(&MainRecord, &mut C, bool, bool), { let width = C::width(); - let num_rows = records.len().max(1); + let dense_records = dense_main_records(records); + let num_rows = dense_records.len().max(1); let height = if let Some(height) = required_height { if height < num_rows { return None; @@ -49,13 +51,10 @@ where num_rows.next_power_of_two().max(1) }; let mut trace = vec![F::ZERO; height * width]; - if records.is_empty() { - return Some(RowMajorMatrix::new(trace, width)); - } let mut prev_proof_idx = usize::MAX; let mut prev_idx = usize::MAX; - for (row_idx, record) in records.iter().enumerate() { + for (row_idx, record) in dense_records.iter().enumerate() { let offset = row_idx * width; let cols_slice = &mut trace[offset..offset + width]; let cols = C::from_bytes(cols_slice); @@ -69,6 +68,39 @@ where Some(RowMajorMatrix::new(trace, width)) } +fn dense_main_records(records: &[MainRecord]) -> Vec { + if records.is_empty() { + return vec![MainRecord::default()]; + } + + let mut dense = Vec::new(); + let mut record_idx = 0; + while record_idx < records.len() { + let proof_idx = records[record_idx].proof_idx; + let proof_start = record_idx; + while record_idx < records.len() && records[record_idx].proof_idx == proof_idx { + record_idx += 1; + } + let proof_records = &records[proof_start..record_idx]; + let max_idx = proof_records.last().map(|record| record.idx).unwrap_or(0); + let mut proof_record_idx = 0; + for idx in 0..=max_idx { + if proof_record_idx < proof_records.len() && proof_records[proof_record_idx].idx == idx + { + dense.push(proof_records[proof_record_idx].clone()); + proof_record_idx += 1; + } else { + dense.push(MainRecord { + proof_idx, + idx, + ..Default::default() + }); + } + } + } + dense +} + trait ColumnAccess: Sized { fn width() -> usize; fn from_bytes(slice: &mut [F]) -> &mut Self; @@ -86,6 +118,7 @@ impl ColumnAccess for MainCols { fn fill_main_cols(record: &MainRecord, cols: &mut MainCols, is_first_idx: bool, is_first: bool) { cols.is_enabled = F::ONE; + cols.is_present = F::from_bool(record.is_present); cols.proof_idx = F::from_usize(record.proof_idx); cols.idx = F::from_usize(record.idx); cols.is_first_idx = F::from_bool(is_first_idx); diff --git a/ceno_recursion_v2/src/main/transcript_bind.rs b/ceno_recursion_v2/src/main/transcript_bind.rs new file mode 100644 index 000000000..91315e3f7 --- /dev/null +++ b/ceno_recursion_v2/src/main/transcript_bind.rs @@ -0,0 +1,114 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::F; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{Field, PrimeCharacteristicRing}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{ForkedTranscriptBus, ForkedTranscriptBusMessage, TranscriptBus, TranscriptBusMessage}, + system::MainTranscriptRecord, + tracegen::RowMajorChip, +}; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct MainTranscriptBindCols { + pub is_enabled: T, + pub proof_idx: T, + pub fork_id: T, + pub is_fork: T, + pub tidx: T, + pub value: T, + pub is_sample: T, +} + +pub struct MainTranscriptBindAir { + pub transcript_bus: TranscriptBus, + pub forked_transcript_bus: ForkedTranscriptBus, +} + +impl BaseAir for MainTranscriptBindAir { + fn width(&self) -> usize { + MainTranscriptBindCols::::width() + } +} + +impl BaseAirWithPublicValues for MainTranscriptBindAir {} +impl PartitionedBaseAir for MainTranscriptBindAir {} + +impl Air for MainTranscriptBindAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &MainTranscriptBindCols = (*local_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_sample); + builder.assert_bool(local.is_fork); + + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.tidx.into(), + value: local.value.into(), + is_sample: local.is_sample.into(), + }, + local.is_enabled * (AB::Expr::ONE - local.is_fork), + ); + self.forked_transcript_bus.receive( + builder, + local.proof_idx, + ForkedTranscriptBusMessage { + fork_id: local.fork_id.into(), + tidx: local.tidx.into(), + value: local.value.into(), + is_sample: local.is_sample.into(), + }, + local.is_enabled * local.is_fork, + ); + } +} + +pub struct MainTranscriptBindTraceGenerator; + +impl RowMajorChip for MainTranscriptBindTraceGenerator { + type Ctx<'a> = &'a [MainTranscriptRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = MainTranscriptBindCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut MainTranscriptBindCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.fork_id = F::from_usize(record.fork_id); + cols.is_fork = F::from_bool(record.is_fork); + cols.tidx = F::from_usize(record.tidx); + cols.value = record.value; + cols.is_sample = F::from_bool(record.is_sample); + } + + Some(RowMajorMatrix::new(trace, width)) + } +} diff --git a/ceno_recursion_v2/src/pcs/mod.rs b/ceno_recursion_v2/src/pcs/mod.rs new file mode 100644 index 000000000..6b69d7fc0 --- /dev/null +++ b/ceno_recursion_v2/src/pcs/mod.rs @@ -0,0 +1,7377 @@ +use core::borrow::{Borrow, BorrowMut}; +use std::sync::Arc; + +use eyre::{Result, bail}; +use ff_ext::{BabyBearExt4, ExtensionField as CenoExtensionField, SmallField}; +use itertools::Itertools; +use mpcs::{ + BasefoldRSParams, + basefold::BasefoldSpec, + jagged::{compute_q_at_assist_point, evaluate_g}, +}; +use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec}; +use openvm_circuit_primitives::utils::assert_array_eq; +use openvm_cpu_backend::CpuBackend; +use openvm_poseidon2_air::POSEIDON2_WIDTH; +use openvm_stark_backend::{ + AirRef, BaseAirWithPublicValues, FiatShamirTranscript, PartitionedBaseAir, StarkProtocolConfig, + TranscriptHistory, interaction::InteractionBuilder, prover::AirProvingContext, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{ + BabyBearPoseidon2Config, D_EF, DIGEST_SIZE, EF, F, poseidon2_compress_with_capacity, + poseidon2_perm, +}; +use p3::{ + dft::{Radix2Dit, TwoAdicSubgroupDft}, + matrix::dense::RowMajorMatrix as P3RowMajorMatrix, + util::reverse_bits_len, +}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{ + BasedVectorSpace, Field, PrimeCharacteristicRing, PrimeField64, TwoAdicField, + extension::BinomiallyExtendable, +}; +use p3_matrix::{Matrix, bitrev::BitReversibleMatrix, dense::RowMajorMatrix}; +use p3_symmetric::Permutation; +use stark_recursion_circuit_derive::AlignedBorrow; +use sumcheck::util::extrapolate_uni_poly; + +use crate::{ + bus::{ + MainEvalBus, MainEvalMessage, MainGlobalPointBus, MainGlobalPointMessage, + PcsBaseInputOpeningBus, PcsBaseInputOpeningMessage, PcsBasefoldEvalBus, + PcsBasefoldEvalMessage, PcsBasefoldFinalExpectedBus, PcsBasefoldFinalExpectedMessage, + PcsBasefoldFinalPointBus, PcsBasefoldFinalPointMessage, PcsBasefoldQueryBus, + PcsBasefoldQueryMessage, PcsBasefoldQueryStage, PcsBatchAlphaBus, PcsBatchAlphaMessage, + PcsBatchCoeffBus, PcsBatchCoeffMessage, PcsCommitHeightBus, PcsCommitHeightMessage, + PcsCommitPhaseLeafBus, PcsCommitPhaseLeafMessage, PcsCommitmentRootBus, + PcsCommitmentRootMessage, PcsEqProductBus, PcsEqProductMessage, PcsFinalMessageBus, + PcsFinalMessageMessage, PcsFoldChallengeBus, PcsFoldChallengeMessage, PcsJaggedAssistHBus, + PcsJaggedAssistHMessage, PcsJaggedAssistQBus, PcsJaggedAssistQMessage, PcsJaggedFEvalBus, + PcsJaggedFEvalMessage, PcsOpeningEvalBus, PcsOpeningEvalMessage, PcsQuerySampleBus, + PcsQuerySampleMessage, PcsSuffixProductBus, PcsSuffixProductMessage, PcsSumcheckInputBus, + PcsSumcheckInputMessage, PcsSumcheckOutputBus, PcsSumcheckOutputMessage, + PcsTranscriptExtBus, PcsTranscriptExtMessage, TranscriptBus, TranscriptBusMessage, + }, + system::{ + AirModule, GlobalCtxCpu, PcsBaseInputLeafHashRecord, PcsBaseInputMerkleRecord, + PcsBasefoldCommitPhaseQueryRecord, PcsBasefoldFinalClaimRecord, + PcsBasefoldFinalCodewordRecord, PcsBasefoldFinalExpectedRecord, + PcsBasefoldFinalPointRecord, PcsBasefoldInitialClaimRecord, PcsBasefoldQueryIndexRecord, + PcsBasefoldQueryOpenRecord, PcsBatchCoeffRecord, PcsCommitPhaseLeafHashRecord, + PcsCommitPhaseMerkleRecord, PcsCommitmentRootRecord, PcsEqProductKind, PcsEqProductRecord, + PcsEqProductSource, PcsJaggedAssistHRecord, PcsJaggedAssistInputRecord, + PcsJaggedAssistQRecord, PcsJaggedAssistRecord, PcsJaggedClaimRecord, PcsJaggedQEvalRecord, + PcsOpeningCommitKind, PcsOpeningEvalRecord, PcsOpeningPointRecord, PcsSuffixProductRecord, + PcsSumcheckInputRecord, PcsSumcheckRoundRecord, PcsTranscriptValueRecord, Preflight, + RecursionField, RecursionPcs, RecursionProof, RecursionVk, TraceGenModule, + }, + tracegen::{ModuleChip, RowMajorChip}, + utils::digests_to_poseidon2_input, +}; +use recursion_circuit::{ + bus::{ + MerkleVerifyBus, MerkleVerifyBusMessage, Poseidon2CompressBus, Poseidon2CompressMessage, + Poseidon2PermuteBus, Poseidon2PermuteMessage, + }, + primitives::bus::{RangeCheckerBus, RangeCheckerBusMessage}, +}; + +pub struct PcsModule { + transcript_bus: TranscriptBus, + main_global_point_bus: MainGlobalPointBus, + main_eval_bus: MainEvalBus, + basefold_query_bus: PcsBasefoldQueryBus, + basefold_eval_bus: PcsBasefoldEvalBus, + basefold_final_expected_bus: PcsBasefoldFinalExpectedBus, + basefold_final_point_bus: PcsBasefoldFinalPointBus, + transcript_ext_bus: PcsTranscriptExtBus, + base_input_opening_bus: PcsBaseInputOpeningBus, + final_message_bus: PcsFinalMessageBus, + query_sample_bus: PcsQuerySampleBus, + commitment_root_bus: PcsCommitmentRootBus, + commit_phase_leaf_bus: PcsCommitPhaseLeafBus, + sumcheck_input_bus: PcsSumcheckInputBus, + sumcheck_output_bus: PcsSumcheckOutputBus, + fold_challenge_bus: PcsFoldChallengeBus, + batch_coeff_bus: PcsBatchCoeffBus, + batch_alpha_bus: PcsBatchAlphaBus, + jagged_f_eval_bus: PcsJaggedFEvalBus, + opening_eval_bus: PcsOpeningEvalBus, + eq_product_bus: PcsEqProductBus, + suffix_product_bus: PcsSuffixProductBus, + jagged_assist_h_bus: PcsJaggedAssistHBus, + jagged_assist_q_bus: PcsJaggedAssistQBus, + commit_height_bus: PcsCommitHeightBus, + merkle_verify_bus: MerkleVerifyBus, + poseidon2_permute_bus: Poseidon2PermuteBus, + poseidon2_compress_bus: Poseidon2CompressBus, + range_checker_bus: RangeCheckerBus, +} + +pub(crate) fn collect_pcs_range_checks(preflights: &[Preflight]) -> Vec { + preflights + .iter() + .flat_map(|preflight| preflight.pcs.basefold_query_indices.iter()) + .flat_map(|record| { + record + .query_bytes + .iter() + .chain(record.high_bytes.iter()) + .map(|&byte| byte as usize) + .collect_vec() + }) + .collect() +} + +type BasefoldOpeningEvals = Vec<(RecursionField, usize)>; +type BasefoldOpening = (usize, (Vec, BasefoldOpeningEvals)); +type BasefoldRound = ( + mpcs::basefold::BasefoldCommitment, + Vec, +); +type JaggedTermSource = (usize, usize, PcsOpeningCommitKind, usize, usize, usize); + +fn ensure_supported_basecode_log(basecode_log: usize) -> Result<()> { + if basecode_log != 0 { + bail!( + "unsupported basefold basecode_log {basecode_log}; only basecode_log == 0 is supported" + ); + } + Ok(()) +} + +fn push_eq_product_records( + preflight: &mut Preflight, + kind: PcsEqProductKind, + source: PcsEqProductSource, + round_idx: usize, + term_idx: usize, + point: &[RecursionField], + point_tidx_base: usize, + sumcheck_idx: usize, + point_round_base: usize, +) -> RecursionField { + let mut acc = RecursionField::ONE; + let mut index_acc = 0usize; + for (bit_idx, &point_value) in point.iter().enumerate() { + let index_bit = ((term_idx >> bit_idx) & 1) == 1; + let factor = if index_bit { + point_value + } else { + RecursionField::ONE - point_value + }; + let acc_in = acc; + acc *= factor; + let index_acc_in = index_acc; + let index_pow2 = 1usize << bit_idx; + if index_bit { + index_acc += index_pow2; + } + preflight.pcs.eq_products.push(PcsEqProductRecord { + proof_idx: 0, + kind, + source, + round_idx, + term_idx, + bit_idx, + is_first: bit_idx == 0, + is_last: bit_idx + 1 == point.len(), + lookup_count: if bit_idx + 1 == point.len() { + if kind == PcsEqProductKind::JaggedClaim { + 2 + } else { + 1 + } + } else { + 0 + }, + point_tidx: point_tidx_base + bit_idx * D_EF, + sumcheck_idx, + point_round: point_round_base + bit_idx, + index_bit, + index_pow2, + index_acc_in, + index_acc_out: index_acc, + point: point_value, + acc_in, + acc_out: acc, + }); + } + acc +} + +fn push_suffix_product_records( + preflight: &mut Preflight, + round_idx: usize, + term_idx: usize, + start_coord_idx: usize, + point: &[RecursionField], +) -> RecursionField { + let mut acc = RecursionField::ONE; + let factor_count = point.len().saturating_sub(start_coord_idx); + let row_count = factor_count.max(1); + for step_idx in 0..row_count { + let has_factor = step_idx < factor_count; + let coord_idx = if has_factor { + start_coord_idx + step_idx + } else { + start_coord_idx + }; + let point_value = point.get(coord_idx).copied().unwrap_or_default(); + let acc_in = acc; + if has_factor { + acc *= RecursionField::ONE - point_value; + } + preflight.pcs.suffix_products.push(PcsSuffixProductRecord { + proof_idx: 0, + round_idx, + term_idx, + coord_idx, + step_idx, + is_first: step_idx == 0, + is_last: step_idx + 1 == row_count, + has_factor, + point: point_value, + acc_in, + acc_out: acc, + }); + } + acc +} + +fn jagged_transition_weights( + z1i: RecursionField, + z2i: RecursionField, + z3i: RecursionField, + z4i: RecursionField, +) -> [(usize, usize, RecursionField, RecursionField, RecursionField); 4] { + let (nz1, nz2, nz3, nz4) = ( + RecursionField::ONE - z1i, + RecursionField::ONE - z2i, + RecursionField::ONE - z3i, + RecursionField::ONE - z4i, + ); + let ab00 = nz1 * nz2; + let ab01 = nz1 * z2i; + let ab10 = z1i * nz2; + let ab11 = z1i * z2i; + let cd00 = nz3 * nz4; + let cd01 = nz3 * z4i; + let cd10 = z3i * nz4; + let cd11 = z3i * z4i; + [ + ( + 0, + 0, + ab00 * cd00 + ab01 * cd11 + ab11 * cd01, + ab00 * cd01, + ab01 * cd10 + ab11 * cd00, + ), + (0, 1, ab10 * cd10, ab10 * cd11, RecursionField::ZERO), + (1, 0, ab01 * cd01, RecursionField::ZERO, ab01 * cd00), + ( + 1, + 1, + ab00 * cd10 + ab10 * cd00 + ab11 * cd11, + ab00 * cd11 + ab10 * cd01, + ab11 * cd10, + ), + ] +} + +fn jagged_backward_step( + z1i: RecursionField, + z2i: RecursionField, + z3i: RecursionField, + z4i: RecursionField, + val: [RecursionField; 4], +) -> [RecursionField; 4] { + let mut new_val = [RecursionField::ZERO; 4]; + for &(ci, co, w_same, w_lt1, w_lt0) in &jagged_transition_weights(z1i, z2i, z3i, z4i) { + let v0 = val[co * 2]; + let v1 = val[co * 2 + 1]; + new_val[ci * 2] += w_same * v0 + w_lt1 * v1 + w_lt0 * v0; + new_val[ci * 2 + 1] += w_same * v1 + w_lt1 * v1 + w_lt0 * v0; + } + new_val +} + +fn push_jagged_assist_h_records( + preflight: &mut Preflight, + round_idx: usize, + point: &[RecursionField], + rho: &[RecursionField], + assist_point: &[RecursionField], + n_robp: usize, +) -> RecursionField { + let sumcheck_idx = round_idx * 2 + 1; + let mut val = [ + RecursionField::ZERO, + RecursionField::ONE, + RecursionField::ZERO, + RecursionField::ZERO, + ]; + for step_idx in 0..n_robp { + let robp_idx = n_robp - 1 - step_idx; + let has_z_row = robp_idx < point.len(); + let has_rho = robp_idx < rho.len(); + let z_row = point.get(robp_idx).copied().unwrap_or_default(); + let rho_value = rho.get(robp_idx).copied().unwrap_or_default(); + let rho_star_c = assist_point[2 * robp_idx]; + let rho_star_d = assist_point[2 * robp_idx + 1]; + let val_in = val; + val = jagged_backward_step(z_row, rho_value, rho_star_c, rho_star_d, val); + preflight.pcs.jagged_assist_h.push(PcsJaggedAssistHRecord { + proof_idx: 0, + round_idx, + sumcheck_idx, + step_idx, + robp_idx, + is_first: step_idx == 0, + is_last: step_idx + 1 == n_robp, + has_z_row, + has_rho, + z_row, + rho: rho_value, + rho_star_c, + rho_star_d, + val_in, + val_out: val, + }); + } + val[0] +} + +fn push_jagged_assist_q_records( + preflight: &mut Preflight, + round_idx: usize, + commitment_kind: usize, + assist_point: &[RecursionField], + eq_col: &[RecursionField], + cumulative_heights: &[usize], + n_robp: usize, +) -> RecursionField { + let sumcheck_idx = round_idx * 2 + 1; + let num_polys = cumulative_heights.len() - 1; + let mut q_acc = RecursionField::ZERO; + for term_idx in 0..num_polys { + let mut term_acc = eq_col[term_idx]; + let mut c_acc = 0usize; + let mut d_acc = 0usize; + for step_idx in 0..n_robp { + let c_bit = ((cumulative_heights[term_idx] >> step_idx) & 1) == 1; + let d_bit = ((cumulative_heights[term_idx + 1] >> step_idx) & 1) == 1; + let bit_pow2 = 1usize << step_idx; + let c_acc_in = c_acc; + let d_acc_in = d_acc; + if c_bit { + c_acc += bit_pow2; + } + if d_bit { + d_acc += bit_pow2; + } + let rho_star_c = assist_point[2 * step_idx]; + let rho_star_d = assist_point[2 * step_idx + 1]; + let c_factor = if c_bit { + rho_star_c + } else { + RecursionField::ONE - rho_star_c + }; + let d_factor = if d_bit { + rho_star_d + } else { + RecursionField::ONE - rho_star_d + }; + let term_acc_in = term_acc; + term_acc *= c_factor * d_factor; + let q_acc_in = q_acc; + let is_last_step = step_idx + 1 == n_robp; + if is_last_step { + q_acc += term_acc; + } + preflight.pcs.jagged_assist_q.push(PcsJaggedAssistQRecord { + proof_idx: 0, + round_idx, + sumcheck_idx, + commitment_kind, + term_idx, + step_idx, + robp_idx: step_idx, + is_first: term_idx == 0 && step_idx == 0, + is_last: term_idx + 1 == num_polys && is_last_step, + is_first_step: step_idx == 0, + is_last_step, + term_is_last: term_idx + 1 == num_polys, + eq_col: eq_col[term_idx], + t_lo: cumulative_heights[term_idx], + t_hi: cumulative_heights[term_idx + 1], + c_bit, + d_bit, + bit_pow2, + c_acc_in, + c_acc_out: c_acc, + d_acc_in, + d_acc_out: d_acc, + rho_star_c, + rho_star_d, + term_acc_in, + term_acc_out: term_acc, + q_acc_in, + q_acc_out: q_acc, + }); + } + } + q_acc +} + +fn push_basefold_final_expected_records( + preflight: &mut Preflight, + rounds: &[BasefoldRound], + final_message: &[Vec], + fold_challenges: &[RecursionField], + basecode_log: usize, + final_tidx: usize, +) -> Result { + let sumcheck_idx = 8_000_000; + let mut acc = RecursionField::ZERO; + let mut term_idx = preflight.pcs.basefold_final_expected.len(); + let mut flat_message_idx = 0usize; + for (point_idx, (round_idx, point)) in final_message + .iter() + .zip( + rounds + .iter() + .enumerate() + .flat_map(|(round_idx, (_, point_evals))| { + point_evals + .iter() + .map(move |(_, (point, _))| (round_idx, point)) + }) + .filter(|(_, point)| point.len() >= basecode_log), + ) + .map(|(_, point)| point) + .enumerate() + { + for (coord_idx, &value) in point.iter().enumerate() { + let source_sumcheck_idx = round_idx * 2; + if let Some(record) = preflight + .pcs + .sumcheck_rounds + .iter_mut() + .find(|record| record.idx == source_sumcheck_idx && record.round == coord_idx) + { + record.fold_challenge_lookup_count += 1; + } + preflight + .pcs + .basefold_final_points + .push(PcsBasefoldFinalPointRecord { + proof_idx: 0, + sumcheck_idx, + point_idx, + coord_idx, + source_sumcheck_idx, + source_round: coord_idx, + value, + }); + } + } + for (point_idx, (message_row, point)) in final_message + .iter() + .zip( + rounds + .iter() + .flat_map(|(_, point_evals)| point_evals.iter()) + .filter(|(_, (point, _))| point.len() >= basecode_log) + .map(|(_, (point, _))| point), + ) + .enumerate() + { + let num_vars_evaluated = point.len() - basecode_log; + if num_vars_evaluated > fold_challenges.len() { + bail!("basefold final expected point exceeds fold challenge count"); + } + for (elem_idx, final_value) in message_row.iter().copied().enumerate() { + let element_tidx = final_tidx + flat_message_idx * D_EF; + let mut coeff = RecursionField::ONE; + let start_term_idx = term_idx; + let factor_count = num_vars_evaluated + basecode_log; + for factor_idx in 0..factor_count { + let is_prefix = factor_idx < num_vars_evaluated; + let point_value = point[factor_idx]; + let (has_challenge, challenge_round, challenge, bit_value, factor) = if is_prefix { + let challenge_round = fold_challenges.len() - num_vars_evaluated + factor_idx; + let challenge = fold_challenges[challenge_round]; + if let Some(record) = preflight.pcs.sumcheck_rounds.iter_mut().find(|record| { + record.idx == sumcheck_idx && record.round == challenge_round + }) { + record.fold_challenge_lookup_count += 1; + } + let xi_yi = point_value * challenge; + ( + true, + challenge_round, + challenge, + false, + xi_yi + xi_yi - point_value - challenge + RecursionField::ONE, + ) + } else { + let bit_idx = factor_idx - num_vars_evaluated; + let bit_value = ((elem_idx >> bit_idx) & 1) == 1; + ( + false, + 0, + RecursionField::ZERO, + bit_value, + if bit_value { + point_value + } else { + RecursionField::ONE - point_value + }, + ) + }; + let coeff_in = coeff; + coeff *= factor; + let is_elem_last = factor_idx + 1 == factor_count; + let acc_in = acc; + if is_elem_last { + acc += coeff * final_value; + } + preflight + .pcs + .basefold_final_expected + .push(PcsBasefoldFinalExpectedRecord { + proof_idx: 0, + sumcheck_idx, + term_idx, + point_idx, + coord_idx: factor_idx, + final_tidx: element_tidx, + is_first: term_idx == 0, + is_last: false, + is_elem_first: term_idx == start_term_idx, + is_elem_last, + has_challenge, + challenge_round, + point_value, + bit_value, + challenge, + final_value, + coeff_in, + coeff_out: coeff, + acc_in, + acc_out: acc, + }); + term_idx += 1; + } + flat_message_idx += 1; + } + } + if let Some(last) = preflight.pcs.basefold_final_expected.last_mut() { + last.is_last = true; + } + Ok(acc) +} + +impl PcsModule { + pub fn new(bus_inventory: crate::system::BusInventory) -> Self { + Self { + transcript_bus: bus_inventory.transcript_bus, + main_global_point_bus: bus_inventory.main_global_point_bus, + main_eval_bus: bus_inventory.main_eval_bus, + basefold_query_bus: bus_inventory.pcs_basefold_query_bus, + basefold_eval_bus: bus_inventory.pcs_basefold_eval_bus, + basefold_final_expected_bus: bus_inventory.pcs_basefold_final_expected_bus, + basefold_final_point_bus: bus_inventory.pcs_basefold_final_point_bus, + transcript_ext_bus: bus_inventory.pcs_transcript_ext_bus, + base_input_opening_bus: bus_inventory.pcs_base_input_opening_bus, + final_message_bus: bus_inventory.pcs_final_message_bus, + query_sample_bus: bus_inventory.pcs_query_sample_bus, + commitment_root_bus: bus_inventory.pcs_commitment_root_bus, + commit_phase_leaf_bus: bus_inventory.pcs_commit_phase_leaf_bus, + sumcheck_input_bus: bus_inventory.pcs_sumcheck_input_bus, + sumcheck_output_bus: bus_inventory.pcs_sumcheck_output_bus, + fold_challenge_bus: bus_inventory.pcs_fold_challenge_bus, + batch_coeff_bus: bus_inventory.pcs_batch_coeff_bus, + batch_alpha_bus: bus_inventory.pcs_batch_alpha_bus, + jagged_f_eval_bus: bus_inventory.pcs_jagged_f_eval_bus, + opening_eval_bus: bus_inventory.pcs_opening_eval_bus, + eq_product_bus: bus_inventory.pcs_eq_product_bus, + suffix_product_bus: bus_inventory.pcs_suffix_product_bus, + jagged_assist_h_bus: bus_inventory.pcs_jagged_assist_h_bus, + jagged_assist_q_bus: bus_inventory.pcs_jagged_assist_q_bus, + commit_height_bus: bus_inventory.pcs_commit_height_bus, + merkle_verify_bus: bus_inventory.merkle_verify_bus, + poseidon2_permute_bus: bus_inventory.poseidon2_permute_bus, + poseidon2_compress_bus: bus_inventory.poseidon2_compress_bus, + range_checker_bus: bus_inventory.range_checker_bus, + } + } + + pub fn run_preflight( + &self, + child_vk: &RecursionVk, + proof: &RecursionProof, + preflight: &mut Preflight, + ts: &mut TS, + ) where + TS: FiatShamirTranscript + + TranscriptHistory, + { + self.replay_preflight(child_vk, proof, preflight, ts) + .unwrap_or_else(|err| panic!("recursion-v2 PCS preflight replay failed: {err:?}")); + } + + fn replay_preflight( + &self, + child_vk: &RecursionVk, + proof: &RecursionProof, + preflight: &mut Preflight, + ts: &mut TS, + ) -> Result<()> + where + TS: FiatShamirTranscript + + TranscriptHistory, + { + let mut rounds = Vec::new(); + let witin_openings = preflight + .pcs + .opening_claims + .iter() + .filter(|claim| !claim.wits_in_evals.is_empty()) + .map(|claim| { + ( + claim.input_opening_point.len(), + ( + claim.input_opening_point.clone(), + claim.wits_in_evals.clone(), + ), + ) + }) + .collect_vec(); + rounds.push((2usize, proof.witin_commit.clone(), witin_openings)); + + let fixed_openings = preflight + .pcs + .opening_claims + .iter() + .filter(|claim| !claim.fixed_in_evals.is_empty()) + .map(|claim| { + ( + claim.input_opening_point.len(), + ( + claim.input_opening_point.clone(), + claim.fixed_in_evals.clone(), + ), + ) + }) + .collect_vec(); + if proof.public_values.shard_id == 0 { + if let Some(fixed_commit) = child_vk.fixed_commit.as_ref() { + rounds.push((0usize, fixed_commit.clone(), fixed_openings)); + } + } else if let Some(fixed_commit) = child_vk.fixed_no_omc_init_commit.as_ref() { + rounds.push((1usize, fixed_commit.clone(), fixed_openings)); + } + + if rounds.len() != proof.opening_proof.rounds.len() { + bail!( + "jagged proof round count mismatch: {} != {}", + rounds.len(), + proof.opening_proof.rounds.len() + ); + } + + let mut inner_rounds = Vec::with_capacity(rounds.len()); + for (round_idx, ((commitment_kind, comm, openings), round_proof)) in rounds + .into_iter() + .zip(proof.opening_proof.rounds.iter()) + .enumerate() + { + let commit_kind = if round_idx == 0 { + PcsOpeningCommitKind::Witin + } else { + PcsOpeningCommitKind::Fixed + }; + let term_sources = preflight + .pcs + .opening_evals + .iter() + .enumerate() + .filter(|(_, record)| record.commit_kind == commit_kind) + .map(|(record_idx, record)| { + ( + record_idx, + record.opening_idx, + record.commit_kind, + record.eval_idx, + record.main_idx, + record.main_eval_idx, + ) + }) + .collect_vec(); + let poly_heights = comm + .cumulative_heights + .windows(2) + .map(|window| window[1] - window[0]) + .collect_vec(); + let openings = openings + .into_iter() + .map(|(_, (point, evals))| (point, evals)) + .collect_vec(); + let (point, evals) = flatten_padded_openings_as_native(&poly_heights, openings)?; + let inner = self.replay_jagged_round( + round_idx, + commitment_kind, + &comm, + &point, + &evals, + &term_sources, + round_proof, + preflight, + ts, + )?; + preflight + .pcs + .commitment_roots + .push(PcsCommitmentRootRecord { + proof_idx: 0, + commit_major: 0, + commit_minor: round_idx, + root: digest_to_array(&inner.0.commit())?, + lookup_count: 0, + }); + inner_rounds.push(inner); + } + self.replay_basefold( + &inner_rounds, + &proof.opening_proof.inner_proof, + preflight, + ts, + )?; + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + fn replay_jagged_round( + &self, + round_idx: usize, + commitment_kind: usize, + comm: &>::Commitment, + point: &[RecursionField], + evals: &[RecursionField], + term_sources: &[JaggedTermSource], + proof: &mpcs::jagged::JaggedBatchOpenProof< + RecursionField, + mpcs::Basefold, + >, + preflight: &mut Preflight, + ts: &mut TS, + ) -> Result<( + mpcs::basefold::BasefoldCommitment, + Vec, + )> + where + TS: FiatShamirTranscript + + TranscriptHistory, + { + let num_polys = comm.cumulative_heights.len() - 1; + if evals.len() != num_polys { + bail!("jagged eval count mismatch: {} != {num_polys}", evals.len()); + } + if term_sources.len() != num_polys { + bail!( + "jagged eval source count mismatch: {} != {num_polys}", + term_sources.len() + ); + } + let total_evals = *comm.cumulative_heights.last().unwrap_or(&0); + let log_h = comm.reshape_log_height; + let h = 1usize << log_h; + let w = total_evals.div_ceil(h); + let padded_total = w * h; + let num_giga_vars = ceil_log2(padded_total); + + let eval_tidx = ts.len(); + for (i, eval) in evals.iter().copied().enumerate() { + ts.observe_ext(eval); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: round_idx * 1_000_000 + i, + tidx: eval_tidx + i * D_EF, + value: eval, + is_sample: false, + is_ext: true, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + } + let num_col_vars = ceil_log2(num_polys).max(1); + observe_label_with_records(ts, b"z_col", preflight, round_idx * 1_000_000 + 4_000); + let z_col_tidx = ts.len(); + let z_col = sample_vec_with_records( + ts, + num_col_vars, + preflight, + 0, + round_idx * 1_000_000 + 5_000, + ); + for bit_idx in 0..num_col_vars { + let tidx = z_col_tidx + bit_idx * D_EF; + if let Some(record) = preflight + .pcs + .transcript_values + .iter_mut() + .find(|record| record.tidx == tidx) + { + record.transcript_ext_lookup_count += num_polys; + } + } + let eq_col = build_eq_x_r_vec(&z_col); + + let max_s = point.len(); + let mut tail_zero_prod = vec![RecursionField::ONE; max_s + 1]; + for j in (0..max_s).rev() { + tail_zero_prod[j] = tail_zero_prod[j + 1] * (RecursionField::ONE - point[j]); + } + let claimed_sum = (0..num_polys) + .map(|i| { + let h_i = comm.cumulative_heights[i + 1] - comm.cumulative_heights[i]; + let s_i = ceil_log2(h_i); + eq_col[i] * tail_zero_prod[s_i] * evals[i] + }) + .sum(); + let mut acc = RecursionField::ZERO; + for i in 0..num_polys { + let h_i = comm.cumulative_heights[i + 1] - comm.cumulative_heights[i]; + let s_i = ceil_log2(h_i); + let eq_product = push_eq_product_records( + preflight, + PcsEqProductKind::JaggedClaim, + PcsEqProductSource::Transcript, + round_idx, + i, + &z_col, + z_col_tidx, + round_idx * 2, + 0, + ); + let tail_product = push_suffix_product_records(preflight, round_idx, i, s_i, point); + if eq_product != eq_col[i] { + bail!("jagged eq product replay mismatch"); + } + if tail_product != tail_zero_prod[s_i] { + bail!("jagged suffix product replay mismatch"); + } + let term = eq_col[i] * tail_zero_prod[s_i] * evals[i]; + let acc_in = acc; + acc += term; + let ( + opening_eval_record_idx, + opening_idx, + commit_kind, + eval_idx, + main_idx, + main_eval_idx, + ) = term_sources[i]; + let raw_eval = tail_zero_prod[s_i] * evals[i]; + preflight.pcs.opening_evals[opening_eval_record_idx].raw_value = raw_eval; + preflight.pcs.jagged_claims.push(PcsJaggedClaimRecord { + proof_idx: 0, + round_idx, + sumcheck_idx: round_idx * 2, + term_idx: i, + is_first: i == 0, + is_last: i + 1 == num_polys, + opening_idx, + commit_kind, + eval_idx, + main_idx, + main_eval_idx, + eval: evals[i], + z_col_tidx, + eq_col: eq_col[i], + tail_zero: tail_zero_prod[s_i], + acc_in, + acc_out: acc, + }); + } + let (rho, expected) = replay_degree2_sumcheck( + 0, + round_idx * 2, + claimed_sum, + &proof.sumcheck_proof, + num_giga_vars, + false, + preflight, + ts, + )?; + + let n_robp = num_giga_vars + usize::from(padded_total.is_power_of_two()); + let rho_row = &rho[..log_h]; + let rho_col = &rho[log_h..]; + let eq_rho_col = build_eq_x_r_vec(rho_col); + let q_eval = eq_rho_col[..w] + .iter() + .zip(&proof.col_evals) + .map(|(e, v)| *e * *v) + .sum::(); + let col_tidx = ts.len(); + for (i, eval) in proof.col_evals.iter().copied().enumerate() { + ts.observe_ext(eval); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: round_idx * 1_000_000 + 10_000 + i, + tidx: col_tidx + i * D_EF, + value: eval, + is_sample: false, + is_ext: true, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: true, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + } + for bit_idx in 0..rho_col.len() { + if let Some(record) = preflight + .pcs + .sumcheck_rounds + .iter_mut() + .find(|record| record.idx == round_idx * 2 && record.round == log_h + bit_idx) + { + record.fold_challenge_lookup_count += w; + } + } + let f_tidx = ts.len(); + ts.observe_ext(proof.f_at_rho); + let assist_sumcheck_idx = round_idx * 2 + 1; + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: assist_sumcheck_idx, + tidx: f_tidx, + value: proof.f_at_rho, + is_sample: false, + is_ext: true, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: true, + }); + preflight + .pcs + .jagged_assist_inputs + .push(PcsJaggedAssistInputRecord { + proof_idx: 0, + round_idx, + sumcheck_idx: assist_sumcheck_idx, + f_tidx, + f_at_rho: proof.f_at_rho, + }); + let mut q_acc = RecursionField::ZERO; + for (i, (&eq_rho_col, &col_eval)) in + eq_rho_col[..w].iter().zip(&proof.col_evals).enumerate() + { + let eq_product = push_eq_product_records( + preflight, + PcsEqProductKind::JaggedQEval, + PcsEqProductSource::FoldChallenge, + round_idx, + i, + rho_col, + 0, + round_idx * 2, + log_h, + ); + if eq_product != eq_rho_col { + bail!("jagged q-eval eq product replay mismatch"); + } + let acc_in = q_acc; + q_acc += eq_rho_col * col_eval; + preflight.pcs.jagged_q_evals.push(PcsJaggedQEvalRecord { + proof_idx: 0, + round_idx, + sumcheck_idx: round_idx * 2, + term_idx: i, + is_first: i == 0, + is_last: i + 1 == w, + col_tidx: col_tidx + i * D_EF, + col_eval, + eq_rho_col, + acc_in, + acc_out: q_acc, + q_eval, + f_at_rho: proof.f_at_rho, + sumcheck_final: expected, + }); + } + + let mut z_row_padded = point.to_vec(); + z_row_padded.resize(n_robp, RecursionField::ZERO); + let mut rho_padded = rho.clone(); + rho_padded.resize(n_robp, RecursionField::ZERO); + let (assist_point, assist_expected) = replay_degree2_sumcheck( + 0, + assist_sumcheck_idx, + proof.f_at_rho, + &proof.assist_proof, + 2 * n_robp, + false, + preflight, + ts, + )?; + let rho_star_c = (0..n_robp).map(|i| assist_point[2 * i]).collect_vec(); + let rho_star_d = (0..n_robp).map(|i| assist_point[2 * i + 1]).collect_vec(); + for robp_idx in 0..n_robp { + if robp_idx < rho.len() + && let Some(record) = preflight + .pcs + .sumcheck_rounds + .iter_mut() + .find(|record| record.idx == round_idx * 2 && record.round == robp_idx) + { + record.fold_challenge_lookup_count += 1; + } + for assist_round in [2 * robp_idx, 2 * robp_idx + 1] { + if let Some(record) = preflight.pcs.sumcheck_rounds.iter_mut().find(|record| { + record.idx == assist_sumcheck_idx && record.round == assist_round + }) { + record.fold_challenge_lookup_count += 1 + num_polys; + } + } + } + let h_at_rho_star = + push_jagged_assist_h_records(preflight, round_idx, point, &rho, &assist_point, n_robp); + let native_h_at_rho_star = evaluate_g(&z_row_padded, &rho_padded, &rho_star_c, &rho_star_d); + if h_at_rho_star != native_h_at_rho_star { + bail!("jagged assist h replay mismatch"); + } + let q_at_rho_star = push_jagged_assist_q_records( + preflight, + round_idx, + commitment_kind, + &assist_point, + &eq_col, + &comm.cumulative_heights, + n_robp, + ); + let native_q_at_rho_star = + compute_q_at_assist_point(&assist_point, &eq_col, &comm.cumulative_heights, n_robp); + if q_at_rho_star != native_q_at_rho_star { + bail!("jagged assist q replay mismatch"); + } + preflight.pcs.jagged_assists.push(PcsJaggedAssistRecord { + proof_idx: 0, + round_idx, + sumcheck_idx: round_idx * 2 + 1, + h_at_rho_star, + q_at_rho_star, + sumcheck_final: assist_expected, + }); + + Ok(( + comm.inner.clone(), + inner_verify_openings_for_col_evals(log_h, rho_row, &proof.col_evals, col_tidx), + )) + } + + fn replay_basefold( + &self, + rounds: &[BasefoldRound], + proof: &mpcs::basefold::structure::BasefoldProof, + preflight: &mut Preflight, + ts: &mut TS, + ) -> Result<()> + where + TS: FiatShamirTranscript + + TranscriptHistory, + { + if proof.sumcheck_proof.is_none() { + bail!("basefold sumcheck proof is missing"); + } + let total_num_polys = rounds + .iter() + .flat_map(|(_, openings)| openings) + .map(|(_, (_, evals))| evals.len()) + .sum::(); + observe_label_with_records(ts, b"batch coeffs", preflight, 2_400_000); + let batch_coeffs = sample_challenge_pows(ts, total_num_polys, preflight); + let max_num_var = rounds + .iter() + .flat_map(|(_, openings)| openings.iter().map(|(num_vars, _)| *num_vars)) + .max() + .unwrap_or(0); + let basecode_log = >::get_basecode_msg_size_log(); + ensure_supported_basecode_log(basecode_log)?; + if max_num_var < basecode_log { + return Ok(()); + } + let num_rounds = max_num_var - basecode_log; + let sumcheck_messages = proof.sumcheck_proof.as_ref().unwrap(); + if sumcheck_messages.len() != num_rounds || proof.commits.len() != num_rounds { + bail!("basefold round count mismatch"); + } + let mut batch_coeffs_iter = batch_coeffs.iter().copied().enumerate(); + let mut expected_sum = RecursionField::ZERO; + let mut term_idx = 0usize; + for (_, openings) in rounds { + for (num_var, (_, evals)) in openings + .iter() + .filter(|(num_var, _)| *num_var >= basecode_log) + { + let scale = RecursionField::from(F::from_usize(1 << (max_num_var - num_var))); + for (eval, eval_tidx) in evals { + let (global_coeff_idx, coeff) = batch_coeffs_iter + .next() + .ok_or_else(|| eyre::eyre!("basefold missing batch coefficient"))?; + if let Some(record) = preflight.pcs.batch_coeffs.get_mut(global_coeff_idx) { + record.lookup_count += 1; + } + let acc_in = expected_sum; + expected_sum += coeff * *eval * scale; + preflight + .pcs + .basefold_initial_claims + .push(PcsBasefoldInitialClaimRecord { + proof_idx: 0, + sumcheck_idx: 8_000_000, + term_idx, + is_first: term_idx == 0, + is_last: false, + global_coeff_idx, + eval_tidx: *eval_tidx, + eval: *eval, + coeff, + scale, + acc_in, + acc_out: expected_sum, + }); + term_idx += 1; + } + } + } + if let Some(last) = preflight.pcs.basefold_initial_claims.last_mut() { + last.is_last = true; + } + + let mut fold_challenges = Vec::with_capacity(num_rounds); + let mut current_claim = expected_sum; + for (round, msg) in sumcheck_messages.iter().enumerate() { + if msg.evaluations.len() != 2 { + bail!( + "basefold degree-2 sumcheck round {round} has {} evals", + msg.evaluations.len() + ); + } + let ev_tidx = ts.len(); + ts.observe_ext(msg.evaluations[0]); + ts.observe_ext(msg.evaluations[1]); + observe_label_with_records(ts, b"commit round", preflight, 3_400_000 + round * 100); + let challenge_tidx = ts.len(); + let challenge = FiatShamirTranscript::::sample_ext(ts); + fold_challenges.push(challenge); + let eval_1 = msg.evaluations[0]; + let eval_0 = current_claim - eval_1; + let claim_out = extrapolate_uni_poly(eval_0, &msg.evaluations, challenge); + preflight.pcs.sumcheck_rounds.push(PcsSumcheckRoundRecord { + proof_idx: 0, + idx: 8_000_000, + round, + is_first: round == 0, + is_last: round + 1 == num_rounds, + ev_tidx, + challenge_tidx, + ev1: msg.evaluations[0], + ev2: msg.evaluations[1], + claim_in: current_claim, + claim_out, + challenge, + fold_challenge_lookup_count: + >::get_number_queries(), + }); + current_claim = claim_out; + observe_digest( + ts, + &proof.commits[round], + preflight, + 3_700_000 + round * 100, + ); + } + let final_tidx = ts.len(); + for value in proof.final_message.iter().flatten().copied() { + ts.observe_ext(value); + } + for (i, value) in proof.final_message.iter().flatten().copied().enumerate() { + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: 4_000_000 + i, + tidx: final_tidx + i * D_EF, + value, + is_sample: false, + is_ext: true, + is_final_message: true, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + } + let pow_bits = child_pow_bits(); + if pow_bits > 0 && !check_witness_with_records(ts, pow_bits, proof.pow_witness, preflight) { + bail!("basefold pow witness check failed"); + } + observe_label_with_records(ts, b"query indices", preflight, 6_900_000); + let query_bits = max_num_var + + >::get_rate_log(); + let mut queries = Vec::new(); + for query_idx in 0 + ..>::get_number_queries( + ) + { + let sample = + sample_bits_with_record(ts, query_bits, preflight, 7_000_000 + query_idx, true); + preflight + .pcs + .basefold_query_indices + .push(PcsBasefoldQueryIndexRecord { + proof_idx: 0, + query_idx, + sample_tidx: sample.tidx, + query_bits, + sampled_value: RecursionField::from(sample.sampled_value), + query_value: sample.query_value, + high_value: sample.high_value, + query_bytes: u32_bytes(sample.query_value), + high_bytes: u32_bytes(sample.high_value), + query_bit_selectors: one_hot_32(query_bits), + }); + queries.push(sample.query_value); + } + record_basefold_query_checks( + rounds, + proof, + preflight, + &batch_coeffs, + &fold_challenges, + &queries, + max_num_var, + final_tidx, + )?; + let final_expected = push_basefold_final_expected_records( + preflight, + rounds, + &proof.final_message, + &fold_challenges, + basecode_log, + final_tidx, + )?; + preflight + .pcs + .basefold_final_claims + .push(PcsBasefoldFinalClaimRecord { + proof_idx: 0, + sumcheck_idx: 8_000_000, + final_claim: current_claim, + expected: final_expected, + }); + Ok(()) + } +} + +impl AirModule for PcsModule { + fn num_airs(&self) -> usize { + 26 + } + + fn airs>(&self) -> Vec> { + vec![ + Arc::new(PcsCommitmentRootAir { + commitment_root_bus: self.commitment_root_bus, + }) as AirRef<_>, + Arc::new(PcsBaseInputLeafHashAir { + poseidon2_permute_bus: self.poseidon2_permute_bus, + base_input_opening_bus: self.base_input_opening_bus, + merkle_verify_bus: self.merkle_verify_bus, + }) as AirRef<_>, + Arc::new(PcsBaseInputMerkleAir { + poseidon2_compress_bus: self.poseidon2_compress_bus, + merkle_verify_bus: self.merkle_verify_bus, + commitment_root_bus: self.commitment_root_bus, + }) as AirRef<_>, + Arc::new(PcsCommitPhaseLeafHashAir { + poseidon2_permute_bus: self.poseidon2_permute_bus, + commit_phase_leaf_bus: self.commit_phase_leaf_bus, + merkle_verify_bus: self.merkle_verify_bus, + }) as AirRef<_>, + Arc::new(PcsCommitPhaseMerkleAir { + poseidon2_compress_bus: self.poseidon2_compress_bus, + merkle_verify_bus: self.merkle_verify_bus, + commitment_root_bus: self.commitment_root_bus, + }) as AirRef<_>, + Arc::new(PcsOpeningPointAir { + main_global_point_bus: self.main_global_point_bus, + }) as AirRef<_>, + Arc::new(PcsBasefoldFinalPointAir { + fold_challenge_bus: self.fold_challenge_bus, + final_point_bus: self.basefold_final_point_bus, + }) as AirRef<_>, + Arc::new(PcsOpeningEvalAir { + main_eval_bus: self.main_eval_bus, + opening_eval_bus: self.opening_eval_bus, + }) as AirRef<_>, + Arc::new(PcsEqProductAir { + transcript_ext_bus: self.transcript_ext_bus, + fold_challenge_bus: self.fold_challenge_bus, + eq_product_bus: self.eq_product_bus, + }) as AirRef<_>, + Arc::new(PcsSuffixProductAir { + main_global_point_bus: self.main_global_point_bus, + suffix_product_bus: self.suffix_product_bus, + }) as AirRef<_>, + Arc::new(PcsJaggedAssistHAir { + main_global_point_bus: self.main_global_point_bus, + fold_challenge_bus: self.fold_challenge_bus, + jagged_assist_h_bus: self.jagged_assist_h_bus, + }) as AirRef<_>, + Arc::new(PcsJaggedAssistQAir { + fold_challenge_bus: self.fold_challenge_bus, + eq_product_bus: self.eq_product_bus, + commit_height_bus: self.commit_height_bus, + jagged_assist_q_bus: self.jagged_assist_q_bus, + }) as AirRef<_>, + Arc::new(PcsBasefoldQueryIndexAir { + query_sample_bus: self.query_sample_bus, + basefold_query_bus: self.basefold_query_bus, + range_checker_bus: self.range_checker_bus, + }) as AirRef<_>, + Arc::new(PcsBasefoldQueryOpenAir { + basefold_query_bus: self.basefold_query_bus, + base_input_opening_bus: self.base_input_opening_bus, + batch_coeff_bus: self.batch_coeff_bus, + }) as AirRef<_>, + Arc::new(PcsBasefoldCommitPhaseQueryAir { + basefold_query_bus: self.basefold_query_bus, + commit_phase_leaf_bus: self.commit_phase_leaf_bus, + fold_challenge_bus: self.fold_challenge_bus, + }) as AirRef<_>, + Arc::new(PcsBasefoldFinalCodewordAir { + basefold_query_bus: self.basefold_query_bus, + final_message_bus: self.final_message_bus, + }) as AirRef<_>, + Arc::new(PcsTranscriptValueAir { + transcript_bus: self.transcript_bus, + final_message_bus: self.final_message_bus, + query_sample_bus: self.query_sample_bus, + batch_alpha_bus: self.batch_alpha_bus, + basefold_eval_bus: self.basefold_eval_bus, + transcript_ext_bus: self.transcript_ext_bus, + jagged_f_eval_bus: self.jagged_f_eval_bus, + final_message_lookup_count: + >::get_number_queries() + 1, + }) as AirRef<_>, + Arc::new(PcsBasefoldInitialClaimAir { + basefold_eval_bus: self.basefold_eval_bus, + batch_coeff_bus: self.batch_coeff_bus, + sumcheck_input_bus: self.sumcheck_input_bus, + }) as AirRef<_>, + Arc::new(PcsJaggedAssistInputAir { + jagged_f_eval_bus: self.jagged_f_eval_bus, + sumcheck_input_bus: self.sumcheck_input_bus, + }) as AirRef<_>, + Arc::new(PcsJaggedClaimAir { + sumcheck_input_bus: self.sumcheck_input_bus, + opening_eval_bus: self.opening_eval_bus, + eq_product_bus: self.eq_product_bus, + suffix_product_bus: self.suffix_product_bus, + }) as AirRef<_>, + Arc::new(PcsSumcheckInputAir { + sumcheck_input_bus: self.sumcheck_input_bus, + }) as AirRef<_>, + Arc::new(PcsSumcheckAir { + sumcheck_input_bus: self.sumcheck_input_bus, + transcript_bus: self.transcript_bus, + sumcheck_output_bus: self.sumcheck_output_bus, + fold_challenge_bus: self.fold_challenge_bus, + }) as AirRef<_>, + Arc::new(PcsBatchCoeffAir { + batch_alpha_bus: self.batch_alpha_bus, + batch_coeff_bus: self.batch_coeff_bus, + }) as AirRef<_>, + Arc::new(PcsJaggedQEvalAir { + basefold_eval_bus: self.basefold_eval_bus, + eq_product_bus: self.eq_product_bus, + sumcheck_output_bus: self.sumcheck_output_bus, + }) as AirRef<_>, + Arc::new(PcsJaggedAssistAir { + sumcheck_output_bus: self.sumcheck_output_bus, + jagged_assist_h_bus: self.jagged_assist_h_bus, + jagged_assist_q_bus: self.jagged_assist_q_bus, + }) as AirRef<_>, + Arc::new(PcsBasefoldFinalExpectedAir { + fold_challenge_bus: self.fold_challenge_bus, + final_message_bus: self.final_message_bus, + final_point_bus: self.basefold_final_point_bus, + final_expected_bus: self.basefold_final_expected_bus, + }) as AirRef<_>, + Arc::new(PcsBasefoldFinalClaimAir { + sumcheck_output_bus: self.sumcheck_output_bus, + final_expected_bus: self.basefold_final_expected_bus, + }) as AirRef<_>, + ] + } +} + +impl> TraceGenModule> for PcsModule { + type ModuleSpecificCtx<'a> = (); + + fn generate_proving_ctxs( + &self, + _child_vk: &RecursionVk, + _proofs: &[RecursionProof], + preflights: &[Preflight], + _ctx: &Self::ModuleSpecificCtx<'_>, + required_heights: Option<&[usize]>, + ) -> Option>>> { + let mut transcript_values = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .transcript_values + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + transcript_values.sort_by_key(|record| (record.proof_idx, record.tidx)); + let mut sumcheck_rounds = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .sumcheck_rounds + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + sumcheck_rounds.sort_by_key(|record| (record.proof_idx, record.idx, record.round)); + let mut sumcheck_inputs = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .sumcheck_inputs + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + sumcheck_inputs.sort_by_key(|record| (record.proof_idx, record.idx)); + let mut eq_products = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs.eq_products.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + eq_products.sort_by_key(|record| { + ( + record.proof_idx, + record.kind.as_usize(), + record.source.as_usize(), + record.round_idx, + record.term_idx, + record.bit_idx, + ) + }); + let mut suffix_products = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .suffix_products + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + suffix_products.sort_by_key(|record| { + ( + record.proof_idx, + record.round_idx, + record.term_idx, + record.step_idx, + ) + }); + let mut jagged_assist_h = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .jagged_assist_h + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + jagged_assist_h.sort_by_key(|record| (record.proof_idx, record.round_idx, record.step_idx)); + let mut jagged_assist_q = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .jagged_assist_q + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + jagged_assist_q.sort_by_key(|record| { + ( + record.proof_idx, + record.round_idx, + record.term_idx, + record.step_idx, + ) + }); + let mut basefold_initial_claims = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_initial_claims + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_initial_claims + .sort_by_key(|record| (record.proof_idx, record.sumcheck_idx, record.term_idx)); + let mut jagged_assist_inputs = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .jagged_assist_inputs + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + jagged_assist_inputs.sort_by_key(|record| (record.proof_idx, record.sumcheck_idx)); + let mut jagged_claims = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs.jagged_claims.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + jagged_claims + .sort_by_key(|record| (record.proof_idx, record.sumcheck_idx, record.term_idx)); + let mut batch_coeffs = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs.batch_coeffs.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + batch_coeffs.sort_by_key(|record| (record.proof_idx, record.global_coeff_idx)); + let mut jagged_q_evals = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs.jagged_q_evals.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + jagged_q_evals.sort_by_key(|record| (record.proof_idx, record.round_idx, record.term_idx)); + let mut jagged_assists = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs.jagged_assists.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + jagged_assists.sort_by_key(|record| (record.proof_idx, record.round_idx)); + let mut basefold_final_expected = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_final_expected + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_final_expected + .sort_by_key(|record| (record.proof_idx, record.sumcheck_idx, record.term_idx)); + let mut basefold_final_points = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_final_points + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_final_points.sort_by_key(|record| { + ( + record.proof_idx, + record.sumcheck_idx, + record.point_idx, + record.coord_idx, + ) + }); + let mut basefold_final_claims = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_final_claims + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_final_claims.sort_by_key(|record| (record.proof_idx, record.sumcheck_idx)); + let mut opening_points = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs.opening_points.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + opening_points.sort_by_key(|record| { + ( + record.proof_idx, + record.opening_idx, + record.coord_idx, + record.global_round_idx, + ) + }); + let mut opening_evals = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs.opening_evals.iter().cloned().map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + opening_evals.sort_by_key(|record| { + ( + record.proof_idx, + record.opening_idx, + record.commit_kind.as_usize(), + record.eval_idx, + ) + }); + let mut basefold_query_indices = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_query_indices + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_query_indices.sort_by_key(|record| (record.proof_idx, record.query_idx)); + let mut basefold_query_opens = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_query_opens + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_query_opens.sort_by_key(|record| { + ( + record.proof_idx, + record.query_idx, + record.log2_height, + record.opening_idx, + record.value_idx, + record.elem_idx, + ) + }); + let mut basefold_commit_phase_queries = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_commit_phase_queries + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_commit_phase_queries + .sort_by_key(|record| (record.proof_idx, record.query_idx, record.round)); + let mut basefold_final_codeword = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .basefold_final_codeword + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + basefold_final_codeword + .sort_by_key(|record| (record.proof_idx, record.query_idx, record.final_tidx)); + let mut commitment_roots = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .commitment_roots + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + commitment_roots + .sort_by_key(|record| (record.proof_idx, record.commit_major, record.commit_minor)); + let mut base_input_leaf_hashes = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .base_input_leaf_hashes + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + base_input_leaf_hashes.sort_by_key(|record| { + ( + record.proof_idx, + record.query_idx, + record.opening_idx, + record.block_idx, + ) + }); + let mut base_input_merkle_rows = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .base_input_merkle_rows + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + base_input_merkle_rows.sort_by_key(|record| { + ( + record.proof_idx, + record.query_idx, + record.opening_idx, + record.step, + ) + }); + let mut commit_phase_leaf_hashes = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .commit_phase_leaf_hashes + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + commit_phase_leaf_hashes + .sort_by_key(|record| (record.proof_idx, record.query_idx, record.round)); + let mut commit_phase_merkle_rows = preflights + .iter() + .enumerate() + .flat_map(|(proof_idx, p)| { + p.pcs + .commit_phase_merkle_rows + .iter() + .cloned() + .map(move |mut record| { + record.proof_idx = proof_idx; + record + }) + }) + .collect_vec(); + commit_phase_merkle_rows.sort_by_key(|record| { + ( + record.proof_idx, + record.query_idx, + record.round, + record.step, + ) + }); + + let ctx = PcsTraceCtx { + commitment_roots: &commitment_roots, + base_input_leaf_hashes: &base_input_leaf_hashes, + base_input_merkle_rows: &base_input_merkle_rows, + commit_phase_leaf_hashes: &commit_phase_leaf_hashes, + commit_phase_merkle_rows: &commit_phase_merkle_rows, + opening_points: &opening_points, + basefold_final_points: &basefold_final_points, + opening_evals: &opening_evals, + basefold_query_indices: &basefold_query_indices, + basefold_query_opens: &basefold_query_opens, + basefold_commit_phase_queries: &basefold_commit_phase_queries, + basefold_final_codeword: &basefold_final_codeword, + transcript_values: &transcript_values, + basefold_initial_claims: &basefold_initial_claims, + jagged_assist_inputs: &jagged_assist_inputs, + eq_products: &eq_products, + suffix_products: &suffix_products, + jagged_assist_h: &jagged_assist_h, + jagged_assist_q: &jagged_assist_q, + jagged_claims: &jagged_claims, + sumcheck_inputs: &sumcheck_inputs, + sumcheck_rounds: &sumcheck_rounds, + batch_coeffs: &batch_coeffs, + jagged_q_evals: &jagged_q_evals, + jagged_assists: &jagged_assists, + basefold_final_expected: &basefold_final_expected, + basefold_final_claims: &basefold_final_claims, + }; + [ + PcsModuleChip::CommitmentRoot, + PcsModuleChip::BaseInputLeafHash, + PcsModuleChip::BaseInputMerkle, + PcsModuleChip::CommitPhaseLeafHash, + PcsModuleChip::CommitPhaseMerkle, + PcsModuleChip::OpeningPoint, + PcsModuleChip::BasefoldFinalPoint, + PcsModuleChip::OpeningEval, + PcsModuleChip::EqProduct, + PcsModuleChip::SuffixProduct, + PcsModuleChip::JaggedAssistH, + PcsModuleChip::JaggedAssistQ, + PcsModuleChip::BasefoldQueryIndex, + PcsModuleChip::BasefoldQueryOpen, + PcsModuleChip::BasefoldCommitPhaseQuery, + PcsModuleChip::BasefoldFinalCodeword, + PcsModuleChip::TranscriptValues, + PcsModuleChip::BasefoldInitialClaim, + PcsModuleChip::JaggedAssistInput, + PcsModuleChip::JaggedClaim, + PcsModuleChip::SumcheckInput, + PcsModuleChip::Sumcheck, + PcsModuleChip::BatchCoeff, + PcsModuleChip::JaggedQEval, + PcsModuleChip::JaggedAssist, + PcsModuleChip::BasefoldFinalExpected, + PcsModuleChip::BasefoldFinalClaim, + ] + .into_iter() + .enumerate() + .map(|(idx, chip)| { + chip.generate_proving_ctx( + &ctx, + required_heights.and_then(|heights| heights.get(idx).copied()), + ) + }) + .collect() + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsCommitmentRootCols { + pub is_enabled: T, + pub proof_idx: T, + pub commit_major: T, + pub commit_minor: T, + pub lookup_count: T, + pub root: [T; DIGEST_SIZE], +} + +pub struct PcsCommitmentRootAir { + pub commitment_root_bus: PcsCommitmentRootBus, +} + +impl BaseAir for PcsCommitmentRootAir { + fn width(&self) -> usize { + PcsCommitmentRootCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsCommitmentRootAir {} +impl PartitionedBaseAir for PcsCommitmentRootAir {} + +impl Air for PcsCommitmentRootAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsCommitmentRootCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + self.commitment_root_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsCommitmentRootMessage { + commit_major: local.commit_major.into(), + commit_minor: local.commit_minor.into(), + root: local.root.map(Into::into), + }, + local.is_enabled * local.lookup_count, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBaseInputLeafHashCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub opening_idx: T, + pub block_idx: T, + pub log2_height: T, + pub reduced_index: T, + pub is_first: T, + pub is_last: T, + pub value_is_present: [T; 8], + pub value_idx: [T; 8], + pub elem_idx: [T; 8], + pub values: [[T; D_EF]; 8], + pub state_in: [T; POSEIDON2_WIDTH], + pub input: [T; POSEIDON2_WIDTH], + pub output_state: [T; POSEIDON2_WIDTH], +} + +pub struct PcsBaseInputLeafHashAir { + pub poseidon2_permute_bus: Poseidon2PermuteBus, + pub base_input_opening_bus: PcsBaseInputOpeningBus, + pub merkle_verify_bus: MerkleVerifyBus, +} + +impl BaseAir for PcsBaseInputLeafHashAir { + fn width(&self) -> usize { + PcsBaseInputLeafHashCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBaseInputLeafHashAir {} +impl PartitionedBaseAir for PcsBaseInputLeafHashAir {} + +impl Air for PcsBaseInputLeafHashAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsBaseInputLeafHashCols = (*local_row).borrow(); + let next: &PcsBaseInputLeafHashCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.block_idx); + + for i in 0..8 { + builder.assert_bool(local.value_is_present[i]); + if i > 0 { + builder.when(local.is_enabled).assert_zero( + local.value_is_present[i] * (AB::Expr::ONE - local.value_is_present[i - 1]), + ); + } + builder.when(local.is_enabled).assert_eq( + local.input[i], + local.value_is_present[i] * local.values[i][0] + + (AB::Expr::ONE - local.value_is_present[i]) * local.state_in[i], + ); + for limb in 1..D_EF { + builder + .when(local.is_enabled * local.value_is_present[i]) + .assert_zero(local.values[i][limb]); + } + self.base_input_opening_bus.receive( + builder, + local.proof_idx, + PcsBaseInputOpeningMessage { + query_idx: local.query_idx.into(), + opening_idx: local.opening_idx.into(), + reduced_index: local.reduced_index.into(), + value_idx: local.value_idx[i].into(), + elem_idx: local.elem_idx[i].into(), + log2_height: local.log2_height.into(), + opened_value: local.values[i].map(Into::into), + }, + local.is_enabled * local.value_is_present[i], + ); + } + for i in 8..POSEIDON2_WIDTH { + builder + .when(local.is_enabled) + .assert_eq(local.input[i], local.state_in[i]); + } + for i in 0..POSEIDON2_WIDTH { + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.state_in[i]); + } + + let same_leaf = local.is_enabled * next.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when(same_leaf.clone()) + .assert_eq(next.query_idx, local.query_idx); + builder + .when(same_leaf.clone()) + .assert_eq(next.opening_idx, local.opening_idx); + builder + .when(same_leaf.clone()) + .assert_eq(next.block_idx, local.block_idx + AB::Expr::ONE); + builder + .when(same_leaf.clone()) + .assert_eq(next.log2_height, local.log2_height); + builder + .when(same_leaf.clone()) + .assert_eq(next.reduced_index, local.reduced_index); + assert_array_eq( + &mut builder.when(same_leaf), + next.state_in, + local.output_state, + ); + + self.poseidon2_permute_bus.lookup_key( + builder, + Poseidon2PermuteMessage { + input: local.input.map(Into::into), + output: local.output_state.map(Into::into), + }, + local.is_enabled, + ); + + let digest = core::array::from_fn(|i| local.output_state[i].into()); + self.merkle_verify_bus.send( + builder, + local.proof_idx, + MerkleVerifyBusMessage { + merkle_idx_bit_src: local.reduced_index.into(), + current_idx_bit_src: local.reduced_index.into(), + total_depth: AB::Expr::ZERO, + height: AB::Expr::ZERO, + is_leaf: AB::Expr::ONE, + leaf_sub_idx: AB::Expr::ZERO, + value: digest, + commit_major: AB::Expr::ZERO, + commit_minor: local.opening_idx.into(), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBaseInputMerkleCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub opening_idx: T, + pub step: T, + pub is_first: T, + pub is_last: T, + pub idx_in: T, + pub idx_bit: T, + pub idx_out: T, + pub current: [T; DIGEST_SIZE], + pub sibling: [T; DIGEST_SIZE], + pub left: [T; DIGEST_SIZE], + pub right: [T; DIGEST_SIZE], + pub output: [T; DIGEST_SIZE], +} + +pub struct PcsBaseInputMerkleAir { + pub poseidon2_compress_bus: Poseidon2CompressBus, + pub merkle_verify_bus: MerkleVerifyBus, + pub commitment_root_bus: PcsCommitmentRootBus, +} + +impl BaseAir for PcsBaseInputMerkleAir { + fn width(&self) -> usize { + PcsBaseInputMerkleCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBaseInputMerkleAir {} +impl PartitionedBaseAir for PcsBaseInputMerkleAir {} + +impl Air for PcsBaseInputMerkleAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsBaseInputMerkleCols = (*local_row).borrow(); + let next: &PcsBaseInputMerkleCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder.assert_bool(local.idx_bit); + + builder + .when(local.is_enabled) + .assert_eq(local.idx_in, local.idx_bit + local.idx_out * AB::Expr::TWO); + for i in 0..DIGEST_SIZE { + let expected_left = local.idx_bit * local.sibling[i] + + (AB::Expr::ONE - local.idx_bit) * local.current[i]; + let expected_right = local.idx_bit * local.current[i] + + (AB::Expr::ONE - local.idx_bit) * local.sibling[i]; + builder + .when(local.is_enabled) + .assert_eq(local.left[i], expected_left); + builder + .when(local.is_enabled) + .assert_eq(local.right[i], expected_right); + } + + let same_path = local.is_enabled * next.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when(same_path.clone()) + .assert_eq(next.query_idx, local.query_idx); + builder + .when(same_path.clone()) + .assert_eq(next.opening_idx, local.opening_idx); + builder + .when(same_path.clone()) + .assert_eq(next.step, local.step + AB::Expr::ONE); + builder + .when(same_path.clone()) + .assert_eq(next.idx_in, local.idx_out); + assert_array_eq(&mut builder.when(same_path), next.current, local.output); + + self.merkle_verify_bus.receive( + builder, + local.proof_idx, + MerkleVerifyBusMessage { + merkle_idx_bit_src: local.idx_in.into(), + current_idx_bit_src: local.idx_in.into(), + total_depth: AB::Expr::ZERO, + height: AB::Expr::ZERO, + is_leaf: AB::Expr::ONE, + leaf_sub_idx: AB::Expr::ZERO, + value: local.current.map(Into::into), + commit_major: AB::Expr::ZERO, + commit_minor: local.opening_idx.into(), + }, + local.is_enabled * local.is_first, + ); + self.poseidon2_compress_bus.lookup_key( + builder, + Poseidon2CompressMessage { + input: digests_to_poseidon2_input( + local.left.map(Into::into), + local.right.map(Into::into), + ), + output: local.output.map(Into::into), + }, + local.is_enabled, + ); + self.commitment_root_bus.lookup_key( + builder, + local.proof_idx, + PcsCommitmentRootMessage { + commit_major: AB::Expr::ZERO, + commit_minor: local.opening_idx.into(), + root: local.output.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsCommitPhaseLeafHashCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub round: T, + pub leaf_idx: T, + pub left: [T; D_EF], + pub right: [T; D_EF], + pub input: [T; POSEIDON2_WIDTH], + pub output_state: [T; POSEIDON2_WIDTH], +} + +pub struct PcsCommitPhaseLeafHashAir { + pub poseidon2_permute_bus: Poseidon2PermuteBus, + pub commit_phase_leaf_bus: PcsCommitPhaseLeafBus, + pub merkle_verify_bus: MerkleVerifyBus, +} + +impl BaseAir for PcsCommitPhaseLeafHashAir { + fn width(&self) -> usize { + PcsCommitPhaseLeafHashCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsCommitPhaseLeafHashAir {} +impl PartitionedBaseAir for PcsCommitPhaseLeafHashAir {} + +impl Air for PcsCommitPhaseLeafHashAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsCommitPhaseLeafHashCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + + for i in 0..D_EF { + builder + .when(local.is_enabled) + .assert_eq(local.input[i], local.left[i]); + builder + .when(local.is_enabled) + .assert_eq(local.input[D_EF + i], local.right[i]); + } + for i in 2 * D_EF..POSEIDON2_WIDTH { + builder.when(local.is_enabled).assert_zero(local.input[i]); + } + + self.commit_phase_leaf_bus.receive( + builder, + local.proof_idx, + PcsCommitPhaseLeafMessage { + query_idx: local.query_idx.into(), + round: local.round.into(), + leaf_idx: local.leaf_idx.into(), + left: local.left.map(Into::into), + right: local.right.map(Into::into), + }, + local.is_enabled, + ); + self.poseidon2_permute_bus.lookup_key( + builder, + Poseidon2PermuteMessage { + input: local.input.map(Into::into), + output: local.output_state.map(Into::into), + }, + local.is_enabled, + ); + let digest = core::array::from_fn(|i| local.output_state[i].into()); + self.merkle_verify_bus.send( + builder, + local.proof_idx, + MerkleVerifyBusMessage { + merkle_idx_bit_src: local.leaf_idx.into(), + current_idx_bit_src: local.leaf_idx.into(), + total_depth: AB::Expr::ZERO, + height: AB::Expr::ZERO, + is_leaf: AB::Expr::ONE, + leaf_sub_idx: AB::Expr::ZERO, + value: digest, + commit_major: AB::Expr::from_usize(1), + commit_minor: local.round.into(), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsCommitPhaseMerkleCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub round: T, + pub step: T, + pub is_first: T, + pub is_last: T, + pub idx_in: T, + pub idx_bit: T, + pub idx_out: T, + pub current: [T; DIGEST_SIZE], + pub sibling: [T; DIGEST_SIZE], + pub left: [T; DIGEST_SIZE], + pub right: [T; DIGEST_SIZE], + pub output: [T; DIGEST_SIZE], +} + +pub struct PcsCommitPhaseMerkleAir { + pub poseidon2_compress_bus: Poseidon2CompressBus, + pub merkle_verify_bus: MerkleVerifyBus, + pub commitment_root_bus: PcsCommitmentRootBus, +} + +impl BaseAir for PcsCommitPhaseMerkleAir { + fn width(&self) -> usize { + PcsCommitPhaseMerkleCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsCommitPhaseMerkleAir {} +impl PartitionedBaseAir for PcsCommitPhaseMerkleAir {} + +impl Air for PcsCommitPhaseMerkleAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsCommitPhaseMerkleCols = (*local_row).borrow(); + let next: &PcsCommitPhaseMerkleCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder.assert_bool(local.idx_bit); + + builder + .when(local.is_enabled) + .assert_eq(local.idx_in, local.idx_bit + local.idx_out * AB::Expr::TWO); + for i in 0..DIGEST_SIZE { + let expected_left = local.idx_bit * local.sibling[i] + + (AB::Expr::ONE - local.idx_bit) * local.current[i]; + let expected_right = local.idx_bit * local.current[i] + + (AB::Expr::ONE - local.idx_bit) * local.sibling[i]; + builder + .when(local.is_enabled) + .assert_eq(local.left[i], expected_left); + builder + .when(local.is_enabled) + .assert_eq(local.right[i], expected_right); + } + + let same_path = local.is_enabled * next.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when(same_path.clone()) + .assert_eq(next.query_idx, local.query_idx); + builder + .when(same_path.clone()) + .assert_eq(next.round, local.round); + builder + .when(same_path.clone()) + .assert_eq(next.step, local.step + AB::Expr::ONE); + builder + .when(same_path.clone()) + .assert_eq(next.idx_in, local.idx_out); + assert_array_eq(&mut builder.when(same_path), next.current, local.output); + + self.merkle_verify_bus.receive( + builder, + local.proof_idx, + MerkleVerifyBusMessage { + merkle_idx_bit_src: local.idx_in.into(), + current_idx_bit_src: local.idx_in.into(), + total_depth: AB::Expr::ZERO, + height: AB::Expr::ZERO, + is_leaf: AB::Expr::ONE, + leaf_sub_idx: AB::Expr::ZERO, + value: local.current.map(Into::into), + commit_major: AB::Expr::from_usize(1), + commit_minor: local.round.into(), + }, + local.is_enabled * local.is_first, + ); + self.poseidon2_compress_bus.lookup_key( + builder, + Poseidon2CompressMessage { + input: digests_to_poseidon2_input( + local.left.map(Into::into), + local.right.map(Into::into), + ), + output: local.output.map(Into::into), + }, + local.is_enabled, + ); + self.commitment_root_bus.lookup_key( + builder, + local.proof_idx, + PcsCommitmentRootMessage { + commit_major: AB::Expr::from_usize(1), + commit_minor: local.round.into(), + root: local.output.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsOpeningPointCols { + pub is_enabled: T, + pub proof_idx: T, + pub opening_idx: T, + pub coord_idx: T, + pub global_round_idx: T, + pub value: [T; D_EF], +} + +pub struct PcsOpeningPointAir { + pub main_global_point_bus: MainGlobalPointBus, +} + +impl BaseAir for PcsOpeningPointAir { + fn width(&self) -> usize { + PcsOpeningPointCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsOpeningPointAir {} +impl PartitionedBaseAir for PcsOpeningPointAir {} + +impl Air for PcsOpeningPointAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsOpeningPointCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + self.main_global_point_bus.lookup_key( + builder, + local.proof_idx, + MainGlobalPointMessage { + round_idx: local.global_round_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldFinalPointCols { + pub is_enabled: T, + pub proof_idx: T, + pub sumcheck_idx: T, + pub point_idx: T, + pub coord_idx: T, + pub source_sumcheck_idx: T, + pub source_round: T, + pub value: [T; D_EF], +} + +pub struct PcsBasefoldFinalPointAir { + pub fold_challenge_bus: PcsFoldChallengeBus, + pub final_point_bus: PcsBasefoldFinalPointBus, +} + +impl BaseAir for PcsBasefoldFinalPointAir { + fn width(&self) -> usize { + PcsBasefoldFinalPointCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldFinalPointAir {} +impl PartitionedBaseAir for PcsBasefoldFinalPointAir {} + +impl Air for PcsBasefoldFinalPointAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsBasefoldFinalPointCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.source_sumcheck_idx.into(), + round: local.source_round.into(), + challenge: local.value.map(Into::into), + }, + local.is_enabled, + ); + self.final_point_bus.send( + builder, + local.proof_idx, + PcsBasefoldFinalPointMessage { + sumcheck_idx: local.sumcheck_idx.into(), + point_idx: local.point_idx.into(), + coord_idx: local.coord_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsOpeningEvalCols { + pub is_enabled: T, + pub proof_idx: T, + pub opening_idx: T, + pub commit_kind: T, + pub eval_idx: T, + pub main_idx: T, + pub main_eval_idx: T, + pub value: [T; D_EF], + pub raw_value: [T; D_EF], +} + +pub struct PcsOpeningEvalAir { + pub main_eval_bus: MainEvalBus, + pub opening_eval_bus: PcsOpeningEvalBus, +} + +impl BaseAir for PcsOpeningEvalAir { + fn width(&self) -> usize { + PcsOpeningEvalCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsOpeningEvalAir {} +impl PartitionedBaseAir for PcsOpeningEvalAir {} + +impl Air for PcsOpeningEvalAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsOpeningEvalCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + self.main_eval_bus.lookup_key( + builder, + local.proof_idx, + MainEvalMessage { + idx: local.main_idx.into(), + eval_idx: local.main_eval_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled, + ); + self.opening_eval_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsOpeningEvalMessage { + opening_idx: local.opening_idx.into(), + commit_kind: local.commit_kind.into(), + eval_idx: local.eval_idx.into(), + value: local.raw_value.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsEqProductCols { + pub is_enabled: T, + pub proof_idx: T, + pub kind: T, + pub source: T, + pub round_idx: T, + pub term_idx: T, + pub bit_idx: T, + pub is_first: T, + pub is_last: T, + pub lookup_count: T, + pub point_tidx: T, + pub sumcheck_idx: T, + pub point_round: T, + pub index_bit: T, + pub index_pow2: T, + pub index_acc_in: T, + pub index_acc_out: T, + pub point: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], +} + +pub struct PcsEqProductAir { + pub transcript_ext_bus: PcsTranscriptExtBus, + pub fold_challenge_bus: PcsFoldChallengeBus, + pub eq_product_bus: PcsEqProductBus, +} + +impl BaseAir for PcsEqProductAir { + fn width(&self) -> usize { + PcsEqProductCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsEqProductAir {} +impl PartitionedBaseAir for PcsEqProductAir {} + +impl Air for PcsEqProductAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsEqProductCols = (*local_row).borrow(); + let next: &PcsEqProductCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.source); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder.assert_bool(local.index_bit); + builder + .when(local.is_enabled * (AB::Expr::ONE - local.is_last)) + .assert_zero(local.lookup_count); + + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.bit_idx); + builder + .when(local.is_enabled * local.is_first) + .assert_one(local.index_pow2); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.index_acc_in); + let one = core::array::from_fn(|i| { + if i == 0 { + AB::Expr::ONE + } else { + AB::Expr::ZERO + } + }); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + one.clone(), + ); + + let expected_index_acc = local.index_acc_in + local.index_bit * local.index_pow2; + builder + .when(local.is_enabled) + .assert_eq(local.index_acc_out, expected_index_acc); + builder + .when(local.is_enabled * local.is_last) + .assert_eq(local.index_acc_out, local.term_idx); + + let factor = core::array::from_fn(|i| { + let point = local.point[i].into(); + let one_minus_point = if i == 0 { + AB::Expr::ONE - point.clone() + } else { + -point.clone() + }; + local.index_bit * point + (AB::Expr::ONE - local.index_bit) * one_minus_point + }); + let expected_acc = recursion_circuit::utils::ext_field_multiply(local.acc_in, factor); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.acc_out, + expected_acc, + ); + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.kind, local.kind); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.source, local.source); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.round_idx, local.round_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.term_idx, local.term_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.bit_idx, local.bit_idx + AB::Expr::ONE); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.index_pow2, local.index_pow2 * AB::Expr::TWO); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.index_acc_in, local.index_acc_out); + assert_array_eq( + &mut builder.when_transition().when(continues), + next.acc_in, + local.acc_out, + ); + + let is_transcript = AB::Expr::ONE - local.source; + self.transcript_ext_bus.lookup_key( + builder, + local.proof_idx, + PcsTranscriptExtMessage { + tidx: local.point_tidx.into(), + value: local.point.map(Into::into), + }, + local.is_enabled * is_transcript, + ); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.sumcheck_idx.into(), + round: local.point_round.into(), + challenge: local.point.map(Into::into), + }, + local.is_enabled * local.source, + ); + self.eq_product_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsEqProductMessage { + kind: local.kind.into(), + round_idx: local.round_idx.into(), + term_idx: local.term_idx.into(), + value: local.acc_out.map(Into::into), + }, + local.is_enabled * local.lookup_count, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsSuffixProductCols { + pub is_enabled: T, + pub proof_idx: T, + pub round_idx: T, + pub term_idx: T, + pub coord_idx: T, + pub step_idx: T, + pub is_first: T, + pub is_last: T, + pub has_factor: T, + pub point: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], +} + +pub struct PcsSuffixProductAir { + pub main_global_point_bus: MainGlobalPointBus, + pub suffix_product_bus: PcsSuffixProductBus, +} + +impl BaseAir for PcsSuffixProductAir { + fn width(&self) -> usize { + PcsSuffixProductCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsSuffixProductAir {} +impl PartitionedBaseAir for PcsSuffixProductAir {} + +impl Air for PcsSuffixProductAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsSuffixProductCols = (*local_row).borrow(); + let next: &PcsSuffixProductCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder.assert_bool(local.has_factor); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.step_idx); + let one = core::array::from_fn(|i| { + if i == 0 { + AB::Expr::ONE + } else { + AB::Expr::ZERO + } + }); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + one.clone(), + ); + + let factor = core::array::from_fn(|i| { + let point = local.point[i].into(); + let one_minus_point = if i == 0 { + AB::Expr::ONE - point.clone() + } else { + -point.clone() + }; + local.has_factor * one_minus_point + (AB::Expr::ONE - local.has_factor) * one[i].clone() + }); + let expected_acc = recursion_circuit::utils::ext_field_multiply(local.acc_in, factor); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.acc_out, + expected_acc, + ); + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.round_idx, local.round_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.term_idx, local.term_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.coord_idx, local.coord_idx + AB::Expr::ONE); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.step_idx, local.step_idx + AB::Expr::ONE); + assert_array_eq( + &mut builder.when_transition().when(continues), + next.acc_in, + local.acc_out, + ); + + self.main_global_point_bus.lookup_key( + builder, + local.proof_idx, + MainGlobalPointMessage { + round_idx: local.coord_idx.into(), + value: local.point.map(Into::into), + }, + local.is_enabled * local.has_factor, + ); + self.suffix_product_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsSuffixProductMessage { + round_idx: local.round_idx.into(), + term_idx: local.term_idx.into(), + value: local.acc_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsJaggedAssistHCols { + pub is_enabled: T, + pub proof_idx: T, + pub round_idx: T, + pub sumcheck_idx: T, + pub step_idx: T, + pub robp_idx: T, + pub is_first: T, + pub is_last: T, + pub has_z_row: T, + pub has_rho: T, + pub z_row: [T; D_EF], + pub rho: [T; D_EF], + pub rho_star_c: [T; D_EF], + pub rho_star_d: [T; D_EF], + pub val_in: [[T; D_EF]; 4], + pub val_out: [[T; D_EF]; 4], +} + +pub struct PcsJaggedAssistHAir { + pub main_global_point_bus: MainGlobalPointBus, + pub fold_challenge_bus: PcsFoldChallengeBus, + pub jagged_assist_h_bus: PcsJaggedAssistHBus, +} + +impl BaseAir for PcsJaggedAssistHAir { + fn width(&self) -> usize { + PcsJaggedAssistHCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsJaggedAssistHAir {} +impl PartitionedBaseAir for PcsJaggedAssistHAir {} + +impl Air for PcsJaggedAssistHAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsJaggedAssistHCols = (*local_row).borrow(); + let next: &PcsJaggedAssistHCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder.assert_bool(local.has_z_row); + builder.assert_bool(local.has_rho); + builder.when(local.is_enabled).assert_eq( + local.sumcheck_idx, + local.round_idx + local.round_idx + AB::Expr::ONE, + ); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.step_idx); + let zero_ext = core::array::from_fn(|_| AB::Expr::ZERO); + let one_ext = core::array::from_fn(|i| { + if i == 0 { + AB::Expr::ONE + } else { + AB::Expr::ZERO + } + }); + assert_array_eq( + &mut builder.when(local.is_enabled * (AB::Expr::ONE - local.has_z_row)), + local.z_row, + zero_ext.clone(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * (AB::Expr::ONE - local.has_rho)), + local.rho, + zero_ext.clone(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.val_in[0], + zero_ext.clone(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.val_in[1], + one_ext.clone(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.val_in[2], + zero_ext.clone(), + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.val_in[3], + zero_ext.clone(), + ); + + let one_minus = |x: [AB::Var; D_EF]| -> [AB::Expr; D_EF] { + core::array::from_fn(|i| { + let limb = x[i].into(); + if i == 0 { AB::Expr::ONE - limb } else { -limb } + }) + }; + let z1 = local.z_row; + let z2 = local.rho; + let z3 = local.rho_star_c; + let z4 = local.rho_star_d; + let nz1 = one_minus(z1); + let nz2 = one_minus(z2); + let nz3 = one_minus(z3); + let nz4 = one_minus(z4); + let ab00: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_multiply(nz1.clone(), nz2.clone()); + let ab01: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply(nz1.clone(), z2); + let ab10: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply(z1, nz2.clone()); + let ab11: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply(z1, z2); + let cd00: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_multiply(nz3.clone(), nz4.clone()); + let cd01: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply(nz3.clone(), z4); + let cd10: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply(z3, nz4.clone()); + let cd11: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply(z3, z4); + + let transitions: [( + usize, + usize, + [AB::Expr; D_EF], + [AB::Expr; D_EF], + [AB::Expr; D_EF], + ); 4] = [ + ( + 0usize, + 0usize, + ext_add::( + ext_add::( + recursion_circuit::utils::ext_field_multiply::( + ab00.clone(), + cd00.clone(), + ), + recursion_circuit::utils::ext_field_multiply::( + ab01.clone(), + cd11.clone(), + ), + ), + recursion_circuit::utils::ext_field_multiply::( + ab11.clone(), + cd01.clone(), + ), + ), + recursion_circuit::utils::ext_field_multiply::( + ab00.clone(), + cd01.clone(), + ), + ext_add::( + recursion_circuit::utils::ext_field_multiply::( + ab01.clone(), + cd10.clone(), + ), + recursion_circuit::utils::ext_field_multiply::( + ab11.clone(), + cd00.clone(), + ), + ), + ), + ( + 0usize, + 1usize, + recursion_circuit::utils::ext_field_multiply::( + ab10.clone(), + cd10.clone(), + ), + recursion_circuit::utils::ext_field_multiply::( + ab10.clone(), + cd11.clone(), + ), + zero_ext.clone(), + ), + ( + 1usize, + 0usize, + recursion_circuit::utils::ext_field_multiply::( + ab01.clone(), + cd01.clone(), + ), + zero_ext.clone(), + recursion_circuit::utils::ext_field_multiply::( + ab01.clone(), + cd00.clone(), + ), + ), + ( + 1usize, + 1usize, + ext_add::( + ext_add::( + recursion_circuit::utils::ext_field_multiply::( + ab00.clone(), + cd10.clone(), + ), + recursion_circuit::utils::ext_field_multiply::( + ab10.clone(), + cd00.clone(), + ), + ), + recursion_circuit::utils::ext_field_multiply::( + ab11.clone(), + cd11.clone(), + ), + ), + ext_add::( + recursion_circuit::utils::ext_field_multiply::(ab00, cd11), + recursion_circuit::utils::ext_field_multiply::(ab10, cd01), + ), + recursion_circuit::utils::ext_field_multiply::(ab11, cd10), + ), + ]; + let mut expected: [[AB::Expr; D_EF]; 4] = core::array::from_fn(|_| zero_ext.clone()); + for (ci, co, w_same, w_lt1, w_lt0) in transitions { + let v0 = local.val_in[co * 2]; + let v1 = local.val_in[co * 2 + 1]; + let same_v0 = + recursion_circuit::utils::ext_field_multiply::(w_same.clone(), v0); + let lt1_v1 = + recursion_circuit::utils::ext_field_multiply::(w_lt1.clone(), v1); + let lt0_v0 = + recursion_circuit::utils::ext_field_multiply::(w_lt0.clone(), v0); + expected[ci * 2] = ext_add::( + expected[ci * 2].clone(), + ext_add::(same_v0, ext_add::(lt1_v1, lt0_v0)), + ); + + let same_v1 = recursion_circuit::utils::ext_field_multiply::(w_same, v1); + let lt1_v1 = recursion_circuit::utils::ext_field_multiply::(w_lt1, v1); + let lt0_v0 = recursion_circuit::utils::ext_field_multiply::(w_lt0, v0); + expected[ci * 2 + 1] = ext_add::( + expected[ci * 2 + 1].clone(), + ext_add::(same_v1, ext_add::(lt1_v1, lt0_v0)), + ); + } + for (actual, expected) in local.val_out.into_iter().zip(expected) { + assert_array_eq(builder, actual, expected); + } + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.round_idx, local.round_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.sumcheck_idx, local.sumcheck_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.step_idx, local.step_idx + AB::Expr::ONE); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.robp_idx + AB::Expr::ONE, local.robp_idx); + for i in 0..4 { + assert_array_eq( + &mut builder.when_transition().when(continues.clone()), + next.val_in[i], + local.val_out[i], + ); + } + + self.main_global_point_bus.lookup_key( + builder, + local.proof_idx, + MainGlobalPointMessage { + round_idx: local.robp_idx.into(), + value: local.z_row.map(Into::into), + }, + local.is_enabled * local.has_z_row, + ); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: (local.round_idx + local.round_idx).into(), + round: local.robp_idx.into(), + challenge: local.rho.map(Into::into), + }, + local.is_enabled * local.has_rho, + ); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.sumcheck_idx.into(), + round: (local.robp_idx + local.robp_idx).into(), + challenge: local.rho_star_c.map(Into::into), + }, + local.is_enabled, + ); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.sumcheck_idx.into(), + round: (local.robp_idx + local.robp_idx + AB::Expr::ONE).into(), + challenge: local.rho_star_d.map(Into::into), + }, + local.is_enabled, + ); + self.jagged_assist_h_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsJaggedAssistHMessage { + round_idx: local.round_idx.into(), + value: local.val_out[0].map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsJaggedAssistQCols { + pub is_enabled: T, + pub proof_idx: T, + pub round_idx: T, + pub sumcheck_idx: T, + pub commitment_kind: T, + pub term_idx: T, + pub step_idx: T, + pub robp_idx: T, + pub is_first: T, + pub is_last: T, + pub is_first_step: T, + pub is_last_step: T, + pub term_is_last: T, + pub eq_col: [T; D_EF], + pub t_lo: T, + pub t_hi: T, + pub c_bit: T, + pub d_bit: T, + pub bit_pow2: T, + pub c_acc_in: T, + pub c_acc_out: T, + pub d_acc_in: T, + pub d_acc_out: T, + pub rho_star_c: [T; D_EF], + pub rho_star_d: [T; D_EF], + pub term_acc_in: [T; D_EF], + pub term_acc_out: [T; D_EF], + pub q_acc_in: [T; D_EF], + pub q_acc_out: [T; D_EF], +} + +pub struct PcsJaggedAssistQAir { + pub fold_challenge_bus: PcsFoldChallengeBus, + pub eq_product_bus: PcsEqProductBus, + pub commit_height_bus: PcsCommitHeightBus, + pub jagged_assist_q_bus: PcsJaggedAssistQBus, +} + +impl BaseAir for PcsJaggedAssistQAir { + fn width(&self) -> usize { + PcsJaggedAssistQCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsJaggedAssistQAir {} +impl PartitionedBaseAir for PcsJaggedAssistQAir {} + +impl Air for PcsJaggedAssistQAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsJaggedAssistQCols = (*local_row).borrow(); + let next: &PcsJaggedAssistQCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder.assert_bool(local.is_first_step); + builder.assert_bool(local.is_last_step); + builder.assert_bool(local.term_is_last); + builder.assert_bool(local.c_bit); + builder.assert_bool(local.d_bit); + builder.when(local.is_enabled).assert_eq( + local.sumcheck_idx, + local.round_idx + local.round_idx + AB::Expr::ONE, + ); + builder + .when(local.is_enabled) + .assert_eq(local.is_last, local.is_last_step * local.term_is_last); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.term_idx); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.step_idx); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.robp_idx); + let zero_ext = core::array::from_fn(|_| AB::Expr::ZERO); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.q_acc_in, + zero_ext.clone(), + ); + builder + .when(local.is_enabled * local.is_first_step) + .assert_zero(local.step_idx); + builder + .when(local.is_enabled * local.is_first_step) + .assert_zero(local.robp_idx); + builder + .when(local.is_enabled * local.is_first_step) + .assert_one(local.bit_pow2); + builder + .when(local.is_enabled * local.is_first_step) + .assert_zero(local.c_acc_in); + builder + .when(local.is_enabled * local.is_first_step) + .assert_zero(local.d_acc_in); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first_step), + local.term_acc_in, + local.eq_col, + ); + + let expected_c_acc = local.c_acc_in + local.c_bit * local.bit_pow2; + let expected_d_acc = local.d_acc_in + local.d_bit * local.bit_pow2; + builder + .when(local.is_enabled) + .assert_eq(local.c_acc_out, expected_c_acc); + builder + .when(local.is_enabled) + .assert_eq(local.d_acc_out, expected_d_acc); + builder + .when(local.is_enabled * local.is_last_step) + .assert_eq(local.c_acc_out, local.t_lo); + builder + .when(local.is_enabled * local.is_last_step) + .assert_eq(local.d_acc_out, local.t_hi); + + let ext_factor = |bit: AB::Var, point: [AB::Var; D_EF]| -> [AB::Expr; D_EF] { + core::array::from_fn(|i| { + let point_limb = point[i].into(); + let one_minus = if i == 0 { + AB::Expr::ONE - point_limb.clone() + } else { + -point_limb.clone() + }; + bit * point_limb + (AB::Expr::ONE - bit) * one_minus + }) + }; + let c_factor = ext_factor(local.c_bit, local.rho_star_c); + let d_factor = ext_factor(local.d_bit, local.rho_star_d); + let factor: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_multiply(c_factor, d_factor); + let expected_term_acc: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_multiply(local.term_acc_in, factor); + assert_array_eq(builder, local.term_acc_out, expected_term_acc); + + let q_delta = local.term_acc_out.map(|limb| local.is_last_step * limb); + let expected_q_acc = ext_add::(local.q_acc_in, q_delta); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.q_acc_out, + expected_q_acc, + ); + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.round_idx, local.round_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.sumcheck_idx, local.sumcheck_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.commitment_kind, local.commitment_kind); + + let same_term = continues.clone() * (AB::Expr::ONE - local.is_last_step); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.term_idx, local.term_idx); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.step_idx, local.step_idx + AB::Expr::ONE); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.robp_idx, local.robp_idx + AB::Expr::ONE); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.bit_pow2, local.bit_pow2 * AB::Expr::TWO); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.c_acc_in, local.c_acc_out); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.d_acc_in, local.d_acc_out); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.t_lo, local.t_lo); + builder + .when_transition() + .when(same_term.clone()) + .assert_eq(next.t_hi, local.t_hi); + assert_array_eq( + &mut builder.when_transition().when(same_term.clone()), + next.term_acc_in, + local.term_acc_out, + ); + assert_array_eq( + &mut builder.when_transition().when(same_term), + next.q_acc_in, + local.q_acc_out, + ); + + let next_term = continues * local.is_last_step; + builder + .when_transition() + .when(next_term.clone()) + .assert_eq(next.term_idx, local.term_idx + AB::Expr::ONE); + builder + .when_transition() + .when(next_term.clone()) + .assert_one(next.is_first_step); + assert_array_eq( + &mut builder.when_transition().when(next_term), + next.q_acc_in, + local.q_acc_out, + ); + + self.eq_product_bus.lookup_key( + builder, + local.proof_idx, + PcsEqProductMessage { + kind: AB::Expr::from_usize(PcsEqProductKind::JaggedClaim.as_usize()), + round_idx: local.round_idx.into(), + term_idx: local.term_idx.into(), + value: local.eq_col.map(Into::into), + }, + local.is_enabled * local.is_first_step, + ); + self.commit_height_bus.lookup_key( + builder, + local.proof_idx, + PcsCommitHeightMessage { + commitment_kind: local.commitment_kind.into(), + height_idx: local.term_idx.into(), + value: local.t_lo.into(), + }, + local.is_enabled * local.is_first_step, + ); + self.commit_height_bus.lookup_key( + builder, + local.proof_idx, + PcsCommitHeightMessage { + commitment_kind: local.commitment_kind.into(), + height_idx: (local.term_idx + AB::Expr::ONE).into(), + value: local.t_hi.into(), + }, + local.is_enabled * local.is_first_step, + ); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.sumcheck_idx.into(), + round: (local.robp_idx + local.robp_idx).into(), + challenge: local.rho_star_c.map(Into::into), + }, + local.is_enabled, + ); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.sumcheck_idx.into(), + round: (local.robp_idx + local.robp_idx + AB::Expr::ONE).into(), + challenge: local.rho_star_d.map(Into::into), + }, + local.is_enabled, + ); + self.jagged_assist_q_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsJaggedAssistQMessage { + round_idx: local.round_idx.into(), + value: local.q_acc_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldQueryIndexCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub sample_tidx: T, + pub query_bits: T, + pub sampled_value: T, + pub query_value: T, + pub high_value: T, + pub query_bytes: [T; 4], + pub high_bytes: [T; 4], + pub query_bit_selectors: [T; 32], +} + +pub struct PcsBasefoldQueryIndexAir { + pub query_sample_bus: PcsQuerySampleBus, + pub basefold_query_bus: PcsBasefoldQueryBus, + pub range_checker_bus: RangeCheckerBus, +} + +impl BaseAir for PcsBasefoldQueryIndexAir { + fn width(&self) -> usize { + PcsBasefoldQueryIndexCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldQueryIndexAir {} +impl PartitionedBaseAir for PcsBasefoldQueryIndexAir {} + +impl Air for PcsBasefoldQueryIndexAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsBasefoldQueryIndexCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + + let mut query_from_bytes = AB::Expr::ZERO; + let mut high_from_bytes = AB::Expr::ZERO; + let mut coeff = AB::Expr::ONE; + for i in 0..4 { + query_from_bytes += local.query_bytes[i] * coeff.clone(); + high_from_bytes += local.high_bytes[i] * coeff.clone(); + self.range_checker_bus.lookup_key( + builder, + RangeCheckerBusMessage { + value: local.query_bytes[i].into(), + max_bits: AB::Expr::from_usize(8), + }, + local.is_enabled, + ); + self.range_checker_bus.lookup_key( + builder, + RangeCheckerBusMessage { + value: local.high_bytes[i].into(), + max_bits: AB::Expr::from_usize(8), + }, + local.is_enabled, + ); + coeff *= AB::Expr::from_usize(256); + } + builder + .when(local.is_enabled) + .assert_eq(local.query_value, query_from_bytes); + builder + .when(local.is_enabled) + .assert_eq(local.high_value, high_from_bytes); + + let mut bit_selector_sum = AB::Expr::ZERO; + let mut query_bits_from_selectors = AB::Expr::ZERO; + let mut two_pow_bits = AB::Expr::ZERO; + for bit in 0..32 { + builder.assert_bool(local.query_bit_selectors[bit]); + bit_selector_sum += local.query_bit_selectors[bit].into(); + query_bits_from_selectors += local.query_bit_selectors[bit] * AB::Expr::from_usize(bit); + two_pow_bits += local.query_bit_selectors[bit] * AB::Expr::from_usize(1usize << bit); + } + builder.when(local.is_enabled).assert_one(bit_selector_sum); + builder + .when(local.is_enabled) + .assert_eq(local.query_bits, query_bits_from_selectors); + builder.when(local.is_enabled).assert_eq( + local.sampled_value, + local.query_value + local.high_value * two_pow_bits, + ); + + self.query_sample_bus.lookup_key( + builder, + local.proof_idx, + PcsQuerySampleMessage { + tidx: local.sample_tidx.into(), + value: local.sampled_value.into(), + }, + local.is_enabled, + ); + + let query_value = core::array::from_fn(|i| { + if i == 0 { + local.query_value.into() + } else { + AB::Expr::ZERO + } + }); + self.basefold_query_bus.send( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::QueryIndex.as_usize()), + round: AB::Expr::ZERO, + value: query_value, + }, + local.is_enabled, + ); + self.basefold_query_bus.send( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::FinalCodewordStart.as_usize()), + round: AB::Expr::ZERO, + value: [AB::Expr::ZERO; D_EF], + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldQueryOpenCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub opening_idx: T, + pub reduced_index: T, + pub global_coeff_idx: T, + pub value_idx: T, + pub elem_idx: T, + pub log2_height: T, + pub is_last_for_height: T, + pub coeff: [T; D_EF], + pub opened_value: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], +} + +pub struct PcsBasefoldQueryOpenAir { + pub basefold_query_bus: PcsBasefoldQueryBus, + pub base_input_opening_bus: PcsBaseInputOpeningBus, + pub batch_coeff_bus: PcsBatchCoeffBus, +} + +impl BaseAir for PcsBasefoldQueryOpenAir { + fn width(&self) -> usize { + PcsBasefoldQueryOpenCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldQueryOpenAir {} +impl PartitionedBaseAir for PcsBasefoldQueryOpenAir {} + +impl Air for PcsBasefoldQueryOpenAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsBasefoldQueryOpenCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_last_for_height); + + self.batch_coeff_bus.lookup_key( + builder, + local.proof_idx, + PcsBatchCoeffMessage { + global_coeff_idx: local.global_coeff_idx.into(), + coeff: local.coeff.map(Into::into), + }, + local.is_enabled, + ); + + let product = recursion_circuit::utils::ext_field_multiply(local.coeff, local.opened_value); + let expected = ext_add(local.acc_in, product); + assert_array_eq(&mut builder.when(local.is_enabled), local.acc_out, expected); + self.basefold_query_bus.send( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::ReducedOpening.as_usize()), + round: local.log2_height.into(), + value: local.acc_out.map(Into::into), + }, + local.is_enabled * local.is_last_for_height, + ); + self.base_input_opening_bus.send( + builder, + local.proof_idx, + PcsBaseInputOpeningMessage { + query_idx: local.query_idx.into(), + opening_idx: local.opening_idx.into(), + reduced_index: local.reduced_index.into(), + value_idx: local.value_idx.into(), + elem_idx: local.elem_idx.into(), + log2_height: local.log2_height.into(), + opened_value: local.opened_value.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldCommitPhaseQueryCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub round: T, + pub query_value: T, + pub idx_in: T, + pub idx_out: T, + pub log2_height: T, + pub is_first: T, + pub has_reduced_opening: T, + pub folded_idx: T, + pub is_last: T, + pub reduced_opening: [T; D_EF], + pub folded_in: [T; D_EF], + pub sibling_value: [T; D_EF], + pub challenge: [T; D_EF], + pub coeff: [T; D_EF], + pub folded_out: [T; D_EF], +} + +pub struct PcsBasefoldCommitPhaseQueryAir { + pub basefold_query_bus: PcsBasefoldQueryBus, + pub commit_phase_leaf_bus: PcsCommitPhaseLeafBus, + pub fold_challenge_bus: PcsFoldChallengeBus, +} + +impl BaseAir for PcsBasefoldCommitPhaseQueryAir { + fn width(&self) -> usize { + PcsBasefoldCommitPhaseQueryCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldCommitPhaseQueryAir {} +impl PartitionedBaseAir for PcsBasefoldCommitPhaseQueryAir {} + +impl Air for PcsBasefoldCommitPhaseQueryAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsBasefoldCommitPhaseQueryCols = (*local_row).borrow(); + let next_row = main.row_slice(1).expect("next row exists"); + let next: &PcsBasefoldCommitPhaseQueryCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.has_reduced_opening); + builder.assert_bool(local.folded_idx); + builder.assert_bool(local.is_last); + + builder.when(local.is_enabled).assert_eq( + local.idx_in, + local.folded_idx + local.idx_out * AB::Expr::TWO, + ); + let same_query_next = local.is_enabled * next.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when(same_query_next.clone()) + .assert_eq(next.query_idx, local.query_idx); + builder + .when(same_query_next.clone()) + .assert_eq(next.idx_in, local.idx_out); + + let query_value = core::array::from_fn(|i| { + if i == 0 { + local.query_value.into() + } else { + AB::Expr::ZERO + } + }); + self.basefold_query_bus.receive( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::QueryIndex.as_usize()), + round: AB::Expr::ZERO, + value: query_value, + }, + local.is_enabled * local.is_first, + ); + builder + .when(local.is_enabled * local.is_first) + .assert_eq(local.idx_in, local.query_value); + + self.basefold_query_bus.receive( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::ReducedOpening.as_usize()), + round: local.log2_height.into(), + value: local.reduced_opening.map(Into::into), + }, + local.is_enabled * local.has_reduced_opening, + ); + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: AB::Expr::from_usize(8_000_000), + round: local.round.into(), + challenge: local.challenge.map(Into::into), + }, + local.is_enabled, + ); + + let selected = ext_add(local.folded_in, local.reduced_opening); + let left = core::array::from_fn(|i| { + local.folded_idx * local.sibling_value[i] + + (AB::Expr::ONE - local.folded_idx) * selected[i].clone() + }); + let right = core::array::from_fn(|i| { + local.folded_idx * selected[i].clone() + + (AB::Expr::ONE - local.folded_idx) * local.sibling_value[i] + }); + let inv_2_scalar = + ::PrimeSubfield::from_usize(2).inverse(); + let mut inv_2 = core::array::from_fn(|_| AB::Expr::ZERO); + inv_2[0] = AB::Expr::from_prime_subfield(inv_2_scalar); + let leaf_left = left.clone(); + let leaf_right = right.clone(); + let sum_lr: [AB::Expr; D_EF] = ext_add(left.clone(), right.clone()); + let diff_lr: [AB::Expr; D_EF] = ext_sub(left, right); + let lo: [AB::Expr; D_EF] = recursion_circuit::utils::ext_field_multiply(sum_lr, inv_2); + let hi: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_multiply(diff_lr, local.coeff); + let diff: [AB::Expr; D_EF] = ext_sub(hi, lo.clone()); + let folded: [AB::Expr; D_EF] = ext_add( + lo, + recursion_circuit::utils::ext_field_multiply(local.challenge, diff), + ); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.folded_out, + folded, + ); + self.commit_phase_leaf_bus.send( + builder, + local.proof_idx, + PcsCommitPhaseLeafMessage { + query_idx: local.query_idx.into(), + round: local.round.into(), + leaf_idx: local.idx_out.into(), + left: leaf_left.map(Into::into), + right: leaf_right.map(Into::into), + }, + local.is_enabled, + ); + self.basefold_query_bus.send( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::FinalFolded.as_usize()), + round: AB::Expr::ZERO, + value: local.folded_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldFinalCodewordCols { + pub is_enabled: T, + pub proof_idx: T, + pub query_idx: T, + pub elem_idx: T, + pub final_tidx: T, + pub is_first: T, + pub is_last: T, + pub final_value: [T; D_EF], + pub coeff: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], + pub folded: [T; D_EF], +} + +pub struct PcsBasefoldFinalCodewordAir { + pub basefold_query_bus: PcsBasefoldQueryBus, + pub final_message_bus: PcsFinalMessageBus, +} + +impl BaseAir for PcsBasefoldFinalCodewordAir { + fn width(&self) -> usize { + PcsBasefoldFinalCodewordCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldFinalCodewordAir {} +impl PartitionedBaseAir for PcsBasefoldFinalCodewordAir {} + +impl Air for PcsBasefoldFinalCodewordAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsBasefoldFinalCodewordCols = (*local_row).borrow(); + let next: &PcsBasefoldFinalCodewordCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + // BasefoldRSParams has basecode_log = 0, so the final codeword is + // reconstructed from one base-code coordinate with coefficient 1. + builder.when(local.is_enabled).assert_zero(local.elem_idx); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.coeff, + [ + AB::Expr::ONE, + AB::Expr::ZERO, + AB::Expr::ZERO, + AB::Expr::ZERO, + ], + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + [AB::Expr::ZERO; D_EF], + ); + self.basefold_query_bus.receive( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::FinalCodewordStart.as_usize()), + round: AB::Expr::ZERO, + value: [AB::Expr::ZERO; D_EF], + }, + local.is_enabled * local.is_first, + ); + + self.final_message_bus.lookup_key( + builder, + local.proof_idx, + PcsFinalMessageMessage { + tidx: local.final_tidx.into(), + value: local.final_value.map(Into::into), + }, + local.is_enabled, + ); + + let product = recursion_circuit::utils::ext_field_multiply(local.coeff, local.final_value); + let expected = ext_add(local.acc_in, product); + assert_array_eq(&mut builder.when(local.is_enabled), local.acc_out, expected); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.acc_out, + local.folded, + ); + let continue_selector = local.is_enabled * (AB::Expr::ONE - local.is_last.into()); + builder + .when(continue_selector.clone()) + .assert_one(next.is_enabled); + builder + .when(continue_selector.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when(continue_selector.clone()) + .assert_eq(next.query_idx, local.query_idx); + builder.when(continue_selector.clone()).assert_eq( + next.final_tidx, + local.final_tidx + AB::Expr::from_usize(D_EF), + ); + assert_array_eq( + &mut builder.when(continue_selector), + next.acc_in, + local.acc_out.map(Into::into), + ); + self.basefold_query_bus.receive( + builder, + local.proof_idx, + PcsBasefoldQueryMessage { + query_idx: local.query_idx.into(), + stage: AB::Expr::from_usize(PcsBasefoldQueryStage::FinalFolded.as_usize()), + round: AB::Expr::ZERO, + value: local.folded.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsTranscriptValueCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub tidx: T, + pub is_sample: T, + pub is_ext: T, + pub is_final_message: T, + pub is_query_sample: T, + pub is_batch_alpha: T, + pub is_basefold_eval: T, + pub is_jagged_f_at_rho: T, + pub transcript_ext_lookup_count: T, + pub value: [T; D_EF], +} + +pub struct PcsTranscriptValueAir { + pub transcript_bus: TranscriptBus, + pub final_message_bus: PcsFinalMessageBus, + pub query_sample_bus: PcsQuerySampleBus, + pub batch_alpha_bus: PcsBatchAlphaBus, + pub basefold_eval_bus: PcsBasefoldEvalBus, + pub transcript_ext_bus: PcsTranscriptExtBus, + pub jagged_f_eval_bus: PcsJaggedFEvalBus, + pub final_message_lookup_count: usize, +} + +impl BaseAir for PcsTranscriptValueAir { + fn width(&self) -> usize { + PcsTranscriptValueCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsTranscriptValueAir {} +impl PartitionedBaseAir for PcsTranscriptValueAir {} + +impl Air for PcsTranscriptValueAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsTranscriptValueCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_sample); + builder.assert_bool(local.is_ext); + builder.assert_bool(local.is_final_message); + builder.assert_bool(local.is_query_sample); + builder.assert_bool(local.is_batch_alpha); + builder.assert_bool(local.is_basefold_eval); + builder.assert_bool(local.is_jagged_f_at_rho); + builder + .when(local.is_final_message) + .assert_one(local.is_enabled); + builder + .when(local.is_final_message) + .assert_zero(local.is_sample); + builder + .when(local.is_final_message) + .assert_one(local.is_ext); + builder + .when(local.is_query_sample) + .assert_one(local.is_enabled); + builder + .when(local.is_query_sample) + .assert_one(local.is_sample); + builder + .when(local.is_query_sample) + .assert_zero(local.is_ext); + for i in 1..D_EF { + builder + .when(local.is_query_sample) + .assert_zero(local.value[i]); + } + for i in 0..D_EF { + let selector = local.is_enabled + * if i == 0 { + AB::Expr::ONE + } else { + local.is_ext.into() + }; + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.tidx + AB::Expr::from_usize(i), + value: local.value[i].into(), + is_sample: local.is_sample.into(), + }, + selector, + ); + } + self.final_message_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsFinalMessageMessage { + tidx: local.tidx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled + * local.is_final_message + * AB::Expr::from_usize(self.final_message_lookup_count), + ); + self.query_sample_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsQuerySampleMessage { + tidx: local.tidx.into(), + value: local.value[0].into(), + }, + local.is_enabled * local.is_query_sample, + ); + self.batch_alpha_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsBatchAlphaMessage { + tidx: local.tidx.into(), + alpha: local.value.map(Into::into), + }, + local.is_enabled * local.is_batch_alpha, + ); + self.basefold_eval_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsBasefoldEvalMessage { + tidx: local.tidx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.is_basefold_eval * AB::Expr::from_usize(2), + ); + self.transcript_ext_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsTranscriptExtMessage { + tidx: local.tidx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.transcript_ext_lookup_count, + ); + self.jagged_f_eval_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsJaggedFEvalMessage { + sumcheck_idx: local.idx.into(), + tidx: local.tidx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.is_jagged_f_at_rho, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsSumcheckCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub round: T, + pub is_first: T, + pub is_last: T, + pub ev_tidx: T, + pub challenge_tidx: T, + pub fold_challenge_lookup_count: T, + pub ev1: [T; D_EF], + pub ev2: [T; D_EF], + pub claim_in: [T; D_EF], + pub claim_out: [T; D_EF], + pub challenge: [T; D_EF], +} + +pub struct PcsSumcheckAir { + pub sumcheck_input_bus: PcsSumcheckInputBus, + pub transcript_bus: TranscriptBus, + pub sumcheck_output_bus: PcsSumcheckOutputBus, + pub fold_challenge_bus: PcsFoldChallengeBus, +} + +impl BaseAir for PcsSumcheckAir { + fn width(&self) -> usize { + PcsSumcheckCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsSumcheckAir {} +impl PartitionedBaseAir for PcsSumcheckAir {} + +impl Air for PcsSumcheckAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local_row, next_row) = ( + main.row_slice(0).expect("main row exists"), + main.row_slice(1).expect("next row exists"), + ); + let local: &PcsSumcheckCols = (*local_row).borrow(); + let next: &PcsSumcheckCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + let same_sumcheck = local.is_enabled + * next.is_enabled + * (AB::Expr::ONE - next.is_first) + * (AB::Expr::ONE - local.is_last); + builder + .when(same_sumcheck.clone()) + .assert_eq(next.idx, local.idx); + builder + .when(same_sumcheck.clone()) + .assert_eq(next.round, local.round + AB::Expr::ONE); + assert_array_eq( + &mut builder.when(same_sumcheck), + local.claim_out, + next.claim_in, + ); + self.sumcheck_input_bus.receive( + builder, + local.proof_idx, + PcsSumcheckInputMessage { + idx: local.idx.into(), + claim: local.claim_in.map(Into::into), + }, + local.is_enabled * local.is_first, + ); + + let ev0 = ext_sub(local.claim_in, local.ev1); + let claim_out = interpolate_quad_at_012(ev0, local.ev1, local.ev2, local.challenge); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.claim_out, + claim_out, + ); + self.sumcheck_output_bus.send( + builder, + local.proof_idx, + PcsSumcheckOutputMessage { + idx: local.idx.into(), + claim: local.claim_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + + for i in 0..D_EF { + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.ev_tidx + AB::Expr::from_usize(i), + value: local.ev1[i].into(), + is_sample: AB::Expr::ZERO, + }, + local.is_enabled, + ); + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.ev_tidx + AB::Expr::from_usize(D_EF + i), + value: local.ev2[i].into(), + is_sample: AB::Expr::ZERO, + }, + local.is_enabled, + ); + self.transcript_bus.receive( + builder, + local.proof_idx, + TranscriptBusMessage { + tidx: local.challenge_tidx + AB::Expr::from_usize(i), + value: local.challenge[i].into(), + is_sample: AB::Expr::ONE, + }, + local.is_enabled, + ); + } + self.fold_challenge_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.idx.into(), + round: local.round.into(), + challenge: local.challenge.map(Into::into), + }, + local.is_enabled * local.fold_challenge_lookup_count, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldInitialClaimCols { + pub is_enabled: T, + pub proof_idx: T, + pub sumcheck_idx: T, + pub term_idx: T, + pub is_first: T, + pub is_last: T, + pub global_coeff_idx: T, + pub eval_tidx: T, + pub eval: [T; D_EF], + pub coeff: [T; D_EF], + pub scale: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], +} + +pub struct PcsBasefoldInitialClaimAir { + pub basefold_eval_bus: PcsBasefoldEvalBus, + pub batch_coeff_bus: PcsBatchCoeffBus, + pub sumcheck_input_bus: PcsSumcheckInputBus, +} + +impl BaseAir for PcsBasefoldInitialClaimAir { + fn width(&self) -> usize { + PcsBasefoldInitialClaimCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldInitialClaimAir {} +impl PartitionedBaseAir for PcsBasefoldInitialClaimAir {} + +impl Air for PcsBasefoldInitialClaimAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsBasefoldInitialClaimCols = (*local_row).borrow(); + let next: &PcsBasefoldInitialClaimCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + let zero = core::array::from_fn(|_| AB::Expr::ZERO); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + zero, + ); + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.sumcheck_idx, local.sumcheck_idx); + builder + .when_transition() + .when(continues) + .assert_eq(next.term_idx, local.term_idx + AB::Expr::ONE); + assert_array_eq( + &mut builder + .when_transition() + .when(local.is_enabled * (AB::Expr::ONE - local.is_last)), + next.acc_in, + local.acc_out, + ); + + let coeff_eval = recursion_circuit::utils::ext_field_multiply(local.coeff, local.eval); + let term = recursion_circuit::utils::ext_field_multiply(coeff_eval, local.scale); + let expected_acc = ext_add(local.acc_in, term); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.acc_out, + expected_acc, + ); + + self.batch_coeff_bus.lookup_key( + builder, + local.proof_idx, + PcsBatchCoeffMessage { + global_coeff_idx: local.global_coeff_idx.into(), + coeff: local.coeff.map(Into::into), + }, + local.is_enabled, + ); + self.basefold_eval_bus.lookup_key( + builder, + local.proof_idx, + PcsBasefoldEvalMessage { + tidx: local.eval_tidx.into(), + value: local.eval.map(Into::into), + }, + local.is_enabled, + ); + self.sumcheck_input_bus.send( + builder, + local.proof_idx, + PcsSumcheckInputMessage { + idx: local.sumcheck_idx.into(), + claim: local.acc_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsJaggedAssistInputCols { + pub is_enabled: T, + pub proof_idx: T, + pub round_idx: T, + pub sumcheck_idx: T, + pub f_tidx: T, + pub f_at_rho: [T; D_EF], +} + +pub struct PcsJaggedAssistInputAir { + pub jagged_f_eval_bus: PcsJaggedFEvalBus, + pub sumcheck_input_bus: PcsSumcheckInputBus, +} + +impl BaseAir for PcsJaggedAssistInputAir { + fn width(&self) -> usize { + PcsJaggedAssistInputCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsJaggedAssistInputAir {} +impl PartitionedBaseAir for PcsJaggedAssistInputAir {} + +impl Air for PcsJaggedAssistInputAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsJaggedAssistInputCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.when(local.is_enabled).assert_eq( + local.sumcheck_idx, + local.round_idx + local.round_idx + AB::Expr::ONE, + ); + self.jagged_f_eval_bus.lookup_key( + builder, + local.proof_idx, + PcsJaggedFEvalMessage { + sumcheck_idx: local.sumcheck_idx.into(), + tidx: local.f_tidx.into(), + value: local.f_at_rho.map(Into::into), + }, + local.is_enabled, + ); + self.sumcheck_input_bus.send( + builder, + local.proof_idx, + PcsSumcheckInputMessage { + idx: local.sumcheck_idx.into(), + claim: local.f_at_rho.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsJaggedClaimCols { + pub is_enabled: T, + pub proof_idx: T, + pub round_idx: T, + pub sumcheck_idx: T, + pub term_idx: T, + pub is_first: T, + pub is_last: T, + pub opening_idx: T, + pub commit_kind: T, + pub eval_idx: T, + pub main_idx: T, + pub main_eval_idx: T, + pub z_col_tidx: T, + pub eval: [T; D_EF], + pub eq_col: [T; D_EF], + pub tail_zero: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], +} + +pub struct PcsJaggedClaimAir { + pub sumcheck_input_bus: PcsSumcheckInputBus, + pub opening_eval_bus: PcsOpeningEvalBus, + pub eq_product_bus: PcsEqProductBus, + pub suffix_product_bus: PcsSuffixProductBus, +} + +impl BaseAir for PcsJaggedClaimAir { + fn width(&self) -> usize { + PcsJaggedClaimCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsJaggedClaimAir {} +impl PartitionedBaseAir for PcsJaggedClaimAir {} + +impl Air for PcsJaggedClaimAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsJaggedClaimCols = (*local_row).borrow(); + let next: &PcsJaggedClaimCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder + .when(local.is_enabled) + .assert_eq(local.sumcheck_idx, local.round_idx + local.round_idx); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.term_idx); + let zero = core::array::from_fn(|_| AB::Expr::ZERO); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + zero, + ); + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.round_idx, local.round_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.sumcheck_idx, local.sumcheck_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.term_idx, local.term_idx + AB::Expr::ONE); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.z_col_tidx, local.z_col_tidx); + assert_array_eq( + &mut builder.when_transition().when(continues), + next.acc_in, + local.acc_out, + ); + + let opened_eval = recursion_circuit::utils::ext_field_multiply(local.tail_zero, local.eval); + let term = recursion_circuit::utils::ext_field_multiply(local.eq_col, opened_eval.clone()); + let expected_acc = ext_add(local.acc_in, term); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.acc_out, + expected_acc, + ); + + self.sumcheck_input_bus.send( + builder, + local.proof_idx, + PcsSumcheckInputMessage { + idx: local.sumcheck_idx.into(), + claim: local.acc_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + self.eq_product_bus.lookup_key( + builder, + local.proof_idx, + PcsEqProductMessage { + kind: AB::Expr::from_usize(PcsEqProductKind::JaggedClaim.as_usize()), + round_idx: local.round_idx.into(), + term_idx: local.term_idx.into(), + value: local.eq_col.map(Into::into), + }, + local.is_enabled, + ); + self.suffix_product_bus.lookup_key( + builder, + local.proof_idx, + PcsSuffixProductMessage { + round_idx: local.round_idx.into(), + term_idx: local.term_idx.into(), + value: local.tail_zero.map(Into::into), + }, + local.is_enabled, + ); + self.opening_eval_bus.lookup_key( + builder, + local.proof_idx, + PcsOpeningEvalMessage { + opening_idx: local.opening_idx.into(), + commit_kind: local.commit_kind.into(), + eval_idx: local.eval_idx.into(), + value: opened_eval.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsSumcheckInputCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub claim: [T; D_EF], +} + +pub struct PcsSumcheckInputAir { + pub sumcheck_input_bus: PcsSumcheckInputBus, +} + +impl BaseAir for PcsSumcheckInputAir { + fn width(&self) -> usize { + PcsSumcheckInputCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsSumcheckInputAir {} +impl PartitionedBaseAir for PcsSumcheckInputAir {} + +impl Air for PcsSumcheckInputAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsSumcheckInputCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + self.sumcheck_input_bus.send( + builder, + local.proof_idx, + PcsSumcheckInputMessage { + idx: local.idx.into(), + claim: local.claim.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBatchCoeffCols { + pub is_enabled: T, + pub proof_idx: T, + pub global_coeff_idx: T, + pub alpha_tidx: T, + pub lookup_count: T, + pub is_first: T, + pub is_last: T, + pub alpha: [T; D_EF], + pub coeff: [T; D_EF], + pub next_coeff: [T; D_EF], +} + +pub struct PcsBatchCoeffAir { + pub batch_alpha_bus: PcsBatchAlphaBus, + pub batch_coeff_bus: PcsBatchCoeffBus, +} + +impl BaseAir for PcsBatchCoeffAir { + fn width(&self) -> usize { + PcsBatchCoeffCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBatchCoeffAir {} +impl PartitionedBaseAir for PcsBatchCoeffAir {} + +impl Air for PcsBatchCoeffAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsBatchCoeffCols = (*local_row).borrow(); + let next: &PcsBatchCoeffCols = (*next_row).borrow(); + + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.global_coeff_idx); + let one = core::array::from_fn(|i| { + if i == 0 { + AB::Expr::ONE + } else { + AB::Expr::ZERO + } + }); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.coeff, + one, + ); + + let expected_next = recursion_circuit::utils::ext_field_multiply(local.coeff, local.alpha); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.next_coeff, + expected_next, + ); + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder.when_transition().when(continues.clone()).assert_eq( + next.global_coeff_idx, + local.global_coeff_idx + AB::Expr::ONE, + ); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.alpha_tidx, local.alpha_tidx); + assert_array_eq( + &mut builder.when_transition().when(continues.clone()), + next.alpha, + local.alpha, + ); + assert_array_eq( + &mut builder.when_transition().when(continues), + next.coeff, + local.next_coeff, + ); + + self.batch_alpha_bus.lookup_key( + builder, + local.proof_idx, + PcsBatchAlphaMessage { + tidx: local.alpha_tidx.into(), + alpha: local.alpha.map(Into::into), + }, + local.is_enabled * local.is_first, + ); + self.batch_coeff_bus.add_key_with_lookups( + builder, + local.proof_idx, + PcsBatchCoeffMessage { + global_coeff_idx: local.global_coeff_idx.into(), + coeff: local.coeff.map(Into::into), + }, + local.is_enabled * local.lookup_count, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsJaggedQEvalCols { + pub is_enabled: T, + pub proof_idx: T, + pub round_idx: T, + pub sumcheck_idx: T, + pub term_idx: T, + pub is_first: T, + pub is_last: T, + pub col_tidx: T, + pub col_eval: [T; D_EF], + pub eq_rho_col: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], + pub q_eval: [T; D_EF], + pub f_at_rho: [T; D_EF], + pub sumcheck_final: [T; D_EF], +} + +pub struct PcsJaggedQEvalAir { + pub basefold_eval_bus: PcsBasefoldEvalBus, + pub eq_product_bus: PcsEqProductBus, + pub sumcheck_output_bus: PcsSumcheckOutputBus, +} + +impl BaseAir for PcsJaggedQEvalAir { + fn width(&self) -> usize { + PcsJaggedQEvalCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsJaggedQEvalAir {} +impl PartitionedBaseAir for PcsJaggedQEvalAir {} + +impl Air for PcsJaggedQEvalAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsJaggedQEvalCols = (*local_row).borrow(); + let next: &PcsJaggedQEvalCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder + .when(local.is_enabled) + .assert_eq(local.sumcheck_idx, local.round_idx + local.round_idx); + builder + .when(local.is_enabled * local.is_first) + .assert_zero(local.term_idx); + let zero = core::array::from_fn(|_| AB::Expr::ZERO); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + zero, + ); + + let term = recursion_circuit::utils::ext_field_multiply(local.eq_rho_col, local.col_eval); + let expected_acc = ext_add(local.acc_in, term); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.acc_out, + expected_acc, + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + local.q_eval, + local.acc_out, + ); + + let continues = local.is_enabled * (AB::Expr::ONE - local.is_last); + builder + .when_transition() + .when(continues.clone()) + .assert_one(next.is_enabled); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.round_idx, local.round_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.sumcheck_idx, local.sumcheck_idx); + builder + .when_transition() + .when(continues.clone()) + .assert_eq(next.term_idx, local.term_idx + AB::Expr::ONE); + assert_array_eq( + &mut builder.when_transition().when(continues), + next.acc_in, + local.acc_out, + ); + + self.basefold_eval_bus.lookup_key( + builder, + local.proof_idx, + PcsBasefoldEvalMessage { + tidx: local.col_tidx.into(), + value: local.col_eval.map(Into::into), + }, + local.is_enabled, + ); + self.eq_product_bus.lookup_key( + builder, + local.proof_idx, + PcsEqProductMessage { + kind: AB::Expr::from_usize(PcsEqProductKind::JaggedQEval.as_usize()), + round_idx: local.round_idx.into(), + term_idx: local.term_idx.into(), + value: local.eq_rho_col.map(Into::into), + }, + local.is_enabled, + ); + let product = recursion_circuit::utils::ext_field_multiply(local.acc_out, local.f_at_rho); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_last), + product, + local.sumcheck_final, + ); + self.sumcheck_output_bus.receive( + builder, + local.proof_idx, + PcsSumcheckOutputMessage { + idx: local.sumcheck_idx.into(), + claim: local.sumcheck_final.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsJaggedAssistCols { + pub is_enabled: T, + pub proof_idx: T, + pub round_idx: T, + pub sumcheck_idx: T, + pub h_at_rho_star: [T; D_EF], + pub q_at_rho_star: [T; D_EF], + pub sumcheck_final: [T; D_EF], +} + +pub struct PcsJaggedAssistAir { + pub sumcheck_output_bus: PcsSumcheckOutputBus, + pub jagged_assist_h_bus: PcsJaggedAssistHBus, + pub jagged_assist_q_bus: PcsJaggedAssistQBus, +} + +impl BaseAir for PcsJaggedAssistAir { + fn width(&self) -> usize { + PcsJaggedAssistCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsJaggedAssistAir {} +impl PartitionedBaseAir for PcsJaggedAssistAir {} + +impl Air for PcsJaggedAssistAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsJaggedAssistCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + let product = + recursion_circuit::utils::ext_field_multiply(local.h_at_rho_star, local.q_at_rho_star); + assert_array_eq( + &mut builder.when(local.is_enabled), + product, + local.sumcheck_final, + ); + self.sumcheck_output_bus.receive( + builder, + local.proof_idx, + PcsSumcheckOutputMessage { + idx: local.sumcheck_idx.into(), + claim: local.sumcheck_final.map(Into::into), + }, + local.is_enabled, + ); + self.jagged_assist_h_bus.lookup_key( + builder, + local.proof_idx, + PcsJaggedAssistHMessage { + round_idx: local.round_idx.into(), + value: local.h_at_rho_star.map(Into::into), + }, + local.is_enabled, + ); + self.jagged_assist_q_bus.lookup_key( + builder, + local.proof_idx, + PcsJaggedAssistQMessage { + round_idx: local.round_idx.into(), + value: local.q_at_rho_star.map(Into::into), + }, + local.is_enabled, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldFinalExpectedCols { + pub is_enabled: T, + pub proof_idx: T, + pub sumcheck_idx: T, + pub term_idx: T, + pub point_idx: T, + pub coord_idx: T, + pub final_tidx: T, + pub is_first: T, + pub is_last: T, + pub is_elem_first: T, + pub is_elem_last: T, + pub has_challenge: T, + pub challenge_round: T, + pub point_value: [T; D_EF], + pub bit_value: T, + pub challenge: [T; D_EF], + pub final_value: [T; D_EF], + pub coeff_in: [T; D_EF], + pub coeff_out: [T; D_EF], + pub acc_in: [T; D_EF], + pub acc_out: [T; D_EF], +} + +pub struct PcsBasefoldFinalExpectedAir { + pub fold_challenge_bus: PcsFoldChallengeBus, + pub final_message_bus: PcsFinalMessageBus, + pub final_point_bus: PcsBasefoldFinalPointBus, + pub final_expected_bus: PcsBasefoldFinalExpectedBus, +} + +impl BaseAir for PcsBasefoldFinalExpectedAir { + fn width(&self) -> usize { + PcsBasefoldFinalExpectedCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldFinalExpectedAir {} +impl PartitionedBaseAir for PcsBasefoldFinalExpectedAir {} + +impl Air for PcsBasefoldFinalExpectedAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let next_row = main.row_slice(1).expect("next row exists"); + let local: &PcsBasefoldFinalExpectedCols = (*local_row).borrow(); + let next: &PcsBasefoldFinalExpectedCols = (*next_row).borrow(); + builder.assert_bool(local.is_enabled); + builder.assert_bool(local.is_first); + builder.assert_bool(local.is_last); + builder.assert_bool(local.is_elem_first); + builder.assert_bool(local.is_elem_last); + builder.assert_bool(local.has_challenge); + builder.assert_bool(local.bit_value); + + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_first), + local.acc_in, + [AB::Expr::ZERO; D_EF], + ); + assert_array_eq( + &mut builder.when(local.is_enabled * local.is_elem_first), + local.coeff_in, + [ + AB::Expr::ONE, + AB::Expr::ZERO, + AB::Expr::ZERO, + AB::Expr::ZERO, + ], + ); + + self.fold_challenge_bus.lookup_key( + builder, + local.proof_idx, + PcsFoldChallengeMessage { + sumcheck_idx: local.sumcheck_idx.into(), + round: local.challenge_round.into(), + challenge: local.challenge.map(Into::into), + }, + local.is_enabled * local.has_challenge, + ); + self.final_point_bus.receive( + builder, + local.proof_idx, + PcsBasefoldFinalPointMessage { + sumcheck_idx: local.sumcheck_idx.into(), + point_idx: local.point_idx.into(), + coord_idx: local.coord_idx.into(), + value: local.point_value.map(Into::into), + }, + local.is_enabled, + ); + + let point = local.point_value.map(Into::into); + let challenge = local.challenge.map(Into::into); + let xi_yi: [AB::Expr; D_EF] = + recursion_circuit::utils::ext_field_multiply(point.clone(), challenge.clone()); + let prefix_factor: [AB::Expr; D_EF] = core::array::from_fn(|i| { + xi_yi[i].clone() + xi_yi[i].clone() - point[i].clone() - challenge[i].clone() + + if i == 0 { + AB::Expr::ONE + } else { + AB::Expr::ZERO + } + }); + let suffix_one_minus: [AB::Expr; D_EF] = ext_sub( + [ + AB::Expr::ONE, + AB::Expr::ZERO, + AB::Expr::ZERO, + AB::Expr::ZERO, + ], + point.clone(), + ); + let suffix_factor: [AB::Expr; D_EF] = core::array::from_fn(|i| { + local.bit_value.into() * point[i].clone() + + (AB::Expr::ONE - local.bit_value.into()) * suffix_one_minus[i].clone() + }); + let factor: [AB::Expr; D_EF] = core::array::from_fn(|i| { + local.has_challenge.into() * prefix_factor[i].clone() + + (AB::Expr::ONE - local.has_challenge.into()) * suffix_factor[i].clone() + }); + let coeff_out = recursion_circuit::utils::ext_field_multiply(local.coeff_in, factor); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.coeff_out, + coeff_out, + ); + + self.final_message_bus.lookup_key( + builder, + local.proof_idx, + PcsFinalMessageMessage { + tidx: local.final_tidx.into(), + value: local.final_value.map(Into::into), + }, + local.is_enabled * local.is_elem_last, + ); + let contribution = + recursion_circuit::utils::ext_field_multiply(local.coeff_out, local.final_value); + let expected_acc: [AB::Expr; D_EF] = core::array::from_fn(|i| { + local.acc_in[i].into() + local.is_elem_last.into() * contribution[i].clone() + }); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.acc_out, + expected_acc, + ); + + let continue_selector = local.is_enabled * (AB::Expr::ONE - local.is_last.into()); + builder + .when(continue_selector.clone()) + .assert_one(next.is_enabled); + builder + .when(continue_selector.clone()) + .assert_eq(next.proof_idx, local.proof_idx); + builder + .when(continue_selector.clone()) + .assert_eq(next.sumcheck_idx, local.sumcheck_idx); + builder + .when(continue_selector.clone()) + .assert_eq(next.term_idx, local.term_idx + AB::Expr::ONE); + assert_array_eq( + &mut builder.when(continue_selector.clone()), + next.acc_in, + local.acc_out.map(Into::into), + ); + assert_array_eq( + &mut builder.when(continue_selector * (AB::Expr::ONE - next.is_elem_first.into())), + next.coeff_in, + local.coeff_out.map(Into::into), + ); + self.final_expected_bus.send( + builder, + local.proof_idx, + PcsBasefoldFinalExpectedMessage { + sumcheck_idx: local.sumcheck_idx.into(), + expected: local.acc_out.map(Into::into), + }, + local.is_enabled * local.is_last, + ); + } +} + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct PcsBasefoldFinalClaimCols { + pub is_enabled: T, + pub proof_idx: T, + pub sumcheck_idx: T, + pub final_claim: [T; D_EF], + pub expected: [T; D_EF], +} + +pub struct PcsBasefoldFinalClaimAir { + pub sumcheck_output_bus: PcsSumcheckOutputBus, + pub final_expected_bus: PcsBasefoldFinalExpectedBus, +} + +impl BaseAir for PcsBasefoldFinalClaimAir { + fn width(&self) -> usize { + PcsBasefoldFinalClaimCols::::width() + } +} + +impl BaseAirWithPublicValues for PcsBasefoldFinalClaimAir {} +impl PartitionedBaseAir for PcsBasefoldFinalClaimAir {} + +impl Air for PcsBasefoldFinalClaimAir +where + AB: AirBuilder + InteractionBuilder, + ::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &PcsBasefoldFinalClaimCols = (*local_row).borrow(); + builder.assert_bool(local.is_enabled); + assert_array_eq( + &mut builder.when(local.is_enabled), + local.final_claim, + local.expected, + ); + self.sumcheck_output_bus.receive( + builder, + local.proof_idx, + PcsSumcheckOutputMessage { + idx: local.sumcheck_idx.into(), + claim: local.final_claim.map(Into::into), + }, + local.is_enabled, + ); + self.final_expected_bus.receive( + builder, + local.proof_idx, + PcsBasefoldFinalExpectedMessage { + sumcheck_idx: local.sumcheck_idx.into(), + expected: local.expected.map(Into::into), + }, + local.is_enabled, + ); + } +} + +struct PcsTraceCtx<'a> { + commitment_roots: &'a [PcsCommitmentRootRecord], + base_input_leaf_hashes: &'a [PcsBaseInputLeafHashRecord], + base_input_merkle_rows: &'a [PcsBaseInputMerkleRecord], + commit_phase_leaf_hashes: &'a [PcsCommitPhaseLeafHashRecord], + commit_phase_merkle_rows: &'a [PcsCommitPhaseMerkleRecord], + opening_points: &'a [PcsOpeningPointRecord], + basefold_final_points: &'a [PcsBasefoldFinalPointRecord], + opening_evals: &'a [PcsOpeningEvalRecord], + eq_products: &'a [PcsEqProductRecord], + suffix_products: &'a [PcsSuffixProductRecord], + jagged_assist_h: &'a [PcsJaggedAssistHRecord], + jagged_assist_q: &'a [PcsJaggedAssistQRecord], + basefold_query_indices: &'a [PcsBasefoldQueryIndexRecord], + basefold_query_opens: &'a [PcsBasefoldQueryOpenRecord], + basefold_commit_phase_queries: &'a [PcsBasefoldCommitPhaseQueryRecord], + basefold_final_codeword: &'a [PcsBasefoldFinalCodewordRecord], + transcript_values: &'a [PcsTranscriptValueRecord], + basefold_initial_claims: &'a [PcsBasefoldInitialClaimRecord], + jagged_assist_inputs: &'a [PcsJaggedAssistInputRecord], + jagged_claims: &'a [PcsJaggedClaimRecord], + sumcheck_inputs: &'a [PcsSumcheckInputRecord], + sumcheck_rounds: &'a [PcsSumcheckRoundRecord], + batch_coeffs: &'a [PcsBatchCoeffRecord], + jagged_q_evals: &'a [PcsJaggedQEvalRecord], + jagged_assists: &'a [PcsJaggedAssistRecord], + basefold_final_expected: &'a [PcsBasefoldFinalExpectedRecord], + basefold_final_claims: &'a [PcsBasefoldFinalClaimRecord], +} + +enum PcsModuleChip { + CommitmentRoot, + BaseInputLeafHash, + BaseInputMerkle, + CommitPhaseLeafHash, + CommitPhaseMerkle, + OpeningPoint, + BasefoldFinalPoint, + OpeningEval, + EqProduct, + SuffixProduct, + JaggedAssistH, + JaggedAssistQ, + BasefoldQueryIndex, + BasefoldQueryOpen, + BasefoldCommitPhaseQuery, + BasefoldFinalCodeword, + TranscriptValues, + BasefoldInitialClaim, + JaggedAssistInput, + JaggedClaim, + SumcheckInput, + Sumcheck, + BatchCoeff, + JaggedQEval, + JaggedAssist, + BasefoldFinalExpected, + BasefoldFinalClaim, +} + +impl RowMajorChip for PcsModuleChip { + type Ctx<'a> = PcsTraceCtx<'a>; + + fn generate_trace( + &self, + ctx: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + match self { + PcsModuleChip::CommitmentRoot => PcsCommitmentRootTraceGenerator + .generate_trace(&ctx.commitment_roots, required_height), + PcsModuleChip::BaseInputLeafHash => PcsBaseInputLeafHashTraceGenerator + .generate_trace(&ctx.base_input_leaf_hashes, required_height), + PcsModuleChip::BaseInputMerkle => PcsBaseInputMerkleTraceGenerator + .generate_trace(&ctx.base_input_merkle_rows, required_height), + PcsModuleChip::CommitPhaseLeafHash => PcsCommitPhaseLeafHashTraceGenerator + .generate_trace(&ctx.commit_phase_leaf_hashes, required_height), + PcsModuleChip::CommitPhaseMerkle => PcsCommitPhaseMerkleTraceGenerator + .generate_trace(&ctx.commit_phase_merkle_rows, required_height), + PcsModuleChip::OpeningPoint => { + PcsOpeningPointTraceGenerator.generate_trace(&ctx.opening_points, required_height) + } + PcsModuleChip::BasefoldFinalPoint => PcsBasefoldFinalPointTraceGenerator + .generate_trace(&ctx.basefold_final_points, required_height), + PcsModuleChip::OpeningEval => { + PcsOpeningEvalTraceGenerator.generate_trace(&ctx.opening_evals, required_height) + } + PcsModuleChip::EqProduct => { + PcsEqProductTraceGenerator.generate_trace(&ctx.eq_products, required_height) + } + PcsModuleChip::SuffixProduct => { + PcsSuffixProductTraceGenerator.generate_trace(&ctx.suffix_products, required_height) + } + PcsModuleChip::JaggedAssistH => { + PcsJaggedAssistHTraceGenerator.generate_trace(&ctx.jagged_assist_h, required_height) + } + PcsModuleChip::JaggedAssistQ => { + PcsJaggedAssistQTraceGenerator.generate_trace(&ctx.jagged_assist_q, required_height) + } + PcsModuleChip::BasefoldQueryIndex => PcsBasefoldQueryIndexTraceGenerator + .generate_trace(&ctx.basefold_query_indices, required_height), + PcsModuleChip::BasefoldQueryOpen => PcsBasefoldQueryOpenTraceGenerator + .generate_trace(&ctx.basefold_query_opens, required_height), + PcsModuleChip::BasefoldCommitPhaseQuery => PcsBasefoldCommitPhaseQueryTraceGenerator + .generate_trace(&ctx.basefold_commit_phase_queries, required_height), + PcsModuleChip::BasefoldFinalCodeword => PcsBasefoldFinalCodewordTraceGenerator + .generate_trace(&ctx.basefold_final_codeword, required_height), + PcsModuleChip::TranscriptValues => PcsTranscriptValueTraceGenerator + .generate_trace(&ctx.transcript_values, required_height), + PcsModuleChip::BasefoldInitialClaim => PcsBasefoldInitialClaimTraceGenerator + .generate_trace(&ctx.basefold_initial_claims, required_height), + PcsModuleChip::JaggedAssistInput => PcsJaggedAssistInputTraceGenerator + .generate_trace(&ctx.jagged_assist_inputs, required_height), + PcsModuleChip::JaggedClaim => { + PcsJaggedClaimTraceGenerator.generate_trace(&ctx.jagged_claims, required_height) + } + PcsModuleChip::SumcheckInput => { + PcsSumcheckInputTraceGenerator.generate_trace(&ctx.sumcheck_inputs, required_height) + } + PcsModuleChip::Sumcheck => { + PcsSumcheckTraceGenerator.generate_trace(&ctx.sumcheck_rounds, required_height) + } + PcsModuleChip::BatchCoeff => { + PcsBatchCoeffTraceGenerator.generate_trace(&ctx.batch_coeffs, required_height) + } + PcsModuleChip::JaggedQEval => { + PcsJaggedQEvalTraceGenerator.generate_trace(&ctx.jagged_q_evals, required_height) + } + PcsModuleChip::JaggedAssist => { + PcsJaggedAssistTraceGenerator.generate_trace(&ctx.jagged_assists, required_height) + } + PcsModuleChip::BasefoldFinalExpected => PcsBasefoldFinalExpectedTraceGenerator + .generate_trace(&ctx.basefold_final_expected, required_height), + PcsModuleChip::BasefoldFinalClaim => PcsBasefoldFinalClaimTraceGenerator + .generate_trace(&ctx.basefold_final_claims, required_height), + } + } +} + +struct PcsOpeningPointTraceGenerator; +struct PcsBasefoldFinalPointTraceGenerator; +struct PcsCommitmentRootTraceGenerator; +struct PcsBaseInputLeafHashTraceGenerator; +struct PcsBaseInputMerkleTraceGenerator; +struct PcsCommitPhaseLeafHashTraceGenerator; +struct PcsCommitPhaseMerkleTraceGenerator; +struct PcsOpeningEvalTraceGenerator; +struct PcsEqProductTraceGenerator; +struct PcsSuffixProductTraceGenerator; +struct PcsJaggedAssistHTraceGenerator; +struct PcsJaggedAssistQTraceGenerator; +struct PcsBasefoldFinalExpectedTraceGenerator; +struct PcsBasefoldQueryIndexTraceGenerator; +struct PcsBasefoldQueryOpenTraceGenerator; +struct PcsBasefoldCommitPhaseQueryTraceGenerator; +struct PcsBasefoldFinalCodewordTraceGenerator; +struct PcsTranscriptValueTraceGenerator; +struct PcsBasefoldInitialClaimTraceGenerator; +struct PcsJaggedAssistInputTraceGenerator; +struct PcsJaggedClaimTraceGenerator; +struct PcsSumcheckInputTraceGenerator; +struct PcsSumcheckTraceGenerator; +struct PcsBatchCoeffTraceGenerator; +struct PcsJaggedQEvalTraceGenerator; +struct PcsJaggedAssistTraceGenerator; +struct PcsBasefoldFinalClaimTraceGenerator; + +impl RowMajorChip for PcsCommitmentRootTraceGenerator { + type Ctx<'a> = &'a [PcsCommitmentRootRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsCommitmentRootCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsCommitmentRootCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.commit_major = F::from_usize(record.commit_major); + cols.commit_minor = F::from_usize(record.commit_minor); + cols.lookup_count = F::from_usize(record.lookup_count); + cols.root = record.root; + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBaseInputLeafHashTraceGenerator { + type Ctx<'a> = &'a [PcsBaseInputLeafHashRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBaseInputLeafHashCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBaseInputLeafHashCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.opening_idx = F::from_usize(record.opening_idx); + cols.block_idx = F::from_usize(record.block_idx); + cols.log2_height = F::from_usize(record.log2_height); + cols.reduced_index = F::from_usize(record.reduced_index); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.value_is_present = record.value_is_present.map(F::from_bool); + cols.value_idx = record.value_idx.map(F::from_usize); + cols.elem_idx = record.elem_idx.map(F::from_usize); + cols.values = record.values.map(ext_limbs); + cols.state_in = record.state_in; + cols.input = record.input; + cols.output_state = record.output_state; + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBaseInputMerkleTraceGenerator { + type Ctx<'a> = &'a [PcsBaseInputMerkleRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBaseInputMerkleCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBaseInputMerkleCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.opening_idx = F::from_usize(record.opening_idx); + cols.step = F::from_usize(record.step); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.idx_in = F::from_usize(record.idx_in); + cols.idx_bit = F::from_usize(record.idx_bit); + cols.idx_out = F::from_usize(record.idx_out); + cols.current = record.current; + cols.sibling = record.sibling; + cols.left = record.left; + cols.right = record.right; + cols.output = record.output; + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsCommitPhaseLeafHashTraceGenerator { + type Ctx<'a> = &'a [PcsCommitPhaseLeafHashRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsCommitPhaseLeafHashCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsCommitPhaseLeafHashCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.round = F::from_usize(record.round); + cols.leaf_idx = F::from_usize(record.leaf_idx); + cols.left = ext_limbs(record.left); + cols.right = ext_limbs(record.right); + cols.input = record.input; + cols.output_state = record.output_state; + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsCommitPhaseMerkleTraceGenerator { + type Ctx<'a> = &'a [PcsCommitPhaseMerkleRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsCommitPhaseMerkleCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsCommitPhaseMerkleCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.round = F::from_usize(record.round); + cols.step = F::from_usize(record.step); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.idx_in = F::from_usize(record.idx_in); + cols.idx_bit = F::from_usize(record.idx_bit); + cols.idx_out = F::from_usize(record.idx_out); + cols.current = record.current; + cols.sibling = record.sibling; + cols.left = record.left; + cols.right = record.right; + cols.output = record.output; + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsOpeningPointTraceGenerator { + type Ctx<'a> = &'a [PcsOpeningPointRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsOpeningPointCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsOpeningPointCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.opening_idx = F::from_usize(record.opening_idx); + cols.coord_idx = F::from_usize(record.coord_idx); + cols.global_round_idx = F::from_usize(record.global_round_idx); + cols.value = ext_limbs(record.value); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldFinalPointTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldFinalPointRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldFinalPointCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldFinalPointCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.point_idx = F::from_usize(record.point_idx); + cols.coord_idx = F::from_usize(record.coord_idx); + cols.source_sumcheck_idx = F::from_usize(record.source_sumcheck_idx); + cols.source_round = F::from_usize(record.source_round); + cols.value = ext_limbs(record.value); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsOpeningEvalTraceGenerator { + type Ctx<'a> = &'a [PcsOpeningEvalRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsOpeningEvalCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsOpeningEvalCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.opening_idx = F::from_usize(record.opening_idx); + cols.commit_kind = F::from_usize(record.commit_kind.as_usize()); + cols.eval_idx = F::from_usize(record.eval_idx); + cols.main_idx = F::from_usize(record.main_idx); + cols.main_eval_idx = F::from_usize(record.main_eval_idx); + cols.value = ext_limbs(record.value); + cols.raw_value = ext_limbs(record.raw_value); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsEqProductTraceGenerator { + type Ctx<'a> = &'a [PcsEqProductRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsEqProductCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsEqProductCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.kind = F::from_usize(record.kind.as_usize()); + cols.source = F::from_usize(record.source.as_usize()); + cols.round_idx = F::from_usize(record.round_idx); + cols.term_idx = F::from_usize(record.term_idx); + cols.bit_idx = F::from_usize(record.bit_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.lookup_count = F::from_usize(record.lookup_count); + cols.point_tidx = F::from_usize(record.point_tidx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.point_round = F::from_usize(record.point_round); + cols.index_bit = F::from_bool(record.index_bit); + cols.index_pow2 = F::from_usize(record.index_pow2); + cols.index_acc_in = F::from_usize(record.index_acc_in); + cols.index_acc_out = F::from_usize(record.index_acc_out); + cols.point = ext_limbs(record.point); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsSuffixProductTraceGenerator { + type Ctx<'a> = &'a [PcsSuffixProductRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsSuffixProductCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsSuffixProductCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.term_idx = F::from_usize(record.term_idx); + cols.coord_idx = F::from_usize(record.coord_idx); + cols.step_idx = F::from_usize(record.step_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.has_factor = F::from_bool(record.has_factor); + cols.point = ext_limbs(record.point); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsJaggedAssistHTraceGenerator { + type Ctx<'a> = &'a [PcsJaggedAssistHRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsJaggedAssistHCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsJaggedAssistHCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.step_idx = F::from_usize(record.step_idx); + cols.robp_idx = F::from_usize(record.robp_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.has_z_row = F::from_bool(record.has_z_row); + cols.has_rho = F::from_bool(record.has_rho); + cols.z_row = ext_limbs(record.z_row); + cols.rho = ext_limbs(record.rho); + cols.rho_star_c = ext_limbs(record.rho_star_c); + cols.rho_star_d = ext_limbs(record.rho_star_d); + cols.val_in = record.val_in.map(ext_limbs); + cols.val_out = record.val_out.map(ext_limbs); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsJaggedAssistQTraceGenerator { + type Ctx<'a> = &'a [PcsJaggedAssistQRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsJaggedAssistQCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsJaggedAssistQCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.commitment_kind = F::from_usize(record.commitment_kind); + cols.term_idx = F::from_usize(record.term_idx); + cols.step_idx = F::from_usize(record.step_idx); + cols.robp_idx = F::from_usize(record.robp_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.is_first_step = F::from_bool(record.is_first_step); + cols.is_last_step = F::from_bool(record.is_last_step); + cols.term_is_last = F::from_bool(record.term_is_last); + cols.eq_col = ext_limbs(record.eq_col); + cols.t_lo = F::from_usize(record.t_lo); + cols.t_hi = F::from_usize(record.t_hi); + cols.c_bit = F::from_bool(record.c_bit); + cols.d_bit = F::from_bool(record.d_bit); + cols.bit_pow2 = F::from_usize(record.bit_pow2); + cols.c_acc_in = F::from_usize(record.c_acc_in); + cols.c_acc_out = F::from_usize(record.c_acc_out); + cols.d_acc_in = F::from_usize(record.d_acc_in); + cols.d_acc_out = F::from_usize(record.d_acc_out); + cols.rho_star_c = ext_limbs(record.rho_star_c); + cols.rho_star_d = ext_limbs(record.rho_star_d); + cols.term_acc_in = ext_limbs(record.term_acc_in); + cols.term_acc_out = ext_limbs(record.term_acc_out); + cols.q_acc_in = ext_limbs(record.q_acc_in); + cols.q_acc_out = ext_limbs(record.q_acc_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldQueryIndexTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldQueryIndexRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldQueryIndexCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldQueryIndexCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.sample_tidx = F::from_usize(record.sample_tidx); + cols.query_bits = F::from_usize(record.query_bits); + cols.sampled_value = record.sampled_value.as_bases()[0]; + cols.query_value = F::from_usize(record.query_value); + cols.high_value = F::from_usize(record.high_value); + cols.query_bytes = record.query_bytes.map(F::from_u8); + cols.high_bytes = record.high_bytes.map(F::from_u8); + cols.query_bit_selectors = record.query_bit_selectors.map(F::from_bool); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldQueryOpenTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldQueryOpenRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldQueryOpenCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldQueryOpenCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.opening_idx = F::from_usize(record.opening_idx); + cols.reduced_index = F::from_usize(record.reduced_index); + cols.global_coeff_idx = F::from_usize(record.global_coeff_idx); + cols.value_idx = F::from_usize(record.value_idx); + cols.elem_idx = F::from_usize(record.elem_idx); + cols.log2_height = F::from_usize(record.log2_height); + cols.is_last_for_height = F::from_bool(record.is_last_for_height); + cols.coeff = ext_limbs(record.coeff); + cols.opened_value = ext_limbs(record.opened_value); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldCommitPhaseQueryTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldCommitPhaseQueryRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldCommitPhaseQueryCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldCommitPhaseQueryCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.round = F::from_usize(record.round); + cols.query_value = F::from_usize(record.query_value); + cols.idx_in = F::from_usize(record.idx_in); + cols.idx_out = F::from_usize(record.idx_out); + cols.log2_height = F::from_usize(record.log2_height); + cols.is_first = F::from_bool(record.is_first); + cols.has_reduced_opening = F::from_bool(record.has_reduced_opening); + cols.folded_idx = F::from_usize(record.folded_idx); + cols.is_last = F::from_bool(record.is_last); + cols.reduced_opening = ext_limbs(record.reduced_opening); + cols.folded_in = ext_limbs(record.folded_in); + cols.sibling_value = ext_limbs(record.sibling_value); + cols.challenge = ext_limbs(record.challenge); + cols.coeff = ext_limbs(record.coeff); + cols.folded_out = ext_limbs(record.folded_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldFinalCodewordTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldFinalCodewordRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldFinalCodewordCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldFinalCodewordCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.query_idx = F::from_usize(record.query_idx); + cols.elem_idx = F::from_usize(record.elem_idx); + cols.final_tidx = F::from_usize(record.final_tidx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.final_value = ext_limbs(record.final_value); + cols.coeff = ext_limbs(record.coeff); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + cols.folded = ext_limbs(record.folded); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldFinalExpectedTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldFinalExpectedRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldFinalExpectedCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldFinalExpectedCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.term_idx = F::from_usize(record.term_idx); + cols.point_idx = F::from_usize(record.point_idx); + cols.coord_idx = F::from_usize(record.coord_idx); + cols.final_tidx = F::from_usize(record.final_tidx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.is_elem_first = F::from_bool(record.is_elem_first); + cols.is_elem_last = F::from_bool(record.is_elem_last); + cols.has_challenge = F::from_bool(record.has_challenge); + cols.challenge_round = F::from_usize(record.challenge_round); + cols.point_value = ext_limbs(record.point_value); + cols.bit_value = F::from_bool(record.bit_value); + cols.challenge = ext_limbs(record.challenge); + cols.final_value = ext_limbs(record.final_value); + cols.coeff_in = ext_limbs(record.coeff_in); + cols.coeff_out = ext_limbs(record.coeff_out); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsTranscriptValueTraceGenerator { + type Ctx<'a> = &'a [PcsTranscriptValueRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsTranscriptValueCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsTranscriptValueCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.tidx = F::from_usize(record.tidx); + cols.is_sample = F::from_bool(record.is_sample); + cols.is_ext = F::from_bool(record.is_ext); + cols.is_final_message = F::from_bool(record.is_final_message); + cols.is_query_sample = F::from_bool(record.is_query_sample); + cols.is_batch_alpha = F::from_bool(record.is_batch_alpha); + cols.is_basefold_eval = F::from_bool(record.is_basefold_eval); + cols.is_jagged_f_at_rho = F::from_bool(record.is_jagged_f_at_rho); + cols.transcript_ext_lookup_count = F::from_usize(record.transcript_ext_lookup_count); + cols.value = ext_limbs(record.value); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldInitialClaimTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldInitialClaimRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldInitialClaimCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldInitialClaimCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.term_idx = F::from_usize(record.term_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.global_coeff_idx = F::from_usize(record.global_coeff_idx); + cols.eval_tidx = F::from_usize(record.eval_tidx); + cols.eval = ext_limbs(record.eval); + cols.coeff = ext_limbs(record.coeff); + cols.scale = ext_limbs(record.scale); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsJaggedAssistInputTraceGenerator { + type Ctx<'a> = &'a [PcsJaggedAssistInputRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsJaggedAssistInputCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsJaggedAssistInputCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.f_tidx = F::from_usize(record.f_tidx); + cols.f_at_rho = ext_limbs(record.f_at_rho); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsJaggedClaimTraceGenerator { + type Ctx<'a> = &'a [PcsJaggedClaimRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsJaggedClaimCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsJaggedClaimCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.term_idx = F::from_usize(record.term_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.opening_idx = F::from_usize(record.opening_idx); + cols.commit_kind = F::from_usize(record.commit_kind.as_usize()); + cols.eval_idx = F::from_usize(record.eval_idx); + cols.main_idx = F::from_usize(record.main_idx); + cols.main_eval_idx = F::from_usize(record.main_eval_idx); + cols.z_col_tidx = F::from_usize(record.z_col_tidx); + cols.eval = ext_limbs(record.eval); + cols.eq_col = ext_limbs(record.eq_col); + cols.tail_zero = ext_limbs(record.tail_zero); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsSumcheckInputTraceGenerator { + type Ctx<'a> = &'a [PcsSumcheckInputRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsSumcheckInputCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsSumcheckInputCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.claim = ext_limbs(record.claim); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsSumcheckTraceGenerator { + type Ctx<'a> = &'a [PcsSumcheckRoundRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsSumcheckCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsSumcheckCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.round = F::from_usize(record.round); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.ev_tidx = F::from_usize(record.ev_tidx); + cols.challenge_tidx = F::from_usize(record.challenge_tidx); + cols.fold_challenge_lookup_count = F::from_usize(record.fold_challenge_lookup_count); + cols.ev1 = ext_limbs(record.ev1); + cols.ev2 = ext_limbs(record.ev2); + cols.claim_in = ext_limbs(record.claim_in); + cols.claim_out = ext_limbs(record.claim_out); + cols.challenge = ext_limbs(record.challenge); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBatchCoeffTraceGenerator { + type Ctx<'a> = &'a [PcsBatchCoeffRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBatchCoeffCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBatchCoeffCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.global_coeff_idx = F::from_usize(record.global_coeff_idx); + cols.alpha_tidx = F::from_usize(record.alpha_tidx); + cols.lookup_count = F::from_usize(record.lookup_count); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.alpha = ext_limbs(record.alpha); + cols.coeff = ext_limbs(record.coeff); + cols.next_coeff = ext_limbs(record.next_coeff); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsJaggedQEvalTraceGenerator { + type Ctx<'a> = &'a [PcsJaggedQEvalRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsJaggedQEvalCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsJaggedQEvalCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.term_idx = F::from_usize(record.term_idx); + cols.is_first = F::from_bool(record.is_first); + cols.is_last = F::from_bool(record.is_last); + cols.col_tidx = F::from_usize(record.col_tidx); + cols.col_eval = ext_limbs(record.col_eval); + cols.eq_rho_col = ext_limbs(record.eq_rho_col); + cols.acc_in = ext_limbs(record.acc_in); + cols.acc_out = ext_limbs(record.acc_out); + cols.q_eval = ext_limbs(record.q_eval); + cols.f_at_rho = ext_limbs(record.f_at_rho); + cols.sumcheck_final = ext_limbs(record.sumcheck_final); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsJaggedAssistTraceGenerator { + type Ctx<'a> = &'a [PcsJaggedAssistRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsJaggedAssistCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsJaggedAssistCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.h_at_rho_star = ext_limbs(record.h_at_rho_star); + cols.q_at_rho_star = ext_limbs(record.q_at_rho_star); + cols.sumcheck_final = ext_limbs(record.sumcheck_final); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +impl RowMajorChip for PcsBasefoldFinalClaimTraceGenerator { + type Ctx<'a> = &'a [PcsBasefoldFinalClaimRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = PcsBasefoldFinalClaimCols::::width(); + let height = trace_height(records.len(), required_height)?; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut PcsBasefoldFinalClaimCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.sumcheck_idx = F::from_usize(record.sumcheck_idx); + cols.final_claim = ext_limbs(record.final_claim); + cols.expected = ext_limbs(record.expected); + } + Some(RowMajorMatrix::new(trace, width)) + } +} + +fn replay_degree2_sumcheck( + proof_idx: usize, + idx: usize, + claimed_sum: RecursionField, + proof: &sumcheck::structs::IOPProof, + num_variables: usize, + emit_input_record: bool, + preflight: &mut Preflight, + ts: &mut TS, +) -> Result<(Vec, RecursionField)> +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + if proof.proofs.len() != num_variables { + bail!( + "sumcheck round count mismatch: {} != {num_variables}", + proof.proofs.len() + ); + } + observe_label_with_records( + ts, + &num_variables.to_le_bytes(), + preflight, + 5_000_000 + idx * 10_000, + ); + observe_label_with_records( + ts, + &2usize.to_le_bytes(), + preflight, + 5_000_000 + idx * 10_000 + 10, + ); + let mut claim = claimed_sum; + let mut challenges = Vec::with_capacity(num_variables); + if emit_input_record { + preflight.pcs.sumcheck_inputs.push(PcsSumcheckInputRecord { + proof_idx, + idx, + claim: claimed_sum, + }); + } + for round in 0..num_variables { + let msg = &proof.proofs[round]; + if msg.evaluations.len() != 2 { + bail!( + "degree-2 sumcheck round {round} has {} evals", + msg.evaluations.len() + ); + } + let ev_tidx = ts.len(); + ts.observe_ext(msg.evaluations[0]); + ts.observe_ext(msg.evaluations[1]); + observe_label_with_records( + ts, + b"Internal round", + preflight, + 5_000_000 + idx * 10_000 + 100 + round * 100, + ); + let challenge_tidx = ts.len(); + let challenge = FiatShamirTranscript::::sample_ext(ts); + let ev0 = claim - msg.evaluations[0]; + let claim_out = extrapolate_uni_poly(ev0, &msg.evaluations, challenge); + challenges.push(challenge); + preflight.pcs.sumcheck_rounds.push(PcsSumcheckRoundRecord { + proof_idx, + idx, + round, + is_first: round == 0, + is_last: round + 1 == num_variables, + ev_tidx, + challenge_tidx, + ev1: msg.evaluations[0], + ev2: msg.evaluations[1], + claim_in: claim, + claim_out, + challenge, + fold_challenge_lookup_count: 0, + }); + claim = claim_out; + } + Ok((challenges, claim)) +} + +fn flatten_padded_openings_as_native( + poly_heights: &[usize], + openings: Vec<(Vec, Vec)>, +) -> Result<(Vec, Vec)> { + let max_native_point_len = poly_heights + .iter() + .map(|&h| ceil_log2(h)) + .max() + .unwrap_or(0); + let mut common_point = vec![None; max_native_point_len]; + let mut evals = Vec::with_capacity(poly_heights.len()); + let mut poly_idx = 0; + for (point, point_evals) in openings { + for value in point_evals { + let height = *poly_heights + .get(poly_idx) + .ok_or_else(|| eyre::eyre!("jagged: too many opening evaluations"))?; + let native_num_vars = ceil_log2(height); + if point.len() < native_num_vars { + bail!("jagged point length is smaller than poly num vars"); + } + for (dst, src) in common_point.iter_mut().zip(point.iter()) { + match dst { + Some(existing) if *existing != *src => { + bail!("jagged opening points are not prefix-compatible"); + } + Some(_) => {} + None => *dst = Some(*src), + } + } + let tail_zero_factor = point[native_num_vars..] + .iter() + .fold(RecursionField::ONE, |acc, r| { + acc * (RecursionField::ONE - *r) + }); + if tail_zero_factor == RecursionField::ZERO { + bail!("jagged padded opening tail factor is zero"); + } + evals.push(value * tail_zero_factor.inverse()); + poly_idx += 1; + } + } + if poly_idx != poly_heights.len() { + bail!( + "jagged opening eval count mismatch: {} != {}", + poly_idx, + poly_heights.len() + ); + } + let point = common_point + .into_iter() + .map(|value| value.ok_or_else(|| eyre::eyre!("jagged missing common opening point"))) + .collect::>>()?; + Ok((point, evals)) +} + +fn inner_verify_openings_for_col_evals( + log_h: usize, + rho_row: &[RecursionField], + col_evals: &[RecursionField], + col_tidx: usize, +) -> Vec { + col_evals + .chunks(8) + .enumerate() + .map(|(chunk_idx, evals)| { + ( + log_h, + ( + rho_row.to_vec(), + evals + .iter() + .copied() + .enumerate() + .map(|(idx, eval)| (eval, col_tidx + (chunk_idx * 8 + idx) * D_EF)) + .collect(), + ), + ) + }) + .collect() +} + +#[allow(clippy::too_many_arguments)] +fn record_basefold_query_checks( + rounds: &[BasefoldRound], + proof: &mpcs::basefold::structure::BasefoldProof, + preflight: &mut Preflight, + batch_coeffs: &[RecursionField], + fold_challenges: &[RecursionField], + queries: &[usize], + max_num_var: usize, + final_tidx: usize, +) -> Result<()> { + let rate_log = >::get_rate_log(); + let basecode_log = + >::get_basecode_msg_size_log(); + ensure_supported_basecode_log(basecode_log)?; + if proof.final_message[0].len() != (1usize << basecode_log) { + bail!("basefold final message width does not match basecode size"); + } + let log2_max_codeword_size = max_num_var + rate_log; + let rounds_len = max_num_var - basecode_log; + if proof.query_opening_proof.len() != queries.len() { + bail!("basefold query opening count mismatch"); + } + if fold_challenges.len() != rounds_len || proof.commits.len() != rounds_len { + bail!("basefold fold challenge count mismatch"); + } + for (round, root) in proof.commits.iter().enumerate() { + preflight + .pcs + .commitment_roots + .push(PcsCommitmentRootRecord { + proof_idx: 0, + commit_major: 1, + commit_minor: round, + root: digest_to_array(root)?, + lookup_count: queries.len(), + }); + } + + let final_codeword = encode_small_final_message(&proof.final_message)?; + for (query_idx, (&query, query_proof)) in queries + .iter() + .zip(proof.query_opening_proof.iter()) + .enumerate() + { + if query_proof.input_proofs.len() != rounds.len() + || query_proof.commit_phase_openings.len() != rounds_len + { + bail!("basefold query proof shape mismatch"); + } + let mut reduced_by_height = + std::collections::BTreeMap::)>::new(); + let mut coeff_iter = batch_coeffs.iter().copied().enumerate(); + + for (opening_idx, ((commit, openings), input_proof)) in rounds + .iter() + .zip(query_proof.input_proofs.iter()) + .enumerate() + { + let bits_reduced = log2_max_codeword_size - commit.log2_max_codeword_size; + let reduced_index = query >> bits_reduced; + if input_proof.opened_values.len() != openings.len() { + bail!("basefold opened-value shape mismatch"); + } + record_base_input_mmcs( + query_idx, + opening_idx, + reduced_index, + openings, + &input_proof.opened_values, + &input_proof.opening_proof, + preflight, + )?; + for (value_idx, ((num_var, (_, evals)), mat)) in openings + .iter() + .zip(input_proof.opened_values.iter()) + .enumerate() + { + if mat.len() != evals.len() { + bail!("basefold opened-value width mismatch"); + } + let log2_height = num_var + rate_log; + for (elem_idx, opened_value) in mat.iter().copied().enumerate() { + let (global_coeff_idx, coeff) = coeff_iter + .next() + .ok_or_else(|| eyre::eyre!("basefold missing batch coefficient"))?; + let (acc_in, last_row) = reduced_by_height + .get(&log2_height) + .copied() + .unwrap_or((RecursionField::ZERO, None)); + if let Some(last_row) = last_row { + preflight.pcs.basefold_query_opens[last_row].is_last_for_height = false; + } + let opened_value = RecursionField::from(opened_value); + let acc_out = acc_in + coeff * opened_value; + let row_idx = preflight.pcs.basefold_query_opens.len(); + preflight + .pcs + .basefold_query_opens + .push(PcsBasefoldQueryOpenRecord { + proof_idx: 0, + query_idx, + opening_idx, + reduced_index, + global_coeff_idx, + value_idx, + elem_idx, + log2_height, + is_last_for_height: true, + coeff, + opened_value, + acc_in, + acc_out, + }); + reduced_by_height.insert(log2_height, (acc_out, Some(row_idx))); + } + } + } + + let mut idx = query; + let mut folded = RecursionField::ZERO; + let mut log2_height = log2_max_codeword_size; + for (round, (challenge, opening)) in fold_challenges + .iter() + .copied() + .zip(query_proof.commit_phase_openings.iter()) + .enumerate() + { + let idx_in = idx; + let folded_idx = idx & 1; + let (reduced_opening, had_reduced) = reduced_by_height + .remove(&log2_height) + .map(|(value, _)| (value, true)) + .unwrap_or((RecursionField::ZERO, false)); + let mut leafs = [opening.sibling_value; 2]; + leafs[folded_idx] = folded + reduced_opening; + let leaf_idx = idx >> 1; + let coeff = verifier_folding_coeff(log2_height, leaf_idx); + let folded_out = fold_codeword_pair(leafs, challenge, coeff); + let leaf_state = poseidon2_hash_ext_pair(leafs[0], leafs[1]); + let leaf_hash = leaf_state[..DIGEST_SIZE].try_into().unwrap(); + preflight + .pcs + .commit_phase_leaf_hashes + .push(PcsCommitPhaseLeafHashRecord { + proof_idx: 0, + query_idx, + round, + leaf_idx, + left: leafs[0], + right: leafs[1], + input: ext_pair_poseidon2_input(leafs[0], leafs[1]), + output_state: leaf_state, + }); + record_commit_phase_merkle_path( + query_idx, + round, + leaf_idx, + leaf_hash, + &opening.opening_proof, + preflight, + ); + preflight + .pcs + .basefold_commit_phase_queries + .push(PcsBasefoldCommitPhaseQueryRecord { + proof_idx: 0, + query_idx, + round, + query_value: query, + idx_in, + idx_out: idx >> 1, + log2_height, + is_first: round == 0, + has_reduced_opening: had_reduced, + reduced_opening, + folded_idx, + folded_in: folded, + sibling_value: opening.sibling_value, + challenge, + coeff: RecursionField::from(coeff), + folded_out, + is_last: round + 1 == rounds_len, + }); + folded = folded_out; + log2_height -= 1; + idx >>= 1; + } + if !reduced_by_height.is_empty() { + bail!("basefold unused reduced openings remain"); + } + let final_index = idx; + let final_coeffs = final_codeword_coeffs(proof.final_message[0].len(), final_index)?; + let mut acc = RecursionField::ZERO; + let mut flat_idx = 0usize; + for row in &proof.final_message { + for (elem_idx, value) in row.iter().copied().enumerate() { + let coeff = final_coeffs[elem_idx]; + let acc_in = acc; + acc += coeff * value; + preflight + .pcs + .basefold_final_codeword + .push(PcsBasefoldFinalCodewordRecord { + proof_idx: 0, + query_idx, + elem_idx, + final_tidx: final_tidx + flat_idx * D_EF, + is_first: flat_idx == 0, + is_last: flat_idx + 1 + == proof.final_message.iter().map(Vec::len).sum::(), + final_value: value, + coeff, + acc_in, + acc_out: acc, + folded, + }); + flat_idx += 1; + } + } + if final_codeword.get(final_index).copied() != Some(acc) { + bail!("basefold final codeword reconstruction mismatch"); + } + } + Ok(()) +} + +fn encode_small_final_message( + final_message: &[Vec], +) -> Result> { + if final_message.is_empty() { + bail!("basefold final message is empty"); + } + let width = final_message[0].len(); + if width == 0 || !width.is_power_of_two() || final_message.iter().any(|row| row.len() != width) + { + bail!("basefold final message shape mismatch"); + } + let mut values = (0..width) + .map(|j| final_message.iter().map(|row| row[j]).sum()) + .collect_vec(); + values.resize( + width * (1 << >::get_rate_log()), + RecursionField::ZERO, + ); + Ok(Radix2Dit::::default() + .dft_batch(P3RowMajorMatrix::new(values, 1)) + .bit_reverse_rows() + .to_row_major_matrix() + .values) +} + +fn final_codeword_coeffs(width: usize, final_index: usize) -> Result> { + let encoded_len = + width * (1 << >::get_rate_log()); + if final_index >= encoded_len { + bail!("basefold final codeword index out of range"); + } + let mut coeffs = Vec::with_capacity(width); + for elem_idx in 0..width { + let mut values = vec![RecursionField::ZERO; encoded_len]; + values[elem_idx] = RecursionField::ONE; + let encoded = Radix2Dit::::default() + .dft_batch(P3RowMajorMatrix::new(values, 1)) + .bit_reverse_rows() + .to_row_major_matrix() + .values; + coeffs.push(encoded[final_index]); + } + Ok(coeffs) +} + +fn verifier_folding_coeff(log2_height: usize, leaf_idx: usize) -> F { + let g_inv = F::two_adic_generator(log2_height).inverse(); + let idx_bit_rev = reverse_bits_len(leaf_idx, log2_height - 1); + g_inv.exp_u64(idx_bit_rev as u64) * F::from_usize(2).inverse() +} + +fn fold_codeword_pair( + leafs: [RecursionField; 2], + challenge: RecursionField, + coeff: F, +) -> RecursionField { + let inv_2 = RecursionField::from(F::from_usize(2).inverse()); + let coeff = RecursionField::from(coeff); + let lo = (leafs[0] + leafs[1]) * inv_2; + let hi = (leafs[0] - leafs[1]) * coeff; + lo + challenge * (hi - lo) +} + +fn record_base_input_mmcs( + query_idx: usize, + opening_idx: usize, + reduced_index: usize, + openings: &[BasefoldOpening], + opened_values: &[Vec], + proof: &[[F; DIGEST_SIZE]], + preflight: &mut Preflight, +) -> Result<()> { + let rate_log = >::get_rate_log(); + let Some((first_num_var, _)) = openings.first() else { + bail!("base input MMCS opening is empty"); + }; + let log2_height = first_num_var + rate_log; + if openings + .iter() + .any(|(num_var, _)| num_var + rate_log != log2_height) + { + bail!("base input mixed-height MMCS opening is not yet AIR-encoded"); + } + + let mut flat_values = Vec::new(); + for (value_idx, ((_, (_, evals)), mat)) in openings.iter().zip(opened_values).enumerate() { + if mat.len() != evals.len() { + bail!("base input opened-value width mismatch"); + } + for (elem_idx, value) in mat.iter().copied().enumerate() { + flat_values.push((value_idx, elem_idx, RecursionField::from(value))); + } + } + if flat_values.is_empty() { + bail!("base input leaf has no opened values"); + } + if proof.is_empty() { + bail!("base input MMCS proof path is empty"); + } + + let mut state = [F::ZERO; POSEIDON2_WIDTH]; + let num_blocks = flat_values.len().div_ceil(8); + for block_idx in 0..num_blocks { + let state_in = state; + let mut input = state_in; + let mut value_is_present = [false; 8]; + let mut value_idx = [0usize; 8]; + let mut elem_idx = [0usize; 8]; + let mut values = [RecursionField::ZERO; 8]; + for (slot, (row_idx, elem, value)) in flat_values + .iter() + .copied() + .skip(block_idx * 8) + .take(8) + .enumerate() + { + value_is_present[slot] = true; + value_idx[slot] = row_idx; + elem_idx[slot] = elem; + values[slot] = value; + input[slot] = value.as_bases()[0]; + } + state = input; + poseidon2_perm().permute_mut(&mut state); + preflight + .pcs + .base_input_leaf_hashes + .push(PcsBaseInputLeafHashRecord { + proof_idx: 0, + query_idx, + opening_idx, + block_idx, + log2_height, + reduced_index, + is_first: block_idx == 0, + is_last: block_idx + 1 == num_blocks, + value_is_present, + value_idx, + elem_idx, + values, + state_in, + input, + output_state: state, + }); + } + + let expected_root = if let Some(root) = preflight + .pcs + .commitment_roots + .iter_mut() + .find(|record| record.commit_major == 0 && record.commit_minor == opening_idx) + { + root.lookup_count += 1; + Some(root.root) + } else { + None + }; + + let leaf_hash = state[..DIGEST_SIZE].try_into().unwrap(); + let replayed_root = record_base_input_merkle_path( + query_idx, + opening_idx, + reduced_index, + leaf_hash, + proof, + preflight, + ); + if let Some(expected_root) = expected_root + && replayed_root != expected_root + { + bail!("base input MMCS root replay mismatch"); + } + Ok(()) +} + +fn record_base_input_merkle_path( + query_idx: usize, + opening_idx: usize, + leaf_idx: usize, + leaf_hash: [F; DIGEST_SIZE], + proof: &[[F; DIGEST_SIZE]], + preflight: &mut Preflight, +) -> [F; DIGEST_SIZE] { + let mut current = leaf_hash; + let mut idx = leaf_idx; + for (step, sibling) in proof.iter().copied().enumerate() { + let idx_bit = idx & 1; + let (left, right) = if idx_bit == 0 { + (current, sibling) + } else { + (sibling, current) + }; + let output = poseidon2_compress_with_capacity(left, right).0; + preflight + .pcs + .base_input_merkle_rows + .push(PcsBaseInputMerkleRecord { + proof_idx: 0, + query_idx, + opening_idx, + step, + is_first: step == 0, + is_last: step + 1 == proof.len(), + idx_in: idx, + idx_bit, + idx_out: idx >> 1, + current, + sibling, + left, + right, + output, + }); + current = output; + idx >>= 1; + } + current +} + +fn digest_to_array(digest: &D) -> Result<[F; DIGEST_SIZE]> +where + D: Clone + IntoIterator, +{ + let values = digest.clone().into_iter().collect_vec(); + values + .try_into() + .map_err(|values: Vec| eyre::eyre!("unexpected digest width {}", values.len())) +} + +fn ext_pair_poseidon2_input(left: RecursionField, right: RecursionField) -> [F; POSEIDON2_WIDTH] { + let mut input = [F::ZERO; POSEIDON2_WIDTH]; + input[..D_EF].copy_from_slice(left.as_bases()); + input[D_EF..2 * D_EF].copy_from_slice(right.as_bases()); + input +} + +fn poseidon2_hash_ext_pair(left: RecursionField, right: RecursionField) -> [F; POSEIDON2_WIDTH] { + let mut state = ext_pair_poseidon2_input(left, right); + poseidon2_perm().permute_mut(&mut state); + state +} + +fn record_commit_phase_merkle_path( + query_idx: usize, + round: usize, + leaf_idx: usize, + leaf_hash: [F; DIGEST_SIZE], + proof: &[[F; DIGEST_SIZE]], + preflight: &mut Preflight, +) { + let mut current = leaf_hash; + let mut idx = leaf_idx; + for (step, sibling) in proof.iter().copied().enumerate() { + let idx_bit = idx & 1; + let (left, right) = if idx_bit == 0 { + (current, sibling) + } else { + (sibling, current) + }; + let output = poseidon2_compress_with_capacity(left, right).0; + preflight + .pcs + .commit_phase_merkle_rows + .push(PcsCommitPhaseMerkleRecord { + proof_idx: 0, + query_idx, + round, + step, + is_first: step == 0, + is_last: step + 1 == proof.len(), + idx_in: idx, + idx_bit, + idx_out: idx >> 1, + current, + sibling, + left, + right, + output, + }); + current = output; + idx >>= 1; + } +} + +fn sample_vec_with_records( + ts: &mut TS, + n: usize, + preflight: &mut Preflight, + proof_idx: usize, + idx_base: usize, +) -> Vec +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + (0..n) + .map(|i| { + let tidx = ts.len(); + let value = FiatShamirTranscript::::sample_ext(ts); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx, + idx: idx_base + i, + tidx, + value, + is_sample: true, + is_ext: true, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + value + }) + .collect() +} + +fn sample_challenge_pows( + ts: &mut TS, + size: usize, + preflight: &mut Preflight, +) -> Vec +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + if size == 0 { + return Vec::new(); + } + let tidx = ts.len(); + let alpha = FiatShamirTranscript::::sample_ext(ts); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: 2_500_000, + tidx, + value: alpha, + is_sample: true, + is_ext: true, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: true, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + let mut out = Vec::with_capacity(size); + let mut acc = RecursionField::ONE; + for global_coeff_idx in 0..size { + out.push(acc); + let next_coeff = acc * alpha; + preflight.pcs.batch_coeffs.push(PcsBatchCoeffRecord { + proof_idx: 0, + global_coeff_idx, + alpha_tidx: tidx, + alpha, + coeff: acc, + next_coeff, + lookup_count: >::get_number_queries(), + is_first: global_coeff_idx == 0, + is_last: global_coeff_idx + 1 == size, + }); + acc = next_coeff; + } + out +} + +fn observe_label_with_records( + ts: &mut TS, + label: &[u8], + preflight: &mut Preflight, + idx_base: usize, +) where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + let label_f = ::BaseField::bytes_to_field_elements(label); + let start_tidx = ts.len(); + for (i, elem) in label_f.into_iter().enumerate() { + ts.observe(elem); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: idx_base + i, + tidx: start_tidx + i, + value: RecursionField::from(elem), + is_sample: false, + is_ext: false, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + } +} + +fn check_witness_with_records( + ts: &mut TS, + bits: usize, + witness: F, + preflight: &mut Preflight, +) -> bool +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + let witness_tidx = ts.len(); + ts.observe(witness); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: 6_000_000, + tidx: witness_tidx, + value: RecursionField::from(witness), + is_sample: false, + is_ext: false, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + sample_bits_with_record(ts, bits, preflight, 6_000_001, false).query_value == 0 +} + +struct SampleBitsRecord { + tidx: usize, + sampled_value: F, + query_value: usize, + high_value: usize, +} + +fn sample_bits_with_record( + ts: &mut TS, + bits: usize, + preflight: &mut Preflight, + idx: usize, + is_query_sample: bool, +) -> SampleBitsRecord +where + TS: FiatShamirTranscript + + TranscriptHistory, +{ + assert!(bits < (u32::BITS as usize)); + assert!((1 << bits) < F::ORDER_U64); + let tidx = ts.len(); + let value = FiatShamirTranscript::::sample(ts); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx, + tidx, + value: RecursionField::from(value), + is_sample: true, + is_ext: false, + is_final_message: false, + is_query_sample, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + let raw = value.as_canonical_u64() as usize; + let mask = (1usize << bits) - 1; + SampleBitsRecord { + tidx, + sampled_value: value, + query_value: raw & mask, + high_value: raw >> bits, + } +} + +fn u32_bytes(value: usize) -> [u8; 4] { + (value as u32).to_le_bytes() +} + +fn one_hot_32(index: usize) -> [bool; 32] { + let mut out = [false; 32]; + out[index] = true; + out +} + +fn observe_digest(ts: &mut TS, digest: &D, preflight: &mut Preflight, idx_base: usize) +where + TS: FiatShamirTranscript + + TranscriptHistory, + D: Clone + IntoIterator, +{ + let start_tidx = ts_len_or_zero(ts); + for (i, value) in digest.clone().into_iter().enumerate() { + ts.observe(value); + preflight + .pcs + .transcript_values + .push(PcsTranscriptValueRecord { + proof_idx: 0, + idx: idx_base + i, + tidx: start_tidx + i, + value: RecursionField::from(value), + is_sample: false, + is_ext: false, + is_final_message: false, + is_query_sample: false, + is_batch_alpha: false, + is_basefold_eval: false, + transcript_ext_lookup_count: 0, + is_jagged_f_at_rho: false, + }); + } +} + +fn ts_len_or_zero(ts: &TS) -> usize +where + TS: TranscriptHistory, +{ + ts.len() +} + +fn child_pow_bits() -> usize { + 16 +} + +fn trace_height(len: usize, required_height: Option) -> Option { + let num_valid_rows = len.max(1); + if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + Some(height) + } else { + Some(num_valid_rows.next_power_of_two()) + } +} + +fn ext_limbs(value: EF) -> [F; D_EF] { + value.as_basis_coefficients_slice().try_into().unwrap() +} + +fn ext_sub(x: [impl Into; D_EF], y: [impl Into; D_EF]) -> [FA; D_EF] +where + FA: PrimeCharacteristicRing, +{ + let [x0, x1, x2, x3] = x.map(Into::into); + let [y0, y1, y2, y3] = y.map(Into::into); + [x0 - y0, x1 - y1, x2 - y2, x3 - y3] +} + +fn ext_add(x: [impl Into; D_EF], y: [impl Into; D_EF]) -> [FA; D_EF] +where + FA: PrimeCharacteristicRing, +{ + let [x0, x1, x2, x3] = x.map(Into::into); + let [y0, y1, y2, y3] = y.map(Into::into); + [x0 + y0, x1 + y1, x2 + y2, x3 + y3] +} + +fn interpolate_quad_at_012( + ev0: [FA; D_EF], + ev1: [F; D_EF], + ev2: [F; D_EF], + x: [F; D_EF], +) -> [FA; D_EF] +where + F: Into + Copy, + FA: PrimeCharacteristicRing, + FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>, +{ + let x = x.map(Into::into); + let ev1 = ev1.map(Into::into); + let ev2 = ev2.map(Into::into); + let one = [FA::ONE, FA::ZERO, FA::ZERO, FA::ZERO]; + let inv_two = FA::from_prime_subfield(FA::PrimeSubfield::from_usize(2).inverse()); + + let slope: [FA; D_EF] = + recursion_circuit::utils::ext_field_subtract::(ev1.clone(), ev0.clone()); + let second_diff: [FA; D_EF] = recursion_circuit::utils::ext_field_add::( + recursion_circuit::utils::ext_field_subtract::(ev2, ev1.map(|v| v.clone() + v)), + ev0.clone(), + ) + .map(|v| v * inv_two.clone()); + let x_minus_one: [FA; D_EF] = + recursion_circuit::utils::ext_field_subtract::(x.clone(), one); + recursion_circuit::utils::ext_field_add::( + recursion_circuit::utils::ext_field_add::( + ev0, + recursion_circuit::utils::ext_field_multiply::(slope, x.clone()), + ), + recursion_circuit::utils::ext_field_multiply::( + second_diff, + recursion_circuit::utils::ext_field_multiply::(x, x_minus_one), + ), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn basefold_basecode_log_zero_is_only_supported_shape() { + ensure_supported_basecode_log(0).expect("basecode_log 0 should be supported"); + let err = ensure_supported_basecode_log(1).expect_err("basecode_log > 0 must reject"); + assert!( + err.to_string().contains("only basecode_log == 0"), + "unexpected error: {err}" + ); + } +} diff --git a/ceno_recursion_v2/src/proof_shape/bus.rs b/ceno_recursion_v2/src/proof_shape/bus.rs index b551dd691..a79cee2ff 100644 --- a/ceno_recursion_v2/src/proof_shape/bus.rs +++ b/ceno_recursion_v2/src/proof_shape/bus.rs @@ -33,9 +33,6 @@ define_typed_per_proof_permutation_bus!(NumPublicValuesBus, NumPublicValuesMessa #[repr(u8)] #[derive(Debug, Copy, Clone)] pub enum AirShapeProperty { - AirId, - NumInteractions, - NeedRot, NumRead, NumWrite, NumLk, diff --git a/ceno_recursion_v2/src/proof_shape/mod.rs b/ceno_recursion_v2/src/proof_shape/mod.rs index 053918f68..442d1dffd 100644 --- a/ceno_recursion_v2/src/proof_shape/mod.rs +++ b/ceno_recursion_v2/src/proof_shape/mod.rs @@ -1,13 +1,16 @@ use std::sync::Arc; +use ceno_zkvm::structs::VK_DIGEST_LEN; +use eyre::{Result, bail, eyre}; use itertools::Itertools; +use multilinear_extensions::Expression; use openvm_circuit_primitives::encoder::Encoder; use openvm_cpu_backend::CpuBackend; use openvm_stark_backend::{ AirRef, FiatShamirTranscript, StarkProtocolConfig, TranscriptHistory, p3_maybe_rayon::prelude::*, prover::AirProvingContext, }; -use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, F}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, D_EF, F}; use p3_field::PrimeCharacteristicRing; use p3_matrix::dense::RowMajorMatrix; use witness::next_pow2_instance_padding; @@ -20,9 +23,10 @@ use crate::{ }, system::{ AirModule, BusIndexManager, BusInventory, GlobalCtxCpu, POW_CHECKER_HEIGHT, Preflight, - RecursionProof, RecursionVk, TraceGenModule, TraceVData, + RecursionField, RecursionProof, RecursionVk, TraceGenModule, TraceVData, }, tracegen::{ModuleChip, RowMajorChip}, + utils::TranscriptLabel, }; use recursion_circuit::primitives::{ bus::RangeCheckerBus, @@ -42,9 +46,6 @@ mod cuda_abi; pub struct AirMetadata { is_required: bool, num_public_values: usize, - num_witin: usize, - num_structural_witin: usize, - num_fixed: usize, num_read_count: usize, num_write_count: usize, num_logup_count: usize, @@ -53,6 +54,25 @@ pub struct AirMetadata { read_op_vars: usize, write_op_vars: usize, logup_op_vars: usize, + selectors: Vec, +} + +#[derive(Clone)] +pub struct SelectorMetadata { + selector_idx: usize, + kind: usize, + point_source: usize, + eval_idx: usize, + context_mode: SelectorContextMode, + ordered_sparse_num_vars: usize, + sparse_indices: Vec, +} + +#[derive(Clone, Copy)] +enum SelectorContextMode { + Total, + Read, + Write, } pub struct ProofShapeModule { @@ -72,6 +92,7 @@ pub struct ProofShapeModule { // Module sends extra public values message for use outside of verifier // sub-circuit if true continuations_enabled: bool, + public_values_start_tidx: usize, } impl ProofShapeModule { @@ -84,6 +105,7 @@ impl ProofShapeModule { let num_airs = child_vk.circuit_vks.len(); let idx_encoder = Arc::new(Encoder::new(num_airs, 2, true)); let per_air = extract_air_metadata_from_vk(child_vk); + let public_values_start_tidx = TranscriptLabel::Riscv.field_len() + VK_DIGEST_LEN * D_EF; let range_bus = bus_inventory.range_checker_bus; Self { @@ -95,6 +117,7 @@ impl ProofShapeModule { num_pvs_bus: NumPublicValuesBus::new(b.new_bus_idx()), idx_encoder, continuations_enabled, + public_values_start_tidx, } } @@ -197,9 +220,6 @@ fn extract_air_metadata_from_vk(child_vk: &RecursionVk) -> Vec { .map(|idx| { let ( num_public_values, - num_witin, - num_structural_witin, - num_fixed, raw_read_count, raw_write_count, raw_logup_count, @@ -214,9 +234,6 @@ fn extract_air_metadata_from_vk(child_vk: &RecursionVk) -> Vec { let css = &cs.zkvm_v1_css; ( css.instance.len(), - css.num_witin as usize, - css.num_structural_witin as usize, - css.num_fixed, cs.num_reads(), cs.num_writes(), cs.num_lks(), @@ -224,7 +241,17 @@ fn extract_air_metadata_from_vk(child_vk: &RecursionVk) -> Vec { usize::from(cs.has_ecc_ops()), ) }) - .unwrap_or((0, 0, 0, 0, 0, 0, 0, 0, 0)); + .unwrap_or((0, 0, 0, 0, 0, 0)); + let selectors = child_vk + .circuit_index_to_name + .get(&idx) + .and_then(|name| child_vk.circuit_vks.get(name)) + .map(|circuit_vk| selector_metadata_from_circuit(circuit_vk.get_cs())) + .transpose() + .unwrap_or_else(|err| { + panic!("failed to extract selector metadata for air {idx}: {err}") + }) + .unwrap_or_default(); let num_read_count = usize::from(raw_read_count > 0); let num_write_count = usize::from(raw_write_count > 0); let num_logup_count = usize::from(raw_logup_count > 0); @@ -232,9 +259,6 @@ fn extract_air_metadata_from_vk(child_vk: &RecursionVk) -> Vec { AirMetadata { is_required: false, num_public_values, - num_witin, - num_structural_witin, - num_fixed, num_read_count, num_write_count, num_logup_count, @@ -243,11 +267,181 @@ fn extract_air_metadata_from_vk(child_vk: &RecursionVk) -> Vec { read_op_vars: grouped_op_vars(raw_read_count), write_op_vars: grouped_op_vars(raw_write_count), logup_op_vars: grouped_op_vars(raw_logup_count), + selectors, } }) .collect_vec() } +fn selector_metadata_from_circuit( + composed_cs: &ceno_zkvm::structs::ComposedConstrainSystem, +) -> Result> { + let Some(circuit) = composed_cs.gkr_circuit.as_ref() else { + return Ok(Vec::new()); + }; + let first_layer = circuit + .layers + .first() + .ok_or_else(|| eyre!("empty gkr circuit layer"))?; + let group_stage_masks = first_layer_output_group_stage_masks(composed_cs, circuit)?; + let mut point_sources = vec![0usize; first_layer.out_sel_and_eval_exprs.len()]; + if let Some([left, right, origin]) = first_layer.rotation_selector_group_indices() { + point_sources[left] = 1; + point_sources[right] = 2; + point_sources[origin] = 3; + } + if let Some([x, y, slope, x3, y3]) = first_layer.ecc_bridge_group_indices() { + point_sources[x] = 4; + point_sources[y] = 4; + point_sources[slope] = 5; + point_sources[x3] = 6; + point_sources[y3] = 6; + } + let cs = &composed_cs.zkvm_v1_css; + let distinct_rw_selectors = + cs.r_selector.is_some() && cs.w_selector.is_some() && cs.r_selector != cs.w_selector; + + first_layer + .out_sel_and_eval_exprs + .iter() + .zip_eq(group_stage_masks.iter()) + .enumerate() + .map(|(selector_idx, ((selector, _), stage_mask))| { + let (kind, ordered_sparse_num_vars, sparse_indices, wit_id) = + selector_shape_metadata(selector)?; + let context_mode = if stage_mask.contains(GkrOutputStageMask::TOWER) + && distinct_rw_selectors + && matches!(selector, gkr_iop::selector::SelectorType::Prefix(_)) + { + if cs.r_selector.as_ref() == Some(selector) { + SelectorContextMode::Read + } else if cs.w_selector.as_ref() == Some(selector) { + SelectorContextMode::Write + } else { + SelectorContextMode::Total + } + } else { + SelectorContextMode::Total + }; + + Ok(SelectorMetadata { + selector_idx, + kind, + point_source: point_sources[selector_idx], + eval_idx: first_layer.n_witin + first_layer.n_fixed + wit_id as usize, + context_mode, + ordered_sparse_num_vars, + sparse_indices, + }) + }) + .collect() +} + +fn selector_shape_metadata( + selector: &gkr_iop::selector::SelectorType, +) -> Result<(usize, usize, Vec, multilinear_extensions::WitnessId)> { + use gkr_iop::selector::SelectorType; + let (kind, ordered_sparse_num_vars, sparse_indices, expr) = match selector { + SelectorType::None => bail!("SelectorType::None is not supported in recursion-v2 main"), + SelectorType::Whole(expr) => (0, 0, Vec::new(), expr), + SelectorType::Prefix(expr) => (1, 0, Vec::new(), expr), + SelectorType::OrderedSparse { + num_vars, + indices, + expression, + } => (2, *num_vars, indices.clone(), expression), + SelectorType::QuarkBinaryTreeLessThan(expr) => (3, 0, Vec::new(), expr), + }; + let Expression::StructuralWitIn(wit_id, _) = expr else { + bail!("selector expression must be StructuralWitIn"); + }; + Ok((kind, ordered_sparse_num_vars, sparse_indices, *wit_id)) +} + +#[derive(Clone, Copy, Default)] +struct GkrOutputStageMask(u8); + +impl GkrOutputStageMask { + const TOWER: Self = Self(1 << 0); + const ECC: Self = Self(1 << 1); + const ROTATION: Self = Self(1 << 2); + const ZERO: Self = Self(1 << 3); + + const fn union(self, other: Self) -> Self { + Self(self.0 | other.0) + } + + const fn contains(self, other: Self) -> bool { + (self.0 & other.0) == other.0 + } +} + +fn first_layer_output_group_stage_masks( + composed_cs: &ceno_zkvm::structs::ComposedConstrainSystem, + circuit: &gkr_iop::gkr::GKRCircuit, +) -> Result> { + let first_layer = circuit + .layers + .first() + .ok_or_else(|| eyre!("empty gkr circuit layer"))?; + let mut group_masks = vec![GkrOutputStageMask::ZERO; first_layer.out_sel_and_eval_exprs.len()]; + + if let Some(rotation_groups) = first_layer.rotation_selector_group_indices() { + for group_idx in rotation_groups { + let Some(mask) = group_masks.get_mut(group_idx) else { + bail!("rotation selector group index {group_idx} out of range"); + }; + *mask = GkrOutputStageMask::ROTATION; + } + } + if let Some(ecc_groups) = first_layer.ecc_bridge_group_indices() { + for group_idx in ecc_groups { + let Some(mask) = group_masks.get_mut(group_idx) else { + bail!("ecc selector group index {group_idx} out of range"); + }; + *mask = GkrOutputStageMask::ECC; + } + } + + let tower_outputs = tower_output_count(composed_cs); + let mut seen_tower_outputs = 0usize; + for (group_mask, (_, outputs)) in group_masks + .iter_mut() + .zip(first_layer.out_sel_and_eval_exprs.iter()) + { + if seen_tower_outputs >= tower_outputs { + break; + } + *group_mask = group_mask.union(GkrOutputStageMask::TOWER); + seen_tower_outputs += outputs.len(); + } + if seen_tower_outputs < tower_outputs { + bail!( + "failed to cover all tower outputs: layer={}, seen_tower_outputs={}, tower_outputs={}", + first_layer.name, + seen_tower_outputs, + tower_outputs + ); + } + + Ok(group_masks) +} + +fn tower_output_count( + composed_cs: &ceno_zkvm::structs::ComposedConstrainSystem, +) -> usize { + let cs = &composed_cs.zkvm_v1_css; + let num_reads = cs.r_expressions.len() + cs.r_table_expressions.len(); + let num_writes = cs.w_expressions.len() + cs.w_table_expressions.len(); + let num_lk_num = cs.lk_table_expressions.len(); + let num_lk_den = if !cs.lk_table_expressions.is_empty() { + cs.lk_table_expressions.len() + } else { + cs.lk_expressions.len() + }; + num_reads + num_writes + num_lk_num + num_lk_den +} + impl AirModule for ProofShapeModule { fn num_airs(&self) -> usize { 3 @@ -261,23 +455,23 @@ impl AirModule for ProofShapeModule { permutation_bus: self.permutation_bus, starting_tidx_bus: self.starting_tidx_bus, lookup_challenge_bus: self.bus_inventory.lookup_challenge_bus, - fraction_folder_input_bus: self.bus_inventory.fraction_folder_input_bus, - expression_claim_n_max_bus: self.bus_inventory.expression_claim_n_max_bus, tower_module_bus: self.bus_inventory.tower_module_bus, + air_presence_bus: self.bus_inventory.air_presence_bus, air_shape_bus: self.bus_inventory.air_shape_bus, - hyperdim_bus: self.bus_inventory.hyperdim_bus, - lifted_heights_bus: self.bus_inventory.lifted_heights_bus, transcript_bus: self.bus_inventory.transcript_bus, forked_transcript_bus: self.bus_inventory.forked_transcript_bus, fork_final_sample_bus: self.bus_inventory.fork_final_sample_bus, - n_lift_bus: self.bus_inventory.n_lift_bus, - tower_prefix_only: crate::system::TOWER_PREFIX_ONLY, + main_selector_shape_bus: self.bus_inventory.main_selector_shape_bus, + main_selector_sparse_index_shape_bus: self + .bus_inventory + .main_selector_sparse_index_shape_bus, }; let pvs_air = PublicValuesAir { public_values_bus: self.bus_inventory.public_values_bus, num_pvs_bus: self.num_pvs_bus, transcript_bus: self.bus_inventory.transcript_bus, continuations_enabled: self.continuations_enabled, + public_values_start_tidx: self.public_values_start_tidx, }; let range_checker = RangeCheckerAir::<8> { bus: self.range_bus, diff --git a/ceno_recursion_v2/src/proof_shape/proof_shape/air.rs b/ceno_recursion_v2/src/proof_shape/proof_shape/air.rs index 7c032240d..369733ffd 100644 --- a/ceno_recursion_v2/src/proof_shape/proof_shape/air.rs +++ b/ceno_recursion_v2/src/proof_shape/proof_shape/air.rs @@ -4,7 +4,7 @@ use itertools::fold; use openvm_circuit_primitives::{ SubAir, encoder::Encoder, - utils::{and, not, or}, + utils::{and, not}, }; use openvm_stark_backend::{ BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, @@ -17,19 +17,19 @@ use stark_recursion_circuit_derive::AlignedBorrow; use crate::{ bus::{ - AirShapeBus, AirShapeBusMessage, ExpressionClaimNMaxBus, ExpressionClaimNMaxMessage, + AirPresenceBus, AirPresenceBusMessage, AirShapeBus, AirShapeBusMessage, ForkFinalSampleBus, ForkFinalSampleMessage, ForkedTranscriptBus, - ForkedTranscriptBusMessage, FractionFolderInputBus, FractionFolderInputMessage, - HyperdimBus, HyperdimBusMessage, LiftedHeightsBus, LiftedHeightsBusMessage, - LookupChallengeBus, LookupChallengeKind, LookupChallengeMessage, NLiftBus, NLiftMessage, - TowerModuleBus, TowerModuleMessage, TranscriptBus, TranscriptBusMessage, + ForkedTranscriptBusMessage, LookupChallengeBus, LookupChallengeKind, + LookupChallengeMessage, MainSelectorShapeBus, MainSelectorShapeMessage, + MainSelectorSparseIndexShapeBus, MainSelectorSparseIndexShapeMessage, TowerModuleBus, + TowerModuleMessage, TranscriptBus, TranscriptBusMessage, }, primitives::bus::{RangeCheckerBus, RangeCheckerBusMessage}, proof_shape::{ - AirMetadata, + AirMetadata, SelectorContextMode, bus::{ AirShapeProperty, ProofShapePermutationBus, ProofShapePermutationMessage, - StartingTidxBus, StartingTidxMessage, + StartingTidxBus, }, }, subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir}, @@ -49,11 +49,10 @@ pub struct ProofShapeCols { /// /// Has a special use on summary row (when `is_last`). pub log_height: F, - /// Whether this AIR needs rotation openings. - pub need_rot: F, - // First possible transcript index of the current AIR. pub starting_tidx: F, + // Fork-local transcript index where the tower verifier starts for this AIR. + pub tower_tidx: F, // First trunk transcript index used by the fork-merge phase. pub fork_start_tidx: F, // Fork id assigned by native chip-proof iteration order. @@ -79,9 +78,6 @@ pub struct ProofShapeCols { pub n_max: F, pub is_n_max_greater: F, - pub num_air_id_lookups: F, - pub num_columns: F, - /// The Poseidon2 sponge state at the fork point (trunk state just before /// forking). Constrained to be identical across all rows within a proof. pub lookup_challenge_alpha: [F; D_EF], @@ -119,16 +115,13 @@ pub struct ProofShapeAir { // Inter-module buses pub tower_module_bus: TowerModuleBus, + pub air_presence_bus: AirPresenceBus, pub air_shape_bus: AirShapeBus, - pub expression_claim_n_max_bus: ExpressionClaimNMaxBus, - pub fraction_folder_input_bus: FractionFolderInputBus, - pub hyperdim_bus: HyperdimBus, - pub lifted_heights_bus: LiftedHeightsBus, pub transcript_bus: TranscriptBus, pub forked_transcript_bus: ForkedTranscriptBus, pub fork_final_sample_bus: ForkFinalSampleBus, - pub n_lift_bus: NLiftBus, - pub tower_prefix_only: bool, + pub main_selector_shape_bus: MainSelectorShapeBus, + pub main_selector_sparse_index_shape_bus: MainSelectorSparseIndexShapeBus, } impl BaseAir @@ -206,13 +199,6 @@ where let mut is_required = AB::Expr::ZERO; let mut air_idx = AB::Expr::ZERO; - // Select values for LiftedHeightsBus - let mut num_witin = AB::Expr::ZERO; - let mut num_structural_witin = AB::Expr::ZERO; - let mut num_fixed = AB::Expr::ZERO; - - // Select values for NumPublicValuesBus - let mut num_pvs = AB::Expr::ZERO; let mut num_read_count = AB::Expr::ZERO; let mut num_write_count = AB::Expr::ZERO; let mut num_logup_count = AB::Expr::ZERO; @@ -230,12 +216,6 @@ where let is_current_air = self.idx_encoder.get_flag_expr::(i, localv.idx_flags); let mut when_current = builder.when(is_current_air.clone()); air_idx += is_current_air.clone() * AB::F::from_usize(i); - num_witin += is_current_air.clone() * AB::F::from_usize(air_data.num_witin); - num_structural_witin += - is_current_air.clone() * AB::F::from_usize(air_data.num_structural_witin); - num_fixed += is_current_air.clone() * AB::F::from_usize(air_data.num_fixed); - - num_pvs += is_current_air.clone() * AB::F::from_usize(air_data.num_public_values); if air_data.is_required { is_required += is_current_air.clone(); @@ -257,6 +237,63 @@ where read_op_vars += is_current_air.clone() * AB::Expr::from_usize(air_data.read_op_vars); write_op_vars += is_current_air.clone() * AB::Expr::from_usize(air_data.write_op_vars); logup_op_vars += is_current_air.clone() * AB::Expr::from_usize(air_data.logup_op_vars); + + let selector_enabled = local.is_present * local.is_valid * is_current_air.clone(); + for selector in &air_data.selectors { + self.air_presence_bus.add_key_with_lookups( + builder, + local.proof_idx, + AirPresenceBusMessage { + air_idx: AB::Expr::from_usize(i), + is_present: AB::Expr::ONE, + }, + selector_enabled.clone(), + ); + let height_1: AB::Expr = local.height_1.into(); + let height_2: AB::Expr = local.height_2.into(); + let (ctx_offset, ctx_num_instances) = match selector.context_mode { + SelectorContextMode::Total => { + (AB::Expr::ZERO, height_1.clone() + height_2.clone()) + } + SelectorContextMode::Read => (AB::Expr::ZERO, height_1.clone()), + SelectorContextMode::Write => (height_1.clone(), height_2.clone()), + }; + let ctx_num_vars = n.clone() + + AB::Expr::from_usize(air_data.rotation_vars) + + AB::Expr::from_usize(air_data.ecc_extra_vars); + self.main_selector_shape_bus.send( + builder, + local.proof_idx, + MainSelectorShapeMessage { + air_idx: AB::Expr::from_usize(i), + selector_idx: AB::Expr::from_usize(selector.selector_idx), + kind: AB::Expr::from_usize(selector.kind), + point_source: AB::Expr::from_usize(selector.point_source), + eval_idx: AB::Expr::from_usize(selector.eval_idx), + ctx_offset, + ctx_num_instances, + ctx_num_vars, + ordered_sparse_num_vars: AB::Expr::from_usize( + selector.ordered_sparse_num_vars, + ), + num_sparse_indices: AB::Expr::from_usize(selector.sparse_indices.len()), + }, + selector_enabled.clone(), + ); + for (sparse_pos, sparse_index) in selector.sparse_indices.iter().enumerate() { + self.main_selector_sparse_index_shape_bus.send( + builder, + local.proof_idx, + MainSelectorSparseIndexShapeMessage { + air_idx: AB::Expr::from_usize(i), + selector_idx: AB::Expr::from_usize(selector.selector_idx), + sparse_pos: AB::Expr::from_usize(sparse_pos), + sparse_index: AB::Expr::from_usize(*sparse_index), + }, + selector_enabled.clone(), + ); + } + } } /////////////////////////////////////////////////////////////////////////////////////////// @@ -337,22 +374,6 @@ where // TRANSCRIPT OBSERVATIONS /////////////////////////////////////////////////////////////////////////////////////////// - let is_first_idx = self.idx_encoder.get_flag_expr::(0, localv.idx_flags); - - self.starting_tidx_bus.receive( - builder, - local.proof_idx, - StartingTidxMessage { - air_idx: air_idx.clone() * local.is_valid - + AB::Expr::from_usize(self.per_air.len()) * local.is_last, - tidx: local.starting_tidx.into(), - }, - or( - local.is_last, - and(local.is_valid, not::(is_first_idx)), - ) * AB::Expr::from_bool(!self.tower_prefix_only), - ); - // All active proof-shape rows for a proof agree on the trunk fork-merge // start. The value is also constrained by the trunk TranscriptBus // receives below; it is not the same concept as proof-shape post_tidx. @@ -377,17 +398,6 @@ where ); } - // constrain next air tid - self.starting_tidx_bus.send( - builder, - local.proof_idx, - StartingTidxMessage { - air_idx: air_idx.clone() + AB::F::ONE, - tidx: next.starting_tidx.into(), - }, - local.is_valid * AB::Expr::from_bool(!self.tower_prefix_only), - ); - /////////////////////////////////////////////////////////////////////////////////////////// // SNAPSHOT STATE CONTINUITY (forked transcript) /////////////////////////////////////////////////////////////////////////////////////////// @@ -508,40 +518,6 @@ where /////////////////////////////////////////////////////////////////////////////////////////// // AIR SHAPE LOOKUP /////////////////////////////////////////////////////////////////////////////////////////// - let downstream_enabled = AB::Expr::from_bool(!self.tower_prefix_only); - - self.air_shape_bus.add_key_with_lookups( - builder, - local.proof_idx, - AirShapeBusMessage { - sort_idx: local.sorted_idx.into(), - property_idx: AirShapeProperty::AirId.to_field(), - value: air_idx.clone(), - }, - local.is_present * local.num_air_id_lookups * downstream_enabled.clone(), - ); - - self.air_shape_bus.add_key_with_lookups( - builder, - local.proof_idx, - AirShapeBusMessage { - sort_idx: local.sorted_idx.into(), - property_idx: AirShapeProperty::NumInteractions.to_field(), - value: AB::Expr::ZERO, - }, - local.is_present * downstream_enabled.clone(), - ); - - self.air_shape_bus.add_key_with_lookups( - builder, - local.proof_idx, - AirShapeBusMessage { - sort_idx: local.sorted_idx.into(), - property_idx: AirShapeProperty::NeedRot.to_field(), - value: local.need_rot.into(), - }, - local.is_present * local.num_columns * downstream_enabled.clone(), - ); let base_tower_vars = n.clone() + rotation_vars + ecc_extra_vars; // TODO(recursion-v2): prove this low-degree by sending the // TowerShapeAir-derived max_layer_count to TowerInputAir. The direct @@ -595,13 +571,6 @@ where /////////////////////////////////////////////////////////////////////////////////////////// // HYPERDIM LOOKUP /////////////////////////////////////////////////////////////////////////////////////////// - builder.assert_bool(local.need_rot); - builder - .when(not(local.is_present)) - .assert_zero(local.need_rot); - builder - .when(not(local.is_present)) - .assert_zero(local.num_columns); // We range check n in [0, 32). self.range_bus.lookup_key( builder, @@ -612,56 +581,26 @@ where local.is_present, ); - self.hyperdim_bus.add_key_with_lookups( - builder, - local.proof_idx, - HyperdimBusMessage { - sort_idx: local.sorted_idx.into(), - n_abs: n.clone(), - n_sign_bit: AB::Expr::ZERO, - }, - local.is_present - * (local.num_air_id_lookups + AB::F::ONE) - * AB::Expr::from_bool(!self.tower_prefix_only), - ); - /////////////////////////////////////////////////////////////////////////////////////////// // LIFTED HEIGHTS LOOKUP + STACKING COMMITMENTS /////////////////////////////////////////////////////////////////////////////////////////// - let _raw_height_1 = fold( + let raw_height_1 = fold( local.height_1_limbs.iter().enumerate(), AB::Expr::ZERO, |acc, (i, limb)| acc + (AB::Expr::from_u32(1 << (i * LIMB_BITS)) * *limb), ); - let _raw_height_2 = fold( + let raw_height_2 = fold( local.height_2_limbs.iter().enumerate(), AB::Expr::ZERO, |acc, (i, limb)| acc + (AB::Expr::from_u32(1 << (i * LIMB_BITS)) * *limb), ); - let combined_height = local.height_1 + local.height_2; - - self.lifted_heights_bus.add_key_with_lookups( - builder, - local.proof_idx, - LiftedHeightsBusMessage { - sort_idx: local.sorted_idx.into(), - part_idx: AB::Expr::ZERO, - commit_idx: AB::Expr::ZERO, - hypercube_dim: n.clone(), - lifted_height: combined_height.into(), - log_lifted_height: local.log_height.into(), - }, - local.is_present - * (num_witin + num_structural_witin + num_fixed) - * AB::Expr::from_bool(!self.tower_prefix_only), - ); - - /////////////////////////////////////////////////////////////////////////////////////////// - // NUM PUBLIC VALUES - /////////////////////////////////////////////////////////////////////////////////////////// - let _ = num_pvs; - + builder + .when(local.is_valid) + .assert_eq(local.height_1, raw_height_1); + builder + .when(local.is_valid) + .assert_eq(local.height_2, raw_height_2); /////////////////////////////////////////////////////////////////////////////////////////// // HEIGHT + GKR MESSAGE /////////////////////////////////////////////////////////////////////////////////////////// @@ -725,51 +664,13 @@ where builder, local.proof_idx, TowerModuleMessage { - idx: if self.tower_prefix_only { - local.sorted_idx.into() - } else { - air_idx.clone() - }, - tidx: local.starting_tidx.into(), - n_logup: if self.tower_prefix_only { - local.tower_n_logup.into() - } else { - n - }, + idx: local.sorted_idx.into(), + tidx: local.tower_tidx.into(), + n_logup: local.tower_n_logup.into(), }, local.is_present * local.is_valid, ); - // Send n_max value to expression claim air - self.expression_claim_n_max_bus.send( - builder, - local.proof_idx, - ExpressionClaimNMaxMessage { - n_max: local.n_max.into(), - }, - local.is_last * downstream_enabled.clone(), - ); - - // Send n_lift to constraint folding air - self.n_lift_bus.send( - builder, - local.proof_idx, - NLiftMessage { - air_idx: air_idx, - n_lift: local.log_height.into(), - }, - local.is_present * downstream_enabled.clone(), - ); - - // Send count of present airs to fraction folder air - self.fraction_folder_input_bus.send( - builder, - local.proof_idx, - FractionFolderInputMessage { - num_present_airs: local.num_present, - }, - local.is_last * downstream_enabled, - ); } } diff --git a/ceno_recursion_v2/src/proof_shape/proof_shape/trace.rs b/ceno_recursion_v2/src/proof_shape/proof_shape/trace.rs index 4ec73433a..c1b928c43 100644 --- a/ceno_recursion_v2/src/proof_shape/proof_shape/trace.rs +++ b/ceno_recursion_v2/src/proof_shape/proof_shape/trace.rs @@ -37,10 +37,12 @@ fn decompose_usize( fn two_instance_heights_from_chip_proof(instance: &impl BorrowNumInstances) -> (usize, usize) { let num_instances = instance.borrow_num_instances(); - ( - num_instances.first().copied().unwrap_or(0), - num_instances.get(1).copied().unwrap_or(0), - ) + assert_eq!( + num_instances.len(), + 2, + "recursion-v2 currently supports exactly two num_instances entries per chip" + ); + (num_instances[0], num_instances[1]) } fn tower_shape_max(metadata: &AirMetadata, log_height: usize) -> (usize, bool, bool, bool) { @@ -164,13 +166,12 @@ impl RowMajorChip .gkr .chips .iter() - .filter_map(|entry| { - proof.chip_proofs.get(&entry.chip_idx).map(|chip_proof| { - (entry.chip_idx, tower_pre_alpha_tidx(chip_proof, entry.tidx)) + .filter_map(|chip| { + proof.chip_proofs.get(&chip.chip_idx).map(|chip_proof| { + (chip.chip_idx, tower_pre_alpha_tidx(chip_proof, chip.tidx)) }) }) .collect::>(); - for (air_idx, vdata) in &preflight.proof_shape.sorted_trace_vdata { let chunk = chunks.next().unwrap(); let (fixed_cols, variable_cols) = chunk.split_at_mut(cols_width); @@ -191,16 +192,10 @@ impl RowMajorChip cols.is_last = F::ZERO; cols.sorted_idx = F::from_usize(sorted_idx); cols.log_height = F::from_usize(log_height); - cols.need_rot = F::ZERO; - let starting_tidx = if crate::system::TOWER_PREFIX_ONLY { - tower_tidx_by_chip - .get(air_idx) - .copied() - .unwrap_or(preflight.proof_shape.starting_tidx[*air_idx]) - } else { - preflight.proof_shape.starting_tidx[*air_idx] - }; + let starting_tidx = preflight.proof_shape.starting_tidx[*air_idx]; cols.starting_tidx = F::from_usize(starting_tidx); + cols.tower_tidx = + F::from_usize(tower_tidx_by_chip.get(air_idx).copied().unwrap_or(0)); cols.fork_start_tidx = F::from_usize(preflight.proof_shape.fork_start_tidx); let fork_id = fork_id_by_chip.get(air_idx).copied().unwrap_or(0); cols.fork_id = F::from_usize(fork_id); @@ -220,8 +215,6 @@ impl RowMajorChip } cols.n_max = F::from_usize(preflight.proof_shape.n_max); cols.is_n_max_greater = F::ZERO; - cols.num_air_id_lookups = F::ZERO; - cols.num_columns = F::ZERO; cols.lookup_challenge_alpha = preflight.proof_shape.lookup_challenge_alpha; cols.lookup_challenge_beta = preflight.proof_shape.lookup_challenge_beta; cols.after_forked_challenge_1 = ef_to_limbs(fork_sample(preflight, fork_id)); @@ -264,8 +257,8 @@ impl RowMajorChip cols.is_last = F::ZERO; cols.sorted_idx = F::from_usize(sorted_idx); cols.log_height = F::ZERO; - cols.need_rot = F::ZERO; cols.starting_tidx = F::from_usize(preflight.proof_shape.starting_tidx[air_idx]); + cols.tower_tidx = F::ZERO; cols.fork_start_tidx = F::from_usize(preflight.proof_shape.fork_start_tidx); cols.fork_id = F::ZERO; cols.is_present = F::ZERO; @@ -279,8 +272,6 @@ impl RowMajorChip } cols.n_max = F::from_usize(preflight.proof_shape.n_max); cols.is_n_max_greater = F::ZERO; - cols.num_air_id_lookups = F::ZERO; - cols.num_columns = F::ZERO; cols.lookup_challenge_alpha = preflight.proof_shape.lookup_challenge_alpha; cols.lookup_challenge_beta = preflight.proof_shape.lookup_challenge_beta; cols.after_forked_challenge_1 = [F::ZERO; D_EF]; @@ -317,8 +308,8 @@ impl RowMajorChip cols.is_last = F::ONE; cols.sorted_idx = F::ZERO; cols.log_height = F::from_usize(preflight.proof_shape.n_logup); - cols.need_rot = F::ZERO; cols.starting_tidx = F::from_usize(preflight.proof_shape.post_tidx); + cols.tower_tidx = F::ZERO; cols.fork_start_tidx = F::from_usize(preflight.proof_shape.fork_start_tidx); cols.fork_id = F::ZERO; cols.is_present = F::ZERO; @@ -336,8 +327,6 @@ impl RowMajorChip .n_max .abs_diff(preflight.proof_shape.n_logup), ); - cols.num_air_id_lookups = F::ZERO; - cols.num_columns = F::ZERO; cols.lookup_challenge_alpha = preflight.proof_shape.lookup_challenge_alpha; cols.lookup_challenge_beta = preflight.proof_shape.lookup_challenge_beta; cols.after_forked_challenge_1 = [F::ZERO; D_EF]; diff --git a/ceno_recursion_v2/src/proof_shape/pvs/air.rs b/ceno_recursion_v2/src/proof_shape/pvs/air.rs index ecdb334a8..7ce16716c 100644 --- a/ceno_recursion_v2/src/proof_shape/pvs/air.rs +++ b/ceno_recursion_v2/src/proof_shape/pvs/air.rs @@ -12,7 +12,6 @@ use crate::{ bus::{PublicValuesBus, PublicValuesBusMessage, TranscriptBus, TranscriptBusMessage}, proof_shape::bus::{NumPublicValuesBus, NumPublicValuesMessage}, subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir}, - utils::TranscriptLabel, }; #[repr(C)] @@ -36,6 +35,7 @@ pub struct PublicValuesAir { pub num_pvs_bus: NumPublicValuesBus, pub transcript_bus: TranscriptBus, pub(crate) continuations_enabled: bool, + pub public_values_start_tidx: usize, } impl BaseAir for PublicValuesAir { @@ -94,7 +94,7 @@ where .when(local.is_valid * local.is_first_in_proof * local.is_first_in_air) .assert_eq( local.tidx, - AB::Expr::from_usize(TranscriptLabel::Riscv.field_len()), + AB::Expr::from_usize(self.public_values_start_tidx), ); // self.num_pvs_bus.receive( @@ -121,20 +121,9 @@ where pv_idx: local.pv_idx, value: local.value, }, - local.is_valid * AB::Expr::from_bool(!crate::system::TOWER_PREFIX_ONLY), + local.is_valid, ); - if self.continuations_enabled && !crate::system::TOWER_PREFIX_ONLY { - self.public_values_bus.send( - builder, - local.proof_idx, - PublicValuesBusMessage { - air_idx: local.air_idx, - pv_idx: local.pv_idx, - value: local.value, - }, - local.is_valid, - ); - } + let _ = self.continuations_enabled; // Receive transcript read of public values self.transcript_bus.receive( diff --git a/ceno_recursion_v2/src/proof_shape/pvs/trace.rs b/ceno_recursion_v2/src/proof_shape/pvs/trace.rs index ce1d47bc8..a131ac313 100644 --- a/ceno_recursion_v2/src/proof_shape/pvs/trace.rs +++ b/ceno_recursion_v2/src/proof_shape/pvs/trace.rs @@ -1,6 +1,7 @@ use core::borrow::BorrowMut; -use openvm_stark_sdk::config::baby_bear_poseidon2::F; +use ceno_zkvm::structs::VK_DIGEST_LEN; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F}; use p3_field::PrimeCharacteristicRing; use p3_matrix::dense::RowMajorMatrix; @@ -44,7 +45,7 @@ impl RowMajorChip for PublicValuesTraceGenerator { for (proof_idx, proof) in proofs.iter().enumerate() { let mut is_first_in_proof = true; - let mut tidx = TranscriptLabel::Riscv.field_len(); + let mut tidx = TranscriptLabel::Riscv.field_len() + VK_DIGEST_LEN * D_EF; for (air_idx, (_, circuit_vk)) in child_vk.circuit_vks.iter().enumerate() { let instance_openings = &circuit_vk.get_cs().zkvm_v1_css.instance; diff --git a/ceno_recursion_v2/src/system/bus_inventory.rs b/ceno_recursion_v2/src/system/bus_inventory.rs index 697bf341f..1dfa0b980 100644 --- a/ceno_recursion_v2/src/system/bus_inventory.rs +++ b/ceno_recursion_v2/src/system/bus_inventory.rs @@ -1,19 +1,27 @@ use recursion_circuit::{ bus::{ - AirShapeBus, FinalTranscriptStateBus, MerkleVerifyBus, Poseidon2CompressBus, - Poseidon2PermuteBus, XiRandomnessBus, + AirPresenceBus, AirShapeBus, FinalTranscriptStateBus, MerkleVerifyBus, + Poseidon2CompressBus, Poseidon2PermuteBus, XiRandomnessBus, }, primitives::bus::{ExpBitsLenBus, PowerCheckerBus, RangeCheckerBus, RightShiftBus}, system::BusIndexManager, }; use crate::bus::{ - CachedCommitBus as LocalCachedCommitBus, ExpressionClaimNMaxBus as LocalExpressionClaimNMaxBus, - ForkFinalSampleBus, ForkedTranscriptBus as LocalForkedTranscriptBus, - FractionFolderInputBus as LocalFractionFolderInputBus, HyperdimBus as LocalHyperdimBus, - LiftedHeightsBus as LocalLiftedHeightsBus, LookupChallengeBus, MainBus, MainExpressionClaimBus, - MainSumcheckInputBus, MainSumcheckOutputBus, NLiftBus as LocalNLiftBus, - PublicValuesBus as LocalPublicValuesBus, TowerModuleBus, TranscriptBus as LocalTranscriptBus, + CachedCommitBus as LocalCachedCommitBus, EccRtBus, ForkFinalSampleBus, + ForkedTranscriptBus as LocalForkedTranscriptBus, LookupChallengeBus, MainBus, MainContributionBus, + MainEccRtChallengeBus, MainEccRtEquationTotalsBus, MainEccRtQuarkFinalBus, + MainEccRtSumcheckFinalBus, MainEvalBus, MainExpressionClaimBus, MainGlobalClaimBus, + MainGlobalPointBus, MainSelectorPointBus, MainSelectorResultBus, MainSelectorShapeBus, + MainSelectorSparseIndexShapeBus, MainSumcheckInputBus, MainSumcheckOutputBus, + PcsBaseInputOpeningBus, PcsBasefoldEvalBus, PcsBasefoldFinalExpectedBus, + PcsBasefoldFinalPointBus, PcsBasefoldQueryBus, PcsBatchAlphaBus, PcsBatchCoeffBus, + PcsCommitHeightBus, PcsCommitPhaseLeafBus, PcsCommitmentRootBus, PcsEqProductBus, + PcsFinalMessageBus, PcsFoldChallengeBus, PcsJaggedAssistHBus, PcsJaggedAssistQBus, + PcsJaggedFEvalBus, PcsOpeningEvalBus, PcsQuerySampleBus, PcsSuffixProductBus, + PcsSumcheckInputBus, PcsSumcheckOutputBus, PcsTranscriptExtBus, + PublicValuesBus as LocalPublicValuesBus, TowerMainPointBus, TowerModuleBus, + TranscriptBus as LocalTranscriptBus, }; #[derive(Clone, Debug)] @@ -23,12 +31,8 @@ pub struct BusInventory { pub poseidon2_compress_bus: Poseidon2CompressBus, pub merkle_verify_bus: MerkleVerifyBus, pub tower_module_bus: TowerModuleBus, - pub expression_claim_n_max_bus: LocalExpressionClaimNMaxBus, - pub fraction_folder_input_bus: LocalFractionFolderInputBus, + pub air_presence_bus: AirPresenceBus, pub air_shape_bus: AirShapeBus, - pub hyperdim_bus: LocalHyperdimBus, - pub lifted_heights_bus: LocalLiftedHeightsBus, - pub n_lift_bus: LocalNLiftBus, pub cached_commit_bus: LocalCachedCommitBus, pub public_values_bus: LocalPublicValuesBus, pub range_checker_bus: RangeCheckerBus, @@ -38,12 +42,48 @@ pub struct BusInventory { pub main_sumcheck_input_bus: MainSumcheckInputBus, pub main_sumcheck_output_bus: MainSumcheckOutputBus, pub main_expression_claim_bus: MainExpressionClaimBus, + pub main_global_claim_bus: MainGlobalClaimBus, + pub main_global_point_bus: MainGlobalPointBus, + pub main_eval_bus: MainEvalBus, + pub main_contribution_bus: MainContributionBus, + pub main_selector_point_bus: MainSelectorPointBus, + pub main_selector_result_bus: MainSelectorResultBus, + pub main_selector_shape_bus: MainSelectorShapeBus, + pub main_selector_sparse_index_shape_bus: MainSelectorSparseIndexShapeBus, + pub main_ecc_rt_challenge_bus: MainEccRtChallengeBus, + pub main_ecc_rt_sumcheck_final_bus: MainEccRtSumcheckFinalBus, + pub main_ecc_rt_equation_totals_bus: MainEccRtEquationTotalsBus, + pub main_ecc_rt_quark_final_bus: MainEccRtQuarkFinalBus, + pub ecc_rt_bus: EccRtBus, + pub tower_main_point_bus: TowerMainPointBus, pub right_shift_bus: RightShiftBus, pub xi_randomness_bus: XiRandomnessBus, pub final_state_bus: FinalTranscriptStateBus, pub forked_transcript_bus: LocalForkedTranscriptBus, pub fork_final_sample_bus: ForkFinalSampleBus, pub lookup_challenge_bus: LookupChallengeBus, + pub pcs_basefold_query_bus: PcsBasefoldQueryBus, + pub pcs_basefold_eval_bus: PcsBasefoldEvalBus, + pub pcs_basefold_final_expected_bus: PcsBasefoldFinalExpectedBus, + pub pcs_basefold_final_point_bus: PcsBasefoldFinalPointBus, + pub pcs_transcript_ext_bus: PcsTranscriptExtBus, + pub pcs_base_input_opening_bus: PcsBaseInputOpeningBus, + pub pcs_final_message_bus: PcsFinalMessageBus, + pub pcs_query_sample_bus: PcsQuerySampleBus, + pub pcs_commitment_root_bus: PcsCommitmentRootBus, + pub pcs_commit_phase_leaf_bus: PcsCommitPhaseLeafBus, + pub pcs_sumcheck_input_bus: PcsSumcheckInputBus, + pub pcs_sumcheck_output_bus: PcsSumcheckOutputBus, + pub pcs_fold_challenge_bus: PcsFoldChallengeBus, + pub pcs_batch_coeff_bus: PcsBatchCoeffBus, + pub pcs_batch_alpha_bus: PcsBatchAlphaBus, + pub pcs_jagged_f_eval_bus: PcsJaggedFEvalBus, + pub pcs_opening_eval_bus: PcsOpeningEvalBus, + pub pcs_eq_product_bus: PcsEqProductBus, + pub pcs_suffix_product_bus: PcsSuffixProductBus, + pub pcs_jagged_assist_h_bus: PcsJaggedAssistHBus, + pub pcs_jagged_assist_q_bus: PcsJaggedAssistQBus, + pub pcs_commit_height_bus: PcsCommitHeightBus, } impl BusInventory { @@ -57,14 +97,10 @@ impl BusInventory { let tower_module_bus = TowerModuleBus::new(gkr_bus_idx); let air_shape_bus = AirShapeBus::new(b.new_bus_idx()); - let hyperdim_bus = LocalHyperdimBus::new(b.new_bus_idx()); - let lifted_heights_bus = LocalLiftedHeightsBus::new(b.new_bus_idx()); let public_values_bus = LocalPublicValuesBus::new(b.new_bus_idx()); let range_checker_bus = RangeCheckerBus::new(b.new_bus_idx()); let power_checker_bus = PowerCheckerBus::new(b.new_bus_idx()); - let expression_claim_n_max_bus = LocalExpressionClaimNMaxBus::new(b.new_bus_idx()); - let fraction_folder_input_bus = LocalFractionFolderInputBus::new(b.new_bus_idx()); - let n_lift_bus = LocalNLiftBus::new(b.new_bus_idx()); + let air_presence_bus = AirPresenceBus::new(b.new_bus_idx()); let xi_randomness_bus = XiRandomnessBus::new(b.new_bus_idx()); @@ -74,12 +110,49 @@ impl BusInventory { let main_sumcheck_input_bus = MainSumcheckInputBus::new(b.new_bus_idx()); let main_sumcheck_output_bus = MainSumcheckOutputBus::new(b.new_bus_idx()); let main_expression_claim_bus = MainExpressionClaimBus::new(b.new_bus_idx()); + let main_global_claim_bus = MainGlobalClaimBus::new(b.new_bus_idx()); + let main_global_point_bus = MainGlobalPointBus::new(b.new_bus_idx()); + let main_eval_bus = MainEvalBus::new(b.new_bus_idx()); + let main_contribution_bus = MainContributionBus::new(b.new_bus_idx()); + let main_selector_point_bus = MainSelectorPointBus::new(b.new_bus_idx()); + let main_selector_result_bus = MainSelectorResultBus::new(b.new_bus_idx()); + let main_selector_shape_bus = MainSelectorShapeBus::new(b.new_bus_idx()); + let main_selector_sparse_index_shape_bus = + MainSelectorSparseIndexShapeBus::new(b.new_bus_idx()); + let main_ecc_rt_challenge_bus = MainEccRtChallengeBus::new(b.new_bus_idx()); + let main_ecc_rt_sumcheck_final_bus = MainEccRtSumcheckFinalBus::new(b.new_bus_idx()); + let main_ecc_rt_equation_totals_bus = MainEccRtEquationTotalsBus::new(b.new_bus_idx()); + let main_ecc_rt_quark_final_bus = MainEccRtQuarkFinalBus::new(b.new_bus_idx()); + let ecc_rt_bus = EccRtBus::new(b.new_bus_idx()); + let tower_main_point_bus = TowerMainPointBus::new(b.new_bus_idx()); let cached_commit_bus = LocalCachedCommitBus::new(b.new_bus_idx()); let final_state_bus = FinalTranscriptStateBus::new(b.new_bus_idx()); let forked_transcript_bus = LocalForkedTranscriptBus::new(b.new_bus_idx()); let fork_final_sample_bus = ForkFinalSampleBus::new(b.new_bus_idx()); let lookup_challenge_bus = LookupChallengeBus::new(b.new_bus_idx()); + let pcs_basefold_query_bus = PcsBasefoldQueryBus::new(b.new_bus_idx()); + let pcs_basefold_eval_bus = PcsBasefoldEvalBus::new(b.new_bus_idx()); + let pcs_basefold_final_expected_bus = PcsBasefoldFinalExpectedBus::new(b.new_bus_idx()); + let pcs_basefold_final_point_bus = PcsBasefoldFinalPointBus::new(b.new_bus_idx()); + let pcs_transcript_ext_bus = PcsTranscriptExtBus::new(b.new_bus_idx()); + let pcs_base_input_opening_bus = PcsBaseInputOpeningBus::new(b.new_bus_idx()); + let pcs_final_message_bus = PcsFinalMessageBus::new(b.new_bus_idx()); + let pcs_query_sample_bus = PcsQuerySampleBus::new(b.new_bus_idx()); + let pcs_commitment_root_bus = PcsCommitmentRootBus::new(b.new_bus_idx()); + let pcs_commit_phase_leaf_bus = PcsCommitPhaseLeafBus::new(b.new_bus_idx()); + let pcs_sumcheck_input_bus = PcsSumcheckInputBus::new(b.new_bus_idx()); + let pcs_sumcheck_output_bus = PcsSumcheckOutputBus::new(b.new_bus_idx()); + let pcs_fold_challenge_bus = PcsFoldChallengeBus::new(b.new_bus_idx()); + let pcs_batch_coeff_bus = PcsBatchCoeffBus::new(b.new_bus_idx()); + let pcs_batch_alpha_bus = PcsBatchAlphaBus::new(b.new_bus_idx()); + let pcs_jagged_f_eval_bus = PcsJaggedFEvalBus::new(b.new_bus_idx()); + let pcs_opening_eval_bus = PcsOpeningEvalBus::new(b.new_bus_idx()); + let pcs_eq_product_bus = PcsEqProductBus::new(b.new_bus_idx()); + let pcs_suffix_product_bus = PcsSuffixProductBus::new(b.new_bus_idx()); + let pcs_jagged_assist_h_bus = PcsJaggedAssistHBus::new(b.new_bus_idx()); + let pcs_jagged_assist_q_bus = PcsJaggedAssistQBus::new(b.new_bus_idx()); + let pcs_commit_height_bus = PcsCommitHeightBus::new(b.new_bus_idx()); Self { transcript_bus, @@ -87,12 +160,8 @@ impl BusInventory { poseidon2_compress_bus, merkle_verify_bus, tower_module_bus, - expression_claim_n_max_bus, - fraction_folder_input_bus, + air_presence_bus, air_shape_bus, - hyperdim_bus, - lifted_heights_bus, - n_lift_bus, cached_commit_bus, public_values_bus, range_checker_bus, @@ -102,12 +171,48 @@ impl BusInventory { main_sumcheck_input_bus, main_sumcheck_output_bus, main_expression_claim_bus, + main_global_claim_bus, + main_global_point_bus, + main_eval_bus, + main_contribution_bus, + main_selector_point_bus, + main_selector_result_bus, + main_selector_shape_bus, + main_selector_sparse_index_shape_bus, + main_ecc_rt_challenge_bus, + main_ecc_rt_sumcheck_final_bus, + main_ecc_rt_equation_totals_bus, + main_ecc_rt_quark_final_bus, + ecc_rt_bus, + tower_main_point_bus, right_shift_bus, xi_randomness_bus, final_state_bus, forked_transcript_bus, fork_final_sample_bus, lookup_challenge_bus, + pcs_basefold_query_bus, + pcs_basefold_eval_bus, + pcs_basefold_final_expected_bus, + pcs_basefold_final_point_bus, + pcs_transcript_ext_bus, + pcs_base_input_opening_bus, + pcs_final_message_bus, + pcs_query_sample_bus, + pcs_commitment_root_bus, + pcs_commit_phase_leaf_bus, + pcs_sumcheck_input_bus, + pcs_sumcheck_output_bus, + pcs_fold_challenge_bus, + pcs_batch_coeff_bus, + pcs_batch_alpha_bus, + pcs_jagged_f_eval_bus, + pcs_opening_eval_bus, + pcs_eq_product_bus, + pcs_suffix_product_bus, + pcs_jagged_assist_h_bus, + pcs_jagged_assist_q_bus, + pcs_commit_height_bus, } } } diff --git a/ceno_recursion_v2/src/system/mod.rs b/ceno_recursion_v2/src/system/mod.rs index 595abc632..c22c69c9b 100644 --- a/ceno_recursion_v2/src/system/mod.rs +++ b/ceno_recursion_v2/src/system/mod.rs @@ -4,8 +4,22 @@ mod types; pub use crate::proof_shape::ProofShapeModule; pub use preflight::{ - BatchConstraintPreflight, ChipTranscriptRange, ForkTranscriptLog, MainPreflight, Preflight, - ProofShapePreflight, TowerChipTranscriptRange, TowerPreflight, TraceVData, + BatchConstraintPreflight, ChipTranscriptRange, EccReplayClaims, ForkTranscriptLog, + MainEccRtRecord, MainEvalRecord, MainFinalClaimRecord, MainFrontloadTermRecord, + MainGlobalSumcheckRecord, MainGlobalSumcheckRoundRecord, MainPreflight, MainSelectorEvalRecord, + MainSelectorKind, MainSelectorPointDeriveKind, MainSelectorPointRecord, + MainSelectorPointSourceKind, MainTowerPointEqRecord, MainTranscriptRecord, + PcsBaseInputLeafHashRecord, PcsBaseInputMerkleRecord, PcsBasefoldCommitPhaseQueryRecord, + PcsBasefoldFinalClaimRecord, PcsBasefoldFinalCodewordRecord, PcsBasefoldFinalExpectedRecord, + PcsBasefoldFinalPointRecord, PcsBasefoldInitialClaimRecord, PcsBasefoldQueryIndexRecord, + PcsBasefoldQueryOpenRecord, PcsBatchCoeffRecord, PcsCommitPhaseLeafHashRecord, + PcsCommitPhaseMerkleRecord, PcsCommitmentRootRecord, PcsEqProductKind, PcsEqProductRecord, + PcsEqProductSource, PcsJaggedAssistHRecord, PcsJaggedAssistInputRecord, PcsJaggedAssistQRecord, + PcsJaggedAssistRecord, PcsJaggedClaimRecord, PcsJaggedQEvalRecord, PcsOpeningClaimRecord, + PcsOpeningCommitKind, PcsOpeningEvalRecord, PcsOpeningPointRecord, PcsPreflight, + PcsSuffixProductRecord, PcsSumcheckInputRecord, PcsSumcheckRoundRecord, + PcsTranscriptValueRecord, Preflight, ProofShapePreflight, RotationReplayClaims, + TowerChipTranscriptRange, TowerMainPointRecord, TowerPreflight, TraceVData, }; pub use recursion_circuit::system::{ AirModule, BusIndexManager, GlobalTraceGenCtx, TraceGenModule, VerifierConfig, @@ -20,15 +34,20 @@ pub use types::{ convert_vk_from_zkvm, }; -use std::{iter, mem, sync::Arc}; +use std::{ + iter, mem, + sync::{Arc, Mutex, OnceLock}, +}; use crate::{ - batch_constraint::{self, BatchConstraintModule}, + batch_constraint::BatchConstraintModule, main::MainModule, + pcs::PcsModule, tower::TowerModule, transcript::TranscriptModule, utils::{TranscriptLabel, transcript_observe_label}, }; +use ceno_zkvm::structs::VK_DIGEST_LEN; use openvm_cpu_backend::CpuBackend; use openvm_poseidon2_air::POSEIDON2_WIDTH; use openvm_stark_backend::{ @@ -42,7 +61,6 @@ use openvm_stark_sdk::{ p3_baby_bear::Poseidon2BabyBear, }; use p3_field::{BasedVectorSpace, PrimeCharacteristicRing}; -use p3_matrix::Matrix; use recursion_circuit::primitives::{ exp_bits_len::{ExpBitsLenAir, ExpBitsLenTraceGenerator}, pow::{PowerCheckerAir, PowerCheckerCpuTraceGenerator}, @@ -50,7 +68,22 @@ use recursion_circuit::primitives::{ use tracing::Span; pub const POW_CHECKER_HEIGHT: usize = 32; -pub(crate) const TOWER_PREFIX_ONLY: bool = true; + +pub(crate) fn child_vk_digest(child_vk: &RecursionVk) -> [RecursionField; VK_DIGEST_LEN] { + static CACHE: OnceLock< + Mutex>, + > = OnceLock::new(); + + let key = child_vk as *const RecursionVk as usize; + let cache = CACHE.get_or_init(|| Mutex::new(std::collections::HashMap::new())); + if let Some(digest) = cache.lock().unwrap().get(&key).copied() { + return digest; + } + + let digest = child_vk.compute_digest(); + cache.lock().unwrap().insert(key, digest); + digest +} /// Local override of the upstream CPU tracegen context so modules accept ZKVM proofs. pub struct GlobalCtxCpu; @@ -79,7 +112,7 @@ pub trait VerifierTraceGen> { &self, engine: &E, child_vk: &RecursionVk, - ) -> CommittedTraceData; + ) -> Option>; #[allow(clippy::ptr_arg)] fn generate_proving_ctxs< @@ -89,7 +122,7 @@ pub trait VerifierTraceGen> { >( &self, child_vk: &RecursionVk, - child_vk_pcs_data: CommittedTraceData, + child_vk_pcs_data: Option>, proofs: &[RecursionProof], external_data: &mut VerifierExternalData<'_>, initial_transcripts: Vec, @@ -103,7 +136,7 @@ pub trait VerifierTraceGen> { >( &self, child_vk: &RecursionVk, - child_vk_pcs_data: CommittedTraceData, + child_vk_pcs_data: Option>, proofs: &[RecursionProof], initial_transcript: TS, ) -> Vec> { @@ -144,6 +177,7 @@ pub struct VerifierSubCircuit { pub(crate) gkr: TowerModule, #[allow(dead_code)] pub(crate) batch_constraint: BatchConstraintModule, + pub(crate) pcs: PcsModule, } #[derive(Copy, Clone)] @@ -154,6 +188,7 @@ enum TraceModuleRef<'a> { Tower(&'a TowerModule), #[allow(dead_code)] BatchConstraint(&'a BatchConstraintModule), + Pcs(&'a PcsModule), } impl<'a> TraceModuleRef<'a> { @@ -164,6 +199,7 @@ impl<'a> TraceModuleRef<'a> { TraceModuleRef::Main(_) => "Main", TraceModuleRef::Tower(_) => "Tower", TraceModuleRef::BatchConstraint(_) => "BatchConstraint", + TraceModuleRef::Pcs(_) => "Pcs", } } @@ -188,6 +224,7 @@ impl<'a> TraceModuleRef<'a> { TraceModuleRef::BatchConstraint(module) => { module.run_preflight(child_vk, proof, preflight, sponge) } + TraceModuleRef::Pcs(module) => module.run_preflight(child_vk, proof, preflight, sponge), TraceModuleRef::Transcript(_) | TraceModuleRef::ProofShape(_) => { panic!( "{} module does not participate in per-module preflight", @@ -204,7 +241,7 @@ impl<'a> TraceModuleRef<'a> { child_vk: &RecursionVk, proofs: &[RecursionProof], preflights: &[Preflight], - child_vk_pcs_data: &CommittedTraceData>, + _child_vk_pcs_data: &Option>>, pow_checker_gen: &Arc>, exp_bits_len_gen: &ExpBitsLenTraceGenerator, external_data: &VerifierExternalData<'_>, @@ -223,12 +260,10 @@ impl<'a> TraceModuleRef<'a> { ), TraceModuleRef::ProofShape(module) => { let mut range_check_inputs = external_data.range_check_inputs.to_vec(); - if TOWER_PREFIX_ONLY { - range_check_inputs.extend( - crate::tower::collect_tower_range_checks(child_vk, proofs, preflights) - .ok()?, - ); - } + range_check_inputs.extend( + crate::tower::collect_tower_range_checks(child_vk, proofs, preflights).ok()?, + ); + range_check_inputs.extend(crate::pcs::collect_pcs_range_checks(preflights)); module.generate_proving_ctxs( child_vk, proofs, @@ -247,8 +282,12 @@ impl<'a> TraceModuleRef<'a> { exp_bits_len_gen, required_heights, ), - TraceModuleRef::BatchConstraint(module) => module - .generate_placeholder_proving_ctxs(child_vk_pcs_data.clone(), required_heights), + TraceModuleRef::BatchConstraint(module) => { + module.generate_proving_ctxs(child_vk, proofs, preflights, &(), required_heights) + } + TraceModuleRef::Pcs(module) => { + module.generate_proving_ctxs(child_vk, proofs, preflights, &(), required_heights) + } } } } @@ -298,6 +337,7 @@ impl VerifierSubCircuit { ); let batch_constraint = BatchConstraintModule::new(&mut bus_idx_manager, bus_inventory.clone(), MAX_NUM_PROOFS); + let pcs = PcsModule::new(bus_inventory.clone()); VerifierSubCircuit { bus_inventory, @@ -307,6 +347,7 @@ impl VerifierSubCircuit { main_module, gkr, batch_constraint, + pcs, } } @@ -320,7 +361,14 @@ impl VerifierSubCircuit { proof .chip_proofs .iter() - .map(|(&chip_idx, chip_proof)| (chip_idx, chip_proof)) + .map(|(&chip_idx, chip_proof)| { + assert_eq!( + chip_proof.num_instances.len(), + 2, + "recursion-v2 currently supports exactly two num_instances entries per chip" + ); + (chip_idx, chip_proof) + }) .collect(); chip_proof_list } @@ -346,24 +394,22 @@ impl VerifierSubCircuit { { let mut preflight = Preflight::default(); - if TOWER_PREFIX_ONLY { - let log = sponge.clone().into_log(); - let values = log.values(); - if values.len() >= 2 * D_EF { - let alpha_start = values.len() - 2 * D_EF; - let beta_start = values.len() - D_EF; - preflight.vm_pvs.lookup_challenge_alpha_tidx = alpha_start; - preflight.vm_pvs.lookup_challenge_beta_tidx = beta_start; - if let Some(alpha) = - EF::from_basis_coefficients_slice(&values[alpha_start..alpha_start + D_EF]) - { - preflight.vm_pvs.lookup_challenge_alpha = alpha; - } - if let Some(beta) = - EF::from_basis_coefficients_slice(&values[beta_start..beta_start + D_EF]) - { - preflight.vm_pvs.lookup_challenge_beta = beta; - } + let log = sponge.clone().into_log(); + let values = log.values(); + if values.len() >= 2 * D_EF { + let alpha_start = values.len() - 2 * D_EF; + let beta_start = values.len() - D_EF; + preflight.vm_pvs.lookup_challenge_alpha_tidx = alpha_start; + preflight.vm_pvs.lookup_challenge_beta_tidx = beta_start; + if let Some(alpha) = + EF::from_basis_coefficients_slice(&values[alpha_start..alpha_start + D_EF]) + { + preflight.vm_pvs.lookup_challenge_alpha = alpha; + } + if let Some(beta) = + EF::from_basis_coefficients_slice(&values[beta_start..beta_start + D_EF]) + { + preflight.vm_pvs.lookup_challenge_beta = beta; } } @@ -411,15 +457,30 @@ impl VerifierSubCircuit { // order from `build_chip_proof_list` (not proof-shape sorted order). // Fork IDs are normalized to 0-based indexing for forked // transcripts (fork_id in 0..num_forks). - fs.observe(F::from_usize(fork_id)); - fs.observe(F::from_u64(chip_idx as u64)); + FiatShamirTranscript::::observe(fs, F::from_usize(fork_id)); + FiatShamirTranscript::::observe( + fs, + F::from_u64(chip_idx as u64), + ); for num_instance in chip_proof.num_instances { - fs.observe(F::from_usize(num_instance)); + FiatShamirTranscript::::observe( + fs, + F::from_usize(num_instance), + ); } let tower_tidx = fs.len(); let tower_replay = crate::tower::record_and_replay_tower_preflight(fs, child_vk, chip_idx, chip_proof); + let (rotation_replay, ecc_replay) = crate::main::replay_chip_pre_main_tail_transcript( + fs, + child_vk, + chip_idx, + chip_proof, + &tower_replay, + [alpha_ext, beta_ext], + ) + .unwrap_or_else(|err| panic!("failed to replay pre-main transcript tail: {err}")); // Record tower entry with fork-local tidx at tower stage start. preflight.gkr.chips.push(TowerChipTranscriptRange { @@ -427,23 +488,9 @@ impl VerifierSubCircuit { tidx: tower_tidx, fork_idx: fork_id, tower_replay, + rotation_replay, + ecc_replay, }); - - if !TOWER_PREFIX_ONLY { - // Main preflight for this chip. - let main_tidx = fs.len(); - crate::main::record_main_transcript( - fs, - chip_idx, - proof.main_constraint_proof.claimed_sum, - ); - - preflight.main.chips.push(ChipTranscriptRange { - chip_idx, - tidx: main_tidx, - fork_idx: fork_id, - }); - } } // Phase 4: Merge — sample 1 ext element from each fork, observe into trunk. @@ -453,12 +500,16 @@ impl VerifierSubCircuit { sponge.observe_ext(sample); } - if !TOWER_PREFIX_ONLY { - // Phase 5: batch-constraint transcript replay on the trunk, after the - // fork merge and before PCS opening verification. - self.batch_constraint - .run_preflight(child_vk, proof, &mut preflight, &mut sponge); - } + self.main_module + .run_preflight(child_vk, proof, &mut preflight, &mut sponge); + + // Phase 5: batch-constraint transcript replay on the trunk, after the + // fork merge and before PCS opening verification. + self.batch_constraint + .run_preflight(child_vk, proof, &mut preflight, &mut sponge); + + self.pcs + .run_preflight(child_vk, proof, &mut preflight, &mut sponge); // The trunk log now contains pre-fork ops (0..fork_offset), merge ops // (fork_offset..merge_end), and batch-constraint ops after that. The @@ -478,21 +529,6 @@ impl VerifierSubCircuit { fork_id, }); } - if std::env::var_os("CENO_REC_V2_DEBUG_TRANSCRIPT").is_some() { - eprintln!( - "rec-v2-debug module=transcript source=preflight proof_idx=? pvs_end={} fork_start={} trunk_len={} num_forks={} fork_lens={:?}", - fork_offset, - preflight.proof_shape.fork_start_tidx, - trunk_log.len(), - preflight.fork_transcripts.len(), - preflight - .fork_transcripts - .iter() - .map(|fork| fork.log.len()) - .collect::>() - ); - } - preflight.proof_shape.after_forked_challenge_1 = preflight .fork_transcripts .first() @@ -526,24 +562,21 @@ impl VerifierSubCircuit { let t_n = self.transcript.num_airs(); let ps_n = self.proof_shape.num_airs(); let gkr_n = self.gkr.num_airs(); - let module_air_counts = if TOWER_PREFIX_ONLY { - vec![t_n, ps_n, gkr_n] - } else { - vec![ - t_n, - ps_n, - gkr_n, - self.main_module.num_airs(), - self.batch_constraint.num_airs(), - ] - }; + let module_air_counts = vec![ + t_n, + ps_n, + gkr_n, + self.main_module.num_airs(), + self.batch_constraint.num_airs(), + self.pcs.num_airs(), + ]; let Some(heights) = required_heights else { return (vec![None; module_air_counts.len()], None, None); }; let total_module_airs: usize = module_air_counts.iter().sum(); - let total = total_module_airs + if TOWER_PREFIX_ONLY { 1 } else { 2 }; + let total = total_module_airs + 1; assert_eq!(heights.len(), total); let mut offset = 0usize; @@ -552,13 +585,70 @@ impl VerifierSubCircuit { per_module.push(Some(&heights[offset..offset + n])); offset += n; } - if TOWER_PREFIX_ONLY { - debug_assert_eq!(heights.len() - offset, 1); - (per_module, None, Some(heights[offset])) - } else { - debug_assert_eq!(heights.len() - offset, 2); - (per_module, Some(heights[offset]), Some(heights[offset + 1])) - } + debug_assert_eq!(heights.len() - offset, 1); + (per_module, None, Some(heights[offset])) + } + + #[cfg(test)] + pub(crate) fn test_run_preflight( + &self, + sponge: TS, + child_vk: &RecursionVk, + proof: &RecursionProof, + ) -> Preflight + where + TS: FiatShamirTranscript + + TranscriptHistory + + From> + + Clone, + { + self.run_preflight(sponge, child_vk, proof) + } + + #[cfg(test)] + pub(crate) fn test_generate_proving_ctxs_from_preflights>( + &self, + child_vk: &RecursionVk, + proofs: &[RecursionProof], + preflights: &[Preflight], + external_data: &VerifierExternalData<'_>, + ) -> Vec>> { + let power_checker_gen = + Arc::new(PowerCheckerCpuTraceGenerator::<2, POW_CHECKER_HEIGHT>::default()); + let exp_bits_len_gen = ExpBitsLenTraceGenerator::default(); + let module_required = vec![None; 6]; + let modules = vec![ + TraceModuleRef::Transcript(&self.transcript), + TraceModuleRef::ProofShape(&self.proof_shape), + TraceModuleRef::Tower(&self.gkr), + TraceModuleRef::Main(&self.main_module), + TraceModuleRef::BatchConstraint(&self.batch_constraint), + TraceModuleRef::Pcs(&self.pcs), + ]; + + let mut ctxs = modules + .into_iter() + .zip(module_required) + .flat_map(|(module, required_heights)| { + module + .generate_cpu_ctxs( + child_vk, + proofs, + preflights, + &None, + &power_checker_gen, + &exp_bits_len_gen, + external_data, + required_heights, + ) + .expect("test trace generation should succeed") + }) + .collect::>(); + let exp_bits_trace = exp_bits_len_gen + .generate_trace_row_major(None) + .expect("exp bits trace should generate"); + ctxs.push(AirProvingContext::simple_no_pis(exp_bits_trace)); + ctxs } } @@ -579,8 +669,9 @@ impl, const MAX_NUM_PROOFS: usize> &self, engine: &E, child_vk: &RecursionVk, - ) -> CommittedTraceData> { - batch_constraint::commit_child_vk(engine, child_vk) + ) -> Option>> { + let _ = (engine, child_vk); + None } #[tracing::instrument(name = "subcircuit_generate_proving_ctxs", skip_all)] @@ -591,7 +682,7 @@ impl, const MAX_NUM_PROOFS: usize> >( &self, child_vk: &RecursionVk, - child_vk_pcs_data: CommittedTraceData>, + child_vk_pcs_data: Option>>, proofs: &[RecursionProof], external_data: &mut VerifierExternalData<'_>, initial_transcripts: Vec, @@ -634,21 +725,14 @@ impl, const MAX_NUM_PROOFS: usize> let (module_required, power_checker_required, exp_bits_len_required) = self.split_required_heights(external_data.required_heights); - let modules = if TOWER_PREFIX_ONLY { - vec![ - TraceModuleRef::Transcript(&self.transcript), - TraceModuleRef::ProofShape(&self.proof_shape), - TraceModuleRef::Tower(&self.gkr), - ] - } else { - vec![ - TraceModuleRef::Transcript(&self.transcript), - TraceModuleRef::ProofShape(&self.proof_shape), - TraceModuleRef::Tower(&self.gkr), - TraceModuleRef::Main(&self.main_module), - TraceModuleRef::BatchConstraint(&self.batch_constraint), - ] - }; + let modules = vec![ + TraceModuleRef::Transcript(&self.transcript), + TraceModuleRef::ProofShape(&self.proof_shape), + TraceModuleRef::Tower(&self.gkr), + TraceModuleRef::Main(&self.main_module), + TraceModuleRef::BatchConstraint(&self.batch_constraint), + TraceModuleRef::Pcs(&self.pcs), + ]; let span = Span::current(); let ctxs_by_module = modules @@ -685,14 +769,7 @@ impl, const MAX_NUM_PROOFS: usize> if power_checker_required.is_some_and(|h| h != POW_CHECKER_HEIGHT) { return None; } - if !TOWER_PREFIX_ONLY { - let power_height = power_checker_required.unwrap_or(POW_CHECKER_HEIGHT); - let power_trace = power_checker_gen.generate_trace_row_major(); - if power_trace.height() != power_height { - return None; - } - ctx_per_trace.push(AirProvingContext::simple_no_pis(power_trace)); - } + let _ = power_checker_required; let exp_bits_height = exp_bits_len_required; let exp_bits_trace = exp_bits_len_gen.generate_trace_row_major(exp_bits_height)?; @@ -712,28 +789,15 @@ impl AggregationSubCircuit for VerifierSubCircuit { - pow_bus: self.bus_inventory.power_checker_bus, - range_bus: self.bus_inventory.range_checker_bus, - }); + let power_checker_air = None::>; iter::empty() .chain(self.transcript.airs()) .chain(self.proof_shape.airs()) .chain(self.gkr.airs()) - .chain( - (!TOWER_PREFIX_ONLY) - .then(|| self.main_module.airs()) - .into_iter() - .flatten(), - ) - .chain( - (!TOWER_PREFIX_ONLY) - .then(|| self.batch_constraint.airs()) - .into_iter() - .flatten(), - ) + .chain(self.main_module.airs()) + .chain(self.batch_constraint.airs()) + .chain(self.pcs.airs()) .chain(power_checker_air.map(|air| Arc::new(air) as AirRef<_>)) .chain([Arc::new(exp_bits_len_air) as AirRef<_>]) .collect() diff --git a/ceno_recursion_v2/src/system/preflight/mod.rs b/ceno_recursion_v2/src/system/preflight/mod.rs index ac9aa7268..1c4436a68 100644 --- a/ceno_recursion_v2/src/system/preflight/mod.rs +++ b/ceno_recursion_v2/src/system/preflight/mod.rs @@ -1,8 +1,9 @@ use openvm_poseidon2_air::POSEIDON2_WIDTH; use openvm_stark_backend::TranscriptLog; -use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, EF, F}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, DIGEST_SIZE, EF, F}; +use p3_field::PrimeCharacteristicRing; -use crate::tower::TowerReplayResult; +use crate::{system::RecursionField, tower::TowerReplayResult}; /// Placeholder types mirroring upstream recursion preflight records. /// These will be populated with real transcript metadata once the @@ -18,6 +19,7 @@ pub struct Preflight { pub main: MainPreflight, pub gkr: TowerPreflight, pub batch_constraint: BatchConstraintPreflight, + pub pcs: PcsPreflight, pub vm_pvs: VmPvsPreflight, } @@ -84,6 +86,842 @@ pub struct TraceVData { #[derive(Clone, Debug, Default)] pub struct MainPreflight { pub chips: Vec, + pub transcript_start: usize, + pub transcript_end: usize, + pub global_sumchecks: Vec, + pub evals: Vec, + pub selector_evals: Vec, + pub selector_points: Vec, + pub tower_point_eqs: Vec, + pub frontload_terms: Vec, + pub final_claims: Vec, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsPreflight { + pub opening_claims: Vec, + pub opening_points: Vec, + pub opening_evals: Vec, + pub basefold_query_indices: Vec, + pub basefold_query_opens: Vec, + pub basefold_commit_phase_queries: Vec, + pub basefold_final_codeword: Vec, + pub transcript_values: Vec, + pub sumcheck_rounds: Vec, + pub sumcheck_inputs: Vec, + pub eq_products: Vec, + pub suffix_products: Vec, + pub jagged_assist_h: Vec, + pub jagged_assist_q: Vec, + pub jagged_claims: Vec, + pub basefold_initial_claims: Vec, + pub jagged_assist_inputs: Vec, + pub batch_coeffs: Vec, + pub jagged_q_evals: Vec, + pub jagged_assists: Vec, + pub basefold_final_points: Vec, + pub basefold_final_expected: Vec, + pub basefold_final_claims: Vec, + pub commitment_roots: Vec, + pub base_input_leaf_hashes: Vec, + pub base_input_merkle_rows: Vec, + pub commit_phase_leaf_hashes: Vec, + pub commit_phase_merkle_rows: Vec, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsOpeningClaimRecord { + pub input_opening_point: Vec, + pub wits_in_evals: Vec, + pub fixed_in_evals: Vec, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsOpeningPointRecord { + pub proof_idx: usize, + pub opening_idx: usize, + pub coord_idx: usize, + pub global_round_idx: usize, + pub value: RecursionField, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum PcsOpeningCommitKind { + Witin, + Fixed, +} + +impl PcsOpeningCommitKind { + pub const fn as_usize(self) -> usize { + match self { + Self::Witin => 0, + Self::Fixed => 1, + } + } +} + +impl Default for PcsOpeningCommitKind { + fn default() -> Self { + Self::Witin + } +} + +#[derive(Clone, Debug)] +pub struct PcsOpeningEvalRecord { + pub proof_idx: usize, + pub opening_idx: usize, + pub commit_kind: PcsOpeningCommitKind, + pub eval_idx: usize, + pub main_idx: usize, + pub main_eval_idx: usize, + pub value: RecursionField, + pub raw_value: RecursionField, +} + +#[derive(Clone, Debug)] +pub struct PcsBasefoldQueryIndexRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub sample_tidx: usize, + pub query_bits: usize, + pub sampled_value: RecursionField, + pub query_value: usize, + pub high_value: usize, + pub query_bytes: [u8; 4], + pub high_bytes: [u8; 4], + pub query_bit_selectors: [bool; 32], +} + +#[derive(Clone, Debug)] +pub struct PcsBasefoldQueryOpenRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub opening_idx: usize, + pub reduced_index: usize, + pub global_coeff_idx: usize, + pub value_idx: usize, + pub elem_idx: usize, + pub log2_height: usize, + pub is_last_for_height: bool, + pub coeff: RecursionField, + pub opened_value: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, +} + +#[derive(Clone, Debug)] +pub struct PcsBaseInputLeafHashRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub opening_idx: usize, + pub block_idx: usize, + pub log2_height: usize, + pub reduced_index: usize, + pub is_first: bool, + pub is_last: bool, + pub value_is_present: [bool; 8], + pub value_idx: [usize; 8], + pub elem_idx: [usize; 8], + pub values: [RecursionField; 8], + pub state_in: [F; POSEIDON2_WIDTH], + pub input: [F; POSEIDON2_WIDTH], + pub output_state: [F; POSEIDON2_WIDTH], +} + +#[derive(Clone, Debug)] +pub struct PcsBaseInputMerkleRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub opening_idx: usize, + pub step: usize, + pub is_first: bool, + pub is_last: bool, + pub idx_in: usize, + pub idx_bit: usize, + pub idx_out: usize, + pub current: [F; DIGEST_SIZE], + pub sibling: [F; DIGEST_SIZE], + pub left: [F; DIGEST_SIZE], + pub right: [F; DIGEST_SIZE], + pub output: [F; DIGEST_SIZE], +} + +#[derive(Clone, Debug)] +pub struct PcsBasefoldCommitPhaseQueryRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub round: usize, + pub query_value: usize, + pub idx_in: usize, + pub idx_out: usize, + pub log2_height: usize, + pub is_first: bool, + pub has_reduced_opening: bool, + pub reduced_opening: RecursionField, + pub folded_idx: usize, + pub folded_in: RecursionField, + pub sibling_value: RecursionField, + pub challenge: RecursionField, + pub coeff: RecursionField, + pub folded_out: RecursionField, + pub is_last: bool, +} + +#[derive(Clone, Debug)] +pub struct PcsCommitmentRootRecord { + pub proof_idx: usize, + pub commit_major: usize, + pub commit_minor: usize, + pub root: [F; DIGEST_SIZE], + pub lookup_count: usize, +} + +#[derive(Clone, Debug)] +pub struct PcsCommitPhaseLeafHashRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub round: usize, + pub leaf_idx: usize, + pub left: RecursionField, + pub right: RecursionField, + pub input: [F; POSEIDON2_WIDTH], + pub output_state: [F; POSEIDON2_WIDTH], +} + +#[derive(Clone, Debug)] +pub struct PcsCommitPhaseMerkleRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub round: usize, + pub step: usize, + pub is_first: bool, + pub is_last: bool, + pub idx_in: usize, + pub idx_bit: usize, + pub idx_out: usize, + pub current: [F; DIGEST_SIZE], + pub sibling: [F; DIGEST_SIZE], + pub left: [F; DIGEST_SIZE], + pub right: [F; DIGEST_SIZE], + pub output: [F; DIGEST_SIZE], +} + +#[derive(Clone, Debug)] +pub struct PcsBasefoldFinalCodewordRecord { + pub proof_idx: usize, + pub query_idx: usize, + pub elem_idx: usize, + pub final_tidx: usize, + pub is_first: bool, + pub is_last: bool, + pub final_value: RecursionField, + pub coeff: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, + pub folded: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsTranscriptValueRecord { + pub proof_idx: usize, + pub idx: usize, + pub tidx: usize, + pub value: RecursionField, + pub is_sample: bool, + pub is_ext: bool, + pub is_final_message: bool, + pub is_query_sample: bool, + pub is_batch_alpha: bool, + pub is_basefold_eval: bool, + pub transcript_ext_lookup_count: usize, + pub is_jagged_f_at_rho: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsSumcheckRoundRecord { + pub proof_idx: usize, + pub idx: usize, + pub round: usize, + pub is_first: bool, + pub is_last: bool, + pub ev_tidx: usize, + pub challenge_tidx: usize, + pub ev1: RecursionField, + pub ev2: RecursionField, + pub claim_in: RecursionField, + pub claim_out: RecursionField, + pub challenge: RecursionField, + pub fold_challenge_lookup_count: usize, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsSumcheckInputRecord { + pub proof_idx: usize, + pub idx: usize, + pub claim: RecursionField, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum PcsEqProductKind { + JaggedClaim = 0, + JaggedQEval = 1, +} + +impl PcsEqProductKind { + pub const fn as_usize(self) -> usize { + self as usize + } +} + +impl Default for PcsEqProductKind { + fn default() -> Self { + Self::JaggedClaim + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum PcsEqProductSource { + Transcript = 0, + FoldChallenge = 1, +} + +impl PcsEqProductSource { + pub const fn as_usize(self) -> usize { + self as usize + } +} + +impl Default for PcsEqProductSource { + fn default() -> Self { + Self::Transcript + } +} + +#[derive(Clone, Debug, Default)] +pub struct PcsEqProductRecord { + pub proof_idx: usize, + pub kind: PcsEqProductKind, + pub source: PcsEqProductSource, + pub round_idx: usize, + pub term_idx: usize, + pub bit_idx: usize, + pub is_first: bool, + pub is_last: bool, + pub lookup_count: usize, + pub point_tidx: usize, + pub sumcheck_idx: usize, + pub point_round: usize, + pub index_bit: bool, + pub index_pow2: usize, + pub index_acc_in: usize, + pub index_acc_out: usize, + pub point: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsSuffixProductRecord { + pub proof_idx: usize, + pub round_idx: usize, + pub term_idx: usize, + pub coord_idx: usize, + pub step_idx: usize, + pub is_first: bool, + pub is_last: bool, + pub has_factor: bool, + pub point: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsJaggedAssistHRecord { + pub proof_idx: usize, + pub round_idx: usize, + pub sumcheck_idx: usize, + pub step_idx: usize, + pub robp_idx: usize, + pub is_first: bool, + pub is_last: bool, + pub has_z_row: bool, + pub has_rho: bool, + pub z_row: RecursionField, + pub rho: RecursionField, + pub rho_star_c: RecursionField, + pub rho_star_d: RecursionField, + pub val_in: [RecursionField; 4], + pub val_out: [RecursionField; 4], +} + +#[derive(Clone, Debug, Default)] +pub struct PcsJaggedAssistQRecord { + pub proof_idx: usize, + pub round_idx: usize, + pub sumcheck_idx: usize, + pub commitment_kind: usize, + pub term_idx: usize, + pub step_idx: usize, + pub robp_idx: usize, + pub is_first: bool, + pub is_last: bool, + pub is_first_step: bool, + pub is_last_step: bool, + pub term_is_last: bool, + pub eq_col: RecursionField, + pub t_lo: usize, + pub t_hi: usize, + pub c_bit: bool, + pub d_bit: bool, + pub bit_pow2: usize, + pub c_acc_in: usize, + pub c_acc_out: usize, + pub d_acc_in: usize, + pub d_acc_out: usize, + pub rho_star_c: RecursionField, + pub rho_star_d: RecursionField, + pub term_acc_in: RecursionField, + pub term_acc_out: RecursionField, + pub q_acc_in: RecursionField, + pub q_acc_out: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsJaggedClaimRecord { + pub proof_idx: usize, + pub round_idx: usize, + pub sumcheck_idx: usize, + pub term_idx: usize, + pub is_first: bool, + pub is_last: bool, + pub opening_idx: usize, + pub commit_kind: PcsOpeningCommitKind, + pub eval_idx: usize, + pub main_idx: usize, + pub main_eval_idx: usize, + pub eval: RecursionField, + pub z_col_tidx: usize, + pub eq_col: RecursionField, + pub tail_zero: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsBasefoldInitialClaimRecord { + pub proof_idx: usize, + pub sumcheck_idx: usize, + pub term_idx: usize, + pub is_first: bool, + pub is_last: bool, + pub global_coeff_idx: usize, + pub eval_tidx: usize, + pub eval: RecursionField, + pub coeff: RecursionField, + pub scale: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsJaggedAssistInputRecord { + pub proof_idx: usize, + pub round_idx: usize, + pub sumcheck_idx: usize, + pub f_tidx: usize, + pub f_at_rho: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsBatchCoeffRecord { + pub proof_idx: usize, + pub global_coeff_idx: usize, + pub alpha_tidx: usize, + pub alpha: RecursionField, + pub coeff: RecursionField, + pub next_coeff: RecursionField, + pub lookup_count: usize, + pub is_first: bool, + pub is_last: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsJaggedQEvalRecord { + pub proof_idx: usize, + pub round_idx: usize, + pub sumcheck_idx: usize, + pub term_idx: usize, + pub is_first: bool, + pub is_last: bool, + pub col_tidx: usize, + pub col_eval: RecursionField, + pub eq_rho_col: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, + pub q_eval: RecursionField, + pub f_at_rho: RecursionField, + pub sumcheck_final: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsJaggedAssistRecord { + pub proof_idx: usize, + pub round_idx: usize, + pub sumcheck_idx: usize, + pub h_at_rho_star: RecursionField, + pub q_at_rho_star: RecursionField, + pub sumcheck_final: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsBasefoldFinalClaimRecord { + pub proof_idx: usize, + pub sumcheck_idx: usize, + pub final_claim: RecursionField, + pub expected: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsBasefoldFinalPointRecord { + pub proof_idx: usize, + pub sumcheck_idx: usize, + pub point_idx: usize, + pub coord_idx: usize, + pub source_sumcheck_idx: usize, + pub source_round: usize, + pub value: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct PcsBasefoldFinalExpectedRecord { + pub proof_idx: usize, + pub sumcheck_idx: usize, + pub term_idx: usize, + pub point_idx: usize, + pub coord_idx: usize, + pub final_tidx: usize, + pub is_first: bool, + pub is_last: bool, + pub is_elem_first: bool, + pub is_elem_last: bool, + pub has_challenge: bool, + pub challenge_round: usize, + pub point_value: RecursionField, + pub bit_value: bool, + pub challenge: RecursionField, + pub final_value: RecursionField, + pub coeff_in: RecursionField, + pub coeff_out: RecursionField, + pub acc_in: RecursionField, + pub acc_out: RecursionField, +} + +#[derive(Clone, Debug, Default)] +pub struct MainTranscriptRecord { + pub proof_idx: usize, + pub fork_id: usize, + pub is_fork: bool, + pub tidx: usize, + pub value: F, + pub is_sample: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct MainFinalClaimRecord { + pub proof_idx: usize, + pub idx: usize, + pub contribution: EF, + pub acc_in: EF, + pub acc_out: EF, + pub expected: EF, +} + +#[derive(Clone, Debug, Default)] +pub struct MainEvalRecord { + pub proof_idx: usize, + pub idx: usize, + pub eval_idx: usize, + pub tidx: usize, + pub value: EF, + pub lookup_count: usize, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum MainSelectorKind { + #[default] + Whole, + Prefix, + OrderedSparse, + QuarkBinaryTreeLessThan, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum MainSelectorPointSourceKind { + #[default] + TowerMain, + RotationLeft, + RotationRight, + RotationOrigin, + EccXY, + EccSlope, + EccX3Y3, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum MainSelectorPointDeriveKind { + #[default] + Identity, + OneMinus, + Zero, + One, +} + +#[derive(Clone, Debug, Default)] +pub struct MainSelectorEvalRecord { + pub proof_idx: usize, + pub idx: usize, + pub tower_idx: usize, + pub air_idx: usize, + pub selector_idx: usize, + pub has_eval: bool, + pub eval_idx: usize, + pub kind: MainSelectorKind, + pub ctx_offset: usize, + pub ctx_num_instances: usize, + pub ctx_num_vars: usize, + pub ordered_sparse_num_vars: usize, + pub sparse_indices: Vec, + pub in_point: Vec, + pub out_point: Vec, + pub point_source: MainSelectorPointSourceKind, + pub fork_id: usize, + pub rotation_cyclic_group_log2: usize, + pub rotation_origin_selector_idx: Option, + pub rotation_origin_tidxs: Vec, + pub ecc_sample_tidx: Option, + pub ecc_rt_tidxs: Vec, + pub ecc_xy_selector_idx: Option, + pub ecc_x3y3_selector_idx: Option, + pub value: EF, +} + +#[derive(Clone, Debug, Default)] +pub struct MainSelectorPointRecord { + pub proof_idx: usize, + pub idx: usize, + pub tower_idx: usize, + pub air_idx: usize, + pub selector_idx: usize, + pub round_idx: usize, + pub value: EF, + pub source_kind: MainSelectorPointSourceKind, + pub lookup_count: usize, + pub fork_id: usize, + pub has_transcript: bool, + pub transcript_tidx: usize, + pub has_ecc_rt: bool, + pub has_source: bool, + pub source_selector_idx: usize, + pub source_source_kind: MainSelectorPointSourceKind, + pub source_round_idx: usize, + pub source_value: EF, + pub derive_kind: MainSelectorPointDeriveKind, +} + +#[derive(Clone, Debug)] +pub struct MainEccRtRecord { + pub proof_idx: usize, + pub idx: usize, + pub fork_id: usize, + pub round_idx: usize, + pub num_rounds: usize, + pub is_first: bool, + pub is_last: bool, + pub tidx: usize, + pub out_tidx: usize, + pub alpha_tidx: usize, + pub value: EF, + pub out_value: EF, + pub alpha: EF, + pub alpha_pows: [EF; 49], + pub ev1: EF, + pub ev2: EF, + pub ev3: EF, + pub claim_in: EF, + pub claim_out: EF, + pub sel_add: EF, + pub sel_bypass: EF, + pub sel_export: EF, + pub s0: [EF; 7], + pub x0: [EF; 7], + pub y0: [EF; 7], + pub x1: [EF; 7], + pub y1: [EF; 7], + pub x3: [EF; 7], + pub y3: [EF; 7], + pub sum_x: [EF; 7], + pub sum_y: [EF; 7], + pub eq_in: EF, + pub eq_out: EF, + pub last_in: EF, + pub last_out: EF, + pub export_out_in: EF, + pub export_out_out: EF, + pub export_rt_in: EF, + pub export_rt_out: EF, + pub quark_in: EF, + pub quark_factor: EF, + pub quark_out: EF, + pub add_eval: EF, + pub bypass_eval: EF, + pub export_eval: EF, + pub lte_out_point: [EF; 32], + pub lte_rt_point: [EF; 32], + pub lte_prefix_acc: [EF; 33], + pub lte_less_acc: [EF; 33], + pub lte_bits: [bool; 32], + pub lte_active: [bool; 32], + pub quark_prefix_count: usize, + pub quark_layer_n: usize, + pub quark_parity: bool, + pub lookup_count: usize, +} + +impl Default for MainEccRtRecord { + fn default() -> Self { + Self { + proof_idx: 0, + idx: 0, + fork_id: 0, + round_idx: 0, + num_rounds: 0, + is_first: false, + is_last: false, + tidx: 0, + out_tidx: 0, + alpha_tidx: 0, + value: EF::ZERO, + out_value: EF::ZERO, + alpha: EF::ZERO, + alpha_pows: [EF::ZERO; 49], + ev1: EF::ZERO, + ev2: EF::ZERO, + ev3: EF::ZERO, + claim_in: EF::ZERO, + claim_out: EF::ZERO, + sel_add: EF::ZERO, + sel_bypass: EF::ZERO, + sel_export: EF::ZERO, + s0: [EF::ZERO; 7], + x0: [EF::ZERO; 7], + y0: [EF::ZERO; 7], + x1: [EF::ZERO; 7], + y1: [EF::ZERO; 7], + x3: [EF::ZERO; 7], + y3: [EF::ZERO; 7], + sum_x: [EF::ZERO; 7], + sum_y: [EF::ZERO; 7], + eq_in: EF::ZERO, + eq_out: EF::ZERO, + last_in: EF::ZERO, + last_out: EF::ZERO, + export_out_in: EF::ZERO, + export_out_out: EF::ZERO, + export_rt_in: EF::ZERO, + export_rt_out: EF::ZERO, + quark_in: EF::ZERO, + quark_factor: EF::ZERO, + quark_out: EF::ZERO, + add_eval: EF::ZERO, + bypass_eval: EF::ZERO, + export_eval: EF::ZERO, + lte_out_point: [EF::ZERO; 32], + lte_rt_point: [EF::ZERO; 32], + lte_prefix_acc: [EF::ZERO; 33], + lte_less_acc: [EF::ZERO; 33], + lte_bits: [false; 32], + lte_active: [false; 32], + quark_prefix_count: 0, + quark_layer_n: 0, + quark_parity: false, + lookup_count: 0, + } + } +} + +#[derive(Clone, Debug, Default)] +pub struct MainTowerPointEqRecord { + pub proof_idx: usize, + pub idx: usize, + pub round_idx: usize, + pub global_value: EF, + pub tower_value: EF, + pub eq_in: EF, + pub eq_out: EF, +} + +#[derive(Clone, Debug, Default)] +pub struct MainFrontloadTermRecord { + pub proof_idx: usize, + pub idx: usize, + pub row_idx: usize, + pub node_idx: usize, + pub eval_idx: usize, + pub has_eval_factor: bool, + pub instance_idx: usize, + pub challenge_idx: usize, + pub global_round_idx: usize, + pub has_global_factor: bool, + pub is_wit: bool, + pub is_const: bool, + pub is_instance: bool, + pub is_challenge: bool, + pub is_add: bool, + pub is_sub: bool, + pub is_neg: bool, + pub is_mul: bool, + pub is_fold: bool, + pub is_tail: bool, + pub constraint_idx: usize, + pub alpha: EF, + pub arg0: EF, + pub arg1: EF, + pub value: EF, + pub chip_acc_in: EF, + pub chip_acc_out: EF, + pub is_last_chip_step: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct TowerMainPointRecord { + pub proof_idx: usize, + pub idx: usize, + pub round_idx: usize, + pub value: EF, + pub lookup_count: usize, +} + +#[derive(Clone, Debug, Default)] +pub struct MainGlobalSumcheckRoundRecord { + pub evaluations: [EF; 4], + pub challenge: EF, + pub challenge_tidx: usize, + pub claim_in: EF, + pub claim_out: EF, + pub point_lookup_count: usize, +} + +#[derive(Clone, Debug, Default)] +pub struct MainGlobalSumcheckRecord { + pub proof_idx: usize, + pub expected: EF, + pub rounds: Vec, +} + +impl MainGlobalSumcheckRecord { + pub fn total_rows(&self) -> usize { + self.rounds.len().max(1) + } } #[derive(Clone, Debug, Default)] @@ -99,13 +937,33 @@ pub struct TowerChipTranscriptRange { /// Index into `Preflight::fork_transcripts`. pub fork_idx: usize, pub tower_replay: TowerReplayResult, + pub rotation_replay: Option, + pub ecc_replay: Option, +} + +#[derive(Clone, Debug, Default)] +pub struct EccReplayClaims { + pub out_rt_tidxs: Vec, + pub alpha_tidx: usize, + pub rt_tidxs: Vec, +} + +#[derive(Clone, Debug, Default)] +pub struct RotationReplayClaims { + pub left_point: Vec, + pub right_point: Vec, + pub origin_point: Vec, + pub origin_tidxs: Vec, + pub left_evals: Vec, + pub right_evals: Vec, + pub target_evals: Vec, } #[derive(Clone, Debug, Default)] pub struct BatchConstraintPreflight { pub lambda_tidx: usize, pub tidx_before_univariate: usize, - pub sumcheck_rnd: Vec, + pub sumcheck_rnd: Vec, pub eq_ns_frontloaded: Vec, pub eq_sharp_ns_frontloaded: Vec, } diff --git a/ceno_recursion_v2/src/tower/input/air.rs b/ceno_recursion_v2/src/tower/input/air.rs index 06138b2ea..ab4f3f716 100644 --- a/ceno_recursion_v2/src/tower/input/air.rs +++ b/ceno_recursion_v2/src/tower/input/air.rs @@ -53,6 +53,7 @@ pub struct TowerInputCols { /// Transcript index pub tidx: T, + pub fork_final_sample_tidx: T, pub r0_claim: [T; D_EF], pub w0_claim: [T; D_EF], @@ -236,7 +237,7 @@ impl Air for TowerInputAir { // layers_cumulative(n) = // n*(SUMCHECK_INIT_LEN + POST_SUMCHECK_LEN) // + n*(n+1)/2*ROUND_LEN - // + (n-1)*ALPHA_LEN, for n > 0. + // + n*ALPHA_LEN, for n > 0, including native's final post-merge alpha. let read_active_layers = local.read_tower_vars - local.has_read_out; let write_active_layers = local.write_tower_vars - local.has_write_out; let logup_active_layers = local.logup_tower_vars - local.has_logup_out; @@ -248,7 +249,7 @@ impl Air for TowerInputAir { let round_span = num_layers.clone() * (num_layers.clone() + AB::Expr::ONE) * AB::Expr::from_usize(ROUND_LEN / 2); - let alpha_span = (num_layers.clone() - AB::Expr::ONE) * AB::Expr::from_usize(ALPHA_LEN); + let alpha_span = num_layers.clone() * AB::Expr::from_usize(ALPHA_LEN); let tidx_after_gkr_layers = tidx_after_alpha_beta.clone() + has_interactions.clone() * (fixed_span + round_span + alpha_span + claim_span); // 1. TowerLayerInputBus @@ -311,7 +312,7 @@ impl Air for TowerInputAir { local.proof_idx, ForkFinalSampleMessage { fork_id: local.fork_id.into(), - tidx: tidx_after_gkr_layers.clone(), + tidx: local.fork_final_sample_tidx.into(), }, local.is_enabled * has_tower_transcript.clone(), ); diff --git a/ceno_recursion_v2/src/tower/input/trace.rs b/ceno_recursion_v2/src/tower/input/trace.rs index 7dce305a2..097191793 100644 --- a/ceno_recursion_v2/src/tower/input/trace.rs +++ b/ceno_recursion_v2/src/tower/input/trace.rs @@ -13,6 +13,7 @@ pub struct TowerInputRecord { pub idx: usize, pub fork_id: usize, pub tidx: usize, + pub fork_final_sample_tidx: usize, pub n_logup: usize, pub alpha_logup: EF, pub beta: EF, @@ -67,10 +68,9 @@ impl RowMajorChip for TowerInputTraceGenerator { let mut prev_proof_idx = usize::MAX; let mut prev_idx = usize::MAX; - for (row_idx, (row_data, (record, q0_claim))) in data_slice + for (row_data, (record, q0_claim)) in data_slice .chunks_exact_mut(width) .zip(gkr_input_records.iter().zip(q0_claims.iter())) - .enumerate() { let cols: &mut TowerInputCols = row_data.borrow_mut(); let is_new_proof_idx = prev_proof_idx != record.proof_idx; @@ -84,6 +84,7 @@ impl RowMajorChip for TowerInputTraceGenerator { cols.fork_id = F::from_usize(record.fork_id); cols.tidx = F::from_usize(record.tidx); + cols.fork_final_sample_tidx = F::from_usize(record.fork_final_sample_tidx); cols.n_logup = F::from_usize(record.n_logup); IsZeroSubAir.generate_subrow( @@ -189,18 +190,6 @@ impl RowMajorChip for TowerInputTraceGenerator { .as_basis_coefficients_slice() .try_into() .unwrap(); - if std::env::var_os("CENO_REC_V2_DEBUG_TOWER_INPUT").is_some() { - eprintln!( - "rec-v2-debug module=tower_input source=trace row={} proof_idx={} idx={} is_first_idx={} is_first={} tidx={} n_logup={}", - row_idx, - record.proof_idx, - record.idx, - is_new_proof_idx, - is_new_idx, - record.tidx, - record.n_logup, - ); - } prev_proof_idx = record.proof_idx; prev_idx = record.idx; } diff --git a/ceno_recursion_v2/src/tower/layer/air.rs b/ceno_recursion_v2/src/tower/layer/air.rs index 1f410dae5..f4d94eb0a 100644 --- a/ceno_recursion_v2/src/tower/layer/air.rs +++ b/ceno_recursion_v2/src/tower/layer/air.rs @@ -49,6 +49,8 @@ pub struct TowerLayerCols { /// Sampled batching challenge pub lambda: [T; D_EF], + /// Native tower samples one final batching challenge after the last merge. + pub final_alpha: [T; D_EF], /// Challenge inherited from previous layer pub lambda_prime: [T; D_EF], /// Reduction point @@ -83,6 +85,7 @@ pub struct TowerLayerCols { pub write_weight: [T; D_EF], pub logup_p_weight: [T; D_EF], pub logup_q_weight: [T; D_EF], + pub weighted_prime_fold: [T; D_EF], /// Received from TowerLayerSumcheckAir pub eq_at_r_prime: [T; D_EF], @@ -361,6 +364,21 @@ where ext_field_multiply::(local.logup_q_weight, local.logup_q_xi), ), ); + let weighted_prime_fold = ext_field_add::( + ext_field_add::( + ext_field_multiply::(local.read_weight, local.read_claim_prime), + ext_field_multiply::(local.write_weight, local.write_claim_prime), + ), + ext_field_add::( + ext_field_multiply::(local.logup_p_weight, logup_p_cross.clone()), + ext_field_multiply::(local.logup_q_weight, logup_q_cross.clone()), + ), + ); + assert_array_eq( + &mut builder.when(enabled_not_dummy.clone()), + local.weighted_prime_fold, + weighted_prime_fold, + ); assert_zeros( &mut builder.when(AB::Expr::ONE - local.read_active), local.read_weight.map(Into::into), @@ -509,7 +527,7 @@ where local.proof_idx, TowerLayerOutputMessage { idx: local.idx.into(), - tidx: tidx_end, + tidx: tidx_end + is_last.clone() * AB::Expr::from_usize(ALPHA_LEN), layer_idx_end: local.layer_idx.into(), input_layer_claim: folded_claim.map(Into::into), lambda: local.lambda.map(Into::into), @@ -538,21 +556,13 @@ where ); // 3. TowerSumcheckOutputBus // 3a. Receive sumcheck results - let weighted_prime_fold = ext_field_add::( - ext_field_add::( - ext_field_multiply::(local.read_weight, local.read_claim_prime), - ext_field_multiply::(local.write_weight, local.write_claim_prime), - ), - ext_field_add::( - ext_field_multiply::(local.logup_p_weight, logup_p_cross), - ext_field_multiply::(local.logup_q_weight, logup_q_cross), - ), - ); let expected_sumcheck_claim_out = - ext_field_multiply::(weighted_prime_fold, local.eq_at_r_prime); - let _ = expected_sumcheck_claim_out; - // TODO(recursion-v2): re-enable the native tower expected-evaluation - // equality here once the full tower transcript/replay oracle is wired. + ext_field_multiply::(local.weighted_prime_fold, local.eq_at_r_prime); + assert_array_eq( + &mut builder.when(enabled_not_dummy.clone()), + local.sumcheck_claim_out, + expected_sumcheck_claim_out, + ); self.sumcheck_output_bus.receive( builder, local.proof_idx, @@ -576,7 +586,7 @@ where sumcheck_round: local.layer_idx + AB::Expr::ONE, challenge: local.mu.map(Into::into), }, - local.is_enabled * (AB::Expr::ONE - is_last) * is_not_dummy.clone(), + local.is_enabled * (AB::Expr::ONE - is_last.clone()) * is_not_dummy.clone(), ); /////////////////////////////////////////////////////////////////////// @@ -699,5 +709,32 @@ where local.is_enabled * is_not_dummy.clone(), ); } + let final_alpha_label_tidx = tidx + AB::Expr::from_usize(D_EF); + for (offset, value) in LABEL_COMBINE_FIELDS.into_iter().enumerate() { + self.forked_transcript_bus.receive( + builder, + local.proof_idx, + ForkedTranscriptBusMessage { + fork_id: local.fork_id.into(), + tidx: final_alpha_label_tidx.clone() + AB::Expr::from_usize(offset), + value: AB::Expr::from_u32(value), + is_sample: AB::Expr::ZERO, + }, + is_last.clone() * is_not_dummy.clone(), + ); + } + for i in 0..D_EF { + self.forked_transcript_bus.receive( + builder, + local.proof_idx, + ForkedTranscriptBusMessage { + fork_id: local.fork_id.into(), + tidx: final_alpha_label_tidx.clone() + AB::Expr::from_usize(LABEL_COMBINE + i), + value: local.final_alpha[i].into(), + is_sample: AB::Expr::ONE, + }, + is_last.clone() * is_not_dummy.clone(), + ); + } } } diff --git a/ceno_recursion_v2/src/tower/layer/trace.rs b/ceno_recursion_v2/src/tower/layer/trace.rs index 6949c6dc0..024debf94 100644 --- a/ceno_recursion_v2/src/tower/layer/trace.rs +++ b/ceno_recursion_v2/src/tower/layer/trace.rs @@ -15,10 +15,7 @@ fn ef_one() -> EF { EF::ONE } -fn weight_bases( - record: &TowerLayerRecord, - layer_idx: usize, -) -> ([F; D_EF], [F; D_EF], [F; D_EF], [F; D_EF]) { +pub(crate) fn weight_values(record: &TowerLayerRecord, layer_idx: usize) -> (EF, EF, EF, EF) { let alpha = record.lambda_at(layer_idx); let mut pow = ef_one(); let has_read = record.read_counts.iter().any(|&count| count != 0); @@ -52,6 +49,15 @@ fn weight_bases( } else { (EF::ZERO, EF::ZERO) }; + (read_weight, write_weight, logup_p_weight, logup_q_weight) +} + +fn weight_bases( + record: &TowerLayerRecord, + layer_idx: usize, +) -> ([F; D_EF], [F; D_EF], [F; D_EF], [F; D_EF]) { + let (read_weight, write_weight, logup_p_weight, logup_q_weight) = + weight_values(record, layer_idx); ( read_weight .as_basis_coefficients_slice() @@ -72,6 +78,52 @@ fn weight_bases( ) } +pub(crate) fn weighted_prime_fold_for_layer( + record: &TowerLayerRecord, + tower: &TowerTowerEvalRecord, + layer_idx: usize, +) -> EF { + let (read_weight, write_weight, logup_p_weight, logup_q_weight) = + weight_values(record, layer_idx); + let read_prime = if record.read_active_at(layer_idx) { + record + .read_prime_claims + .get(layer_idx) + .copied() + .unwrap_or(EF::ZERO) + } else { + EF::ZERO + }; + let write_prime = if record.write_active_at(layer_idx) { + record + .write_prime_claims + .get(layer_idx) + .copied() + .unwrap_or(EF::ZERO) + } else { + EF::ZERO + }; + let (logup_p_cross, logup_q_cross) = if record.logup_active_at(layer_idx) { + let logup_quad = tower + .logup_layers + .get(layer_idx) + .and_then(|rows| rows.first()) + .copied() + .unwrap_or([EF::ZERO; 4]); + ( + logup_quad[0] * logup_quad[3] + logup_quad[1] * logup_quad[2], + logup_quad[2] * logup_quad[3], + ) + } else { + (EF::ZERO, EF::ZERO) + }; + + read_weight * read_prime + + write_weight * write_prime + + logup_p_weight * logup_p_cross + + logup_q_weight * logup_q_cross +} + /// Minimal record for parallel tower layer trace generation #[derive(Debug, Clone, Default)] pub struct TowerLayerRecord { @@ -82,6 +134,7 @@ pub struct TowerLayerRecord { pub tidx: usize, pub layer_claims: Vec<[EF; 4]>, pub lambdas: Vec, + pub final_alpha: EF, pub eq_at_r_primes: Vec, pub read_counts: Vec, pub write_counts: Vec, @@ -305,6 +358,7 @@ impl RowMajorChip for TowerLayerTraceGenerator { cols.layer_idx = F::ZERO; cols.tidx = F::from_usize(record.tidx); cols.lambda = [F::ZERO; D_EF]; + cols.final_alpha = [F::ZERO; D_EF]; let mut lambda_prime_one = [F::ZERO; D_EF]; lambda_prime_one[0] = F::ONE; cols.lambda_prime = lambda_prime_one; @@ -336,6 +390,7 @@ impl RowMajorChip for TowerLayerTraceGenerator { cols.write_weight = [F::ZERO; D_EF]; cols.logup_p_weight = [F::ZERO; D_EF]; cols.logup_q_weight = [F::ZERO; D_EF]; + cols.weighted_prime_fold = [F::ZERO; D_EF]; cols.eq_at_r_prime = [F::ZERO; D_EF]; cols.r0_claim.copy_from_slice(q0_basis); cols.w0_claim.copy_from_slice(q0_basis); @@ -363,6 +418,11 @@ impl RowMajorChip for TowerLayerTraceGenerator { .as_basis_coefficients_slice() .try_into() .unwrap(); + cols.final_alpha = record + .final_alpha + .as_basis_coefficients_slice() + .try_into() + .unwrap(); cols.lambda_prime = record .lambda_prime_at(layer_idx) .as_basis_coefficients_slice() @@ -470,8 +530,14 @@ impl RowMajorChip for TowerLayerTraceGenerator { cols.write_weight = write_weight; cols.logup_p_weight = logup_p_weight; cols.logup_q_weight = logup_q_weight; - cols.eq_at_r_prime = record - .eq_at(layer_idx) + let eq_at_r_prime = record.eq_at(layer_idx); + cols.eq_at_r_prime = eq_at_r_prime + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + let weighted_prime_fold = + weighted_prime_fold_for_layer(record, tower, layer_idx); + cols.weighted_prime_fold = weighted_prime_fold .as_basis_coefficients_slice() .try_into() .unwrap(); diff --git a/ceno_recursion_v2/src/tower/main_point.rs b/ceno_recursion_v2/src/tower/main_point.rs new file mode 100644 index 000000000..5b79568db --- /dev/null +++ b/ceno_recursion_v2/src/tower/main_point.rs @@ -0,0 +1,99 @@ +use core::borrow::{Borrow, BorrowMut}; + +use openvm_stark_backend::{ + BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder, +}; +use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing}; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use stark_recursion_circuit_derive::AlignedBorrow; + +use crate::{ + bus::{TowerMainPointBus, TowerMainPointMessage}, + system::TowerMainPointRecord, + tracegen::RowMajorChip, +}; + +#[repr(C)] +#[derive(AlignedBorrow, Debug)] +pub struct TowerMainPointCols { + pub is_enabled: T, + pub proof_idx: T, + pub idx: T, + pub round_idx: T, + pub lookup_count: T, + pub value: [T; D_EF], +} + +pub struct TowerMainPointAir { + pub tower_point_bus: TowerMainPointBus, +} + +impl BaseAir for TowerMainPointAir { + fn width(&self) -> usize { + TowerMainPointCols::::width() + } +} + +impl BaseAirWithPublicValues for TowerMainPointAir {} +impl PartitionedBaseAir for TowerMainPointAir {} + +impl Air for TowerMainPointAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local_row = main.row_slice(0).expect("main row exists"); + let local: &TowerMainPointCols = (*local_row).borrow(); + + builder.assert_bool(local.is_enabled); + self.tower_point_bus.add_key_with_lookups( + builder, + local.proof_idx, + TowerMainPointMessage { + idx: local.idx.into(), + round_idx: local.round_idx.into(), + value: local.value.map(Into::into), + }, + local.is_enabled * local.lookup_count, + ); + } +} + +pub struct TowerMainPointTraceGenerator; + +impl RowMajorChip for TowerMainPointTraceGenerator { + type Ctx<'a> = &'a [TowerMainPointRecord]; + + fn generate_trace( + &self, + records: &Self::Ctx<'_>, + required_height: Option, + ) -> Option> { + let width = TowerMainPointCols::::width(); + let num_valid_rows = records.len().max(1); + let height = if let Some(height) = required_height { + if height < num_valid_rows { + return None; + } + height + } else { + num_valid_rows.next_power_of_two() + }; + let mut trace = vec![F::ZERO; height * width]; + for (row_idx, record) in records.iter().enumerate() { + let row = &mut trace[row_idx * width..(row_idx + 1) * width]; + let cols: &mut TowerMainPointCols = row.borrow_mut(); + cols.is_enabled = F::ONE; + cols.proof_idx = F::from_usize(record.proof_idx); + cols.idx = F::from_usize(record.idx); + cols.round_idx = F::from_usize(record.round_idx); + cols.lookup_count = F::from_usize(record.lookup_count); + cols.value = record + .value + .as_basis_coefficients_slice() + .try_into() + .unwrap(); + } + Some(RowMajorMatrix::new(trace, width)) + } +} diff --git a/ceno_recursion_v2/src/tower/mod.rs b/ceno_recursion_v2/src/tower/mod.rs index 2ad7b3cc9..5b1bc4b1b 100644 --- a/ceno_recursion_v2/src/tower/mod.rs +++ b/ceno_recursion_v2/src/tower/mod.rs @@ -64,9 +64,11 @@ use strum::EnumCount; use tracing::error; use crate::{ + main::selector::selector_formula_point_lookup_counts, system::{ - AirModule, BusIndexManager, BusInventory, GlobalCtxCpu, Preflight, RecursionField, - RecursionProof, RecursionVk, TowerChipTranscriptRange, TraceGenModule, + AirModule, BusIndexManager, BusInventory, GlobalCtxCpu, MainSelectorPointSourceKind, + Preflight, RecursionField, RecursionProof, RecursionVk, TowerChipTranscriptRange, + TraceGenModule, }, tower::{ alpha_pow::{TowerAlphaPowAir, TowerAlphaPowTraceGenerator}, @@ -74,8 +76,9 @@ use crate::{ TowerActivityBus, TowerInputShapeBus, TowerLayerInputBus, TowerLayerOutputBus, TowerShapeBus, }, - input::{TowerInputAir, TowerInputRecord, TowerInputTraceGenerator}, + input::{TowerInputAir, TowerInputTraceGenerator}, layer::{TowerLayerAir, TowerLayerRecord, TowerLayerTraceGenerator}, + main_point::{TowerMainPointAir, TowerMainPointTraceGenerator}, shape::{ TowerActivityAir, TowerActivityTraceGenerator, TowerShapeAir, TowerShapeRecord, TowerShapeTraceGenerator, @@ -232,11 +235,13 @@ pub mod tower_transcript_len { pub mod alpha_pow; pub mod input; pub mod layer; +pub mod main_point; pub mod shape; pub mod sumcheck; pub(crate) use shape::{TOWER_ACTIVITY_LOGUP, TOWER_ACTIVITY_READ, TOWER_ACTIVITY_WRITE}; #[allow(clippy::module_inception)] mod tower; +pub(crate) use input::TowerInputRecord; pub(crate) use tower::{TowerReplayResult, replay_tower_proof_precomputed}; pub struct TowerModule { // Global bus inventory @@ -262,13 +267,14 @@ pub(crate) struct TowerTowerEvalRecord { pub(crate) struct TowerBlobCpu { shape_records: Vec, - input_records: Vec, + pub(crate) input_records: Vec, /// Per-proof q0 claims matching input_records (one per proof). proof_q0_claims: Vec, layer_records: Vec, tower_records: Vec, sumcheck_records: Vec, mus_records: Vec>, + main_point_records: Vec, /// Per-chip q0 claims matching layer_records. q0_claims: Vec, } @@ -278,6 +284,7 @@ pub(crate) struct TowerTranscriptSchedule { pub(crate) alpha_logup: EF, pub(crate) beta: EF, pub(crate) lambdas: Vec, + pub(crate) final_alpha: EF, pub(crate) mus: Vec, pub(crate) ris: Vec, } @@ -320,6 +327,8 @@ impl TowerModule { tidx, fork_idx: 0, // unused in forked flow tower_replay, + rotation_replay: None, + ecc_replay: None, }); } } @@ -411,19 +420,20 @@ pub(crate) fn record_and_replay_tower_preflight( where TS: FiatShamirTranscript, { - let schedule = record_gkr_transcript(ts, chip_idx, chip_proof); match circuit_vk_for_idx(child_vk, chip_idx) { - Some(circuit_vk) => match replay_tower_proof_precomputed(chip_proof, circuit_vk, &schedule) - { - Ok(replay) => replay, - Err(err) => { - error!( - ?err, - chip_idx, "failed to replay tower proof during preflight" - ); - TowerReplayResult::default() + Some(circuit_vk) => { + let schedule = record_gkr_transcript(ts, circuit_vk, chip_proof); + match replay_tower_proof_precomputed(chip_proof, circuit_vk, &schedule) { + Ok(replay) => replay, + Err(err) => { + error!( + ?err, + chip_idx, "failed to replay tower proof during preflight" + ); + TowerReplayResult::default() + } } - }, + } None => TowerReplayResult::default(), } } @@ -492,6 +502,39 @@ fn build_tower_shape_record( } } +fn tower_num_variables( + chip_proof: &ZKVMChipProof, + circuit_vk: &VerifyingKey, +) -> Vec { + let cs = &circuit_vk.cs; + let num_instances: usize = chip_proof.num_instances.iter().copied().sum(); + let next_pow2_instance = next_pow2_instance_padding(num_instances); + let mut num_var_with_rotation = next_pow2_instance.ilog2() as usize; + if cs.has_ecc_ops() { + num_var_with_rotation += 1; + } + num_var_with_rotation += cs.rotation_vars().unwrap_or(0); + + let group_num_vars = |op_count: usize| num_var_with_rotation + grouped_op_vars(op_count); + chip_proof + .r_out_evals + .iter() + .map(|_| group_num_vars(cs.num_reads())) + .chain( + chip_proof + .w_out_evals + .iter() + .map(|_| group_num_vars(cs.num_writes())), + ) + .chain( + chip_proof + .lk_out_evals + .iter() + .map(|_| group_num_vars(cs.num_lks())), + ) + .collect() +} + #[allow(clippy::too_many_arguments)] fn build_chip_records( proof_idx: usize, @@ -504,6 +547,7 @@ fn build_chip_records( replay: &TowerReplayResult, schedule: &TowerTranscriptSchedule, tidx: usize, + fork_final_sample_tidx: usize, ) -> Result<( TowerShapeRecord, TowerInputRecord, @@ -511,6 +555,7 @@ fn build_chip_records( TowerTowerEvalRecord, TowerSumcheckRecord, Vec, + Vec, EF, )> { let shape_record = build_tower_shape_record(proof_idx, idx, air_idx, chip_proof, circuit_vk); @@ -570,6 +615,7 @@ fn build_chip_records( tidx: tidx + tower_transcript_len::ALPHA_BETA_LEN, layer_claims: Vec::with_capacity(layer_count), lambdas: vec![EF::ZERO; layer_count], + final_alpha: schedule.final_alpha, eq_at_r_primes: vec![EF::ZERO; layer_count], read_counts: vec![1; layer_count], write_counts: vec![1; layer_count], @@ -787,6 +833,39 @@ fn build_chip_records( } else { schedule.mus.last().copied().unwrap_or(EF::ZERO) }; + let main_point_records = if layer_count == 0 || shape_record.num_vars == 0 { + Vec::new() + } else { + let final_layer = layer_count - 1; + let start = TowerSumcheckRecord::layer_start_index(final_layer); + let end = start + TowerSumcheckRecord::layer_rounds(final_layer); + eyre::ensure!( + sumcheck_record.ris.len() >= end, + "tower point source length {} is shorter than final layer end {}", + sumcheck_record.ris.len(), + end + ); + let mut rt_tower = sumcheck_record.ris[start..end].to_vec(); + rt_tower.push(layer_output_mu); + eyre::ensure!( + rt_tower.len() >= shape_record.num_vars, + "tower point length {} is shorter than main point length {}", + rt_tower.len(), + shape_record.num_vars + ); + rt_tower[rt_tower.len() - shape_record.num_vars..] + .iter() + .copied() + .enumerate() + .map(|(round_idx, value)| crate::system::TowerMainPointRecord { + proof_idx, + idx, + round_idx, + value, + lookup_count: 1, + }) + .collect() + }; let mut read_out_evals = [EF::ZERO; 2]; if let Some(values) = chip_proof.r_out_evals.first() { for (dst, src) in read_out_evals.iter_mut().zip(values.iter().take(2)) { @@ -810,6 +889,7 @@ fn build_chip_records( idx, fork_id, tidx, + fork_final_sample_tidx, n_logup: layer_count, alpha_logup: schedule.alpha_logup, beta: schedule.beta, @@ -831,92 +911,201 @@ fn build_chip_records( layer_output_mu, }; - if std::env::var_os("CENO_TOWER_DEBUG").is_some() { + Ok(( + shape_record, + input_record, + layer_record, + tower_record, + sumcheck_record, + mus_record, + main_point_records, + q0_claim, + )) +} + +fn build_chip_input_record( + proof_idx: usize, + idx: usize, + fork_id: usize, + chip_proof: &ZKVMChipProof, + circuit_vk: &VerifyingKey, + schedule: &TowerTranscriptSchedule, + tidx: usize, + fork_final_sample_tidx: usize, +) -> Result { + let shape_record = build_tower_shape_record(proof_idx, idx, idx, chip_proof, circuit_vk); + let layer_count = shape_record.max_layer_count; + + let read_count = usize::from(shape_record.has_read); + let write_count = usize::from(shape_record.has_write); + let logup_count = usize::from(shape_record.has_logup); + let mut read_layers = vec![Vec::with_capacity(read_count); layer_count]; + let mut write_layers = vec![Vec::with_capacity(write_count); layer_count]; + let mut logup_layers = vec![Vec::with_capacity(logup_count); layer_count]; + + for (spec_idx, rounds) in chip_proof.tower_proof.prod_specs_eval.iter().enumerate() { for layer_idx in 0..layer_count { - let alpha = layer_record.lambda_at(layer_idx); - let mut pow = EF::ONE; - let read_weight = if shape_record.has_read { - let weight = pow; - pow *= alpha; - if layer_record.read_active_at(layer_idx) { - weight - } else { - EF::ZERO + let mut pair = [EF::ZERO; 2]; + if let Some(values) = rounds.get(layer_idx) { + for (dst, src) in pair.iter_mut().zip(values.iter().take(2)) { + *dst = *src; } + } + if spec_idx < read_count && layer_idx + 1 < shape_record.read_tower_vars { + read_layers[layer_idx].push(pair); + } else if layer_idx + 1 < shape_record.write_tower_vars { + write_layers[layer_idx].push(pair); + } + } + } + + for rounds in &chip_proof.tower_proof.logup_specs_eval { + for layer_idx in 0..layer_count { + let mut quad = [EF::ZERO; 4]; + if let Some(values) = rounds.get(layer_idx) { + for (dst, src) in quad.iter_mut().zip(values.iter().take(4)) { + *dst = *src; + } + } + if layer_idx + 1 < shape_record.logup_tower_vars { + logup_layers[layer_idx].push(quad); + } + } + } + + for layer_idx in 0..layer_count { + let read_len = read_layers.get(layer_idx).map(Vec::len).unwrap_or(0); + let write_len = write_layers.get(layer_idx).map(Vec::len).unwrap_or(0); + let logup_len = logup_layers.get(layer_idx).map(Vec::len).unwrap_or(0); + eyre::ensure!( + read_len <= 1 && write_len <= 1 && logup_len <= 1, + "grouped tower layer {layer_idx} has too many specs: read={read_len}, write={write_len}, logup={logup_len}" + ); + } + + let layer_output_claim = if layer_count == 0 { + EF::ZERO + } else { + let layer_idx = layer_count - 1; + let alpha = schedule.lambdas.get(layer_idx).copied().unwrap_or(EF::ZERO); + let mut pow = EF::ONE; + let read_weight = if shape_record.has_read { + let weight = pow; + pow *= alpha; + if read_layers + .get(layer_idx) + .is_some_and(|rows| !rows.is_empty()) + { + weight } else { EF::ZERO - }; - let write_weight = if shape_record.has_write { - let weight = pow; - pow *= alpha; - if layer_record.write_active_at(layer_idx) { - weight - } else { - EF::ZERO - } + } + } else { + EF::ZERO + }; + let write_weight = if shape_record.has_write { + let weight = pow; + pow *= alpha; + if write_layers + .get(layer_idx) + .is_some_and(|rows| !rows.is_empty()) + { + weight } else { EF::ZERO - }; - let (logup_p_weight, logup_q_weight) = if shape_record.has_logup { - let p_weight = pow; - let q_weight = pow * alpha; - if layer_record.logup_active_at(layer_idx) { - (p_weight, q_weight) - } else { - (EF::ZERO, EF::ZERO) - } + } + } else { + EF::ZERO + }; + let (logup_p_weight, logup_q_weight) = if shape_record.has_logup { + let p_weight = pow; + let q_weight = pow * alpha; + if logup_layers + .get(layer_idx) + .is_some_and(|rows| !rows.is_empty()) + { + (p_weight, q_weight) } else { (EF::ZERO, EF::ZERO) - }; - let read_prime = layer_record.read_prime_claims[layer_idx]; - let write_prime = layer_record.write_prime_claims[layer_idx]; - let logup_quad = tower_record - .logup_layers - .get(layer_idx) - .and_then(|rows| rows.first()) - .copied() - .unwrap_or([EF::ZERO; 4]); - let logup_p_cross = logup_quad[0] * logup_quad[3] + logup_quad[1] * logup_quad[2]; - let logup_q_cross = logup_quad[2] * logup_quad[3]; - let expected = layer_record.eq_at_r_primes[layer_idx] - * (read_weight * read_prime - + write_weight * write_prime - + logup_p_weight * logup_p_cross - + logup_q_weight * logup_q_cross); - let replay_claim_out = replay - .layers - .get(layer_idx) - .map(|layer| layer.claim_out) - .unwrap_or(EF::ZERO); - let replay_eq = replay - .layers - .get(layer_idx) - .map(|layer| layer.eq_at_r) - .unwrap_or(EF::ZERO); - if expected != replay_claim_out { - panic!( - "tower debug mismatch proof_idx={proof_idx} idx={idx} layer={layer_idx} expected={expected:?} replay_claim_out={replay_claim_out:?} alpha={alpha:?} eq={:?} replay_eq={replay_eq:?} weights=({read_weight:?},{write_weight:?},{logup_p_weight:?},{logup_q_weight:?}) primes=({read_prime:?},{write_prime:?},{logup_p_cross:?},{logup_q_cross:?}) active=({},{},{}) shape=({},{},{})", - layer_record.eq_at_r_primes[layer_idx], - layer_record.read_active_at(layer_idx), - layer_record.write_active_at(layer_idx), - layer_record.logup_active_at(layer_idx), - shape_record.has_read, - shape_record.has_write, - shape_record.has_logup, - ); } + } else { + (EF::ZERO, EF::ZERO) + }; + let mu = schedule.mus.get(layer_idx).copied().unwrap_or(EF::ZERO); + let read_claim = read_layers + .get(layer_idx) + .map(|rows| accumulate_prod_claims(rows, EF::ZERO, EF::ZERO, mu).0) + .unwrap_or(EF::ZERO); + let write_claim = write_layers + .get(layer_idx) + .map(|rows| accumulate_prod_claims(rows, EF::ZERO, EF::ZERO, mu).0) + .unwrap_or(EF::ZERO); + let logup_quad = logup_layers + .get(layer_idx) + .and_then(|rows| rows.first()) + .copied() + .unwrap_or([EF::ZERO; 4]); + read_weight * read_claim + + write_weight * write_claim + + logup_p_weight * interpolate_pair([logup_quad[0], logup_quad[1]], mu) + + logup_q_weight * interpolate_pair([logup_quad[2], logup_quad[3]], mu) + }; + + let layer_output_lambda = if layer_count == 0 { + EF::ZERO + } else { + schedule.lambdas.last().copied().unwrap_or(EF::ZERO) + }; + let layer_output_mu = if layer_count == 0 { + EF::ZERO + } else { + schedule.mus.last().copied().unwrap_or(EF::ZERO) + }; + let mut read_out_evals = [EF::ZERO; 2]; + if let Some(values) = chip_proof.r_out_evals.first() { + for (dst, src) in read_out_evals.iter_mut().zip(values.iter().take(2)) { + *dst = *src; + } + } + let mut write_out_evals = [EF::ZERO; 2]; + if let Some(values) = chip_proof.w_out_evals.first() { + for (dst, src) in write_out_evals.iter_mut().zip(values.iter().take(2)) { + *dst = *src; + } + } + let mut logup_out_evals = [EF::ZERO; 4]; + if let Some(values) = chip_proof.lk_out_evals.first() { + for (dst, src) in logup_out_evals.iter_mut().zip(values.iter().take(4)) { + *dst = *src; } } - Ok(( - shape_record, - input_record, - layer_record, - tower_record, - sumcheck_record, - mus_record, - q0_claim, - )) + Ok(TowerInputRecord { + proof_idx, + idx, + fork_id, + tidx, + fork_final_sample_tidx, + n_logup: layer_count, + alpha_logup: schedule.alpha_logup, + beta: schedule.beta, + read_out_evals, + write_out_evals, + logup_out_evals, + has_read_out: !chip_proof.r_out_evals.is_empty(), + has_write_out: !chip_proof.w_out_evals.is_empty(), + has_logup_out: !chip_proof.lk_out_evals.is_empty(), + has_read: shape_record.has_read, + has_write: shape_record.has_write, + has_logup: shape_record.has_logup, + read_tower_vars: shape_record.read_tower_vars, + write_tower_vars: shape_record.write_tower_vars, + logup_tower_vars: shape_record.logup_tower_vars, + max_layer_count: shape_record.max_layer_count, + input_layer_claim: layer_output_claim, + layer_output_lambda, + layer_output_mu, + }) } impl AirModule for TowerModule { @@ -945,7 +1134,7 @@ impl AirModule for TowerModule { layer_input_bus: self.layer_input_bus, layer_output_bus: self.layer_output_bus, sumcheck_challenge_bus: self.sumcheck_challenge_bus, - send_main: !crate::system::TOWER_PREFIX_ONLY, + send_main: true, }; let gkr_layer_air = TowerLayerAir { @@ -969,6 +1158,9 @@ impl AirModule for TowerModule { let tower_alpha_pow_air = TowerAlphaPowAir { alpha_pow_bus: self.alpha_pow_bus, }; + let tower_main_point_air = TowerMainPointAir { + tower_point_bus: self.bus_inventory.tower_main_point_bus, + }; vec![ Arc::new(tower_shape_air) as AirRef<_>, @@ -977,6 +1169,7 @@ impl AirModule for TowerModule { Arc::new(tower_alpha_pow_air) as AirRef<_>, Arc::new(gkr_layer_air) as AirRef<_>, Arc::new(gkr_sumcheck_air) as AirRef<_>, + Arc::new(tower_main_point_air) as AirRef<_>, ] } } @@ -1007,6 +1200,7 @@ pub(crate) fn build_gkr_blob( let mut tower_records = Vec::new(); let mut sumcheck_records = Vec::new(); let mut mus_records = Vec::new(); + let mut main_point_records = Vec::new(); let mut q0_claims = Vec::new(); eyre::ensure!( @@ -1034,6 +1228,44 @@ pub(crate) fn build_gkr_blob( entry.chip_idx, ) }); + let tower_idx_by_chip = sorted_pf_entries + .iter() + .enumerate() + .map(|(tower_idx, entry)| (entry.chip_idx, tower_idx)) + .collect::>(); + let mut selector_eval_records = preflight.main.selector_evals.clone(); + for record in &mut selector_eval_records { + record.proof_idx = proof_idx; + } + let selector_point_lookup_counts = + selector_formula_point_lookup_counts(&selector_eval_records); + let mut selector_tower_point_counts = + std::collections::BTreeMap::<(usize, usize), usize>::new(); + for record in &selector_eval_records { + if record.point_source != MainSelectorPointSourceKind::TowerMain { + continue; + } + let Some(tower_idx) = tower_idx_by_chip.get(&record.air_idx).copied() else { + continue; + }; + for round_idx in 0..record.out_point.len() { + let count = selector_point_lookup_counts + .get(&( + record.proof_idx, + record.idx, + record.air_idx, + record.selector_idx, + round_idx, + )) + .copied() + .unwrap_or(0); + if count != 0 { + *selector_tower_point_counts + .entry((tower_idx, round_idx)) + .or_default() += 1; + } + } + } for (entry_idx, pf_entry) in sorted_pf_entries.into_iter().enumerate() { let chip_idx = pf_entry.chip_idx; let chip_proof = proof @@ -1046,10 +1278,9 @@ pub(crate) fn build_gkr_blob( let fork_log = preflight.fork_log(pf_entry.fork_idx); ReadOnlyTranscript::new(fork_log, pf_entry.tidx) }; - let transcript_schedule = record_gkr_transcript(&mut ts, chip_idx, chip_proof); - let circuit_vk = circuit_vk_for_idx(child_vk, chip_idx) .ok_or_else(|| eyre::eyre!("missing circuit verifying key for index {chip_idx}"))?; + let transcript_schedule = record_gkr_transcript(&mut ts, circuit_vk, chip_proof); let replay = &pf_entry.tower_replay; let schedule = &transcript_schedule; @@ -1058,6 +1289,10 @@ pub(crate) fn build_gkr_blob( // by 0 or 1 within each proof_idx group). let idx = entry_idx; let tower_air_tidx = tower_pre_alpha_tidx(chip_proof, pf_entry.tidx); + let fork_final_sample_tidx = preflight + .fork_log(pf_entry.fork_idx) + .len() + .saturating_sub(D_EF); let ( shape_record, chip_input_record, @@ -1065,6 +1300,7 @@ pub(crate) fn build_gkr_blob( tower_record, sumcheck_record, mus_record, + mut chip_main_point_records, q0_claim, ) = build_chip_records( proof_idx, @@ -1077,54 +1313,14 @@ pub(crate) fn build_gkr_blob( replay, schedule, tower_air_tidx, + fork_final_sample_tidx, )?; - if std::env::var_os("CENO_REC_V2_DEBUG_TOWER").is_some() { - let out_eval_span = tower_air_tidx.saturating_sub(pf_entry.tidx); - let compact_layer_span = (0..shape_record.max_layer_count) - .map(|layer_idx| { - tower_transcript_len::compact_layer_span( - layer_idx, - layer_record - .read_counts - .get(layer_idx) - .copied() - .unwrap_or(0) - != 0, - layer_record - .write_counts - .get(layer_idx) - .copied() - .unwrap_or(0) - != 0, - layer_record - .logup_counts - .get(layer_idx) - .copied() - .unwrap_or(0) - != 0, - ) - }) - .sum::(); - let computed_tower_span = - out_eval_span + tower_transcript_len::ALPHA_BETA_LEN + compact_layer_span; - let fork_len = preflight.fork_log(pf_entry.fork_idx).len(); - eprintln!( - "rec-v2-debug module=tower source=trace proof_idx={proof_idx} fork_id={} chip_idx={} fork_len={} tower_start={} tower_alpha_tidx={} computed_span={} final_sample_tidx={} expected_final_tidx={} has_read_out={} has_write_out={} has_logup_out={} n_logup={}", - pf_entry.fork_idx, - chip_idx, - fork_len, - pf_entry.tidx, - tower_air_tidx, - computed_tower_span, - pf_entry.tidx + computed_tower_span, - fork_len.saturating_sub(D_EF), - chip_input_record.has_read_out, - chip_input_record.has_write_out, - chip_input_record.has_logup_out, - chip_input_record.n_logup, - ); + for record in &mut chip_main_point_records { + record.lookup_count += selector_tower_point_counts + .get(&(idx, record.round_idx)) + .copied() + .unwrap_or(0); } - input_records.push(chip_input_record); proof_q0_claims.push(q0_claim); shape_records.push(shape_record); @@ -1132,6 +1328,7 @@ pub(crate) fn build_gkr_blob( tower_records.push(tower_record); sumcheck_records.push(sumcheck_record); mus_records.push(mus_record); + main_point_records.extend(chip_main_point_records); q0_claims.push(q0_claim); } @@ -1184,10 +1381,83 @@ pub(crate) fn build_gkr_blob( tower_records, sumcheck_records, mus_records, + main_point_records, q0_claims, }) } +pub(crate) fn build_tower_input_records( + child_vk: &RecursionVk, + proofs: &[RecursionProof], + preflights: &[Preflight], +) -> Result> { + eyre::ensure!( + proofs.len() == preflights.len(), + "proof/preflight length mismatch" + ); + + let mut input_records = Vec::new(); + for (proof_idx, (proof, preflight)) in proofs.iter().zip(preflights).enumerate() { + let sorted_idx_by_chip: std::collections::BTreeMap = preflight + .proof_shape + .sorted_trace_vdata + .iter() + .enumerate() + .map(|(sorted_idx, (chip_idx, _))| (*chip_idx, sorted_idx)) + .collect(); + let mut sorted_pf_entries: Vec<_> = preflight.gkr.chips.iter().collect(); + sorted_pf_entries.sort_by_key(|entry| { + ( + sorted_idx_by_chip + .get(&entry.chip_idx) + .copied() + .unwrap_or(usize::MAX), + entry.chip_idx, + ) + }); + + for (entry_idx, pf_entry) in sorted_pf_entries.into_iter().enumerate() { + let chip_idx = pf_entry.chip_idx; + let chip_proof = proof + .chip_proofs + .get(&chip_idx) + .ok_or_else(|| eyre::eyre!("missing chip proof for chip {chip_idx}"))?; + let mut ts = { + let fork_log = preflight.fork_log(pf_entry.fork_idx); + ReadOnlyTranscript::new(fork_log, pf_entry.tidx) + }; + let circuit_vk = circuit_vk_for_idx(child_vk, chip_idx) + .ok_or_else(|| eyre::eyre!("missing circuit verifying key for index {chip_idx}"))?; + let transcript_schedule = record_gkr_transcript(&mut ts, circuit_vk, chip_proof); + let tower_air_tidx = tower_pre_alpha_tidx(chip_proof, pf_entry.tidx); + let fork_final_sample_tidx = preflight + .fork_log(pf_entry.fork_idx) + .len() + .saturating_sub(D_EF); + input_records.push(build_chip_input_record( + proof_idx, + entry_idx, + pf_entry.fork_idx, + chip_proof, + circuit_vk, + &transcript_schedule, + tower_air_tidx, + fork_final_sample_tidx, + )?); + } + } + + Ok(input_records) +} + +pub(crate) fn build_tower_main_point_records( + child_vk: &RecursionVk, + proofs: &[RecursionProof], + preflights: &[Preflight], +) -> Result> { + Ok(build_gkr_blob(child_vk, proofs, preflights)?.main_point_records) +} + pub(crate) fn collect_tower_range_checks( child_vk: &RecursionVk, proofs: &[RecursionProof], @@ -1263,7 +1533,7 @@ pub(crate) fn collect_tower_range_checks( pub(crate) fn record_gkr_transcript( ts: &mut TS, - _chip_idx: usize, + circuit_vk: &VerifyingKey, chip_proof: &ZKVMChipProof, ) -> TowerTranscriptSchedule where @@ -1297,15 +1567,14 @@ where // Reconstruct the transcript events consumed by tower-related AIRs. // This keeps preflight transcript history aligned with TowerLayer/Sumcheck/ // ProdClaim/LogupClaim transcript bus interactions. - let read_count = chip_proof.r_out_evals.len(); - let layer_count = chip_proof - .tower_proof - .logup_specs_eval + let num_variables = tower_num_variables(chip_proof, circuit_vk); + let num_prod_spec = chip_proof.r_out_evals.len() + chip_proof.w_out_evals.len(); + let layer_count = num_variables .iter() - .map(Vec::len) - .chain(chip_proof.tower_proof.prod_specs_eval.iter().map(Vec::len)) + .copied() .max() - .unwrap_or(0); + .unwrap_or(0) + .saturating_sub(1); let log2_num_fanin: usize = 1; // ceil_log2(NUM_FANIN=2) = 1 @@ -1347,19 +1616,25 @@ where } } - observe_active_tower_eval_round(ts, chip_proof, layer_idx); + observe_active_tower_eval_round(ts, chip_proof, &num_variables, num_prod_spec, layer_idx); // Mirror native: sample_and_append_vec(b"merge", log2_num_fanin) transcript_observe_label(ts, b"merge"); let mu = FiatShamirTranscript::::sample_ext(ts); mus.push(mu); } + let final_alpha = if layer_count > 0 { + transcript_observe_label(ts, b"combine subset evals"); + FiatShamirTranscript::::sample_ext(ts) + } else { + EF::ZERO + }; - let _ = read_count; TowerTranscriptSchedule { alpha_logup, beta, lambdas, + final_alpha, mus, ris, } @@ -1368,21 +1643,39 @@ where fn observe_active_tower_eval_round( ts: &mut TS, chip_proof: &ZKVMChipProof, + num_variables: &[usize], + num_prod_spec: usize, round: usize, ) where TS: FiatShamirTranscript, { - for spec_rounds in &chip_proof.tower_proof.prod_specs_eval { - if let Some(evals) = spec_rounds.get(round) { - for eval in evals { - ts.observe_ext(*eval); + for (spec_idx, spec_rounds) in chip_proof.tower_proof.prod_specs_eval.iter().enumerate() { + if round + < num_variables + .get(spec_idx) + .copied() + .unwrap_or(0) + .saturating_sub(1) + { + if let Some(evals) = spec_rounds.get(round) { + for eval in evals { + ts.observe_ext(*eval); + } } } } - for spec_rounds in &chip_proof.tower_proof.logup_specs_eval { - if let Some(evals) = spec_rounds.get(round) { - for eval in evals { - ts.observe_ext(*eval); + for (spec_idx, spec_rounds) in chip_proof.tower_proof.logup_specs_eval.iter().enumerate() { + if round + < num_variables + .get(num_prod_spec + spec_idx) + .copied() + .unwrap_or(0) + .saturating_sub(1) + { + if let Some(evals) = spec_rounds.get(round) { + for eval in evals { + ts.observe_ext(*eval); + } } } } @@ -1414,6 +1707,7 @@ impl> TraceGenModule TowerModuleChip::AlphaPow, TowerModuleChip::Layer, TowerModuleChip::LayerSumcheck, + TowerModuleChip::MainPoint, ]; let span = tracing::Span::current(); @@ -1444,6 +1738,7 @@ enum TowerModuleChip { AlphaPow, Layer, LayerSumcheck, + MainPoint, } impl TowerModuleChip { @@ -1491,6 +1786,8 @@ impl RowMajorChip for TowerModuleChip { &(&blob.sumcheck_records, &blob.mus_records), required_height, ), + MainPoint => TowerMainPointTraceGenerator + .generate_trace(&blob.main_point_records.as_slice(), required_height), } } } diff --git a/ceno_recursion_v2/src/tower/sumcheck/air.rs b/ceno_recursion_v2/src/tower/sumcheck/air.rs index 33983ff27..034ef7aa8 100644 --- a/ceno_recursion_v2/src/tower/sumcheck/air.rs +++ b/ceno_recursion_v2/src/tower/sumcheck/air.rs @@ -447,10 +447,7 @@ where idx: local.round + AB::Expr::ONE, xi: local.challenge.map(Into::into), }, - local.is_enabled - * local.is_last_layer - * is_not_dummy.clone() - * AB::Expr::from_bool(!crate::system::TOWER_PREFIX_ONLY), + local.is_enabled * local.is_last_layer * is_not_dummy.clone() * AB::Expr::ZERO, ); } } diff --git a/ceno_recursion_v2/src/tower/tower.rs b/ceno_recursion_v2/src/tower/tower.rs index 19a7937a8..c8a78b0b0 100644 --- a/ceno_recursion_v2/src/tower/tower.rs +++ b/ceno_recursion_v2/src/tower/tower.rs @@ -76,10 +76,7 @@ impl PrecomputedTranscript { if k + 1 < schedule.lambdas.len() { challenges.push_back(schedule.lambdas[k + 1]); } else { - // Native still samples the next alpha after the final merge. - // TODO(recursion-v2): record and constrain this final unused - // challenge in the tower schedule instead of using a sentinel. - challenges.push_back(RecursionField::ZERO); + challenges.push_back(schedule.final_alpha); } } Self { challenges } @@ -349,6 +346,7 @@ fn build_tower_replay_result( round, &rt_prime, &coeffs, + &num_variables, &mut prod_spec_point_n_eval, &mut logup_spec_p_point_n_eval, &mut logup_spec_q_point_n_eval, @@ -426,12 +424,20 @@ fn update_point_evals( round: usize, rt_prime: &Point, coeffs: &[RecursionField], + num_variables: &[usize], prod_point_eval: &mut [PointAndEval], logup_p_point_eval: &mut [PointAndEval], logup_q_point_eval: &mut [PointAndEval], ) { + let num_prod_spec = prod_point_eval.len(); for (spec_idx, point_eval) in prod_point_eval.iter_mut().enumerate() { - if round < tower_proof.prod_specs_eval[spec_idx].len() { + if round + < num_variables + .get(spec_idx) + .copied() + .unwrap_or(0) + .saturating_sub(1) + { let evals = &tower_proof.prod_specs_eval[spec_idx][round]; let merged = izip!(evals.iter(), coeffs.iter()) .map(|(a, b)| *a * *b) @@ -445,7 +451,13 @@ fn update_point_evals( .zip(logup_q_point_eval.iter_mut()) .enumerate() { - if round < tower_proof.logup_specs_eval[spec_idx].len() { + if round + < num_variables + .get(num_prod_spec + spec_idx) + .copied() + .unwrap_or(0) + .saturating_sub(1) + { let evals = &tower_proof.logup_specs_eval[spec_idx][round]; let (p_slice, q_slice) = evals.split_at(2); let merged_p = izip!(p_slice.iter(), coeffs.iter()) @@ -476,7 +488,7 @@ fn aggregate_next_eval( .zip(alpha_pows.iter()) .zip(num_variables.iter()) { - if round + 1 < *max_round { + if round + 1 < max_round.saturating_sub(1) { total += *alpha * point_eval.eval; } } @@ -487,7 +499,7 @@ fn aggregate_next_eval( .zip(alpha_pows[num_prod_spec..].chunks(2)) .zip(num_variables[num_prod_spec..].iter()) { - if round + 1 < *max_round { + if round + 1 < max_round.saturating_sub(1) { total += alpha_chunk[0] * p_eval.eval + alpha_chunk[1] * q_eval.eval; } } diff --git a/ceno_recursion_v2/src/transcript/mod.rs b/ceno_recursion_v2/src/transcript/mod.rs index c7321281e..76a22f36a 100644 --- a/ceno_recursion_v2/src/transcript/mod.rs +++ b/ceno_recursion_v2/src/transcript/mod.rs @@ -13,8 +13,12 @@ use p3_field::{PrimeCharacteristicRing, PrimeField32}; use p3_matrix::dense::RowMajorMatrix; use p3_symmetric::Permutation; -use crate::system::{ - AirModule, BusInventory, GlobalCtxCpu, Preflight, RecursionProof, RecursionVk, TraceGenModule, +use crate::{ + system::{ + AirModule, BusInventory, GlobalCtxCpu, Preflight, RecursionProof, RecursionVk, + TraceGenModule, + }, + utils::digests_to_poseidon2_input, }; use recursion_circuit::transcript::poseidon2::{CHUNK, Poseidon2Air, Poseidon2Cols}; @@ -377,6 +381,36 @@ impl> TraceGenModule poseidon2_perm_inputs.extend_from_slice(ctx.0); poseidon2_compress_inputs.extend_from_slice(ctx.1); + for preflight in preflights { + poseidon2_perm_inputs.extend( + preflight + .pcs + .base_input_leaf_hashes + .iter() + .map(|record| record.input), + ); + poseidon2_perm_inputs.extend( + preflight + .pcs + .commit_phase_leaf_hashes + .iter() + .map(|record| record.input), + ); + poseidon2_compress_inputs.extend( + preflight + .pcs + .base_input_merkle_rows + .iter() + .map(|record| digests_to_poseidon2_input(record.left, record.right)), + ); + poseidon2_compress_inputs.extend( + preflight + .pcs + .commit_phase_merkle_rows + .iter() + .map(|record| digests_to_poseidon2_input(record.left, record.right)), + ); + } let poseidon2_trace = { let (mut poseidon_states, poseidon_counts) = diff --git a/ceno_zkvm/src/instructions/riscv/ecall/keccak.rs b/ceno_zkvm/src/instructions/riscv/ecall/keccak.rs index 1f65df079..22091f948 100644 --- a/ceno_zkvm/src/instructions/riscv/ecall/keccak.rs +++ b/ceno_zkvm/src/instructions/riscv/ecall/keccak.rs @@ -446,7 +446,7 @@ mod tests { e2e::ShardContext, structs::ProgramParams, }; - use ff_ext::BabyBearExt4; + use ff_ext::{BabyBearExt4, FieldFrom}; use multilinear_extensions::utils::eval_by_expr_with_instance; type E = BabyBearExt4; @@ -475,7 +475,7 @@ mod tests { ) -> E { let wit_row = wit_row.iter().copied().map(E::from).collect_vec(); let structural_row = structural_row.iter().copied().map(E::from).collect_vec(); - let challenges = [E::from_canonical_u64(7), E::from_canonical_u64(11)]; + let challenges = [E::from_v(7), E::from_v(11)]; eval_by_expr_with_instance::(&[], &wit_row, &structural_row, &[], &challenges, expr) .unwrap_right() } diff --git a/ceno_zkvm/src/scheme/verifier.rs b/ceno_zkvm/src/scheme/verifier.rs index 41ebe6a06..94b536c3c 100644 --- a/ceno_zkvm/src/scheme/verifier.rs +++ b/ceno_zkvm/src/scheme/verifier.rs @@ -1552,13 +1552,13 @@ impl TowerVerify { round, ); - let expected_evaluation: E = (0..num_prod_spec) + let weighted_prime_fold: E = (0..num_prod_spec) .zip(alpha_pows.iter()) .zip(num_variables.iter()) .map(|((spec_index, alpha), max_round)| { // prod'[b] = prod[0,b] * prod[1,b] // prod'[out_rt] = \sum_b eq(out_rt,b) * prod'[b] = \sum_b eq(out_rt,b) * prod[0,b] * prod[1,b] - eq * *alpha + *alpha * if round < *max_round - 1 { tower_proofs.prod_specs_eval[spec_index][round].iter().copied().product() } else { E::ZERO } @@ -1573,7 +1573,7 @@ impl TowerVerify { // logup_p'[out_rt] = \sum_b eq(out_rt,b) * (logup_p[0,b] * logup_q[1,b] + logup_p[1,b] * logup_q[0,b]) // logup_q'[out_rt] = \sum_b eq(out_rt,b) * logup_q[0,b] * logup_q[1,b] let (alpha_numerator, alpha_denominator) = (&alpha[0], &alpha[1]); - eq * if round < *max_round - 1 { + if round < *max_round - 1 { let evals = &tower_proofs.logup_specs_eval[spec_index][round]; let (p1, p2, q1, q2) = (evals[0], evals[1], evals[2], evals[3]); @@ -1584,6 +1584,7 @@ impl TowerVerify { } }) .sum::(); + let expected_evaluation = eq * weighted_prime_fold; if expected_evaluation != sumcheck_claim.expected_evaluation { return Err(ZKVMError::VerifyError("mismatch tower evaluation".into())); diff --git a/ceno_zkvm/src/structs.rs b/ceno_zkvm/src/structs.rs index 7f4ab504a..d64870914 100644 --- a/ceno_zkvm/src/structs.rs +++ b/ceno_zkvm/src/structs.rs @@ -1086,16 +1086,16 @@ pub struct ZKVMVerifyingKey< /// Number of extension-field elements squeezed from the digest sponge. pub const VK_DIGEST_LEN: usize = 2; -/// Convert bytes into base-field elements using ≤ 3-byte chunks. -fn bytes_to_felts_safe(bytes: &[u8]) -> Vec { - bytes - .chunks(3) - .map(|chunk| { - let mut arr = [0u8; 4]; - arr[..chunk.len()].copy_from_slice(chunk); - F::from_u32(u32::from_le_bytes(arr)) - }) - .collect() +fn observe_bytes_safe(challenger: &mut C, bytes: &[u8]) +where + F: FieldAlgebra, + C: CanObserve, +{ + for chunk in bytes.chunks(3) { + let mut arr = [0u8; 4]; + arr[..chunk.len()].copy_from_slice(chunk); + challenger.observe(F::from_u32(u32::from_le_bytes(arr))); + } } /// Hash the supplied VK preimage into `VK_DIGEST_LEN` extension elements. @@ -1128,7 +1128,7 @@ where .expect("vk digest serialization is infallible for a valid VK"); let mut sponge = DefaultChallenger::::new_poseidon_default(); - sponge.observe_slice(&bytes_to_felts_safe::(&bytes)); + observe_bytes_safe::(&mut sponge, &bytes); std::array::from_fn(|_| sponge.sample_algebra_element()) } diff --git a/ceno_zkvm/src/tables/shard_ram.rs b/ceno_zkvm/src/tables/shard_ram.rs index f440df29f..319c4ac40 100644 --- a/ceno_zkvm/src/tables/shard_ram.rs +++ b/ceno_zkvm/src/tables/shard_ram.rs @@ -1031,13 +1031,13 @@ impl ShardRamCircuit { #[cfg(test)] mod tests { use either::Either; - use ff_ext::{BabyBearExt4, FromUniformBytes, PoseidonField}; + use ff_ext::{BabyBearExt4, FieldFrom, FromUniformBytes, PoseidonField}; use gkr_iop::cpu::{CpuBackend, CpuProver}; use itertools::Itertools; use mpcs::{BasefoldDefault, PolynomialCommitmentScheme, SecurityLevel}; use p3::{ babybear::BabyBear, - field::{FieldAlgebra, PrimeField32}, + field::{PrimeCharacteristicRing, PrimeField32}, }; use rand::thread_rng; use std::sync::Arc; @@ -1256,7 +1256,7 @@ mod tests { build_main_witness::, CpuProver>>( &leaf_composed, &leaf_proof_input, - &[E::ONE, E::from_canonical_u32(7)], + &[E::ONE, E::from_v(7)], WitnessBuildStage::Tower, ); let leaf_r_len = leaf_composed.zkvm_v1_css.r_expressions.len() @@ -1331,7 +1331,7 @@ mod tests { build_main_witness::, CpuProver>>( &ec_tree_composed, &ec_tree_proof_input, - &[E::ONE, E::from_canonical_u32(7)], + &[E::ONE, E::from_v(7)], WitnessBuildStage::Tower, ); let ec_tree_r_len = ec_tree_composed.zkvm_v1_css.r_expressions.len()