Generalize auxiliary trace building

.github/workflows/ci.yml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   light-checks:
-    name: Ligth checks
+    name: Light checks
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@main

README.md

@@ -290,7 +290,7 @@ impl Prover for WorkProver {
     fn new_evaluator<'a, E: FieldElement<BaseField = BaseElement>>(
         &self,
         air: &'a WorkAir,
-        aux_rand_elements: winterfell::AuxTraceRandElements<E>,
+        aux_rand_elements: Option<Self::AuxRandElements<E>>,
         composition_coefficients: winterfell::ConstraintCompositionCoefficients<E>,
     ) -> Self::ConstraintEvaluator<'a, E> {
         DefaultConstraintEvaluator::new(air, aux_rand_elements, composition_coefficients)
@@ -307,10 +307,10 @@ Now, we are finally ready to generate a STARK proof. The function below, will ex
 ```Rust
 use winterfell::{
     math::{fields::f128::BaseElement, FieldElement},
-    FieldExtension, HashFunction, ProofOptions, StarkProof,
+    FieldExtension, HashFunction, ProofOptions, Proof,
 };
 
-pub fn prove_work() -> (BaseElement, StarkProof) {
+pub fn prove_work() -> (BaseElement, Proof) {
     // We'll just hard-code the parameters here for this example.
     let start = BaseElement::new(3);
     let n = 1_048_576;
@@ -343,12 +343,12 @@ We can then give this proof (together with the public inputs) to anyone, and the
 use winterfell::{
     crypto::{hashers::Blake3_256, DefaultRandomCoin},
     math::fields::f128::BaseElement,
-    verify, AcceptableOptions, StarkProof,
+    verify, AcceptableOptions, Proof,
 };
 
 type Blake3 = Blake3_256<BaseElement>;
 
-pub fn verify_work(start: BaseElement, result: BaseElement, proof: StarkProof) {
+pub fn verify_work(start: BaseElement, result: BaseElement, proof: Proof) {
     // The verifier will accept proofs with parameters which guarantee 95 bits or more of
     // conjectured security
     let min_opts = AcceptableOptions::MinConjecturedSecurity(95);

air/README.md

@@ -81,7 +81,7 @@ To describe Randomized AIR, you will need to do the following when implementing
 4. Blowup factor - higher values increase proof security, but also increase proof generation time and proof size. However, higher blowup factors require fewer queries for the same security level. Thus, it is frequently possible to increase blowup factor and at the same time decrease the number of queries in such a way that the proofs become smaller.
 5. Grinding factor - higher values increase proof security, but also may increase proof generation time.
 
-See [options.rs](src/options.rs) for more info on currently available options and their meaning. Additionally, security level of a proof can be estimated using `StarkProof::security_level()` function.
+See [options.rs](src/options.rs) for more info on currently available options and their meaning. Additionally, security level of a proof can be estimated using `Proof::security_level()` function.
 
 ## Crate features
 This crate can be compiled with the following features:

air/src/air/aux.rs

@@ -0,0 +1,90 @@
+use alloc::{string::ToString, vec::Vec};
+use crypto::{ElementHasher, RandomCoin, RandomCoinError};
+use math::FieldElement;
+use utils::Deserializable;
+
+use super::lagrange::LagrangeKernelRandElements;
+
+/// Holds the randomly generated elements necessary to build the auxiliary trace.
+///
+/// Specifically, [`AuxRandElements`] currently supports 2 types of random elements:
+/// - the ones needed to build the Lagrange kernel column (when using GKR to accelerate LogUp),
+/// - the ones needed to build all the other auxiliary columns
+#[derive(Debug, Clone)]
+pub struct AuxRandElements<E> {
+    rand_elements: Vec<E>,
+    lagrange: Option<LagrangeKernelRandElements<E>>,
+}
+
+impl<E> AuxRandElements<E> {
+    /// Creates a new [`AuxRandElements`], where the auxiliary trace doesn't contain a Lagrange
+    /// kernel column.
+    pub fn new(rand_elements: Vec<E>) -> Self {
+        Self {
+            rand_elements,
+            lagrange: None,
+        }
+    }
+
+    /// Creates a new [`AuxRandElements`], where the auxiliary trace contains a Lagrange kernel
+    /// column.
+    pub fn new_with_lagrange(
+        rand_elements: Vec<E>,
+        lagrange: Option<LagrangeKernelRandElements<E>>,
+    ) -> Self {
+        Self {
+            rand_elements,
+            lagrange,
+        }
+    }
+
+    /// Returns the random elements needed to build all columns other than the Lagrange kernel one.
+    pub fn rand_elements(&self) -> &[E] {
+        &self.rand_elements
+    }
+
+    /// Returns the random elements needed to build the Lagrange kernel column.
+    pub fn lagrange(&self) -> Option<&LagrangeKernelRandElements<E>> {
+        self.lagrange.as_ref()
+    }
+}
+
+/// A trait for verifying a GKR proof.
+///
+/// Specifically, the use case in mind is proving the constraints of a LogUp bus using GKR, as
+/// described in [Improving logarithmic derivative lookups using
+/// GKR](https://eprint.iacr.org/2023/1284.pdf).
+pub trait GkrVerifier {
+    /// The GKR proof.
+    type GkrProof: Deserializable;
+    /// The error that can occur during GKR proof verification.
+    type Error: ToString;
+
+    /// Verifies the GKR proof, and returns the random elements that were used in building
+    /// the Lagrange kernel auxiliary column.
+    fn verify<E, Hasher>(
+        &self,
+        gkr_proof: Self::GkrProof,
+        public_coin: &mut impl RandomCoin<BaseField = E::BaseField, Hasher = Hasher>,
+    ) -> Result<LagrangeKernelRandElements<E>, Self::Error>
+    where
+        E: FieldElement,
+        Hasher: ElementHasher<BaseField = E::BaseField>;
+}
+
+impl GkrVerifier for () {
+    type GkrProof = ();
+    type Error = RandomCoinError;
+
+    fn verify<E, Hasher>(
+        &self,
+        _gkr_proof: Self::GkrProof,
+        _public_coin: &mut impl RandomCoin<BaseField = E::BaseField, Hasher = Hasher>,
+    ) -> Result<LagrangeKernelRandElements<E>, Self::Error>
+    where
+        E: FieldElement,
+        Hasher: ElementHasher<BaseField = E::BaseField>,
+    {
+        Ok(LagrangeKernelRandElements::new(Vec::new()))
+    }
+}

air/src/air/coefficients.rs

@@ -6,42 +6,6 @@
 use alloc::vec::Vec;
 use math::FieldElement;
 
-// AUXILIARY TRACE SEGMENT RANDOMNESS
-// ================================================================================================
-
-/// 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<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 the auxiliary segment.
-    pub fn get_segment_elements(&self) -> &[E] {
-        &self.0
-    }
-
-    /// Sets the random elements associated with the auxiliary segment.
-    pub fn set_segment_elements(&mut self, rand_elements: Vec<E>) {
-        self.0 = rand_elements;
-    }
-}
-
-impl<E: FieldElement> Default for AuxTraceRandElements<E> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 // CONSTRAINT COMPOSITION COEFFICIENTS
 // ================================================================================================
 /// Coefficients used in construction of constraint composition polynomial.

air/src/air/lagrange/boundary.rs

@@ -5,7 +5,7 @@
 
 use math::FieldElement;
 
-use crate::LagrangeKernelEvaluationFrame;
+use crate::{LagrangeKernelEvaluationFrame, LagrangeKernelRandElements};
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct LagrangeKernelBoundaryConstraint<E>
@@ -21,7 +21,10 @@ where
     E: FieldElement,
 {
     /// Creates a new Lagrange kernel boundary constraint.
-    pub fn new(composition_coefficient: E, lagrange_kernel_rand_elements: &[E]) -> Self {
+    pub fn new(
+        composition_coefficient: E,
+        lagrange_kernel_rand_elements: &LagrangeKernelRandElements<E>,
+    ) -> Self {
         Self {
             assertion_value: Self::assertion_value(lagrange_kernel_rand_elements),
             composition_coefficient,
@@ -57,9 +60,9 @@ where
     }
 
     /// Computes the assertion value given the provided random elements.
-    pub fn assertion_value(lagrange_kernel_rand_elements: &[E]) -> E {
+    pub fn assertion_value(lagrange_kernel_rand_elements: &LagrangeKernelRandElements<E>) -> E {
         let mut assertion_value = E::ONE;
-        for &rand_ele in lagrange_kernel_rand_elements {
+        for &rand_ele in lagrange_kernel_rand_elements.as_ref() {
             assertion_value *= E::ONE - rand_ele;
         }
 

air/src/air/lagrange/mod.rs

@@ -4,6 +4,9 @@
 // LICENSE file in the root directory of this source tree.
 
 mod boundary;
+use core::ops::Deref;
+
+use alloc::vec::Vec;
 pub use boundary::LagrangeKernelBoundaryConstraint;
 
 mod frame;
@@ -26,7 +29,7 @@ impl<E: FieldElement> LagrangeKernelConstraints<E> {
     /// Constructs a new [`LagrangeKernelConstraints`].
     pub fn new(
         lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients<E>,
-        lagrange_kernel_rand_elements: &[E],
+        lagrange_kernel_rand_elements: &LagrangeKernelRandElements<E>,
         lagrange_kernel_col_idx: usize,
     ) -> Self {
         Self {
@@ -41,3 +44,36 @@ impl<E: FieldElement> LagrangeKernelConstraints<E> {
         }
     }
 }
+
+/// Holds the randomly generated elements needed to build the Lagrange kernel auxiliary column.
+#[derive(Debug, Clone)]
+pub struct LagrangeKernelRandElements<E> {
+    elements: Vec<E>,
+}
+
+impl<E> LagrangeKernelRandElements<E> {
+    /// Creates a new [`LagrangeKernelRandElements`].
+    pub fn new(elements: Vec<E>) -> Self {
+        Self { elements }
+    }
+}
+
+impl<E> Deref for LagrangeKernelRandElements<E> {
+    type Target = Vec<E>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.elements
+    }
+}
+
+impl<E> AsRef<[E]> for LagrangeKernelRandElements<E> {
+    fn as_ref(&self) -> &[E] {
+        &self.elements
+    }
+}
+
+impl<E> From<LagrangeKernelRandElements<E>> for Vec<E> {
+    fn from(lagrange_rand_elements: LagrangeKernelRandElements<E>) -> Self {
+        lagrange_rand_elements.elements
+    }
+}