add documentation

Author: nicole-graus

crates/provers/gkr/README.md

@@ -4,6 +4,8 @@ An implementation of the Goldwasser-Kalai-Rothblum (GKR) Non-Interactive Protoco
 
 To help with the understanding of this implementation, we recommend reading our [blog post](https://blog.lambdaclass.com/gkr-protocol-a-step-by-step-example/).
 
+**Warning:** This GKR implementation is for educational purposes and should not be used in production. It uses the Fiat-Shamir transform, which is vulnerable to practical attacks in this context (see ["How to Prove False Statments"](https://eprint.iacr.org/2025/118.pdf)). 
+
 ## Overview
 
 The GKR Protocol allows a Prover to convince a Verifier that he correctly evaluated an arithmetic circuit on a given input without the Verifier having to perform the entire computation.   
@@ -40,13 +42,13 @@ Each layer must have a power-of-two number of gates.
 ### Main Functions
 
 - `gkr_prove(circuit, input)` - Generate a GKR proof for a circuit evaluation.
-- `gkr_verify(proof, circuit, input)` - Verify a GKR proof with complete input checking.
+- `gkr_verify(proof, circuit, input)` - Verify a GKR proof.
 
 ### Circuit Construction
 
-- `Circuit::new(layers, num_inputs)` - Create a new circuit with specified layers and input count
-- `CircuitLayer::new(gates)` - Create a layer with the given gates
-- `Gate::new(gate_type, inputs)` - Create a gate with type (Add/Mul) and input connections
+- `Circuit::new(layers, num_inputs)` - Create a new circuit with specified layers, gates and number of inputs.
+- `CircuitLayer::new(gates)` - Create a layer with the given gates.
+- `Gate::new(gate_type, inputs_idx)` - Create a gate with type (Add/Mul) and certain input indeces.
 
 ## Example
 
@@ -57,29 +59,29 @@ use lambdaworks_math::field::fields::u64_prime_field::U64PrimeField;
 use lambdaworks_math::field::element::FieldElement;
 use lambdaworks_gkr::{gkr_prove, gkr_verify, circuit_from_lambda};
 
-// Define the field (We use modulus 23 as in the example of our blog posr)
+// Define the field (We use modulus 23 as in the example of our blog post).
 const MODULUS23: u64 = 23;
 type F23 = U64PrimeField<MODULUS23>;
 type F23E = FieldElement<F23>;
 
-// Create the circuit from our blog post.
-// This creates a 2-layer circuit.
+// Create the circuit of our blog post.
+// This creates a 2-layer circuit plus the input layer.
 let circuit = lambda_post_circuit.unwrap();
 
-// Define input values (from the post example)
+// Define input values (from the post example).
 let input = [F23E::from(3), F23E::from(1)];
 
-// Generate proof
+// Generate proof.
 let proof_result = gkr_prove(&circuit, &input);
 assert!(proof_result.is_ok());
 let proof = proof_result.unwrap();
 
-// Verify proof
+// Verify proof.
 let verification_result = gkr_verify(&proof, &circuit, &input);
 assert!(verification_result.is_ok() && verification_result.unwrap());
 println!("GKR verification successful!");
 
-// You can also check the actual output
+// You can also check the actual output.
 let evaluation = circuit.evaluate(&input);
 println!("Circuit output: {:?}", evaluation.layers[0]);
 ```
@@ -109,10 +111,10 @@ let custom_circuit = Circuit::new(
 
 // Test with inputs [2, 3, 4, 5]
 let input = [F23E::from(2), F23E::from(3), F23E::from(4), F23E::from(5)];
-// This computes: Layer 1: [2*3, 4*5] = [6, 20], Layer 0: [6+20] = [26 mod 23] = [3]
 
 let proof = gkr_prove(&custom_circuit, &input).unwrap();
 let is_valid = gkr_verify(&proof, &custom_circuit, &input).unwrap();
+
 assert!(is_valid);
 ```
 

crates/provers/gkr/src/circuit.rs

@@ -56,19 +56,6 @@ impl CircuitLayer {
     }
 }
 
-/// An evaluation of a `Circuit` on some input.
-pub struct CircuitEvaluation<F> {
-    /// Evaluations on per-layer basis.
-    pub layers: Vec<Vec<F>>,
-}
-
-impl<F: Copy> CircuitEvaluation<F> {
-    /// Takes a gate label and outputs the corresponding gate's value at layer `layer`.
-    pub fn w(&self, layer: usize, label: usize) -> F {
-        self.layers[layer][label]
-    }
-}
-
 #[derive(Debug, Clone)]
 pub enum CircuitError {
     InputsNotPowerOfTwo,
@@ -99,15 +86,20 @@ pub enum CircuitError {
 /// - The circuit is evaluated from inputs upward, but layers are stored from output to input.
 #[derive(Clone)]
 pub struct Circuit {
-    /// First layer being the output layer, last layer being
-    /// the input layer.
+    /// First layer is the output layer. It doesn't include the input layer.
     layers: Vec<CircuitLayer>,
 
     /// Number of inputs
     num_inputs: usize,
     input_num_vars: usize, // log2 of number of inputs
 }
 
+/// An evaluation of a `Circuit` on some input.
+pub struct CircuitEvaluation<F> {
+    /// Evaluations on per-layer basis. First layer is the output and last layer is the input.
+    pub layers: Vec<Vec<F>>,
+}
+
 impl Circuit {
     pub fn new(layers: Vec<CircuitLayer>, num_inputs: usize) -> Result<Self, CircuitError> {
         if layers.is_empty() {
@@ -196,22 +188,11 @@ impl Circuit {
             current_input = temp_layer;
         }
 
+        // Reverse the order so that the first layer is the output layer.
         layers.reverse();
         CircuitEvaluation { layers }
     }
 
-    /// The $\text{add}_i(a, b, c)$ polynomial value at layer $i$.
-    pub fn add_i(&self, i: usize, a: usize, b: usize, c: usize) -> bool {
-        let gate = &self.layers[i].gates[a];
-        gate.gate_type == GateType::Add && gate.inputs_idx[0] == b && gate.inputs_idx[1] == c
-    }
-
-    /// The $\text{mul}_i(a, b, c)$ polynomial value at layer $i$.
-    pub fn mul_i(&self, i: usize, a: usize, b: usize, c: usize) -> bool {
-        let gate = &self.layers[i].gates[a];
-        gate.gate_type == GateType::Mul && gate.inputs_idx[0] == b && gate.inputs_idx[1] == c
-    }
-
     pub fn layers(&self) -> &[CircuitLayer] {
         &self.layers
     }
@@ -224,6 +205,7 @@ impl Circuit {
         self.num_inputs
     }
 
+    /// The multilinear polynomial extension of the function `add_i(a, b, c)`, where `a` is fixed at `r_i`.
     pub fn add_i_ext<F: IsField>(
         &self,
         r_i: &[FieldElement<F>],
@@ -240,6 +222,8 @@ impl Circuit {
         };
         let total_vars = num_vars_current + 2 * num_vars_next;
         let mut add_i_evals = vec![FieldElement::zero(); 1 << total_vars];
+
+        // For each Add gate, we set the corresponding index `a || b || c` in the evaluation vector to one.
         for (a, gate) in self.layers[i].gates.iter().enumerate() {
             if gate.gate_type == GateType::Add {
                 let b = gate.inputs_idx[0];
@@ -248,13 +232,15 @@ impl Circuit {
                 add_i_evals[idx] = FieldElement::one();
             }
         }
-        let mut p = DenseMultilinearPolynomial::new(add_i_evals);
+
+        let mut add_i_poly = DenseMultilinearPolynomial::new(add_i_evals);
         for val in r_i.iter() {
-            p = p.fix_first_variable(val);
+            add_i_poly = add_i_poly.fix_first_variable(val);
         }
-        p
+        add_i_poly
     }
 
+    /// The multilinear polynomial extension of the function `mul_i(a, b, c)`, where `a` is fixed at `r_i`.
     pub fn mul_i_ext<F: IsField>(
         &self,
         r_i: &[FieldElement<F>],
@@ -271,6 +257,8 @@ impl Circuit {
         };
         let total_vars = num_vars_current + 2 * num_vars_next;
         let mut mul_i_evals = vec![FieldElement::zero(); 1 << total_vars];
+
+        // For each Mul gate, we set the corresponding index `a || b || c` in the evaluation vector to one.
         for (a, gate) in self.layers[i].gates.iter().enumerate() {
             if gate.gate_type == GateType::Mul {
                 let b = gate.inputs_idx[0];
@@ -279,10 +267,11 @@ impl Circuit {
                 mul_i_evals[idx] = FieldElement::one();
             }
         }
-        let mut p = DenseMultilinearPolynomial::new(mul_i_evals);
+
+        let mut mul_i_poly = DenseMultilinearPolynomial::new(mul_i_evals);
         for val in r_i.iter() {
-            p = p.fix_first_variable(val);
+            mul_i_poly = mul_i_poly.fix_first_variable(val);
         }
-        p
+        mul_i_poly
     }
 }

crates/provers/gkr/src/lib.rs

@@ -125,7 +125,7 @@ mod tests {
             2,
         )
     }
-    /// Create a circuit with four layers (without counting the inputs).
+    /// Create a circuit with four layers (without counting the input layer).
     /// To picture this circuit, imagine a tree structure where each layer has twice the number of gates as the layer above.
     pub fn four_layer_circuit() -> Result<Circuit, CircuitError> {
         use crate::circuit::{CircuitLayer, Gate, GateType};

crates/provers/gkr/src/sumcheck.rs

@@ -11,13 +11,17 @@ use lambdaworks_sumcheck::{Channel, Prover};
 
 use crate::prover::ProverError;
 
+/// GKR-specific sumcheck proof, which contains each round polynomial `g_j` and the challenges used.
 #[derive(Debug, Clone)]
 pub struct GKRSumcheckProof<F: IsField> {
     pub round_polynomials: Vec<Polynomial<FieldElement<F>>>,
     pub challenges: Vec<FieldElement<F>>,
 }
 
-/// GKR-specific sumcheck prover
+/// GKR-specific sumcheck prover.
+/// This function will recieve a vector of two terms. Each term contains two multilinear polynomials.
+/// This separation of terms is necessary because the classic/original sumcheck only accepts a product of multilinear polynomials.
+/// In this way, we apply the sumcheck to two products, each consisting of two factors.
 pub fn gkr_sumcheck_prove<F, T>(
     terms: Vec<Vec<DenseMultilinearPolynomial<F>>>,
     transcript: &mut T,
@@ -35,9 +39,11 @@ where
     let factors_term_1 = terms[0].clone();
     let factors_term_2 = terms[1].clone();
 
+    // Create two separate sumcheck provers for each term.
     let mut prover_term_1 = Prover::new(factors_term_1).map_err(|_| ProverError::SumcheckError)?;
     let mut prover_term_2 = Prover::new(factors_term_2).map_err(|_| ProverError::SumcheckError)?;
 
+    // Both terms have the same number of variables.
     let num_vars = prover_term_1.num_vars();
 
     let claimed_sum_term_1 = prover_term_1
@@ -58,7 +64,7 @@ where
     let mut challenges = Vec::new();
     let mut current_challenge: Option<FieldElement<F>> = None;
 
-    // Execute rounds
+    // Execute sumcheck rounds
     for j in 0..num_vars {
         let g_j_term_1 = prover_term_1
             .round(current_challenge.as_ref())
@@ -101,6 +107,7 @@ where
     Ok(sumcheck_proof)
 }
 
+/// GKR-specific sumcheck Verifier.
 pub fn gkr_sumcheck_verify<F, T>(
     claimed_sum: FieldElement<F>,
     sumcheck_proof: &GKRSumcheckProof<F>,
@@ -126,12 +133,11 @@ where
 
     let mut challenges = Vec::new();
 
-    // Process each round polynomial
+    // Verify each round polynomial.
     for (j, g_j) in proof_polys.iter().enumerate() {
         // Add polynomial info to transcript
         let round_label = format!("round_{}_poly", j);
         transcript.append_bytes(round_label.as_bytes());
-
         let coeffs = g_j.coefficients();
         transcript.append_bytes(&(coeffs.len() as u64).to_be_bytes());
         if coeffs.is_empty() {
@@ -142,11 +148,12 @@ where
             }
         }
 
+        // Verify `g_j(0) + g_j(1) = m_{j-1}`, where:
+        // `m_{j-1} = g_{j-1} (s_{j-1})`, the previous claimed sum.
         let g_j_0 = g_j.evaluate::<F>(&FieldElement::zero());
         let g_j_1 = g_j.evaluate::<F>(&FieldElement::one());
         let sum_evals = &g_j_0 + &g_j_1;
 
-        // We need to manually track the expected sum
         let expected_sum = if j == 0 {
             claimed_sum.clone()
         } else {