Add comments and fix some tiny bugs

Cargo.toml

@@ -9,7 +9,7 @@ members = [
   "mpcs",
   "multilinear_extensions",
   "poseidon",
-  "protocol",
+  "gkr_iop",
   "subprotocols",
   "sumcheck",
   "transcript",

protocol/Cargo.toml → gkr_iop/Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 categories.workspace = true
-description = "Protocol"
+description = "GKR IOP protocol implementation"
 edition.workspace = true
 keywords.workspace = true
 license.workspace = true
-name = "protocol"
+name = "gkr_iop"
 readme.workspace = true
 repository.workspace = true
 version.workspace = true

protocol/examples/multi_layer_logup.rs → gkr_iop/examples/multi_layer_logup.rs

@@ -1,9 +1,7 @@
 use std::{marker::PhantomData, sync::Arc};
 
 use ff_ext::ExtensionField;
-use goldilocks::GoldilocksExt2;
-use itertools::{Itertools, izip};
-use protocol::{
+use gkr_iop::{
     ProtocolBuilder, ProtocolWitnessGenerator,
     chip::Chip,
     evaluation::{EvalExpression, PointAndEval},
@@ -13,6 +11,8 @@ use protocol::{
         mock::MockProver,
     },
 };
+use goldilocks::GoldilocksExt2;
+use itertools::{Itertools, izip};
 use rand::{Rng, rngs::OsRng};
 use subprotocols::expression::{Constant, Expression, VectorType};
 use transcript::{BasicTranscript, Transcript};

gkr_iop/src/chip.rs

@@ -0,0 +1,26 @@
+use crate::{evaluation::EvalExpression, gkr::layer::Layer};
+
+pub mod builder;
+pub mod protocol;
+
+#[derive(Clone, Debug, Default)]
+pub struct Chip {
+    /// The number of base inputs committed in the whole protocol.
+    pub n_committed_bases: usize,
+    /// The number of ext inputs committed in the whole protocol.
+    pub n_committed_exts: usize,
+
+    /// The number of challenges generated through the whole protocols
+    /// (except the ones inside sumcheck protocols).
+    pub n_challenges: usize,
+    /// All input evaluations generated at the end of layer protocols will be stored
+    /// in a vector and this is the length.
+    pub n_evaluations: usize,
+    /// The layers of the GKR circuit, in the reverse order.
+    pub layers: Vec<Layer>,
+
+    /// The polynomial index and evaluation expressions of the base inputs.
+    pub base_openings: Vec<(usize, EvalExpression)>,
+    /// The polynomial index and evaluation expressions of the ext inputs.
+    pub ext_openings: Vec<(usize, EvalExpression)>,
+}

protocol/src/chip/builder.rs → gkr_iop/src/chip/builder.rs

@@ -10,16 +10,22 @@ use crate::{
 use super::Chip;
 
 impl Chip {
+    /// Allocate indices for committing base field polynomials.
     pub fn allocate_committed_base<const N: usize>(&mut self) -> [usize; N] {
         self.n_committed_bases += N;
         array::from_fn(|i| i + self.n_committed_bases - N)
     }
 
+    /// Allocate indices for committing extension field polynomials.
     pub fn allocate_committed_ext<const N: usize>(&mut self) -> [usize; N] {
         self.n_committed_exts += N;
         array::from_fn(|i| i + self.n_committed_exts - N)
     }
 
+    /// Allocate `Witness` and `EvalExpression` for the input polynomials in a layer.
+    /// Where `Witness` denotes the index and `EvalExpression` denotes the position
+    /// to place the evaluation of the polynomial after processing the layer prover
+    /// for each polynomial.
     #[allow(clippy::type_complexity)]
     pub fn allocate_wits_in_layer<const M: usize, const N: usize>(
         &mut self,
@@ -44,24 +50,31 @@ impl Chip {
         (bases, exts)
     }
 
+    /// Generate the evaluation expression for each output.
     pub fn allocate_output_evals<const N: usize>(&mut self) -> [EvalExpression; N] {
         self.n_evaluations += N;
         array::from_fn(|i| EvalExpression::Single(i + self.n_evaluations - N))
     }
 
+    /// Allocate challenges.
     pub fn allocate_challenges<const N: usize>(&mut self) -> [Constant; N] {
         self.n_challenges += N;
         array::from_fn(|i| Constant::Challenge(i + self.n_challenges - N))
     }
 
+    /// Allocate a PCS opening action to a base polynomial with index `wit_index`.
+    /// The `EvalExpression` represents the expression to compute the evaluation.
     pub fn allocate_base_opening(&mut self, wit_index: usize, eval: EvalExpression) {
         self.base_openings.push((wit_index, eval));
     }
 
+    /// Allocate a PCS opening action to an ext polynomial with index `wit_index`.
+    /// The `EvalExpression` represents the expression to compute the evaluation.
     pub fn allocate_ext_opening(&mut self, wit_index: usize, eval: EvalExpression) {
         self.ext_openings.push((wit_index, eval));
     }
 
+    /// Add a layer to the circuit.
     pub fn add_layer(&mut self, layer: Layer) {
         assert_eq!(layer.outs.len(), layer.exprs.len());
         match layer.ty {

protocol/src/chip/protocol.rs → gkr_iop/src/chip/protocol.rs

@@ -3,6 +3,7 @@ use crate::gkr::GKRCircuit;
 use super::Chip;
 
 impl Chip {
+    /// Extract information for the GKR protocol.
     pub fn gkr_circuit(&'_ self) -> GKRCircuit<'_> {
         GKRCircuit {
             layers: &self.layers,

protocol/src/evaluation.rs → gkr_iop/src/evaluation.rs

@@ -5,10 +5,18 @@ use itertools::{Itertools, izip};
 use multilinear_extensions::virtual_poly::build_eq_x_r_vec_sequential;
 use subprotocols::expression::{Constant, Point};
 
+/// Evaluation expression for the gkr layer reduction and PCS opening preparation.
 #[derive(Clone, Debug)]
 pub enum EvalExpression {
+    /// Single entry in the evaluation vector.
     Single(usize),
+    /// Linear expression of an entry with the scalar and offset.
     Linear(usize, Constant, Constant),
+    /// Merging multiple evaluations which denotes a partition of the original
+    /// polynomial. `(usize, Constant)` denote the modification of the point.
+    /// For example, when it receive a point `(p0, p1, p2, p3)` from a succeeding
+    /// layer, `vec![(2, c0), (4, c1)]` will modify the point to `(p0, p1, c0, p2, c1, p3)`.
+    /// where the indices specify how the partition applied to the original polynomial.
     Partition(Vec<Box<EvalExpression>>, Vec<(usize, Constant)>),
 }
 

protocol/src/gkr/layer.rs → gkr_iop/src/gkr/layer.rs

@@ -31,10 +31,21 @@ pub enum LayerType {
 pub struct Layer {
     pub name: String,
     pub ty: LayerType,
+    /// Challenges generated at the beginning of the layer protocol.
     pub challenges: Vec<Constant>,
+    /// Expressions to prove in this layer. For zerocheck and linear layers, each
+    /// expression corresponds to an output. While in sumcheck, there is only 1
+    /// expression, which corresponds to the sum of all outputs. This design is
+    /// for the convenience when building the following expression:
+    ///     `e_0 + beta * e_1 = sum_x (eq(p_0, x) + beta * eq(p_1, x)) expr(x)`.
+    /// where `vec![e_0, beta * e_1]` will be the output evaluation expressions.
     pub exprs: Vec<Expression>,
+    /// Positions to place the evaluations of the base inputs of this layer.
     pub in_bases: Vec<EvalExpression>,
+    /// Positions to place the evaluations of the ext inputs of this layer.
     pub in_exts: Vec<EvalExpression>,
+    /// The expressions of the evaluations from the succeeding layers, which are
+    /// connected to the outputs of this layer.
     pub outs: Vec<EvalExpression>,
 }
 
@@ -128,17 +139,14 @@ impl Layer {
             base_mle_evals,
             ext_mle_evals,
         } = match self.ty {
-            LayerType::Sumcheck => {
-                assert_eq!(sigmas.len(), 1);
-                <Layer as SumcheckLayer<E>>::verify(
-                    self,
-                    proof,
-                    &sigmas.iter().sum(),
-                    points.slice_vector(),
-                    challenges,
-                    transcript,
-                )?
-            }
+            LayerType::Sumcheck => <Layer as SumcheckLayer<E>>::verify(
+                self,
+                proof,
+                &sigmas.iter().sum(),
+                points.slice_vector(),
+                challenges,
+                transcript,
+            )?,
             LayerType::Zerocheck => <Layer as ZerocheckLayer<E>>::verify(
                 self,
                 proof,

protocol/src/gkr/mock.rs → gkr_iop/src/gkr/mock.rs

@@ -41,7 +41,7 @@ impl<E: ExtensionField> MockProver<E> {
         mut challenges: Vec<E>,
     ) -> Result<(), MockProverError<E>> {
         evaluations.resize(circuit.n_evaluations, VectorType::Base(vec![]));
-        challenges.resize(circuit.n_challenges, E::ONE + E::ONE);
+        challenges.resize_with(circuit.n_challenges, || E::random(OsRng));
         for (layer, layer_wit) in izip!(circuit.layers, &circuit_wit.layers) {
             let num_vars = layer_wit.num_vars;
             let points = (0..layer.outs.len())

protocol/src/lib.rs → gkr_iop/src/lib.rs

@@ -14,6 +14,7 @@ pub trait ProtocolBuilder: Sized {
 
     fn init(params: Self::Params) -> Self;
 
+    /// Build the protocol for GKR IOP.
     fn build(params: Self::Params) -> (Self, Chip) {
         let mut chip_spec = Self::init(params);
         let mut chip = Chip::default();
@@ -24,8 +25,15 @@ pub trait ProtocolBuilder: Sized {
         (chip_spec, chip)
     }
 
+    /// Specify the polynomials and challenges to be committed and generated in
+    /// Phase 1.
     fn build_commit_phase1(&mut self, spec: &mut Chip);
+    /// Specify the polynomials and challenges to be committed and generated in
+    /// Phase 2.
     fn build_commit_phase2(&mut self, _spec: &mut Chip) {}
+    /// Create the GKR layers in the reverse order. For each layer, specify the
+    /// polynomial expressions, evaluation expressions of outputs and evaluation
+    /// positions of the inputs.
     fn build_gkr_phase(&mut self, spec: &mut Chip);
 }