Explicitly only support 1 auxiliary segment

air/README.md

@@ -65,12 +65,12 @@ To define such columns for your computation, you can override `get_periodic_colu
 ### Randomized AIR
 Randomized AIR is a powerful extension of AIR which enables, among other things, multiset and permutation checks similar to the ones available in PLONKish systems. These, in turn, allow efficient descriptions of "non-local" constraints which can be used to build such components as efficient range checks, random access memory, and many others.
 
-With Randomized AIR, construction of the execution trace is split into multiple stages. During the first stage, the *main trace segment* is built in a manner similar to how the trace is built for regular AIR. In the subsequent stages, *auxiliary trace segments* are built. When building auxiliary trace segments, the prover has access to extra randomness sent by the verifier (in the non-interactive version of the protocol, this randomness is derived from the previous trace segment commitments). Currently, the number of auxiliary trace segments is limited to one.
+With Randomized AIR, construction of the execution trace is split into multiple stages. During the first stage, the *main trace segment* is built in a manner similar to how the trace is built for regular AIR. In the subsequent stages, *auxiliary trace segment* is built. When building the auxiliary trace segment, the prover has access to extra randomness sent by the verifier (in the non-interactive version of the protocol, this randomness is derived from the previous trace segment commitments).
 
 To describe Randomized AIR, you will need to do the following when implementing the `Air` trait:
 * The `AirContext` struct returned from `Air::context()` method must be instantiated using `AirContext::new_multi_segment()` constructor. When building AIR context in this way, you will need to provide a `TraceLayout` which describes the shape of a multi-segment execution trace.
-* Override `Air::evaluate_aux_transition()` method. This method is similar to the `Air::evaluate_transition()` method but it also accepts two extra parameters: `aux_evaluation_frame` and `aux_rand_elements`. These parameters are needed for evaluating transition constraints over the auxiliary trace segments.
-* Override `Air::get_aux_assertions()` method. This method is similar to the `Air::get_assertions()` method, but it should return assertions against columns of the auxiliary trace segments.
+* Override `Air::evaluate_aux_transition()` method. This method is similar to the `Air::evaluate_transition()` method but it also accepts two extra parameters: `aux_evaluation_frame` and `aux_rand_elements`. These parameters are needed for evaluating transition constraints over the auxiliary trace segment.
+* Override `Air::get_aux_assertions()` method. This method is similar to the `Air::get_assertions()` method, but it should return assertions against columns of the auxiliary trace segment.
 
 ## Protocol parameters
 `ProofOptions` struct defines a set of options which are used during STARK proof generation and verification. These options have a direct impact on the security of the generated proofs as well as the proof generation time. Specifically, security of STARK proofs depends on:
@@ -94,4 +94,4 @@ To compile with `no_std`, disable default features via `--no-default-features` f
 License
 -------
 
-This project is [MIT licensed](../LICENSE).
+This project is [MIT licensed](../LICENSE).

air/src/air/boundary/constraint_group.rs

@@ -30,7 +30,7 @@ use alloc::{collections::BTreeMap, vec::Vec};
 /// segments of the execution trace. Specifically:
 /// * For the constraints against columns of the main execution trace, `F` is set to the base field
 ///   of the protocol, and `E` is set to the extension field.
-/// * For the constraints against columns of auxiliary trace segments, both `F` and `E` are set to
+/// * For the constraints against columns of the auxiliary trace segment, both `F` and `E` are set to
 ///   the extension field.
 #[derive(Debug, Clone)]
 pub struct BoundaryConstraintGroup<F, E>

air/src/air/boundary/mod.rs

@@ -25,7 +25,7 @@ mod tests;
 /// Boundary constraints for a computation.
 ///
 /// Boundary constraints are arranged into two categories: constraints against columns of the main
-/// trace segment, and constraints against columns of auxiliary trace segments. Within each
+/// trace segment, and constraints against columns of the auxiliary trace segment. Within each
 /// category, the constraints are grouped by their divisor (see [BoundaryConstraintGroup] for
 /// more info on boundary constraint structure).
 ///
@@ -74,7 +74,7 @@ impl<E: FieldElement> BoundaryConstraints<E> {
         assert_eq!(
             aux_assertions.len(),
             context.num_aux_assertions,
-            "expected {} assertions against auxiliary trace segments, but received {}",
+            "expected {} assertions against the auxiliary trace segment, but received {}",
             context.num_aux_assertions,
             aux_assertions.len(),
         );
@@ -87,7 +87,7 @@ impl<E: FieldElement> BoundaryConstraints<E> {
 
         let trace_length = context.trace_info.length();
         let main_trace_width = context.trace_info.main_trace_width();
-        let aux_trace_width = context.trace_info.aux_trace_width();
+        let aux_trace_width = context.trace_info.aux_segment_width();
 
         // make sure the assertions are valid in the context of their respective trace segments;
         // also, sort the assertions in the deterministic order so that changing the order of
@@ -116,7 +116,7 @@ impl<E: FieldElement> BoundaryConstraints<E> {
             &mut twiddle_map,
         );
 
-        // build constraints for the assertions against auxiliary trace segments
+        // build constraints for the assertions against the auxiliary trace segment
         let aux_constraints = group_constraints(
             aux_assertions,
             context,

air/src/air/coefficients.rs

@@ -9,30 +9,30 @@ use math::FieldElement;
 // AUXILIARY TRACE SEGMENT RANDOMNESS
 // ================================================================================================
 
-/// Random elements used in construction of auxiliary trace segments.
+/// Random elements used in construction of the auxiliary trace segment.
 ///
 /// These elements are generated by the
 /// [Air::get_aux_trace_segment_random_elements()](crate::Air::get_aux_trace_segment_random_elements)
 /// function for each auxiliary trace segment. In the interactive version of the protocol, the
 /// verifier draws these elements uniformly at random from the extension field of the protocol
 /// after the prover commits to a previous trace segment.
 #[derive(Debug, Clone)]
-pub struct AuxTraceRandElements<E: FieldElement>(Vec<Vec<E>>);
+pub struct AuxTraceRandElements<E: FieldElement>(Vec<E>);
 
 impl<E: FieldElement> AuxTraceRandElements<E> {
     /// Instantiates and returns an empty set of random elements.
     pub fn new() -> Self {
         Self(Vec::new())
     }
 
-    /// Returns a list of random elements for an auxiliary segment with the specified index.
-    pub fn get_segment_elements(&self, aux_segment_idx: usize) -> &[E] {
-        &self.0[aux_segment_idx]
+    /// Returns a list of random elements for the auxiliary segment.
+    pub fn get_segment_elements(&self) -> &[E] {
+        &self.0
     }
 
-    /// Adds random elements for a new auxiliary segment to this set of random elements.
-    pub fn add_segment_elements(&mut self, rand_elements: Vec<E>) {
-        self.0.push(rand_elements);
+    /// Sets the random elements associated with the auxiliary segment.
+    pub fn set_segment_elements(&mut self, rand_elements: Vec<E>) {
+        self.0 = rand_elements;
     }
 }
 

air/src/air/context.rs

@@ -105,11 +105,11 @@ impl<B: StarkField> AirContext<B> {
         if trace_info.is_multi_segment() {
             assert!(
                 !aux_transition_constraint_degrees.is_empty(),
-                "at least one transition constraint degree must be specified for auxiliary trace segments"
+                "at least one transition constraint degree must be specified for the auxiliary trace segment"
                 );
             assert!(
                 num_aux_assertions > 0,
-                "at least one assertion must be specified against auxiliary trace segments"
+                "at least one assertion must be specified against the auxiliary trace segment"
             );
         } else {
             assert!(
@@ -124,7 +124,7 @@ impl<B: StarkField> AirContext<B> {
 
         // validate Lagrange kernel aux column, if any
         if let Some(lagrange_kernel_aux_column_idx) = lagrange_kernel_aux_column_idx {
-            assert!(lagrange_kernel_aux_column_idx < trace_info.get_aux_segment_width(0), "Lagrange kernel column index out of bounds: index={}, but only {} columns in segment", lagrange_kernel_aux_column_idx, trace_info.get_aux_segment_width(0));
+            assert!(lagrange_kernel_aux_column_idx < trace_info.get_aux_segment_width(), "Lagrange kernel column index out of bounds: index={}, but only {} columns in segment", lagrange_kernel_aux_column_idx, trace_info.get_aux_segment_width());
         }
 
         // determine minimum blowup factor needed to evaluate transition constraints by taking
@@ -219,7 +219,7 @@ impl<B: StarkField> AirContext<B> {
         self.main_transition_constraint_degrees.len()
     }
 
-    /// Returns the number of transition constraints placed against all auxiliary trace segments.
+    /// Returns the number of transition constraints placed against the auxiliary trace segment.
     pub fn num_aux_transition_constraints(&self) -> usize {
         self.aux_transition_constraint_degrees.len()
     }
@@ -238,7 +238,7 @@ impl<B: StarkField> AirContext<B> {
     /// kernel assertion, which is managed separately.
     ///
     /// The number of assertions consists of the assertions placed against the main segment of an
-    /// execution trace as well as assertions placed against all auxiliary trace segments.
+    /// execution trace as well as assertions placed against the auxiliary trace segment.
     pub fn num_assertions(&self) -> usize {
         self.num_main_assertions + self.num_aux_assertions
     }

air/src/air/mod.rs

@@ -160,11 +160,10 @@ const MIN_CYCLE_LENGTH: usize = 2;
 ///
 /// With Randomized AIR, construction of the execution trace is split into multiple stages. During
 /// the first stage, the *main trace segment* is built in a manner similar to how the trace is
-/// built for regular AIR. In the subsequent stages, *auxiliary trace segments* are built. When
-/// building auxiliary trace segments, the prover has access to extra randomness sent by the
+/// built for regular AIR. In the subsequent stages, the *auxiliary trace segment* is built. When
+/// building the auxiliary trace segment, the prover has access to extra randomness sent by the
 /// verifier (in the non-interactive version of the protocol, this randomness is derived from the
-/// previous trace segment commitments). Currently, the number of auxiliary trace segments is
-/// limited to one.
+/// previous trace segment commitments).
 ///
 /// To describe Randomized AIR, you will need to do the following when implementing the [Air]
 /// trait:
@@ -175,10 +174,10 @@ const MIN_CYCLE_LENGTH: usize = 2;
 /// * Override [Air::evaluate_aux_transition()] method. This method is similar to the
 ///   [Air::evaluate_transition()] method but it also accepts two extra parameters:
 ///   `aux_evaluation_frame` and `aux_rand_elements`. These parameters are needed for evaluating
-///   transition constraints over the auxiliary trace segments.
+///   transition constraints over the auxiliary trace segment.
 /// * Override [Air::get_aux_assertions()] method. This method is similar to the
 ///   [Air::get_assertions()] method, but it should return assertions against columns of the
-///   auxiliary trace segments.
+///   auxiliary trace segment.
 pub trait Air: Send + Sync {
     /// Base field for the computation described by this AIR. STARK protocol for this computation
     /// may be executed in the base field, or in an extension of the base fields as specified
@@ -229,7 +228,7 @@ pub trait Air: Send + Sync {
     // --------------------------------------------------------------------------------------------
 
     /// Evaluates transition constraints over the specified evaluation frames for the main and
-    /// auxiliary trace segments.
+    /// auxiliary trace segment.
     ///
     /// The evaluations should be written into the `results` slice in the same order as the order
     /// of auxiliary transition constraint degree descriptors used to instantiate [AirContext] for
@@ -267,15 +266,15 @@ pub trait Air: Send + Sync {
         unimplemented!("evaluation of auxiliary transition constraints has not been implemented");
     }
 
-    /// Returns a set of assertions placed against auxiliary trace segments.
+    /// Returns a set of assertions placed against the auxiliary trace segment.
     ///
     /// The default implementation of this function returns an empty vector. It should be
-    /// overridden only if the computation relies on auxiliary trace segments. In such a case,
+    /// overridden only if the computation relies on the auxiliary trace segment. In such a case,
     /// the vector returned from this function must contain at least one assertion.
     ///
     /// The column index for assertions is expected to be zero-based across all auxiliary trace
     /// segments. That is, assertion against column 0, is an assertion against the first column
-    /// of the auxiliary trace segments.
+    /// of auxiliary trace segment.
     ///
     /// When the protocol is executed using an extension field, auxiliary assertions are defined
     /// over the extension field. This is in contrast with the assertions returned from
@@ -480,20 +479,18 @@ pub trait Air: Send + Sync {
     // TRACE SEGMENT RANDOMNESS
     // --------------------------------------------------------------------------------------------
 
-    /// Returns a vector of field elements required for construction of an auxiliary trace segment
-    /// with the specified index.
+    /// Returns a vector of field elements required for construction of the auxiliary trace segment.
     ///
     /// The elements are drawn uniformly at random from the provided public coin.
     fn get_aux_trace_segment_random_elements<E, R>(
         &self,
-        aux_segment_idx: usize,
         public_coin: &mut R,
     ) -> Result<Vec<E>, RandomCoinError>
     where
         E: FieldElement<BaseField = Self::BaseField>,
         R: RandomCoin<BaseField = Self::BaseField>,
     {
-        let num_elements = self.trace_info().get_aux_segment_rand_elements(aux_segment_idx);
+        let num_elements = self.trace_info().get_num_aux_segment_rand_elements();
         let mut result = Vec::with_capacity(num_elements);
         for _ in 0..num_elements {
             result.push(public_coin.draw()?);