chore: address feedback

air/src/air/mod.rs

@@ -292,7 +292,10 @@ pub trait Air: Send + Sync {
     /// 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_rand_elements<E, R>(&self, public_coin: &mut R) -> Result<Vec<E>, RandomCoinError>
+    fn get_aux_rand_elements<E, R>(
+        &self,
+        public_coin: &mut R,
+    ) -> Result<AuxRandElements<E>, RandomCoinError>
     where
         E: FieldElement<BaseField = Self::BaseField>,
         R: RandomCoin<BaseField = Self::BaseField>,
@@ -302,7 +305,7 @@ pub trait Air: Send + Sync {
         for _ in 0..num_elements {
             rand_elements.push(public_coin.draw()?);
         }
-        Ok(rand_elements)
+        Ok(AuxRandElements::new(rand_elements))
     }
 
     /// Returns values for all periodic columns used in the computation.

air/src/proof/ood_frame.rs

@@ -64,7 +64,7 @@ impl OodFrame {
     {
         assert!(self.trace_states.is_empty(), "trace sates have already been set");
 
-        // save the evaluations with the current and next evaluations interleaved for each polynomial
+        // save the evaluations of the current and then next evaluations for each polynomial
         let main_and_aux_trace_states = trace_ood_frame.to_trace_states();
 
         // there are 2 frames: current and next
@@ -246,8 +246,8 @@ impl<E: FieldElement> TraceOodFrame<E> {
         self.current_row.len() > self.main_trace_width
     }
 
-    /// Returns the main/aux frames as element vectors. Specifically, the main and auxiliary frames
-    /// are interleaved, as described in [`OodFrame::set_trace_states`].
+    /// Returns the main/aux frame and Lagrange kernel frame as element vectors as described in
+    /// [`OodFrame::set_trace_states`].
     fn to_trace_states(&self) -> Vec<E> {
         let mut main_and_aux_frame_states = Vec::new();
         main_and_aux_frame_states.extend_from_slice(&self.current_row);

prover/src/lib.rs

@@ -321,12 +321,10 @@ pub trait Prover {
         // build the auxiliary trace segment, and append the resulting segments to trace commitment
         // and trace polynomial table structs
         let aux_trace_with_metadata = if air.trace_info().is_multi_segment() {
-            let rand_elements = air
+            let aux_rand_elements = air
                 .get_aux_rand_elements(channel.public_coin())
                 .expect("failed to draw random elements for the auxiliary trace segment");
 
-            let aux_rand_elements = AuxRandElements::new(rand_elements);
-
             let aux_trace = maybe_await!(self.build_aux_trace(&trace, &aux_rand_elements));
 
             // commit to the auxiliary trace segment

verifier/src/channel.rs

@@ -126,11 +126,11 @@ where
             .partition_size::<E>(air.context().num_constraint_composition_columns());
 
         Ok(VerifierChannel {
-            // trace queries
             trace_commitments,
+            // trace queries
             trace_queries: Some(trace_queries),
-            // constraint queries
             constraint_commitment,
+            // constraint queries
             constraint_queries: Some(constraint_queries),
             // num partitions used in commitment
             partition_size_main,

verifier/src/lib.rs

@@ -34,7 +34,6 @@ extern crate alloc;
 use alloc::vec::Vec;
 use core::cmp;
 
-use air::AuxRandElements;
 pub use air::{
     proof::Proof, Air, AirContext, Assertion, BoundaryConstraint, BoundaryConstraintGroup,
     ConstraintCompositionCoefficients, ConstraintDivisor, DeepCompositionCoefficients,
@@ -176,13 +175,13 @@ where
 
     // process auxiliary trace segments (if any), to build a set of random elements for each segment
     let aux_trace_rand_elements = if air.trace_info().is_multi_segment() {
-        let rand_elements = air
+        let aux_rand_elements = air
             .get_aux_rand_elements(&mut public_coin)
             .expect("failed to generate the random elements needed to build the auxiliary trace");
 
         public_coin.reseed(trace_commitments[AUX_TRACE_IDX]);
 
-        Some(AuxRandElements::new(rand_elements))
+        Some(aux_rand_elements)
     } else {
         None
     };