air AIR return (point, values) instead of a vec of evaluations sharing the same point

Author: TomWambsgans

crates/air/src/prove.rs

@@ -9,7 +9,6 @@ use utils::PF;
 use utils::{FSProver, add_multilinears, multilinears_linear_combination};
 use whir_p3::fiat_shamir::FSChallenger;
 use whir_p3::poly::evals::{eval_eq, fold_multilinear, scale_poly};
-use whir_p3::poly::multilinear::Evaluation;
 use whir_p3::poly::{evals::EvaluationsList, multilinear::MultilinearPoint};
 
 use crate::{NormalAir, PackedAir};
@@ -38,7 +37,7 @@ fn prove_air<
     univariate_skips: usize,
     table: &AirTable<EF, A, AP>,
     witness: &[&'a [WF]],
-) -> Vec<Evaluation<EF>> {
+) -> (MultilinearPoint<EF>, Vec<EF>) {
     let n_rows = witness[0].len();
     assert!(witness.iter().all(|col| col.len() == n_rows));
     let log_n_rows = log2_strict_usize(n_rows);
@@ -119,7 +118,7 @@ impl<EF: ExtensionField<PF<EF>>, A: NormalAir<EF>, AP: PackedAir<EF>> AirTable<E
         prover_state: &mut FSProver<EF, impl FSChallenger<EF>>,
         univariate_skips: usize,
         witness: &[&[PF<EF>]],
-    ) -> Vec<Evaluation<EF>> {
+    ) -> (MultilinearPoint<EF>, Vec<EF>) {
         prove_air::<PF<EF>, EF, A, AP>(prover_state, univariate_skips, self, witness)
     }
 
@@ -129,7 +128,7 @@ impl<EF: ExtensionField<PF<EF>>, A: NormalAir<EF>, AP: PackedAir<EF>> AirTable<E
         prover_state: &mut FSProver<EF, impl FSChallenger<EF>>,
         univariate_skips: usize,
         witness: &[&[EF]],
-    ) -> Vec<Evaluation<EF>> {
+    ) -> (MultilinearPoint<EF>, Vec<EF>) {
         prove_air::<EF, EF, A, AP>(prover_state, univariate_skips, self, witness)
     }
 }
@@ -143,7 +142,7 @@ fn open_unstructured_columns<
     univariate_skips: usize,
     witness: &[&[WF]],
     outer_sumcheck_challenge: &[EF],
-) -> Vec<Evaluation<EF>> {
+) -> (MultilinearPoint<EF>, Vec<EF>) {
     let log_n_rows = log2_strict_usize(witness[0].len());
 
     let columns_batching_scalars = prover_state.sample_vec(log2_ceil_usize(witness.len()));
@@ -170,20 +169,13 @@ fn open_unstructured_columns<
         .concat(),
     );
 
-    let mut evaluations_remaining_to_prove = vec![];
-    assert_eq!(sub_evals.len(), 1 << epsilons.len());
-
-    for col in witness {
-        // TODO compute oe time eq(.) then inner product with everything
-        let value = col.evaluate(&common_point);
-        prover_state.add_extension_scalars(&[value]);
-        evaluations_remaining_to_prove.push(Evaluation {
-            point: common_point.clone(),
-            value,
-        });
-    }
+    let evaluations_remaining_to_prove = witness
+        .iter()
+        .map(|col| col.evaluate(&common_point))
+        .collect::<Vec<_>>();
+    prover_state.add_extension_scalars(&evaluations_remaining_to_prove);
 
-    evaluations_remaining_to_prove
+    (common_point, evaluations_remaining_to_prove)
 }
 
 #[instrument(skip_all)]
@@ -192,7 +184,7 @@ fn open_structured_columns<EF: ExtensionField<PF<EF>> + ExtensionField<IF>, IF:
     univariate_skips: usize,
     witness: &[&[IF]],
     outer_sumcheck_challenge: &[EF],
-) -> Vec<Evaluation<EF>> {
+) -> (MultilinearPoint<EF>, Vec<EF>) {
     let n_columns = witness.len();
     let n_rows = witness[0].len();
     let log_n_rows = log2_strict_usize(n_rows);
@@ -246,17 +238,11 @@ fn open_structured_columns<EF: ExtensionField<PF<EF>> + ExtensionField<IF>, IF:
         None,
     );
 
-    // TODO using inner_evals[1], we can avoid 1 of the evaluations below (the last one)
-
-    let mut evaluations_remaining_to_prove = vec![];
-    for col in witness {
-        let value = col.evaluate(&inner_challenges);
-        prover_state.add_extension_scalar(value);
-        evaluations_remaining_to_prove.push(Evaluation {
-            point: inner_challenges.clone(),
-            value,
-        });
-    }
+    let evaluations_remaining_to_prove = witness
+        .iter()
+        .map(|col| col.evaluate(&inner_challenges))
+        .collect::<Vec<_>>();
+    prover_state.add_extension_scalars(&evaluations_remaining_to_prove);
 
-    evaluations_remaining_to_prove
+    (inner_challenges, evaluations_remaining_to_prove)
 }

crates/air/src/test.rs

@@ -181,20 +181,21 @@ fn test_structured_air() {
         ExampleStructuredAir::<N_COLUMNS, N_PREPROCESSED_COLUMNS>,
     );
     table.check_trace_validity(&columns_ref).unwrap();
-    let evaluations_remaining_to_prove =
+    let (point_prover, evaluations_remaining_to_prove) =
         table.prove_base(&mut prover_state, UNIVARIATE_SKIPS, &columns_ref);
     let mut verifier_state = build_verifier_state(&prover_state);
-    let evaluations_remaining_to_verify = table
+    let (point_verifier, evaluations_remaining_to_verify) = table
         .verify(&mut verifier_state, UNIVARIATE_SKIPS, log_n_rows)
         .unwrap();
+    assert_eq!(point_prover, point_verifier);
     assert_eq!(
         &evaluations_remaining_to_prove,
         &evaluations_remaining_to_verify
     );
     for i in 0..N_COLUMNS {
         assert_eq!(
-            columns[i].evaluate(&evaluations_remaining_to_verify[i].point),
-            evaluations_remaining_to_verify[i].value
+            columns[i].evaluate(&point_prover),
+            evaluations_remaining_to_verify[i]
         );
     }
 }
@@ -214,20 +215,21 @@ fn test_unstructured_air() {
         ExampleUnstructuredAir::<N_COLUMNS, N_PREPROCESSED_COLUMNS>,
     );
     table.check_trace_validity(&columns_ref).unwrap();
-    let evaluations_remaining_to_prove =
+    let (point_prover, evaluations_remaining_to_prove) =
         table.prove_base(&mut prover_state, UNIVARIATE_SKIPS, &columns_ref);
     let mut verifier_state = build_verifier_state(&prover_state);
-    let evaluations_remaining_to_verify = table
+    let (point_verifier, evaluations_remaining_to_verify) = table
         .verify(&mut verifier_state, UNIVARIATE_SKIPS, log_n_rows)
         .unwrap();
+    assert_eq!(point_prover, point_verifier);
     assert_eq!(
         &evaluations_remaining_to_prove,
         &evaluations_remaining_to_verify
     );
     for i in 0..N_COLUMNS {
         assert_eq!(
-            columns[i].evaluate(&evaluations_remaining_to_verify[i].point),
-            evaluations_remaining_to_verify[i].value
+            columns[i].evaluate(&point_prover),
+            evaluations_remaining_to_verify[i]
         );
     }
 }

crates/air/src/verify.rs

@@ -22,7 +22,7 @@ fn verify_air<EF: ExtensionField<PF<EF>>, A: NormalAir<EF>, AP: PackedAir<EF>>(
     table: &AirTable<EF, A, AP>,
     univariate_skips: usize,
     log_n_rows: usize,
-) -> Result<Vec<Evaluation<EF>>, ProofError> {
+) -> Result<(MultilinearPoint<EF>, Vec<EF>), ProofError> {
     let constraints_batching_scalar = verifier_state.sample();
 
     let n_zerocheck_challenges = log_n_rows + 1 - univariate_skips;
@@ -109,7 +109,7 @@ impl<EF: ExtensionField<PF<EF>>, A: NormalAir<EF>, AP: PackedAir<EF>> AirTable<E
         verifier_state: &mut FSVerifier<EF, impl FSChallenger<EF>>,
         univariate_skips: usize,
         log_n_rows: usize,
-    ) -> Result<Vec<Evaluation<EF>>, ProofError> {
+    ) -> Result<(MultilinearPoint<EF>, Vec<EF>), ProofError> {
         verify_air::<EF, A, AP>(verifier_state, self, univariate_skips, log_n_rows)
     }
 }
@@ -121,7 +121,7 @@ fn verify_unstructured_columns<EF: ExtensionField<PF<EF>>>(
     outer_sumcheck_point: &MultilinearPoint<EF>,
     outer_selector_evals: &[EF],
     log_n_rows: usize,
-) -> Result<Vec<Evaluation<EF>>, ProofError> {
+) -> Result<(MultilinearPoint<EF>, Vec<EF>), ProofError> {
     let n_columns = inner_sums.len();
     let columns_batching_scalars = verifier_state.sample_vec(log2_ceil_usize(n_columns));
 
@@ -146,25 +146,18 @@ fn verify_unstructured_columns<EF: ExtensionField<PF<EF>>>(
         .concat(),
     );
 
-    let mut evaluations_remaining_to_verify = vec![];
-    for _ in 0..n_columns {
-        let value = verifier_state.next_extension_scalar()?;
-        evaluations_remaining_to_verify.push(Evaluation {
-            point: common_point.clone(),
-            value,
-        });
-    }
+    let evaluations_remaining_to_verify = verifier_state.next_extension_scalars_vec(n_columns)?;
 
     if sub_evals.evaluate(&epsilons)
         != dot_product(
             eval_eq(&columns_batching_scalars).into_iter(),
-            evaluations_remaining_to_verify.iter().map(|e| e.value),
+            evaluations_remaining_to_verify.iter().copied(),
         )
     {
         return Err(ProofError::InvalidProof);
     }
 
-    Ok(evaluations_remaining_to_verify)
+    Ok((common_point, evaluations_remaining_to_verify))
 }
 
 #[allow(clippy::too_many_arguments)] // TODO
@@ -176,7 +169,7 @@ fn verify_structured_columns<EF: ExtensionField<PF<EF>>>(
     outer_sumcheck_challenge: &Evaluation<EF>,
     outer_selector_evals: &[EF],
     log_n_rows: usize,
-) -> Result<Vec<Evaluation<EF>>, ProofError> {
+) -> Result<(MultilinearPoint<EF>, Vec<EF>), ProofError> {
     let columns_batching_scalars = verifier_state.sample_vec(log2_ceil_usize(n_columns));
     let alpha = verifier_state.sample();
 
@@ -220,22 +213,18 @@ fn verify_structured_columns<EF: ExtensionField<PF<EF>>>(
 
     let final_value = inner_sumcheck_stement.value / (up + alpha * down);
 
-    let mut evaluations_remaining_to_verify = vec![];
-    for _ in 0..n_columns {
-        let value = verifier_state.next_extension_scalar()?;
-        evaluations_remaining_to_verify.push(Evaluation {
-            point: inner_sumcheck_stement.point.clone(),
-            value,
-        });
-    }
+    let evaluations_remaining_to_verify = verifier_state.next_extension_scalars_vec(n_columns)?;
 
     if final_value
         != dot_product(
             eval_eq(&columns_batching_scalars).into_iter(),
-            evaluations_remaining_to_verify.iter().map(|e| e.value),
+            evaluations_remaining_to_verify.iter().copied(),
         )
     {
         return Err(ProofError::InvalidProof);
     }
-    Ok(evaluations_remaining_to_verify)
+    Ok((
+        inner_sumcheck_stement.point.clone(),
+        evaluations_remaining_to_verify,
+    ))
 }