enabled domains of size bigger than 2^32

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>
@@ -63,17 +63,11 @@ where
         // deg(composition) + deg(divisor) - deg(trace)
         let target_degree = composition_degree + divisor.degree();
         let degree_adjustment = (target_degree - trace_poly_degree) as u64;
-        assert!(
-            degree_adjustment <= u32::MAX as u64,
-            "boundary constraint degree adjustment cannot exceed {}, but was {}",
-            u32::MAX,
-            degree_adjustment
-        );
 
         BoundaryConstraintGroup {
             constraints: Vec::new(),
             divisor,
-            degree_adjustment: degree_adjustment as u32,
+            degree_adjustment,
         }
     }
 
@@ -91,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

@@ -201,7 +201,7 @@ 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)>,
@@ -223,20 +223,14 @@ impl<E: FieldElement> TransitionConstraintGroup<E> {
         let target_degree = composition_degree + divisor_degree;
         let evaluation_degree = degree.get_evaluation_degree(trace_length);
         let degree_adjustment = (target_degree - evaluation_degree) as u64;
-        assert!(
-            degree_adjustment <= u32::MAX as u64,
-            "transition constraint degree adjustment cannot exceed {}, but was {}",
-            u32::MAX,
-            degree_adjustment
-        );
 
         // 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: degree_adjustment as u32,
+            degree_adjustment,
             domain_offset_exp,
             indexes: vec![],
             coefficients: vec![],
@@ -257,7 +251,7 @@ impl<E: FieldElement> TransitionConstraintGroup<E> {
     }
 
     /// Returns degree adjustment factor for this constraint group.
-    pub fn degree_adjustment(&self) -> u32 {
+    pub fn degree_adjustment(&self) -> u64 {
         self.degree_adjustment
     }
 

prover/src/constraints/boundary.rs

@@ -136,7 +136,7 @@ impl<E: FieldElement> BoundaryConstraints<E> {
 /// 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>>,
@@ -156,7 +156,7 @@ impl<E: FieldElement> BoundaryConstraintGroup<E> {
     /// degree adjustment factor.
     fn new(
         divisor: ConstraintDivisor<E::BaseField>,
-        degree_adjustment: u32,
+        degree_adjustment: u64,
         domain_offset: E::BaseField,
     ) -> Self {
         Self {

prover/src/constraints/evaluation_table.rs

@@ -404,7 +404,7 @@ fn get_inv_evaluation<B: StarkField>(
         128, // min batch size
         |batch: &mut [B], batch_offset: usize| {
             for (i, evaluation) in batch.iter_mut().enumerate() {
-                let x = domain.get_ce_x_power_at(batch_offset + i, a as u32, domain_offset_exp);
+                let x = domain.get_ce_x_power_at(batch_offset + i, a, domain_offset_exp);
                 *evaluation = x - b;
             }
         }

prover/src/domain.rs

@@ -15,14 +15,19 @@ pub struct StarkDomain<B: StarkField> {
     /// vector is half the length of the trace domain size.
     trace_twiddles: Vec<B>,
 
+    /// [g^i for i in (0..ce_domain_size)] where g is the constraint evaluation domain generator.
+    ce_domain: Vec<B>,
+
     /// LDE domain size / constraint evaluation domain size
     ce_to_lde_blowup: usize,
 
+    /// A mask which can be used to compute (x % ce_domain_size) via binary AND. This takes
+    /// advantage of the fact that ce_domain_size is a power of two. The mask is then simply
+    /// ce_domain_size - 1.
+    ce_domain_mod_mask: usize,
+
     /// Offset of the low-degree extension domain.
     domain_offset: B,
-
-    /// [g^i for i in (0..ce_domain_size)] where g is the constraint evaluation domain generator.
-    ce_domain: Vec<B>,
 }
 
 // STARK DOMAIN IMPLEMENTATION
@@ -31,15 +36,6 @@ pub struct StarkDomain<B: StarkField> {
 impl<B: StarkField> StarkDomain<B> {
     /// Returns a new STARK domain initialized with the provided `context`.
     pub fn new<A: Air<BaseField = B>>(air: &A) -> Self {
-        // this is needed to ensure that step * power does not overflow in get_ce_x_power_at().
-        // in the future, we should revisit this to enable domain sizes greater than 2^32
-        assert!(
-            air.ce_domain_size() <= 2_usize.pow(32),
-            "constraint evaluation domain size cannot exceed {}, but was {}",
-            u32::MAX,
-            2_usize.pow(32)
-        );
-
         let trace_twiddles = fft::get_twiddles(air.trace_length());
 
         // build constraint evaluation domain
@@ -48,9 +44,10 @@ impl<B: StarkField> StarkDomain<B> {
 
         StarkDomain {
             trace_twiddles,
+            ce_domain,
             ce_to_lde_blowup: air.lde_domain_size() / air.ce_domain_size(),
+            ce_domain_mod_mask: air.ce_domain_size() - 1,
             domain_offset: air.domain_offset(),
-            ce_domain,
         }
     }
 
@@ -109,9 +106,13 @@ impl<B: StarkField> StarkDomain<B> {
     /// The computation is performed without doing exponentiations. offset_exp is assumed to be
     /// s^power which is pre-computed elsewhere.
     #[inline(always)]
-    pub fn get_ce_x_power_at(&self, step: usize, power: u32, offset_exp: B) -> B {
-        let index: usize = step * power as usize;
-        let index = index % self.ce_domain_size();
+    pub fn get_ce_x_power_at(&self, step: usize, power: u64, offset_exp: B) -> B {
+        debug_assert_eq!(offset_exp, self.offset().exp(power.into()));
+        // this computes (step * power) % ce_domain_size. even though both step and power could be
+        // 64-bit values, we are not concerned about overflow here because we are modding by a
+        // power of two. this is also the reason why we can do & ce_domain_mod_mask instead of
+        // performing the actual modulus operation.
+        let index = step.wrapping_mul(power as usize) & self.ce_domain_mod_mask;
         self.ce_domain[index] * offset_exp
     }