From be265626f064ac1fcc82b1bf13e28f83023a505a Mon Sep 17 00:00:00 2001
From: Andrew Milson <andrewmilson@users.noreply.github.com>
Date: Thu, 18 Jul 2024 10:17:27 -0400
Subject: [PATCH] Generate periodic vanishing polynomial values (#745)

---
 crates/prover/src/examples/poseidon/mod.rs | 29 +++++++++++-----------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs
index fb38786db..7241c288e 100644
--- a/crates/prover/src/examples/poseidon/mod.rs
+++ b/crates/prover/src/examples/poseidon/mod.rs
@@ -1,5 +1,6 @@
 //! AIR for Poseidon2 hash function from <https://eprint.iacr.org/2023/323.pdf>.
 
+use std::array;
 use std::ops::{Add, AddAssign, Mul, Sub};
 
 use itertools::Itertools;
@@ -9,20 +10,20 @@ use crate::constraint_framework::{EvalAtRow, PointEvaluator, SimdDomainEvaluator
 use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use crate::core::air::mask::fixed_mask_points;
 use crate::core::air::{Air, AirProver, Component, ComponentProver, ComponentTrace};
-use crate::core::backend::simd::column::BaseFieldVec;
-use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES};
+use crate::core::backend::simd::m31::{PackedBaseField, PackedM31, LOG_N_LANES};
 use crate::core::backend::simd::SimdBackend;
-use crate::core::backend::{Col, Column, ColumnOps};
+use crate::core::backend::{Col, Column};
 use crate::core::channel::Blake2sChannel;
 use crate::core::circle::CirclePoint;
 use crate::core::constraints::coset_vanishing;
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
-use crate::core::fields::{FieldExpOps, FieldOps};
+use crate::core::fields::FieldExpOps;
 use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps};
 use crate::core::poly::BitReversedOrder;
 use crate::core::prover::VerificationError;
+use crate::core::utils::bit_reverse;
 use crate::core::{ColumnVec, InteractionElements, LookupValues};
 use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier, ComponentTraceGenerator};
 
@@ -385,11 +386,10 @@ impl ComponentProver<SimdBackend> for PoseidonComponent {
         // Denoms.
         let span = span!(Level::INFO, "Constraint eval denominators").entered();
         let zero_domain = CanonicCoset::new(self.log_column_size()).coset;
-        let mut denoms =
-            BaseFieldVec::from_iter(eval_domain.iter().map(|p| coset_vanishing(zero_domain, p)));
-        <SimdBackend as ColumnOps<BaseField>>::bit_reverse_column(&mut denoms);
-        let mut denom_inverses = BaseFieldVec::zeros(denoms.len());
-        <SimdBackend as FieldOps<BaseField>>::batch_inverse(&denoms, &mut denom_inverses);
+        let denoms_inv: [BaseField; 1 << LOG_EXPAND] =
+            array::from_fn(|i| coset_vanishing(zero_domain, eval_domain.at(i)).inverse());
+        let mut packed_denoms_inv = denoms_inv.map(PackedM31::broadcast);
+        bit_reverse(&mut packed_denoms_inv);
         span.exit();
 
         let _span = span!(Level::INFO, "Constraint pointwise eval").entered();
@@ -414,13 +414,14 @@ impl ComponentProver<SimdBackend> for PoseidonComponent {
             for _ in 0..N_INSTANCES_PER_ROW {
                 evaluator.eval();
             }
-            let row_res = evaluator.eval.row_res;
 
+            let packed_denom_inv =
+                packed_denoms_inv[vec_row >> (zero_domain.log_size() - LOG_N_LANES)];
+            let quotient = evaluator.eval.row_res * packed_denom_inv;
             unsafe {
-                accum.col.set_packed(
-                    vec_row,
-                    accum.col.packed_at(vec_row) + row_res * denom_inverses.data[vec_row],
-                )
+                accum
+                    .col
+                    .set_packed(vec_row, accum.col.packed_at(vec_row) + quotient)
             }
             assert_eq!(evaluator.eval.constraint_index, n_constraints);
         }