diff --git a/air/README.md b/air/README.md index 30004b482..5a8ef5fa3 100644 --- a/air/README.md +++ b/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). \ No newline at end of file +This project is [MIT licensed](../LICENSE). diff --git a/air/src/air/boundary/constraint_group.rs b/air/src/air/boundary/constraint_group.rs index e50ec1153..ab71f0836 100644 --- a/air/src/air/boundary/constraint_group.rs +++ b/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 diff --git a/air/src/air/boundary/mod.rs b/air/src/air/boundary/mod.rs index d6db05165..effd5653e 100644 --- a/air/src/air/boundary/mod.rs +++ b/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 BoundaryConstraints { 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 BoundaryConstraints { 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 BoundaryConstraints { &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, diff --git a/air/src/air/coefficients.rs b/air/src/air/coefficients.rs index 144c5846a..11d89d1f4 100644 --- a/air/src/air/coefficients.rs +++ b/air/src/air/coefficients.rs @@ -9,7 +9,7 @@ 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) @@ -17,7 +17,7 @@ use math::FieldElement; /// 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(Vec>); +pub struct AuxTraceRandElements(Vec); impl AuxTraceRandElements { /// Instantiates and returns an empty set of random elements. @@ -25,14 +25,14 @@ impl AuxTraceRandElements { 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) { - self.0.push(rand_elements); + /// Sets the random elements associated with the auxiliary segment. + pub fn set_segment_elements(&mut self, rand_elements: Vec) { + self.0 = rand_elements; } } diff --git a/air/src/air/context.rs b/air/src/air/context.rs index 67117cea9..389eca022 100644 --- a/air/src/air/context.rs +++ b/air/src/air/context.rs @@ -105,11 +105,11 @@ impl AirContext { 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 AirContext { // 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 AirContext { 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 AirContext { /// 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 } diff --git a/air/src/air/mod.rs b/air/src/air/mod.rs index 1fb7e6b46..6fea14414 100644 --- a/air/src/air/mod.rs +++ b/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( &self, - aux_segment_idx: usize, public_coin: &mut R, ) -> Result, RandomCoinError> where E: FieldElement, R: RandomCoin, { - 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()?); diff --git a/air/src/air/trace_info.rs b/air/src/air/trace_info.rs index 29d0d69dd..48d6f56d6 100644 --- a/air/src/air/trace_info.rs +++ b/air/src/air/trace_info.rs @@ -10,23 +10,19 @@ use utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serial // CONSTANTS // ================================================================================================ -/// Number of allowed auxiliary trace segments. -const NUM_AUX_SEGMENTS: usize = 1; - // TRACE INFO // ================================================================================================ /// Information about a specific execution trace. /// -/// Trace info consists of trace layout info, length, and optional custom metadata. Trace layout -/// specifies the number of columns for all trace segments. Currently, a trace can consist of at -/// most two segments. Metadata is just a vector of bytes and can store any values up to 64KB in +/// Trace info consists of the number of columns for all trace segments, trace length, and optional +/// custom metadata. Currently, a trace can consist of at most two segments: the main segment and +/// one auxiliary segment. Metadata is just a vector of bytes and can store any values up to 64KB in /// size. #[derive(Debug, Clone, Eq, PartialEq)] pub struct TraceInfo { main_segment_width: usize, - aux_segment_widths: [usize; NUM_AUX_SEGMENTS], - aux_segment_rands: [usize; NUM_AUX_SEGMENTS], - num_aux_segments: usize, + aux_segment_width: usize, + num_aux_segment_rands: usize, trace_length: usize, trace_meta: Vec, } @@ -38,7 +34,7 @@ impl TraceInfo { pub const MAX_TRACE_WIDTH: usize = 255; /// Maximum number of bytes in trace metadata; currently set at 65535. pub const MAX_META_LENGTH: usize = 65535; - /// Maximum number of random elements per auxiliary trace segment; currently set to 255. + /// Maximum number of random elements in the auxiliary trace segment; currently set to 255. pub const MAX_RAND_SEGMENT_ELEMENTS: usize = 255; // CONSTRUCTORS @@ -67,7 +63,7 @@ impl TraceInfo { /// * Length of `meta` is greater than 65535; pub fn with_meta(width: usize, length: usize, meta: Vec) -> Self { assert!(width > 0, "trace width must be greater than 0"); - Self::new_multi_segment(width, [0], [0], length, meta) + Self::new_multi_segment(width, 0, 0, length, meta) } /// Creates a new [TraceInfo] with main and auxiliary segments. @@ -78,13 +74,13 @@ impl TraceInfo { /// * Total width of all trace segments is greater than 255. /// * Trace length is smaller than 8 or is not a power of two. /// * A zero entry in auxiliary segment width array is followed by a non-zero entry. - /// * Number of random elements for an auxiliary trace segment of non-zero width is set to zero. - /// * Number of random elements for an auxiliary trace segment of zero width is set to non-zero. + /// * Number of random elements for the auxiliary trace segment of non-zero width is set to zero. + /// * Number of random elements for the auxiliary trace segment of zero width is set to non-zero. /// * Number of random elements for any auxiliary trace segment is greater than 255. pub fn new_multi_segment( main_segment_width: usize, - aux_segment_widths: [usize; NUM_AUX_SEGMENTS], - aux_segment_rands: [usize; NUM_AUX_SEGMENTS], + aux_segment_width: usize, + num_aux_segment_rands: usize, trace_length: usize, trace_meta: Vec, ) -> Self { @@ -107,7 +103,7 @@ impl TraceInfo { // validate trace segment widths assert!(main_segment_width > 0, "main trace segment must consist of at least one column"); - let full_width = main_segment_width + aux_segment_widths.iter().sum::(); + let full_width = main_segment_width + aux_segment_width; assert!( full_width <= TraceInfo::MAX_TRACE_WIDTH, "total number of columns in the trace cannot be greater than {}, but was {}", @@ -115,41 +111,29 @@ impl TraceInfo { full_width ); - // validate number of random elements required by each segment - let mut was_zero_width = false; - let mut num_aux_segments = 0; - for (&width, &num_rand_elements) in aux_segment_widths.iter().zip(aux_segment_rands.iter()) - { - if width != 0 { - assert!( - !was_zero_width, - "a non-empty trace segment cannot follow an empty segment" - ); - assert!( - num_rand_elements > 0, - "number of random elements for a non-empty trace segment must be greater than zero" - ); - num_aux_segments += 1; - } else { - assert!( - num_rand_elements == 0, - "number of random elements for an empty trace segment must be zero" - ); - was_zero_width = true; - } + // validate number of random elements required by the auxiliary segment + if aux_segment_width != 0 { assert!( - num_rand_elements <= TraceInfo::MAX_RAND_SEGMENT_ELEMENTS, - "number of random elements required by a segment cannot exceed {}, but was {}", - TraceInfo::MAX_RAND_SEGMENT_ELEMENTS, - num_rand_elements + num_aux_segment_rands > 0, + "number of random elements for a non-empty auxiliary trace segment must be greater than zero" + ); + } else { + assert!( + num_aux_segment_rands == 0, + "number of random elements for an empty auxiliary trace segment must be zero" ); } + assert!( + num_aux_segment_rands <= TraceInfo::MAX_RAND_SEGMENT_ELEMENTS, + "number of random elements required by a segment cannot exceed {}, but was {}", + TraceInfo::MAX_RAND_SEGMENT_ELEMENTS, + num_aux_segment_rands + ); TraceInfo { main_segment_width, - aux_segment_widths, - aux_segment_rands, - num_aux_segments, + aux_segment_width, + num_aux_segment_rands, trace_length, trace_meta, } @@ -162,7 +146,7 @@ impl TraceInfo { /// /// This is guaranteed to be between 1 and 255. pub fn width(&self) -> usize { - self.main_segment_width + self.aux_segment_widths[0] + self.main_segment_width + self.aux_segment_width } /// Returns execution trace length. @@ -177,9 +161,9 @@ impl TraceInfo { &self.trace_meta } - /// Returns true if an execution trace contains more than one segment. + /// Returns true if an execution trace contains the auxiliary trace segment. pub fn is_multi_segment(&self) -> bool { - self.num_aux_segments > 0 + self.aux_segment_width > 0 } /// Returns the number of columns in the main segment of an execution trace. @@ -189,37 +173,37 @@ impl TraceInfo { self.main_segment_width } - /// Returns the number of columns in all auxiliary segments of an execution trace. - pub fn aux_trace_width(&self) -> usize { - self.aux_segment_widths.iter().sum() + /// Returns the number of columns in the auxiliary segment of an execution trace. + pub fn aux_segment_width(&self) -> usize { + self.aux_segment_width } /// Returns the total number of segments in an execution trace. pub fn num_segments(&self) -> usize { - self.num_aux_segments + 1 + if self.is_multi_segment() { + 2 + } else { + 1 + } } /// Returns the number of auxiliary trace segments in an execution trace. pub fn num_aux_segments(&self) -> usize { - self.num_aux_segments + if self.is_multi_segment() { + 1 + } else { + 0 + } } - /// Returns the number of columns in the auxiliary trace segment at the specified index. - pub fn get_aux_segment_width(&self, segment_idx: usize) -> usize { - assert!( - segment_idx < self.num_aux_segments, - "attempted to access segment index {segment_idx}, but there are only {} segments", - self.num_aux_segments - ); - - self.aux_segment_widths[segment_idx] + /// Returns the number of columns in the auxiliary trace segment. + pub fn get_aux_segment_width(&self) -> usize { + self.aux_segment_width } - /// Returns the number of random elements required by the auxiliary trace segment at the - /// specified index. - pub fn get_aux_segment_rand_elements(&self, segment_idx: usize) -> usize { - // TODO: panic if segment_idx is not within num_aux_segments - self.aux_segment_rands[segment_idx] + /// Returns the number of random elements required by the auxiliary trace segment. + pub fn get_num_aux_segment_rand_elements(&self) -> usize { + self.num_aux_segment_rands } } @@ -231,20 +215,13 @@ impl ToElements for TraceInfo { // segment (if present) go into the first field element; we assume that each parameter can // be encoded in 8 bits (which is enforced by the constructor) let mut buf = self.main_segment_width as u32; - buf = (buf << 8) | self.num_aux_segments as u32; - if self.num_aux_segments == 1 { - buf = (buf << 8) | self.aux_segment_widths[0] as u32; - buf = (buf << 8) | self.aux_segment_rands[0] as u32; + buf = (buf << 8) | self.num_aux_segments() as u32; + if self.num_aux_segments() == 1 { + buf = (buf << 8) | self.aux_segment_width as u32; + buf = (buf << 8) | self.num_aux_segment_rands as u32; } result.push(E::from(buf)); - // parameters of all subsequent auxiliary segments go into additional elements - for i in 1..self.num_aux_segments { - buf = self.aux_segment_widths[i] as u32; - buf = (buf << 8) | self.aux_segment_rands[i] as u32; - result.push(E::from(buf)); - } - // We assume here that the trace length is never greater than 2^32. result.push(E::from(self.trace_length as u32)); @@ -266,17 +243,17 @@ impl Serializable for TraceInfo { fn write_into(&self, target: &mut W) { // store segments target.write_u8(self.main_segment_width as u8); - for &w in self.aux_segment_widths.iter() { - debug_assert!(w <= u8::MAX as usize, "aux segment width does not fit into u8 value"); - target.write_u8(w as u8); - } - for &rc in self.aux_segment_rands.iter() { - debug_assert!( - rc <= u8::MAX as usize, - "aux segment random element count does not fit into u8 value" - ); - target.write_u8(rc as u8); - } + + debug_assert!( + self.aux_segment_width <= u8::MAX as usize, + "aux segment width does not fit into u8 value" + ); + target.write_u8(self.aux_segment_width as u8); + debug_assert!( + self.num_aux_segment_rands <= u8::MAX as usize, + "aux segment random element count does not fit into u8 value" + ); + target.write_u8(self.num_aux_segment_rands as u8); // store trace length as power of two target.write_u8(self.trace_length.ilog2() as u8); @@ -301,23 +278,10 @@ impl Deserializable for TraceInfo { )); } - // read and validate auxiliary trace segment widths - let mut was_zero_width = false; - let mut aux_segment_widths = [0; NUM_AUX_SEGMENTS]; - for width in aux_segment_widths.iter_mut() { - *width = source.read_u8()? as usize; - if *width != 0 { - if was_zero_width { - return Err(DeserializationError::InvalidValue( - "a non-empty trace segment cannot follow an empty segment".to_string(), - )); - } - } else { - was_zero_width = true; - } - } + // read auxiliary trace segment width + let aux_segment_width = source.read_u8()? as usize; - let full_trace_width = main_segment_width + aux_segment_widths.iter().sum::(); + let full_trace_width = main_segment_width + aux_segment_width; if full_trace_width >= TraceInfo::MAX_TRACE_WIDTH { return Err(DeserializationError::InvalidValue(format!( "full trace width cannot be greater than {}, but was {}", @@ -326,28 +290,18 @@ impl Deserializable for TraceInfo { ))); } - // read and validate number of random elements for each auxiliary trace segment - let mut aux_segment_rands = [0; NUM_AUX_SEGMENTS]; - for (num_rand_elements, &width) in - aux_segment_rands.iter_mut().zip(aux_segment_widths.iter()) - { - *num_rand_elements = source.read_u8()? as usize; - if width == 0 && *num_rand_elements != 0 { - return Err(DeserializationError::InvalidValue( - "an empty trace segment cannot require random elements".to_string(), - )); - } else if width != 0 && *num_rand_elements == 0 { - return Err(DeserializationError::InvalidValue( - "a non-empty trace segment must require at least one random element" - .to_string(), - )); - } else if *num_rand_elements > TraceInfo::MAX_RAND_SEGMENT_ELEMENTS { - return Err(DeserializationError::InvalidValue(format!( - "number of random elements required by a segment cannot exceed {}, but was {}", - TraceInfo::MAX_RAND_SEGMENT_ELEMENTS, - *num_rand_elements - ))); - } + // read and validate number of random elements for the auxiliary trace segment + let num_aux_segment_rands = source.read_u8()? as usize; + if aux_segment_width != 0 && num_aux_segment_rands == 0 { + return Err(DeserializationError::InvalidValue( + "a non-empty trace segment must require at least one random element".to_string(), + )); + } else if num_aux_segment_rands > TraceInfo::MAX_RAND_SEGMENT_ELEMENTS { + return Err(DeserializationError::InvalidValue(format!( + "number of random elements required by a segment cannot exceed {}, but was {}", + TraceInfo::MAX_RAND_SEGMENT_ELEMENTS, + num_aux_segment_rands + ))); } // read and validate trace length (which was stored as a power of two) @@ -371,8 +325,8 @@ impl Deserializable for TraceInfo { Ok(Self::new_multi_segment( main_segment_width, - aux_segment_widths, - aux_segment_rands, + aux_segment_width, + num_aux_segment_rands, trace_length, trace_meta, )) @@ -425,8 +379,8 @@ mod tests { let info = TraceInfo::new_multi_segment( main_width as usize, - [aux_width as usize], - [aux_rands as usize], + aux_width as usize, + aux_rands, trace_length as usize, trace_meta, ); diff --git a/air/src/air/transition/mod.rs b/air/src/air/transition/mod.rs index c3f065dcc..28be07f38 100644 --- a/air/src/air/transition/mod.rs +++ b/air/src/air/transition/mod.rs @@ -96,7 +96,7 @@ impl TransitionConstraints { self.main_constraint_coef.clone() } - /// Returns a list of transition constraint degree descriptors for auxiliary trace segments of + /// Returns a list of transition constraint degree descriptors for the auxiliary trace segment of /// a computation. /// /// This list will be identical to the list passed into [AirContext::new_multi_segment()] @@ -105,13 +105,13 @@ impl TransitionConstraints { &self.aux_constraint_degrees } - /// Returns the number of constraints applied against auxiliary trace segments of a + /// Returns the number of constraints applied against the auxiliary trace segment of a /// computation. pub fn num_aux_constraints(&self) -> usize { self.aux_constraint_degrees.len() } - /// Returns the random coefficients for constraints applied against auxiliary trace segments of a + /// Returns the random coefficients for constraints applied against the auxiliary trace segment of a /// computation. pub fn aux_constraint_coef(&self) -> Vec { self.aux_constraint_coef.clone() diff --git a/air/src/proof/context.rs b/air/src/proof/context.rs index b3fa289fc..1f489ca70 100644 --- a/air/src/proof/context.rs +++ b/air/src/proof/context.rs @@ -187,8 +187,8 @@ mod tests { let expected = { let trace_info = TraceInfo::new_multi_segment( main_width, - [aux_width], - [aux_rands], + aux_width, + aux_rands, trace_length, vec![], ); @@ -214,13 +214,8 @@ mod tests { fri_folding_factor as usize, fri_remainder_max_degree as usize, ); - let trace_info = TraceInfo::new_multi_segment( - main_width, - [aux_width], - [aux_rands], - trace_length, - vec![], - ); + let trace_info = + TraceInfo::new_multi_segment(main_width, aux_width, aux_rands, trace_length, vec![]); let context = Context::new::(trace_info, options); assert_eq!(expected, context.to_elements()); } diff --git a/examples/src/rescue_raps/air.rs b/examples/src/rescue_raps/air.rs index 5a392c92c..5ea1e71f0 100644 --- a/examples/src/rescue_raps/air.rs +++ b/examples/src/rescue_raps/air.rs @@ -171,7 +171,7 @@ impl Air for RescueRapsAir { let aux_current = aux_frame.current(); let aux_next = aux_frame.next(); - let random_elements = aux_rand_elements.get_segment_elements(0); + let random_elements = aux_rand_elements.get_segment_elements(); let absorption_flag = periodic_values[1]; diff --git a/examples/src/rescue_raps/custom_trace_table.rs b/examples/src/rescue_raps/custom_trace_table.rs index 1bf8fcf14..58b8f330d 100644 --- a/examples/src/rescue_raps/custom_trace_table.rs +++ b/examples/src/rescue_raps/custom_trace_table.rs @@ -93,7 +93,7 @@ impl RapTraceTable { let columns = unsafe { (0..width).map(|_| uninit_vector(length)).collect() }; Self { - info: TraceInfo::new_multi_segment(width, [3], [3], length, meta), + info: TraceInfo::new_multi_segment(width, 3, 3, length, meta), trace: ColMatrix::new(columns), } } @@ -173,22 +173,17 @@ impl Trace for RapTraceTable { fn build_aux_segment( &mut self, - aux_segments: &[ColMatrix], rand_elements: &[E], _lagrange_rand_elements: Option<&[E]>, ) -> Option> where E: FieldElement, { - // We only have one auxiliary segment for this example - if !aux_segments.is_empty() { - return None; - } - let mut current_row = unsafe { uninit_vector(self.width()) }; let mut next_row = unsafe { uninit_vector(self.width()) }; self.read_row_into(0, &mut current_row); - let mut aux_columns = vec![vec![E::ZERO; self.info.length()]; self.info.aux_trace_width()]; + let mut aux_columns = + vec![vec![E::ZERO; self.info.length()]; self.info.aux_segment_width()]; // Columns storing the copied values for the permutation argument are not necessary, but // help understanding the construction of RAPs and are kept for illustrative purposes. diff --git a/prover/benches/lagrange_kernel.rs b/prover/benches/lagrange_kernel.rs index 2da94d6dd..69ee979b0 100644 --- a/prover/benches/lagrange_kernel.rs +++ b/prover/benches/lagrange_kernel.rs @@ -63,8 +63,8 @@ impl LagrangeTrace { main_trace: ColMatrix::new(vec![main_trace_col]), info: TraceInfo::new_multi_segment( 1, - [aux_segment_width], - [num_aux_segment_rands], + aux_segment_width, + num_aux_segment_rands, trace_len, vec![], ), @@ -91,12 +91,9 @@ impl Trace for LagrangeTrace { /// by the main column fn build_aux_segment>( &mut self, - aux_segments: &[ColMatrix], rand_elements: &[E], lagrange_kernel_rand_elements: Option<&[E]>, ) -> Option> { - assert!(aux_segments.is_empty()); - let mut columns = Vec::new(); // first build the Lagrange kernel column diff --git a/prover/src/channel.rs b/prover/src/channel.rs index 4d074d9e1..ec1d68b72 100644 --- a/prover/src/channel.rs +++ b/prover/src/channel.rs @@ -100,14 +100,13 @@ where // PUBLIC COIN METHODS // -------------------------------------------------------------------------------------------- - /// Returns a set of random elements required for constructing an auxiliary trace segment with - /// the specified index. + /// Returns a set of random elements required for constructing the auxiliary trace segment. /// /// The elements are drawn from the public coin uniformly at random. - pub fn get_aux_trace_segment_rand_elements(&mut self, aux_segment_idx: usize) -> Vec { + pub fn get_aux_trace_segment_rand_elements(&mut self) -> Vec { self.air - .get_aux_trace_segment_random_elements(aux_segment_idx, &mut self.public_coin) - .expect("failed to draw random elements for an auxiliary trace segment") + .get_aux_trace_segment_random_elements(&mut self.public_coin) + .expect("failed to draw random elements for the auxiliary trace segment") } /// Returns a set of coefficients for constructing a constraint composition polynomial. diff --git a/prover/src/composer/mod.rs b/prover/src/composer/mod.rs index 729d605ea..f6be2c74a 100644 --- a/prover/src/composer/mod.rs +++ b/prover/src/composer/mod.rs @@ -105,7 +105,7 @@ impl DeepCompositionPoly { i += 1; } - // --- merge polynomials of the auxiliary trace segments ---------------------------------- + // --- merge polynomials of the auxiliary trace segment ---------------------------------- for poly in trace_polys.aux_trace_polys() { // compute T'(x) = T(x) - T(z), multiply it by a pseudo-random coefficient, // and add the result into composition polynomial diff --git a/prover/src/constraints/evaluation_table.rs b/prover/src/constraints/evaluation_table.rs index b57897f4a..bc4fe20b6 100644 --- a/prover/src/constraints/evaluation_table.rs +++ b/prover/src/constraints/evaluation_table.rs @@ -204,7 +204,7 @@ impl<'a, E: FieldElement> ConstraintEvaluationTable<'a, E> { max_degree = core::cmp::max(max_degree, degree); } - // then process transition constraint evaluations for auxiliary trace segments + // then process transition constraint evaluations for the auxiliary trace segment for evaluations in self.aux_transition_evaluations.iter() { let degree = get_transition_poly_degree(evaluations, &inv_twiddles, &div_values); actual_degrees.push(degree); diff --git a/prover/src/constraints/evaluator/boundary.rs b/prover/src/constraints/evaluator/boundary.rs index d2ca0bde0..575b0c94e 100644 --- a/prover/src/constraints/evaluator/boundary.rs +++ b/prover/src/constraints/evaluator/boundary.rs @@ -20,7 +20,7 @@ const SMALL_POLY_DEGREE: usize = 63; /// Contains all boundary constraints defined for an instance of a computation. This includes /// constraints against the main segment of the execution trace as well as constraints against -/// auxiliary trace segments (if any). +/// the auxiliary trace segment (if any). /// /// We transform the constraints defined in the [air] crate into specialized constraints here /// to make evaluation of these constraints more efficient in the prover context. @@ -55,7 +55,7 @@ impl BoundaryConstraints { }) .collect::>>(); - // transform constraints against auxiliary trace segments (if any) into specialized + // transform constraints against the auxiliary trace segment (if any) into specialized // constraints. this also checks if a group with the same divisor has already been // transformed (when processing constraints against the main trace above), and if so, // appends constraints to that group rather than creating a new group. this ensures diff --git a/prover/src/constraints/evaluator/default.rs b/prover/src/constraints/evaluator/default.rs index 6d4607d00..9e40114db 100644 --- a/prover/src/constraints/evaluator/default.rs +++ b/prover/src/constraints/evaluator/default.rs @@ -153,7 +153,7 @@ where Some(LagrangeKernelConstraintsBatchEvaluator::new( air, - aux_rand_elements.get_segment_elements(0)[0..num_rand_elements].to_vec(), + aux_rand_elements.get_segment_elements()[0..num_rand_elements].to_vec(), composition_coefficients .lagrange .expect("expected Lagrange kernel composition coefficients to be present"), @@ -237,7 +237,7 @@ where ) { // initialize buffers to hold trace values and evaluation results at each step let mut main_frame = EvaluationFrame::new(trace.trace_info().main_trace_width()); - let mut aux_frame = EvaluationFrame::new(trace.trace_info().aux_trace_width()); + let mut aux_frame = EvaluationFrame::new(trace.trace_info().aux_segment_width()); let mut tm_evaluations = vec![E::BaseField::ZERO; self.num_main_transition_constraints()]; let mut ta_evaluations = vec![E::ZERO; self.num_aux_transition_constraints()]; let mut evaluations = vec![E::ZERO; fragment.num_columns()]; @@ -337,7 +337,7 @@ where .fold(E::ZERO, |acc, (&const_eval, &coef)| acc + coef.mul_base(const_eval)) } - /// Evaluates all transition constraints (i.e., for main and auxiliary trace segments) at the + /// Evaluates all transition constraints (i.e., for main and the auxiliary trace segment) at the /// specified step of the constraint evaluation domain. /// /// `x` is the corresponding domain value at the specified step. That is, x = s * g^step, @@ -355,7 +355,7 @@ where // get periodic values at the evaluation step let periodic_values = self.periodic_values.get_row(step); - // evaluate transition constraints over auxiliary trace segments and save the results into + // evaluate transition constraints over the auxiliary trace segment and save the results into // evaluations buffer self.air.evaluate_aux_transition( main_frame, @@ -382,7 +382,7 @@ where self.transition_constraints.num_main_constraints() } - /// Returns the number of transition constraints applied against all auxiliary trace segments. + /// Returns the number of transition constraints applied against the auxiliary trace segment. fn num_aux_transition_constraints(&self) -> usize { self.transition_constraints.num_aux_constraints() } diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 69073145c..b3ee6abcf 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -54,7 +54,6 @@ pub use utils::{ SliceReader, }; -use alloc::vec::Vec; use fri::FriProver; pub use math; @@ -278,17 +277,17 @@ pub trait Prover { (trace_lde, trace_polys) }; - // build auxiliary trace segments (if any), and append the resulting segments to trace + // build the auxiliary trace segment (if any), and append the resulting segments to trace // commitment and trace polynomial table structs - let mut aux_trace_segments = Vec::new(); + let mut aux_trace_segment: Option> = None; let mut aux_trace_rand_elements = AuxTraceRandElements::new(); - for i in 0..trace.info().num_aux_segments() { - let num_columns = trace.info().get_aux_segment_width(i); + if trace.info().is_multi_segment() { + let num_columns = trace.info().get_aux_segment_width(); let (aux_segment, rand_elements) = { let span = info_span!("build_aux_trace_segment", num_columns).entered(); - // draw a set of random elements required to build an auxiliary trace segment - let rand_elements = channel.get_aux_trace_segment_rand_elements(i); + // draw a set of random elements required to build the auxiliary trace segment + let rand_elements = channel.get_aux_trace_segment_rand_elements(); let lagrange_rand_elements = if air.context().has_lagrange_kernel_aux_column() { Some(rand_elements.as_ref()) } else { @@ -297,7 +296,7 @@ pub trait Prover { // build the trace segment let aux_segment = trace - .build_aux_segment(&aux_trace_segments, &rand_elements, lagrange_rand_elements) + .build_aux_segment(&rand_elements, lagrange_rand_elements) .expect("failed build auxiliary trace segment"); drop(span); @@ -312,7 +311,7 @@ pub trait Prover { // trace let span = info_span!("commit_to_aux_trace_segment").entered(); let (aux_segment_polys, aux_segment_root) = - trace_lde.add_aux_segment(&aux_segment, &domain); + trace_lde.set_aux_trace(&aux_segment, &domain); // commit to the LDE of the extended auxiliary trace segment by writing the root of // its Merkle tree into the channel @@ -324,19 +323,19 @@ pub trait Prover { trace_polys .add_aux_segment(aux_segment_polys, air.context().lagrange_kernel_aux_column_idx()); - aux_trace_rand_elements.add_segment_elements(rand_elements); - aux_trace_segments.push(aux_segment); + aux_trace_rand_elements.set_segment_elements(rand_elements); + aux_trace_segment = Some(aux_segment); } // make sure the specified trace (including auxiliary segments) is valid against the AIR. // This checks validity of both, assertions and state transitions. We do this in debug // mode only because this is a very expensive operation. #[cfg(debug_assertions)] - trace.validate(&air, &aux_trace_segments, &aux_trace_rand_elements); + trace.validate(&air, aux_trace_segment.as_ref(), &aux_trace_rand_elements); // drop the main trace and aux trace segments as they are no longer needed drop(trace); - drop(aux_trace_segments); + drop(aux_trace_segment); // 2 ----- evaluate constraints ----------------------------------------------------------- // evaluate constraints specified by the AIR over the constraint evaluation domain, and diff --git a/prover/src/matrix/col_matrix.rs b/prover/src/matrix/col_matrix.rs index 8f49b12ab..8c7093a27 100644 --- a/prover/src/matrix/col_matrix.rs +++ b/prover/src/matrix/col_matrix.rs @@ -292,13 +292,23 @@ impl ColMatrix { /// Iterator over columns of [ColMatrix]. pub struct ColumnIter<'a, E: FieldElement> { - matrix: &'a ColMatrix, + matrix: Option<&'a ColMatrix>, cursor: usize, } impl<'a, E: FieldElement> ColumnIter<'a, E> { pub fn new(matrix: &'a ColMatrix) -> Self { - Self { matrix, cursor: 0 } + Self { + matrix: Some(matrix), + cursor: 0, + } + } + + pub fn empty() -> Self { + Self { + matrix: None, + cursor: 0, + } } } @@ -306,25 +316,34 @@ impl<'a, E: FieldElement> Iterator for ColumnIter<'a, E> { type Item = &'a [E]; fn next(&mut self) -> Option { - match self.matrix.num_cols() - self.cursor { - 0 => None, - _ => { - let column = self.matrix.get_column(self.cursor); - self.cursor += 1; - Some(column) - } + match self.matrix { + Some(matrix) => match matrix.num_cols() - self.cursor { + 0 => None, + _ => { + let column = matrix.get_column(self.cursor); + self.cursor += 1; + Some(column) + } + }, + None => None, } } } impl<'a, E: FieldElement> ExactSizeIterator for ColumnIter<'a, E> { fn len(&self) -> usize { - self.matrix.num_cols() + self.matrix.map(|matrix| matrix.num_cols()).unwrap_or_default() } } impl<'a, E: FieldElement> FusedIterator for ColumnIter<'a, E> {} +impl<'a, E: FieldElement> Default for ColumnIter<'a, E> { + fn default() -> Self { + Self::empty() + } +} + // MUTABLE COLUMN ITERATOR // ================================================================================================ @@ -367,66 +386,3 @@ impl<'a, E: FieldElement> ExactSizeIterator for ColumnIterMut<'a, E> { } impl<'a, E: FieldElement> FusedIterator for ColumnIterMut<'a, E> {} - -// MULTI-MATRIX COLUMN ITERATOR -// ================================================================================================ - -/// Iterator over columns several [ColMatrix]'s. -pub struct MultiColumnIter<'a, E: FieldElement> { - matrixes: &'a [ColMatrix], - m_cursor: usize, - c_cursor: usize, -} - -impl<'a, E: FieldElement> MultiColumnIter<'a, E> { - pub fn new(matrixes: &'a [ColMatrix]) -> Self { - // make sure all matrixes have the same number of rows - if !matrixes.is_empty() { - let num_rows = matrixes[0].num_rows(); - for matrix in matrixes.iter().skip(1) { - assert_eq!( - matrix.num_rows(), - num_rows, - "all matrixes must have the same number of rows" - ); - } - } - - Self { - matrixes, - m_cursor: 0, - c_cursor: 0, - } - } -} - -impl<'a, E: FieldElement> Iterator for MultiColumnIter<'a, E> { - type Item = &'a [E]; - - fn next(&mut self) -> Option { - if self.matrixes.is_empty() { - return None; - } - let matrix = &self.matrixes[self.m_cursor]; - match matrix.num_cols() - self.c_cursor { - 0 => None, - _ => { - let column = matrix.get_column(self.c_cursor); - self.c_cursor += 1; - if self.c_cursor == matrix.num_cols() && self.m_cursor < self.matrixes.len() - 1 { - self.m_cursor += 1; - self.c_cursor = 0; - } - Some(column) - } - } - } -} - -impl<'a, E: FieldElement> ExactSizeIterator for MultiColumnIter<'a, E> { - fn len(&self) -> usize { - self.matrixes.iter().fold(0, |s, m| s + m.num_cols()) - } -} - -impl<'a, E: FieldElement> FusedIterator for MultiColumnIter<'a, E> {} diff --git a/prover/src/matrix/mod.rs b/prover/src/matrix/mod.rs index a6e8d9c78..75ad2b551 100644 --- a/prover/src/matrix/mod.rs +++ b/prover/src/matrix/mod.rs @@ -9,7 +9,7 @@ mod row_matrix; pub use row_matrix::{build_segments, get_evaluation_offsets, RowMatrix}; mod col_matrix; -pub use col_matrix::{ColMatrix, ColumnIter, MultiColumnIter}; +pub use col_matrix::{ColMatrix, ColumnIter}; mod segments; pub use segments::Segment; diff --git a/prover/src/trace/mod.rs b/prover/src/trace/mod.rs index 837025cd6..7d059d6d0 100644 --- a/prover/src/trace/mod.rs +++ b/prover/src/trace/mod.rs @@ -3,7 +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::{matrix::MultiColumnIter, ColMatrix}; +use super::ColMatrix; use air::{ Air, AuxTraceRandElements, EvaluationFrame, LagrangeKernelBoundaryConstraint, TraceInfo, }; @@ -51,19 +51,13 @@ pub trait Trace: Sized { /// Returns a reference to a [Matrix] describing the main segment of this trace. fn main_segment(&self) -> &ColMatrix; - /// Builds and returns the next auxiliary trace segment. If there are no more segments to build - /// (i.e., the trace is complete), None is returned. - /// - /// The `aux_segments` slice contains a list of auxiliary trace segments built as a result of - /// prior invocations of this function. Thus, for example, on the first invocation, - /// `aux_segments` will be empty; on the second invocation, it will contain a single matrix (the - /// one built during the first invocation) etc. + /// Builds and returns the auxiliary trace segment. If the trace does not require an auxiliary + /// segment, None is returned. /// /// The `rand_elements` slice contains the random elements to use to build the aux segment. If a /// Lagrange kernel column is present, the `lagrange_kernel_rand_elements` should be used. fn build_aux_segment>( &mut self, - aux_segments: &[ColMatrix], rand_elements: &[E], lagrange_kernel_rand_elements: Option<&[E]>, ) -> Option>; @@ -84,9 +78,9 @@ pub trait Trace: Sized { self.info().main_trace_width() } - /// Returns the number of columns in all auxiliary trace segments. + /// Returns the number of columns in the auxiliary trace segment. fn aux_trace_width(&self) -> usize { - self.info().aux_trace_width() + self.info().aux_segment_width() } /// Checks if this trace is valid against the specified AIR, and panics if not. @@ -95,7 +89,7 @@ pub trait Trace: Sized { fn validate( &self, air: &A, - aux_segments: &[ColMatrix], + aux_segment: Option<&ColMatrix>, aux_rand_elements: &AuxTraceRandElements, ) where A: Air, @@ -125,44 +119,33 @@ pub trait Trace: Sized { }); } - // then, check assertions against auxiliary trace segments - for assertion in air.get_aux_assertions(aux_rand_elements) { - // find which segment the assertion is for and remap assertion column index to the - // column index in the context of this segment - let mut column_idx = assertion.column(); - let mut segment_idx = 0; - for i in 0..self.info().num_aux_segments() { - let segment_width = self.info().get_aux_segment_width(i); - if column_idx < segment_width { - segment_idx = i; - break; - } - column_idx -= segment_width; + // then, check assertions against the auxiliary trace segment + if let Some(aux_segment) = aux_segment { + for assertion in air.get_aux_assertions(aux_rand_elements) { + // get the matrix and verify the assertion against it + assertion.apply(self.length(), |step, value| { + assert!( + value == aux_segment.get(assertion.column(), step), + "trace does not satisfy assertion aux_trace({}, {}) == {}", + assertion.column(), + step, + value + ); + }); } - // get the matrix and verify the assertion against it - assertion.apply(self.length(), |step, value| { - assert!( - value == aux_segments[segment_idx].get(column_idx, step), - "trace does not satisfy assertion aux_trace({}, {}) == {}", - assertion.column(), - step, - value - ); - }); - } + // then, check the Lagrange kernel assertion, if any + if let Some(lagrange_kernel_col_idx) = air.context().lagrange_kernel_aux_column_idx() { + let boundary_constraint_assertion_value = + LagrangeKernelBoundaryConstraint::assertion_value( + aux_rand_elements.get_segment_elements(), + ); - // then, check the Lagrange kernel assertion, if any - if let Some(lagrange_kernel_col_idx) = air.context().lagrange_kernel_aux_column_idx() { - let boundary_constraint_assertion_value = - LagrangeKernelBoundaryConstraint::assertion_value( - aux_rand_elements.get_segment_elements(0), + assert_eq!( + boundary_constraint_assertion_value, + aux_segment.get(lagrange_kernel_col_idx, 0) ); - - assert_eq!( - boundary_constraint_assertion_value, - aux_segments[0].get(lagrange_kernel_col_idx, 0) - ); + } } // --- 2. make sure this trace satisfies all transition constraints ----------------------- @@ -205,10 +188,11 @@ pub trait Trace: Sized { ); } - // evaluate transition constraints for auxiliary trace segments (if any) and make + // evaluate transition constraints for the auxiliary trace segment (if any) and make // sure they all evaluate to zeros if let Some(ref mut aux_frame) = aux_frame { - read_aux_frame(aux_segments, step, aux_frame); + let aux_segment = aux_segment.expect("expected aux segment to be present"); + read_aux_frame(aux_segment, step, aux_frame); air.evaluate_aux_transition( &main_frame, aux_frame, @@ -231,9 +215,11 @@ pub trait Trace: Sized { // evaluate transition constraints for Lagrange kernel column (if any) and make sure // they all evaluate to zeros if let Some(col_idx) = air.context().lagrange_kernel_aux_column_idx() { - let c = aux_segments[0].get_column(col_idx); + let aux_segment = aux_segment + .expect("expected aux segment to be present when the Lagrange kernel column is"); + let c = aux_segment.get_column(col_idx); let v = self.length().ilog2() as usize; - let r = aux_rand_elements.get_segment_elements(0); + let r = aux_rand_elements.get_segment_elements(); // Loop over every constraint for constraint_idx in 1..v + 1 { @@ -262,22 +248,24 @@ pub trait Trace: Sized { // HELPER FUNCTIONS // ================================================================================================ -/// Reads an evaluation frame from the set of provided auxiliary segments. This expects that -/// `aux_segments` contains at least one entry. +/// Reads an evaluation frame from the provided auxiliary segment. /// /// This is probably not the most efficient implementation, but since we call this function only /// for trace validation purposes (which is done in debug mode only), we don't care all that much /// about its performance. -fn read_aux_frame(aux_segments: &[ColMatrix], row_idx: usize, frame: &mut EvaluationFrame) +fn read_aux_frame(aux_segment: &ColMatrix, row_idx: usize, frame: &mut EvaluationFrame) where E: FieldElement, { - for (column, current_value) in MultiColumnIter::new(aux_segments).zip(frame.current_mut()) { - *current_value = column[row_idx]; + for (current_frame_cell, aux_segment_col) in + frame.current_mut().iter_mut().zip(aux_segment.columns()) + { + *current_frame_cell = aux_segment_col[row_idx]; } - let next_row_idx = (row_idx + 1) % aux_segments[0].num_rows(); - for (column, next_value) in MultiColumnIter::new(aux_segments).zip(frame.next_mut()) { - *next_value = column[next_row_idx]; + let next_row_idx = (row_idx + 1) % aux_segment.num_rows(); + for (next_frame_cell, aux_segment_col) in frame.next_mut().iter_mut().zip(aux_segment.columns()) + { + *next_frame_cell = aux_segment_col[next_row_idx]; } } diff --git a/prover/src/trace/poly_table.rs b/prover/src/trace/poly_table.rs index 1aeb53d3d..0a53df51a 100644 --- a/prover/src/trace/poly_table.rs +++ b/prover/src/trace/poly_table.rs @@ -3,10 +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 crate::{ - matrix::{ColumnIter, MultiColumnIter}, - ColMatrix, -}; +use crate::{matrix::ColumnIter, ColMatrix}; use air::{proof::TraceOodFrame, LagrangeKernelEvaluationFrame}; use alloc::vec::Vec; use math::{FieldElement, StarkField}; @@ -17,11 +14,11 @@ use math::{FieldElement, StarkField}; /// Trace polynomials in coefficient from for all segments of the execution trace. /// /// Coefficients of the polynomials for the main trace segment are always in the base field. -/// However, coefficients of the polynomials for the auxiliary trace segments may be either in the +/// However, coefficients of the polynomials for the auxiliary trace segment may be either in the /// base field, or in the extension field, depending on whether extension field is being used. pub struct TracePolyTable { - main_segment_polys: ColMatrix, - aux_segment_polys: Vec>, + main_trace_polys: ColMatrix, + aux_trace_polys: Option>, lagrange_kernel_column_idx: Option, } @@ -31,8 +28,8 @@ impl TracePolyTable { /// Creates a new table of trace polynomials from the provided main trace segment polynomials. pub fn new(main_trace_polys: ColMatrix) -> Self { Self { - main_segment_polys: main_trace_polys, - aux_segment_polys: Vec::new(), + main_trace_polys, + aux_trace_polys: None, lagrange_kernel_column_idx: None, } } @@ -43,15 +40,16 @@ impl TracePolyTable { /// Adds the provided auxiliary segment polynomials to this polynomial table. pub fn add_aux_segment( &mut self, - aux_segment_polys: ColMatrix, + aux_trace_polys: ColMatrix, lagrange_kernel_column_idx: Option, ) { + assert!(self.aux_trace_polys.is_none()); assert_eq!( - self.main_segment_polys.num_rows(), - aux_segment_polys.num_rows(), + self.main_trace_polys.num_rows(), + aux_trace_polys.num_rows(), "polynomials in auxiliary segment must be of the same size as in the main segment" ); - self.aux_segment_polys.push(aux_segment_polys); + self.aux_trace_polys = Some(aux_trace_polys); self.lagrange_kernel_column_idx = lagrange_kernel_column_idx; } @@ -60,13 +58,13 @@ impl TracePolyTable { /// Returns the size of each polynomial - i.e. size of a vector needed to hold a polynomial. pub fn poly_size(&self) -> usize { - self.main_segment_polys.num_rows() + self.main_trace_polys.num_rows() } /// Evaluates all trace polynomials (across all trace segments) at the specified point `x`. pub fn evaluate_at(&self, x: E) -> Vec { - let mut result = self.main_segment_polys.evaluate_columns_at(x); - for aux_polys in self.aux_segment_polys.iter() { + let mut result = self.main_trace_polys.evaluate_columns_at(x); + for aux_polys in self.aux_trace_polys.iter() { result.append(&mut aux_polys.evaluate_columns_at(x)); } result @@ -84,7 +82,11 @@ impl TracePolyTable { let next_row = self.evaluate_at(z * g); let lagrange_kernel_frame = self.lagrange_kernel_column_idx.map(|col_idx| { - let lagrange_kernel_col_poly = self.aux_segment_polys[0].get_column(col_idx); + let aux_segment_poly = self + .aux_trace_polys + .as_ref() + .expect("aux segment poly and lagrange kernel column idx are set together"); + let lagrange_kernel_col_poly = aux_segment_poly.get_column(col_idx); LagrangeKernelEvaluationFrame::from_lagrange_kernel_column_poly( lagrange_kernel_col_poly, @@ -92,19 +94,22 @@ impl TracePolyTable { ) }); - let main_trace_width = self.main_segment_polys.num_cols(); + let main_trace_width = self.main_trace_polys.num_cols(); TraceOodFrame::new(current_row, next_row, main_trace_width, lagrange_kernel_frame) } /// Returns an iterator over the polynomials of the main trace segment. - pub fn main_trace_polys(&self) -> ColumnIter { - self.main_segment_polys.columns() + pub fn main_trace_polys(&self) -> impl Iterator { + self.main_trace_polys.columns() } - /// Returns an iterator over the polynomials of all auxiliary trace segments. - pub fn aux_trace_polys(&self) -> MultiColumnIter { - MultiColumnIter::new(self.aux_segment_polys.as_slice()) + /// Returns an iterator over the polynomials of the auxiliary trace segment. + pub fn aux_trace_polys(&self) -> impl Iterator { + match self.aux_trace_polys { + Some(ref aux_segment_polys) => aux_segment_polys.columns(), + None => ColumnIter::empty(), + } } // TEST HELPERS @@ -113,12 +118,12 @@ impl TracePolyTable { /// Returns the number of polynomials in the main segment of the trace. #[cfg(test)] pub fn num_main_trace_polys(&self) -> usize { - self.main_segment_polys.num_cols() + self.main_trace_polys.num_cols() } /// Returns a polynomial from the main segment of the trace at the specified index. #[cfg(test)] pub fn get_main_trace_poly(&self, idx: usize) -> &[E::BaseField] { - self.main_segment_polys.get_column(idx) + self.main_trace_polys.get_column(idx) } } diff --git a/prover/src/trace/trace_lde/default/mod.rs b/prover/src/trace/trace_lde/default/mod.rs index 8d7fa9c92..0f693146f 100644 --- a/prover/src/trace/trace_lde/default/mod.rs +++ b/prover/src/trace/trace_lde/default/mod.rs @@ -31,10 +31,10 @@ pub struct DefaultTraceLde, // commitment to the main segment of the trace main_segment_tree: MerkleTree, - // low-degree extensions of the auxiliary segments of the trace - aux_segment_ldes: Vec>, - // commitment to the auxiliary segments of the trace - aux_segment_trees: Vec>, + // low-degree extensions of the auxiliary segment of the trace + aux_segment_lde: Option>, + // commitment to the auxiliary segment of the trace + aux_segment_tree: Option>, blowup: usize, trace_info: TraceInfo, } @@ -60,8 +60,8 @@ impl> DefaultTraceLd let trace_lde = DefaultTraceLde { main_segment_lde, main_segment_tree, - aux_segment_ldes: Vec::new(), - aux_segment_trees: Vec::new(), + aux_segment_lde: None, + aux_segment_tree: None, blowup: domain.trace_to_lde_blowup(), trace_info: trace_info.clone(), }; @@ -117,8 +117,8 @@ where /// /// This function will panic if any of the following are true: /// - the number of rows in the provided `aux_trace` does not match the main trace. - /// - this segment would exceed the number of segments specified by the trace layout. - fn add_aux_segment( + /// - the auxiliary trace has been previously set already. + fn set_aux_trace( &mut self, aux_trace: &ColMatrix, domain: &StarkDomain, @@ -129,8 +129,8 @@ where // check errors assert!( - self.aux_segment_ldes.len() < self.trace_info.num_aux_segments(), - "the specified number of auxiliary segments has already been added" + usize::from(self.aux_segment_lde.is_some()) < self.trace_info.num_aux_segments(), + "the auxiliary trace has already been added" ); assert_eq!( self.main_segment_lde.num_rows(), @@ -139,9 +139,9 @@ where ); // save the lde and commitment - self.aux_segment_ldes.push(aux_segment_lde); + self.aux_segment_lde = Some(aux_segment_lde); let root_hash = *aux_segment_tree.root(); - self.aux_segment_trees.push(aux_segment_tree); + self.aux_segment_tree = Some(aux_segment_tree); (aux_segment_polys, root_hash) } @@ -170,7 +170,7 @@ where let next_lde_step = (lde_step + self.blowup()) % self.trace_len(); // copy auxiliary trace segment values into the frame - let segment = &self.aux_segment_ldes[0]; + let segment = self.aux_segment_lde.as_ref().expect("expected aux segment to be present"); frame.current_mut().copy_from_slice(segment.row(lde_step)); frame.next_mut().copy_from_slice(segment.row(next_lde_step)); } @@ -184,7 +184,8 @@ where let frame = frame.frame_mut(); frame.truncate(0); - let aux_segment = &self.aux_segment_ldes[0]; + let aux_segment = + self.aux_segment_lde.as_ref().expect("expected aux segment to be present"); frame.push(aux_segment.get(lagrange_kernel_aux_column_idx, lde_step)); @@ -207,9 +208,10 @@ where positions, )]; - // build queries for auxiliary trace segments - for (i, segment_tree) in self.aux_segment_trees.iter().enumerate() { - let segment_lde = &self.aux_segment_ldes[i]; + // build queries for the auxiliary trace segment + if let Some(ref segment_tree) = self.aux_segment_tree { + let segment_lde = + self.aux_segment_lde.as_ref().expect("expected aux segment to be present"); result.push(build_segment_queries(segment_lde, segment_tree, positions)); } diff --git a/prover/src/trace/trace_lde/mod.rs b/prover/src/trace/trace_lde/mod.rs index 5243c695c..2081371fd 100644 --- a/prover/src/trace/trace_lde/mod.rs +++ b/prover/src/trace/trace_lde/mod.rs @@ -40,7 +40,7 @@ pub trait TraceLde: Sync { /// This function is expected to panic if any of the following are true: /// - the number of rows in the provided `aux_trace` does not match the main trace. /// - this segment would exceed the number of segments specified by the trace layout. - fn add_aux_segment( + fn set_aux_trace( &mut self, aux_trace: &ColMatrix, domain: &StarkDomain, diff --git a/prover/src/trace/trace_table.rs b/prover/src/trace/trace_table.rs index 87aec2a19..8bf726d95 100644 --- a/prover/src/trace/trace_table.rs +++ b/prover/src/trace/trace_table.rs @@ -301,7 +301,6 @@ impl Trace for TraceTable { fn build_aux_segment( &mut self, - _aux_segments: &[ColMatrix], _rand_elements: &[E], _lagrange_rand_elements: Option<&[E]>, ) -> Option> diff --git a/verifier/src/channel.rs b/verifier/src/channel.rs index 038ec1c6c..9f23c566c 100644 --- a/verifier/src/channel.rs +++ b/verifier/src/channel.rs @@ -68,7 +68,7 @@ impl> VerifierChanne let num_trace_segments = air.trace_info().num_segments(); let main_trace_width = air.trace_info().main_trace_width(); - let aux_trace_width = air.trace_info().aux_trace_width(); + let aux_trace_width = air.trace_info().aux_segment_width(); let lde_domain_size = air.lde_domain_size(); let fri_options = air.options().to_fri_options(); @@ -272,19 +272,18 @@ impl> TraceQueries(air.lde_domain_size(), num_queries, segment_width) - .map_err(|err| { - VerifierError::ProofDeserializationError(format!( - "auxiliary trace segment query deserialization failed: {err}" - )) - })?; - - query_proofs.push(segment_query_proof); - aux_trace_states.push(segment_trace_states); - } + let segment_queries = queries.remove(0); + let segment_width = air.trace_info().get_aux_segment_width(); + let (segment_query_proof, segment_trace_states) = segment_queries + .parse::(air.lde_domain_size(), num_queries, segment_width) + .map_err(|err| { + VerifierError::ProofDeserializationError(format!( + "auxiliary trace segment query deserialization failed: {err}" + )) + })?; + + query_proofs.push(segment_query_proof); + aux_trace_states.push(segment_trace_states); // merge tables for each auxiliary segment into a single table Some(Table::merge(aux_trace_states)) diff --git a/verifier/src/evaluator.rs b/verifier/src/evaluator.rs index d1b60e872..d8dfeccd0 100644 --- a/verifier/src/evaluator.rs +++ b/verifier/src/evaluator.rs @@ -43,7 +43,7 @@ pub fn evaluate_constraints>( let mut t_evaluations1 = E::zeroed_vector(t_constraints.num_main_constraints()); air.evaluate_transition(main_trace_frame, &periodic_values, &mut t_evaluations1); - // evaluate transition constraints for auxiliary trace segments (if any) + // evaluate transition constraints for the auxiliary trace segment (if any) let mut t_evaluations2 = E::zeroed_vector(t_constraints.num_aux_constraints()); if let Some(aux_trace_frame) = aux_trace_frame { air.evaluate_aux_transition( @@ -73,7 +73,7 @@ pub fn evaluate_constraints>( result += group.evaluate_at(main_trace_frame.current(), x); } - // iterate over boundary constraint groups for auxiliary trace segments (each group has a + // iterate over boundary constraint groups for the auxiliary trace segment (each group has a // distinct divisor), evaluate constraints in each group and add their combination to the // result if let Some(aux_trace_frame) = aux_trace_frame { @@ -91,7 +91,7 @@ pub fn evaluate_constraints>( let lagrange_kernel_aux_rand_elements = { let num_rand_elements = air.context().trace_len().ilog2() as usize; - &aux_rand_elements.get_segment_elements(0)[0..num_rand_elements] + &aux_rand_elements.get_segment_elements()[0..num_rand_elements] }; let lagrange_constraints = air diff --git a/verifier/src/lib.rs b/verifier/src/lib.rs index 85bcf71d3..328748095 100644 --- a/verifier/src/lib.rs +++ b/verifier/src/lib.rs @@ -165,14 +165,15 @@ where // reseed the coin with the commitment to the main trace segment public_coin.reseed(trace_commitments[0]); - // process auxiliary trace segments (if any), to build a set of random elements for each segment + // process the auxiliary trace segment (if any), to build a set of random elements let mut aux_trace_rand_elements = AuxTraceRandElements::::new(); - for (i, commitment) in trace_commitments.iter().skip(1).enumerate() { + if trace_commitments.len() > 1 { + let aux_segment_commitment = trace_commitments[1]; let rand_elements = air - .get_aux_trace_segment_random_elements(i, &mut public_coin) + .get_aux_trace_segment_random_elements(&mut public_coin) .map_err(|_| VerifierError::RandomCoinError)?; - aux_trace_rand_elements.add_segment_elements(rand_elements); - public_coin.reseed(*commitment); + aux_trace_rand_elements.set_segment_elements(rand_elements); + public_coin.reseed(aux_segment_commitment); } // build random coefficients for the composition polynomial diff --git a/winterfell/src/tests.rs b/winterfell/src/tests.rs index de5984410..49426f763 100644 --- a/winterfell/src/tests.rs +++ b/winterfell/src/tests.rs @@ -53,7 +53,7 @@ impl LagrangeSimpleTrace { Self { main_trace: ColMatrix::new(vec![col]), - info: TraceInfo::new_multi_segment(1, [2], [3], Self::TRACE_LENGTH, vec![]), + info: TraceInfo::new_multi_segment(1, 2, 3, Self::TRACE_LENGTH, vec![]), } } } @@ -71,12 +71,9 @@ impl Trace for LagrangeSimpleTrace { fn build_aux_segment>( &mut self, - aux_segments: &[ColMatrix], _rand_elements: &[E], lagrange_rand_elements: Option<&[E]>, ) -> Option> { - assert!(aux_segments.is_empty()); - let lagrange_rand_elements = lagrange_rand_elements.unwrap(); let r0 = lagrange_rand_elements[0]; @@ -273,8 +270,8 @@ impl LagrangeComplexTrace { main_trace: ColMatrix::new(vec![main_trace_col]), info: TraceInfo::new_multi_segment( 1, - [aux_segment_width], - [num_aux_segment_rands], + aux_segment_width, + num_aux_segment_rands, trace_len, vec![], ), @@ -301,12 +298,9 @@ impl Trace for LagrangeComplexTrace { /// by the main column fn build_aux_segment>( &mut self, - aux_segments: &[ColMatrix], rand_elements: &[E], lagrange_kernel_rand_elements: Option<&[E]>, ) -> Option> { - assert!(aux_segments.is_empty()); - let mut columns = Vec::new(); // first build the Lagrange kernel column