Merge pull request #118 from novifinancial/exp-elimination

Eliminate exponentiation from prover constraint evaluation

air/src/air/boundary/constraint_group.rs

@@ -41,7 +41,7 @@ where
 {
     constraints: Vec<BoundaryConstraint<F, E>>,
     divisor: ConstraintDivisor<F::BaseField>,
-    degree_adjustment: u32,
+    degree_adjustment: u64,
 }
 
 impl<F, E> BoundaryConstraintGroup<F, E>
@@ -62,7 +62,7 @@ where
         // Boundary constraint degree is always deg(trace). So, the degree adjustment is simply:
         // deg(composition) + deg(divisor) - deg(trace)
         let target_degree = composition_degree + divisor.degree();
-        let degree_adjustment = (target_degree - trace_poly_degree) as u32;
+        let degree_adjustment = (target_degree - trace_poly_degree) as u64;
 
         BoundaryConstraintGroup {
             constraints: Vec::new(),
@@ -85,7 +85,7 @@ where
     }
 
     /// Returns a degree adjustment factor for all boundary constraints in this group.
-    pub fn degree_adjustment(&self) -> u32 {
+    pub fn degree_adjustment(&self) -> u64 {
         self.degree_adjustment
     }
 

air/src/air/divisor.rs

@@ -171,7 +171,7 @@ impl<B: StarkField> Display for ConstraintDivisor<B> {
 // ================================================================================================
 
 /// Returns g^step, where g is the generator of trace domain.
-pub fn get_trace_domain_value_at<B: StarkField>(trace_length: usize, step: usize) -> B {
+fn get_trace_domain_value_at<B: StarkField>(trace_length: usize, step: usize) -> B {
     debug_assert!(
         step < trace_length,
         "step must be in the trace domain [0, {})",

air/src/air/transition/mod.rs

@@ -158,7 +158,7 @@ impl<E: FieldElement> TransitionConstraints<E> {
     /// * $z(x)$ is the constraint divisor polynomial.
     ///
     /// Thus, this function computes a linear combination of $C(x)$ evaluations. For more detail on
-    ///  how this linear combination is computed refer to [TransitionConstraintGroup::merge_evaluations].
+    /// how this linear combination is computed refer to [TransitionConstraintGroup::merge_evaluations].
     ///
     /// Since, the divisor polynomial is the same for all transition constraints (see
     /// [ConstraintDivisor::from_transition]), we can divide the linear combination by the
@@ -171,13 +171,15 @@ impl<E: FieldElement> TransitionConstraints<E> {
     {
         // merge constraint evaluations for the main trace segment
         let mut result = self.main_constraints().iter().fold(E::ZERO, |acc, group| {
-            acc + group.merge_evaluations::<F, F>(main_evaluations, x)
+            let xp = x.exp(group.degree_adjustment.into());
+            acc + group.merge_evaluations::<F, F>(main_evaluations, xp)
         });
 
         // merge constraint evaluations for auxiliary trace segments (if any)
         if self.num_aux_constraints() > 0 {
             result += self.aux_constraints().iter().fold(E::ZERO, |acc, group| {
-                acc + group.merge_evaluations::<F, E>(aux_evaluations, x)
+                let xp = x.exp(group.degree_adjustment.into());
+                acc + group.merge_evaluations::<F, E>(aux_evaluations, xp)
             });
         }
 
@@ -199,7 +201,8 @@ impl<E: FieldElement> TransitionConstraints<E> {
 #[derive(Clone, Debug)]
 pub struct TransitionConstraintGroup<E: FieldElement> {
     degree: TransitionConstraintDegree,
-    degree_adjustment: u32,
+    degree_adjustment: u64,
+    domain_offset_exp: E::BaseField,
     indexes: Vec<usize>,
     coefficients: Vec<(E, E)>,
 }
@@ -213,15 +216,22 @@ impl<E: FieldElement> TransitionConstraintGroup<E> {
         trace_length: usize,
         composition_degree: usize,
         divisor_degree: usize,
+        domain_offset: E::BaseField,
     ) -> Self {
         // We want to make sure that once we divide a constraint polynomial by its divisor, the
         // degree of the resulting polynomial will be exactly equal to the composition_degree.
         let target_degree = composition_degree + divisor_degree;
         let evaluation_degree = degree.get_evaluation_degree(trace_length);
-        let degree_adjustment = (target_degree - evaluation_degree) as u32;
+        let degree_adjustment = (target_degree - evaluation_degree) as u64;
+
+        // pre-compute domain offset exponent; this is used only by the prover and is not relevant
+        // for the verifier
+        let domain_offset_exp = domain_offset.exp(degree_adjustment.into());
+
         TransitionConstraintGroup {
             degree,
             degree_adjustment,
+            domain_offset_exp,
             indexes: vec![],
             coefficients: vec![],
         }
@@ -240,6 +250,16 @@ impl<E: FieldElement> TransitionConstraintGroup<E> {
         &self.degree
     }
 
+    /// Returns degree adjustment factor for this constraint group.
+    pub fn degree_adjustment(&self) -> u64 {
+        self.degree_adjustment
+    }
+
+    /// Returns c^degree_adjustment where c is the coset offset.
+    pub fn domain_offset_exp(&self) -> E::BaseField {
+        self.domain_offset_exp
+    }
+
     /// Adds a new constraint to the group. The constraint is identified by an index in the
     /// evaluation table.
     pub fn add(&mut self, constraint_idx: usize, coefficients: (E, E)) {
@@ -253,10 +273,13 @@ impl<E: FieldElement> TransitionConstraintGroup<E> {
     ///
     /// The linear combination is computed as follows:
     /// $$
-    /// \sum_{i=0}^{k-1}{C_i(x) \cdot (\alpha_i + \beta_i \cdot x^d)}
+    /// \sum_{i=0}^{k-1}{C_i(x) \cdot (\alpha_i + \beta_i \cdot xp)}
     /// $$
     /// where:
     /// * $C_i(x)$ is the evaluation of the $i$th constraint at `x` (same as `evaluations[i]`).
+    /// * $xp = x^d$ where $d$ is the degree adjustment factor computed as $D + (n - 1) - deg(C_i(x))$,
+    ///   where $D$ is the degree of the composition polynomial, $n$ is the length of the execution
+    ///   trace, and $deg(C_i(x))$ is the evaluation degree of the $i$th constraint.
     /// * $\alpha$ and $\beta$ are random field elements. In the interactive version of the
     ///   protocol, these are provided by the verifier.
     /// * $d$ is the degree adjustment factor computed as $D + (n - 1) - deg(C_i(x))$, where
@@ -271,15 +294,12 @@ impl<E: FieldElement> TransitionConstraintGroup<E> {
     /// them by the divisor later on. The degree of the divisor for transition constraints is
     /// always $n - 1$. Thus, once we divide out the divisor, the evaluations will represent a
     /// polynomial of degree $D$.
-    pub fn merge_evaluations<B, F>(&self, evaluations: &[F], x: B) -> E
+    pub fn merge_evaluations<B, F>(&self, evaluations: &[F], xp: B) -> E
     where
         B: FieldElement,
         F: FieldElement<BaseField = B::BaseField> + ExtensionOf<B>,
         E: FieldElement<BaseField = B::BaseField> + ExtensionOf<B> + ExtensionOf<F>,
     {
-        // compute degree adjustment factor for this group
-        let xp = x.exp(self.degree_adjustment.into());
-
         // compute linear combination of evaluations as D(x) * (cc_0 + cc_1 * x^p), where D(x)
         // is an evaluation of a particular constraint, and x^p is the degree adjustment factor
         let mut result = E::ZERO;
@@ -312,6 +332,7 @@ fn group_constraints<E: FieldElement>(
                 context.trace_len(),
                 context.composition_degree(),
                 divisor_degree,
+                context.options.domain_offset(),
             )
         });
         group.add(i, coefficients[i]);

prover/src/constraints/boundary.rs

@@ -3,6 +3,7 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
+use super::StarkDomain;
 use air::{Air, AuxTraceRandElements, ConstraintDivisor};
 use math::{fft, ExtensionOf, FieldElement};
 use utils::collections::{BTreeMap, Vec};
@@ -86,59 +87,37 @@ impl<E: FieldElement> BoundaryConstraints<E> {
 
     /// Evaluates boundary constraints against the main segment of an execution trace at the
     /// specified step of constraint evaluation domain.
-    ///
-    /// Specifically, `step` is the step in the constraint evaluation domain, and `x` is the
-    /// corresponding domain value. That is, x = s * g^step, where g is the generator of the
-    /// constraint evaluation domain, and s is the domain offset.
     pub fn evaluate_main(
         &self,
         main_state: &[E::BaseField],
-        x: E::BaseField,
+        domain: &StarkDomain<E::BaseField>,
         step: usize,
         result: &mut [E],
     ) {
-        // compute the adjustment degree outside of the group so that we can re-use
-        // it for groups which have the same adjustment degree
-        let mut degree_adjustment = self.0[0].degree_adjustment;
-        let mut xp: E::BaseField = x.exp(degree_adjustment.into());
-
+        let x = domain.get_ce_x_at(step);
         for (group, result) in self.0.iter().zip(result.iter_mut()) {
-            // recompute adjustment degree only when it has changed
-            if group.degree_adjustment != degree_adjustment {
-                degree_adjustment = group.degree_adjustment;
-                xp = x.exp(degree_adjustment.into());
-            }
             // evaluate the group and save the result
+            let (power, offset_exp) = (group.degree_adjustment, group.domain_offset_exp);
+            let xp = domain.get_ce_x_power_at(step, power, offset_exp);
             *result = group.evaluate_main(main_state, step, x, xp);
         }
     }
 
     /// Evaluates boundary constraints against all segments of an execution trace at the
     /// specified step of constraint evaluation domain.
-    ///
-    /// Specifically, `step` is the step in the constraint evaluation domain, and `x` is the
-    /// corresponding domain value. That is, x = s * g^step, where g is the generator of the
-    /// constraint evaluation domain, and s is the domain offset.
     pub fn evaluate_all(
         &self,
         main_state: &[E::BaseField],
         aux_state: &[E],
-        x: E::BaseField,
+        domain: &StarkDomain<E::BaseField>,
         step: usize,
         result: &mut [E],
     ) {
-        // compute the adjustment degree outside of the group so that we can re-use
-        // it for groups which have the same adjustment degree
-        let mut degree_adjustment = self.0[0].degree_adjustment;
-        let mut xp: E::BaseField = x.exp(degree_adjustment.into());
-
+        let x = domain.get_ce_x_at(step);
         for (group, result) in self.0.iter().zip(result.iter_mut()) {
-            // recompute adjustment degree only when it has changed
-            if group.degree_adjustment != degree_adjustment {
-                degree_adjustment = group.degree_adjustment;
-                xp = x.exp(degree_adjustment.into());
-            }
             // evaluate the group and save the result
+            let (power, offset_exp) = (group.degree_adjustment, group.domain_offset_exp);
+            let xp = domain.get_ce_x_power_at(step, power, offset_exp);
             *result = group.evaluate_all(main_state, aux_state, step, x, xp);
         }
     }
@@ -152,9 +131,13 @@ impl<E: FieldElement> BoundaryConstraints<E> {
 ///
 /// The constraints are also separated into constraints against the main segment of the execution
 /// and the constraints against auxiliary segments of the execution trace (if any).
+///
+/// Domain offset exponent is pre-computed to be used later during constraint evaluation process,
+/// and thus, to help avoid exponentiations.
 pub struct BoundaryConstraintGroup<E: FieldElement> {
     divisor: ConstraintDivisor<E::BaseField>,
-    degree_adjustment: u32,
+    degree_adjustment: u64,
+    domain_offset_exp: E::BaseField,
     // main trace constraints
     main_single_value: Vec<SingleValueConstraint<E::BaseField, E>>,
     main_small_poly: Vec<SmallPolyConstraint<E::BaseField, E>>,
@@ -171,10 +154,15 @@ impl<E: FieldElement> BoundaryConstraintGroup<E> {
 
     /// Returns an empty [BoundaryConstraintGroup] instantiated with the specified divisor and
     /// degree adjustment factor.
-    fn new(divisor: ConstraintDivisor<E::BaseField>, degree_adjustment: u32) -> Self {
+    fn new(
+        divisor: ConstraintDivisor<E::BaseField>,
+        degree_adjustment: u64,
+        domain_offset: E::BaseField,
+    ) -> Self {
         Self {
             divisor,
             degree_adjustment,
+            domain_offset_exp: domain_offset.exp(degree_adjustment.into()),
             main_single_value: Vec::new(),
             main_small_poly: Vec::new(),
             main_large_poly: Vec::new(),
@@ -195,7 +183,11 @@ impl<E: FieldElement> BoundaryConstraintGroup<E> {
         air: &A,
         twiddle_map: &mut BTreeMap<usize, Vec<E::BaseField>>,
     ) -> Self {
-        let mut result = Self::new(source.divisor().clone(), source.degree_adjustment());
+        let mut result = Self::new(
+            source.divisor().clone(),
+            source.degree_adjustment(),
+            air.domain_offset(),
+        );
 
         for constraint in source.constraints() {
             if constraint.poly().len() == 1 {
@@ -224,7 +216,11 @@ impl<E: FieldElement> BoundaryConstraintGroup<E> {
         air: &A,
         twiddle_map: &mut BTreeMap<usize, Vec<E::BaseField>>,
     ) -> Self {
-        let mut result = Self::new(group.divisor().clone(), group.degree_adjustment());
+        let mut result = Self::new(
+            group.divisor().clone(),
+            group.degree_adjustment(),
+            air.domain_offset(),
+        );
         result.add_aux_constraints(group, air, twiddle_map);
         result
     }