From 32c879aa5275f85bfb865e989e206371f6b047b1 Mon Sep 17 00:00:00 2001 From: Alon Haramati Date: Mon, 27 May 2024 15:32:35 +0300 Subject: [PATCH] Extract constraint code into functions. --- .../wide_fibonacci/constraint_eval.rs | 102 ++++++++++++------ 1 file changed, 71 insertions(+), 31 deletions(-) diff --git a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs index 41b9ddd96..8d4188850 100644 --- a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs +++ b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs @@ -3,14 +3,17 @@ use num_traits::Zero; use super::component::{Input, WideFibAir, WideFibComponent, ALPHA_ID, Z_ID}; use super::trace_gen::write_trace_row; -use crate::core::air::accumulation::DomainEvaluationAccumulator; +use crate::core::air::accumulation::{ColumnAccumulator, DomainEvaluationAccumulator}; use crate::core::air::{AirProver, Component, ComponentProver, ComponentTrace}; use crate::core::backend::CpuBackend; +use crate::core::circle::Coset; use crate::core::constraints::{coset_vanishing, point_vanishing}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::FieldExpOps; -use crate::core::poly::circle::{CanonicCoset, SecureCirclePoly}; +use crate::core::pcs::TreeVec; +use crate::core::poly::circle::{CanonicCoset, CircleDomain, CircleEvaluation, SecureCirclePoly}; +use crate::core::poly::BitReversedOrder; use crate::core::utils::{bit_reverse, shifted_secure_combination}; use crate::core::{ColumnVec, InteractionElements}; use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; @@ -23,50 +26,64 @@ impl AirProver for WideFibAir { } } -impl ComponentProver for WideFibComponent { - fn evaluate_constraint_quotients_on_domain( +impl WideFibComponent { + fn evaluate_trace_step_constraints( &self, - trace: &ComponentTrace<'_, CpuBackend>, - evaluation_accumulator: &mut DomainEvaluationAccumulator, - interaction_elements: &InteractionElements, + trace_evals: &TreeVec>>, + trace_eval_domain: CircleDomain, + zero_domain: Coset, + accum: &mut ColumnAccumulator<'_, CpuBackend>, ) { let max_constraint_degree = self.max_constraint_log_degree_bound(); - let trace_eval_domain = CanonicCoset::new(max_constraint_degree).circle_domain(); - let trace_evals = &trace.evals; - let zero_domain = CanonicCoset::new(self.log_column_size()).coset; let mut denoms = vec![]; - let mut lookup_denoms = vec![]; for point in trace_eval_domain.iter() { denoms.push(coset_vanishing(zero_domain, point)); - lookup_denoms.push(point_vanishing(zero_domain.at(0), point)); } bit_reverse(&mut denoms); let mut denom_inverses = vec![BaseField::zero(); 1 << (max_constraint_degree)]; BaseField::batch_inverse(&denoms, &mut denom_inverses); - bit_reverse(&mut lookup_denoms); - let mut lookup_denom_inverses = vec![BaseField::zero(); 1 << (max_constraint_degree)]; - BaseField::batch_inverse(&lookup_denoms, &mut lookup_denom_inverses); let mut numerators = vec![SecureField::zero(); 1 << (max_constraint_degree)]; - let mut lookup_numerators = vec![SecureField::zero(); 1 << (max_constraint_degree)]; - let [mut accum] = - evaluation_accumulator.columns([(max_constraint_degree, self.n_constraints())]); - let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); + #[allow(clippy::needless_range_loop)] for i in 0..trace_eval_domain.size() { - // Step constraints. for j in 0..self.n_columns() - 2 { numerators[i] += accum.random_coeff_powers[self.n_columns() - 3 - j] * (trace_evals[0][j][i].square() + trace_evals[0][j + 1][i].square() - trace_evals[0][j + 2][i]); } + } + for (i, (num, denom_inverse)) in numerators.iter().zip(denom_inverses.iter()).enumerate() { + accum.accumulate(i, *num * *denom_inverse); + } + } - // Lookup constraints. - let lookup_value = + fn evaluate_lookup_boundary_constraint( + &self, + trace_evals: &TreeVec>>, + trace_eval_domain: CircleDomain, + zero_domain: Coset, + accum: &mut ColumnAccumulator<'_, CpuBackend>, + interaction_elements: &InteractionElements, + ) { + let max_constraint_degree = self.max_constraint_log_degree_bound(); + let mut denoms = vec![]; + for point in trace_eval_domain.iter() { + denoms.push(point_vanishing(zero_domain.at(0), point)); + } + bit_reverse(&mut denoms); + let mut denom_inverses = vec![BaseField::zero(); 1 << (max_constraint_degree)]; + BaseField::batch_inverse(&denoms, &mut denom_inverses); + let mut numerators = vec![SecureField::zero(); 1 << (max_constraint_degree)]; + let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); + + #[allow(clippy::needless_range_loop)] + for i in 0..trace_eval_domain.size() { + let value = SecureCirclePoly::::eval_from_partial_evals(std::array::from_fn(|j| { trace_evals[1][j][i].into() })); - lookup_numerators[i] = accum.random_coeff_powers[self.n_columns() - 2] - * ((lookup_value + numerators[i] = accum.random_coeff_powers[self.n_columns() - 2] + * ((value * shifted_secure_combination( &[ trace_evals[0][self.n_columns() - 2][i], @@ -84,13 +101,36 @@ impl ComponentProver for WideFibComponent { for (i, (num, denom_inverse)) in numerators.iter().zip(denom_inverses.iter()).enumerate() { accum.accumulate(i, *num * *denom_inverse); } - for (i, (num, denom_inverse)) in lookup_numerators - .iter() - .zip(lookup_denom_inverses.iter()) - .enumerate() - { - accum.accumulate(i, *num * *denom_inverse); - } + } +} + +impl ComponentProver for WideFibComponent { + fn evaluate_constraint_quotients_on_domain( + &self, + trace: &ComponentTrace<'_, CpuBackend>, + evaluation_accumulator: &mut DomainEvaluationAccumulator, + interaction_elements: &InteractionElements, + ) { + let max_constraint_degree = self.max_constraint_log_degree_bound(); + let trace_eval_domain = CanonicCoset::new(max_constraint_degree).circle_domain(); + let trace_evals = &trace.evals; + let zero_domain = CanonicCoset::new(self.log_column_size()).coset; + let [mut accum] = + evaluation_accumulator.columns([(max_constraint_degree, self.n_constraints())]); + + self.evaluate_trace_step_constraints( + trace_evals, + trace_eval_domain, + zero_domain, + &mut accum, + ); + self.evaluate_lookup_boundary_constraint( + trace_evals, + trace_eval_domain, + zero_domain, + &mut accum, + interaction_elements, + ); } }