Refactor STARK for small fields (#945)

* add enum for evaluate transitions constraints * change constraints evaluator fn * fix fib example * fix simple fib and dummy air examples * fix examples * fix mem example * crate new methods for evaluation_context --------- Co-authored-by: jotabulacios <[email protected]>
lambdaclass · Dec 20, 2024 · 2ce8b2e · 2ce8b2e
1 parent 7b5a638
commit 2ce8b2e
Show file tree

Hide file tree

Showing 14 changed files with 310 additions and 178 deletions.
diff --git a/provers/stark/src/constraints/evaluator.rs b/provers/stark/src/constraints/evaluator.rs
@@ -3,7 +3,7 @@ use super::boundary::BoundaryConstraints;
 use crate::debug::check_boundary_polys_divisibility;
 use crate::domain::Domain;
 use crate::trace::LDETraceTable;
-use crate::traits::AIR;
+use crate::traits::{TransitionEvaluationContext, AIR};
 use crate::{frame::Frame, prover::evaluate_polynomial_on_lde_domain};
 use itertools::Itertools;
 #[cfg(not(feature = "parallel"))]
@@ -14,6 +14,7 @@ use rayon::{
     iter::IndexedParallelIterator,
     prelude::{IntoParallelIterator, ParallelIterator},
 };
+
 #[cfg(feature = "instruments")]
 use std::time::Instant;
 
@@ -183,8 +184,12 @@ impl<A: AIR> ConstraintEvaluator<A> {
                     .collect();
 
                 // Compute all the transition constraints at this point of the LDE domain.
-                let evaluations_transition =
-                    air.compute_transition_prover(&frame, &periodic_values, rap_challenges);
+                let transition_evaluation_context = TransitionEvaluationContext::new_prover(
+                    &frame,
+                    &periodic_values,
+                    rap_challenges,
+                );
+                let evaluations_transition = air.compute_transition(&transition_evaluation_context);
 
                 #[cfg(all(debug_assertions, not(feature = "parallel")))]
                 transition_evaluations.push(evaluations_transition.clone());

diff --git a/provers/stark/src/constraints/transition.rs b/provers/stark/src/constraints/transition.rs
@@ -1,8 +1,8 @@
 use std::ops::Div;
 
 use crate::domain::Domain;
-use crate::frame::Frame;
 use crate::prover::evaluate_polynomial_on_lde_domain;
+use crate::traits::TransitionEvaluationContext;
 use itertools::Itertools;
 use lambdaworks_math::field::element::FieldElement;
 use lambdaworks_math::field::traits::{IsFFTField, IsField, IsSubFieldOf};
@@ -33,10 +33,8 @@ where
     /// vector, in the index corresponding to the constraint as given by `constraint_idx()`.
     fn evaluate(
         &self,
-        frame: &Frame<F, E>,
+        evaluation_context: &TransitionEvaluationContext<F, E>,
         transition_evaluations: &mut [FieldElement<E>],
-        periodic_values: &[FieldElement<F>],
-        rap_challenges: &[FieldElement<E>],
     );
 
     /// The periodicity the constraint is applied over the trace.

diff --git a/provers/stark/src/debug.rs b/provers/stark/src/debug.rs
@@ -1,5 +1,5 @@
 use super::domain::Domain;
-use super::traits::AIR;
+use super::traits::{TransitionEvaluationContext, AIR};
 use crate::{frame::Frame, trace::LDETraceTable};
 use lambdaworks_math::{
     field::{
@@ -93,7 +93,9 @@ pub fn validate_trace<A: AIR>(
             .iter()
             .map(|col| col[step].clone())
             .collect();
-        let evaluations = air.compute_transition_prover(&frame, &periodic_values, rap_challenges);
+        let transition_evaluation_context =
+            TransitionEvaluationContext::new_prover(&frame, &periodic_values, rap_challenges);
+        let evaluations = air.compute_transition(&transition_evaluation_context);
 
         // Iterate over each transition evaluation. When the evaluated step is not from
         // the exemption steps corresponding to the transition, it should have zero as a

diff --git a/provers/stark/src/examples/bit_flags.rs b/provers/stark/src/examples/bit_flags.rs
@@ -1,10 +1,9 @@
 use crate::{
     constraints::{boundary::BoundaryConstraints, transition::TransitionConstraint},
     context::AirContext,
-    frame::Frame,
     proof::options::ProofOptions,
     trace::TraceTable,
-    traits::AIR,
+    traits::{TransitionEvaluationContext, AIR},
     Felt252,
 };
 use lambdaworks_math::field::{
@@ -45,11 +44,22 @@ impl TransitionConstraint<StarkField, StarkField> for BitConstraint {
 
     fn evaluate(
         &self,
-        frame: &Frame<StarkField, StarkField>,
-        transition_evaluations: &mut [Felt252],
-        _periodic_values: &[Felt252],
-        _rap_challenges: &[Felt252],
+        evaluation_context: &TransitionEvaluationContext<StarkField, StarkField>,
+        transition_evaluations: &mut [FieldElement<StarkField>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let step = frame.get_evaluation_step(0);
 
         let prefix_flag = step.get_main_evaluation_element(0, 0);
@@ -92,11 +102,22 @@ impl TransitionConstraint<StarkField, StarkField> for ZeroFlagConstraint {
 
     fn evaluate(
         &self,
-        frame: &Frame<StarkField, StarkField>,
-        transition_evaluations: &mut [FieldElement<Stark252PrimeField>],
-        _periodic_values: &[FieldElement<Stark252PrimeField>],
-        _rap_challenges: &[FieldElement<Stark252PrimeField>],
+        evaluation_context: &TransitionEvaluationContext<StarkField, StarkField>,
+        transition_evaluations: &mut [FieldElement<StarkField>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let step = frame.get_evaluation_step(0);
         let zero_flag = step.get_main_evaluation_element(15, 0);
 
@@ -149,15 +170,6 @@ impl AIR for BitFlagsAIR {
         &self.constraints
     }
 
-    fn compute_transition_verifier(
-        &self,
-        frame: &Frame<Self::FieldExtension, Self::FieldExtension>,
-        periodic_values: &[FieldElement<Self::FieldExtension>],
-        rap_challenges: &[FieldElement<Self::FieldExtension>],
-    ) -> Vec<FieldElement<Self::FieldExtension>> {
-        self.compute_transition_prover(frame, periodic_values, rap_challenges)
-    }
-
     fn boundary_constraints(
         &self,
         _rap_challenges: &[FieldElement<Self::FieldExtension>],

diff --git a/provers/stark/src/examples/dummy_air.rs b/provers/stark/src/examples/dummy_air.rs
@@ -6,10 +6,9 @@ use crate::{
         transition::TransitionConstraint,
     },
     context::AirContext,
-    frame::Frame,
     proof::options::ProofOptions,
     trace::TraceTable,
-    traits::AIR,
+    traits::{TransitionEvaluationContext, AIR},
 };
 use lambdaworks_math::field::{
     element::FieldElement, fields::fft_friendly::stark_252_prime_field::Stark252PrimeField,
@@ -48,11 +47,22 @@ where
 
     fn evaluate(
         &self,
-        frame: &Frame<F, F>,
+        evaluation_context: &TransitionEvaluationContext<F, F>,
         transition_evaluations: &mut [FieldElement<F>],
-        _periodic_values: &[FieldElement<F>],
-        _rap_challenges: &[FieldElement<F>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let first_step = frame.get_evaluation_step(0);
         let second_step = frame.get_evaluation_step(1);
         let third_step = frame.get_evaluation_step(2);
@@ -97,11 +107,22 @@ where
 
     fn evaluate(
         &self,
-        frame: &Frame<F, F>,
+        evaluation_context: &TransitionEvaluationContext<F, F>,
         transition_evaluations: &mut [FieldElement<F>],
-        _periodic_values: &[FieldElement<F>],
-        _rap_challenges: &[FieldElement<F>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let first_step = frame.get_evaluation_step(0);
 
         let bit = first_step.get_main_evaluation_element(0, 0);
@@ -186,15 +207,6 @@ impl AIR for DummyAIR {
     fn pub_inputs(&self) -> &Self::PublicInputs {
         &()
     }
-
-    fn compute_transition_verifier(
-        &self,
-        frame: &Frame<Self::FieldExtension, Self::FieldExtension>,
-        periodic_values: &[FieldElement<Self::FieldExtension>],
-        rap_challenges: &[FieldElement<Self::FieldExtension>],
-    ) -> Vec<FieldElement<Self::Field>> {
-        self.compute_transition_prover(frame, periodic_values, rap_challenges)
-    }
 }
 
 pub fn dummy_trace<F: IsFFTField>(trace_length: usize) -> TraceTable<F, F> {

diff --git a/provers/stark/src/examples/fibonacci_2_cols_shifted.rs b/provers/stark/src/examples/fibonacci_2_cols_shifted.rs
@@ -4,10 +4,9 @@ use crate::{
         transition::TransitionConstraint,
     },
     context::AirContext,
-    frame::Frame,
     proof::options::ProofOptions,
     trace::TraceTable,
-    traits::AIR,
+    traits::{TransitionEvaluationContext, AIR},
 };
 use lambdaworks_math::{
     field::{element::FieldElement, traits::IsFFTField},
@@ -46,11 +45,22 @@ where
 
     fn evaluate(
         &self,
-        frame: &Frame<F, F>,
+        evaluation_context: &TransitionEvaluationContext<F, F>,
         transition_evaluations: &mut [FieldElement<F>],
-        _periodic_values: &[FieldElement<F>],
-        _rap_challenges: &[FieldElement<F>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let first_row = frame.get_evaluation_step(0);
         let second_row = frame.get_evaluation_step(1);
 
@@ -94,11 +104,22 @@ where
 
     fn evaluate(
         &self,
-        frame: &Frame<F, F>,
+        evaluation_context: &TransitionEvaluationContext<F, F>,
         transition_evaluations: &mut [FieldElement<F>],
-        _periodic_values: &[FieldElement<F>],
-        _rap_challenges: &[FieldElement<F>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let first_row = frame.get_evaluation_step(0);
         let second_row = frame.get_evaluation_step(1);
 
@@ -223,15 +244,6 @@ where
     fn pub_inputs(&self) -> &Self::PublicInputs {
         &self.pub_inputs
     }
-
-    fn compute_transition_verifier(
-        &self,
-        frame: &Frame<Self::FieldExtension, Self::FieldExtension>,
-        periodic_values: &[FieldElement<Self::FieldExtension>],
-        rap_challenges: &[FieldElement<Self::FieldExtension>],
-    ) -> Vec<FieldElement<Self::Field>> {
-        self.compute_transition_prover(frame, periodic_values, rap_challenges)
-    }
 }
 
 pub fn compute_trace<F: IsFFTField>(

diff --git a/provers/stark/src/examples/fibonacci_2_columns.rs b/provers/stark/src/examples/fibonacci_2_columns.rs
@@ -7,10 +7,9 @@ use crate::{
         transition::TransitionConstraint,
     },
     context::AirContext,
-    frame::Frame,
     proof::options::ProofOptions,
     trace::TraceTable,
-    traits::AIR,
+    traits::{TransitionEvaluationContext, AIR},
 };
 use lambdaworks_math::field::{element::FieldElement, traits::IsFFTField};
 
@@ -45,11 +44,22 @@ where
 
     fn evaluate(
         &self,
-        frame: &Frame<F, F>,
+        evaluation_context: &TransitionEvaluationContext<F, F>,
         transition_evaluations: &mut [FieldElement<F>],
-        _periodic_values: &[FieldElement<F>],
-        _rap_challenges: &[FieldElement<F>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let first_step = frame.get_evaluation_step(0);
         let second_step = frame.get_evaluation_step(1);
 
@@ -95,11 +105,22 @@ where
 
     fn evaluate(
         &self,
-        frame: &Frame<F, F>,
+        evaluation_context: &TransitionEvaluationContext<F, F>,
         transition_evaluations: &mut [FieldElement<F>],
-        _periodic_values: &[FieldElement<F>],
-        _rap_challenges: &[FieldElement<F>],
     ) {
+        let (frame, _periodic_values, _rap_challenges) = match evaluation_context {
+            TransitionEvaluationContext::Prover {
+                frame,
+                periodic_values,
+                rap_challenges,
+            }
+            | TransitionEvaluationContext::Verifier {
+                frame,
+                periodic_values,
+                rap_challenges,
+            } => (frame, periodic_values, rap_challenges),
+        };
+
         let first_step = frame.get_evaluation_step(0);
         let second_step = frame.get_evaluation_step(1);
 
@@ -194,15 +215,6 @@ where
     fn pub_inputs(&self) -> &Self::PublicInputs {
         &self.pub_inputs
     }
-
-    fn compute_transition_verifier(
-        &self,
-        frame: &Frame<Self::FieldExtension, Self::FieldExtension>,
-        periodic_values: &[FieldElement<Self::FieldExtension>],
-        rap_challenges: &[FieldElement<Self::FieldExtension>],
-    ) -> Vec<FieldElement<Self::Field>> {
-        self.compute_transition_prover(frame, periodic_values, rap_challenges)
-    }
 }
 
 pub fn compute_trace<F: IsFFTField>(