abstractions and refactoring, extract out tensor-iop-based-pcs trait

tonyfloatersu · tonyfloatersu · commit 19fbd6cb9c5a · 2024-11-16T23:47:24.000-05:00
diff --git a/pcs/src/orion.rs b/pcs/src/orion.rs
@@ -10,7 +10,7 @@ use transcript::Transcript;
 use crate::PolynomialCommitmentScheme;
 
 mod utils;
-pub use utils::{OrionPCSError, OrionResult};
+pub use utils::{OrionPCSError, OrionResult, SubsetSumLUTs, TensorIOPPCS};
 
 mod linear_code;
 pub use linear_code::{OrionCodeParameter, ORION_CODE_PARAMETER_INSTANCE};
diff --git a/pcs/src/orion/pcs_impl.rs b/pcs/src/orion/pcs_impl.rs
@@ -8,7 +8,7 @@ use crate::{orion::utils::SubsetSumLUTs, PCS_SOUNDNESS_BITS};
 
 use super::{
     linear_code::{OrionCode, OrionCodeParameter},
-    utils::{transpose_in_place, OrionPCSError, OrionResult},
+    utils::{transpose_in_place, OrionPCSError, OrionResult, TensorIOPPCS},
 };
 
 /**********************************************************
@@ -21,6 +21,16 @@ pub struct OrionPublicParams {
     pub code_instance: OrionCode,
 }
 
+impl TensorIOPPCS for OrionPublicParams {
+    fn codeword_len(&self) -> usize {
+        self.code_instance.code_len()
+    }
+
+    fn hamming_weight(&self) -> f64 {
+        self.code_instance.hamming_weight()
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct OrionCommitmentWithData<F, ComPackF>
 where
@@ -53,17 +63,6 @@ pub struct OrionProof<EvalF: Field + FieldSerde> {
 }
 
 impl OrionPublicParams {
-    pub(crate) fn row_col_from_variables<F: Field>(num_variables: usize) -> (usize, usize) {
-        let poly_variables: usize = num_variables;
-
-        let elems_for_smallest_tree = tree::leaf_adic::<F>() * 2;
-
-        let row_num: usize = elems_for_smallest_tree;
-        let msg_size: usize = (1 << poly_variables) / row_num;
-
-        (row_num, msg_size)
-    }
-
     pub fn new<F: Field>(num_variables: usize, code_instance: OrionCode) -> OrionResult<Self> {
         let (_, msg_size) = Self::row_col_from_variables::<F>(num_variables);
         if msg_size != code_instance.msg_len() {
@@ -91,28 +90,6 @@ impl OrionPublicParams {
         }
     }
 
-    pub fn code_len(&self) -> usize {
-        self.code_instance.code_len()
-    }
-
-    pub fn query_complexity(&self, soundness_bits: usize) -> usize {
-        // NOTE: use Ligero (AHIV22) or Avg-case dist to a code (BKS18)
-        // version of avg case dist in unique decoding technique.
-        let avg_case_dist = self.code_instance.hamming_weight() / 3f64;
-        let sec_bits = -(1f64 - avg_case_dist).log2();
-
-        (soundness_bits as f64 / sec_bits).ceil() as usize
-    }
-
-    pub fn proximity_repetition_num(&self, soundness_bits: usize, field_size_bits: usize) -> usize {
-        // NOTE: use Ligero (AHIV22) or Avg-case dist to a code (BKS18)
-        // version of avg case dist in unique decoding technique.
-        // Here is the probability union bound
-        let code_len_over_f_bits = field_size_bits - self.code_instance.code_len().ilog2() as usize;
-
-        (soundness_bits as f64 / code_len_over_f_bits as f64).ceil() as usize
-    }
-
     pub fn commit<F, ComPackF>(
         &self,
         poly: &MultiLinearPoly<F>,
@@ -146,17 +123,18 @@ impl OrionPublicParams {
         drop(scratch);
 
         // NOTE: packed codeword buffer and encode over packed field
-        let mut packed_interleaved_codewords = vec![ComPackF::ZERO; packed_rows * self.code_len()];
+        let mut packed_interleaved_codewords =
+            vec![ComPackF::ZERO; packed_rows * self.codeword_len()];
         packed_evals
             .chunks(msg_size)
-            .zip(packed_interleaved_codewords.chunks_mut(self.code_len()))
+            .zip(packed_interleaved_codewords.chunks_mut(self.codeword_len()))
             .try_for_each(|(evals, codeword)| {
                 self.code_instance.encode_in_place(evals, codeword)
             })?;
         drop(packed_evals);
 
         // NOTE: transpose codeword s.t., the matrix has codewords being columns
-        let mut scratch = vec![ComPackF::ZERO; packed_rows * self.code_len()];
+        let mut scratch = vec![ComPackF::ZERO; packed_rows * self.codeword_len()];
         transpose_in_place(&mut packed_interleaved_codewords, &mut scratch, packed_rows);
         drop(scratch);
 
@@ -225,9 +203,8 @@ impl OrionPublicParams {
 
         // NOTE: draw random linear combination out
         // and compose proximity response(s) of tensor code IOP based PCS
-        let proximity_repetitions =
-            self.proximity_repetition_num(PCS_SOUNDNESS_BITS, EvalF::FIELD_SIZE);
-        let mut proximity_rows = vec![vec![EvalF::ZERO; msg_size]; proximity_repetitions];
+        let proximity_test_num = self.proximity_repetitions::<EvalF>(PCS_SOUNDNESS_BITS);
+        let mut proximity_rows = vec![vec![EvalF::ZERO; msg_size]; proximity_test_num];
 
         proximity_rows.iter_mut().for_each(|row_buffer| {
             let random_coeffs = transcript.generate_challenge_field_elements(row_num);
@@ -238,6 +215,7 @@ impl OrionPublicParams {
                 .zip(row_buffer.iter_mut())
                 .for_each(|(p_col, res)| *res = luts.lookup_and_sum(p_col));
         });
+        drop(luts);
 
         // NOTE: working on evaluation on top of evaluation response
         let mut scratch = vec![EvalF::ZERO; msg_size];
@@ -246,6 +224,7 @@ impl OrionPublicParams {
             &point[..num_of_vars_in_msg],
             &mut scratch,
         );
+        drop(scratch);
 
         // NOTE: MT opening for point queries
         let leaf_range = row_num / tree::leaf_adic::<F>();
@@ -254,7 +233,7 @@ impl OrionPublicParams {
         let query_openings = query_indices
             .iter()
             .map(|qi| {
-                let index = *qi % self.code_len();
+                let index = *qi % self.codeword_len();
                 let left = index * leaf_range;
                 let right = left + leaf_range - 1;
 
@@ -305,8 +284,7 @@ impl OrionPublicParams {
 
         // NOTE: working on proximity responses, draw random linear combinations
         // then draw query points from fiat shamir transcripts
-        let proximity_test_num =
-            self.proximity_repetition_num(PCS_SOUNDNESS_BITS, EvalF::FIELD_SIZE);
+        let proximity_test_num = self.proximity_repetitions::<EvalF>(PCS_SOUNDNESS_BITS);
         let random_linear_combinations: Vec<Vec<EvalF>> = (0..proximity_test_num)
             .map(|_| transcript.generate_challenge_field_elements(row_num))
             .collect();
@@ -320,7 +298,7 @@ impl OrionPublicParams {
                 .iter()
                 .zip(proof.query_openings.iter())
                 .all(|(&qi, range_path)| {
-                    let index = qi % self.code_len();
+                    let index = qi % self.codeword_len();
                     range_path.verify(commitment) && index == range_path.left / leaf_range
                 });
         if !mt_consistency {
@@ -363,7 +341,7 @@ impl OrionPublicParams {
                     .iter()
                     .zip(packed_interleaved_alphabets.iter())
                     .all(|(&qi, interleaved_alphabet)| {
-                        let index = qi % self.code_len();
+                        let index = qi % self.codeword_len();
                         let alphabet = luts.lookup_and_sum(interleaved_alphabet);
                         alphabet == codeword[index]
                     })
diff --git a/pcs/src/orion/tests.rs b/pcs/src/orion/tests.rs
@@ -8,7 +8,7 @@ use polynomials::MultiLinearPoly;
 use crate::orion::{
     linear_code::{OrionCode, ORION_CODE_PARAMETER_INSTANCE},
     pcs_impl::{OrionCommitment, OrionCommitmentWithData, OrionPublicParams},
-    utils::{transpose_in_place, SubsetSumLUTs},
+    utils::{transpose_in_place, SubsetSumLUTs, TensorIOPPCS},
 };
 
 fn column_combination<F, PackF>(mat: &[F], combination: &[F]) -> Vec<F>
@@ -95,7 +95,7 @@ where
         .flat_map(|msg| orion_pp.code_instance.encode(&msg).unwrap())
         .collect();
 
-    let mut scratch = vec![F::ZERO; row_num * orion_pp.code_len()];
+    let mut scratch = vec![F::ZERO; row_num * orion_pp.codeword_len()];
     transpose_in_place(&mut interleaved_codewords, &mut scratch, row_num);
     drop(scratch);
 
diff --git a/pcs/src/orion/utils.rs b/pcs/src/orion/utils.rs
@@ -1,5 +1,3 @@
-use std::ops::Mul;
-
 use arith::{Field, FieldSerdeError, SimdField};
 use thiserror::Error;
 
@@ -48,44 +46,48 @@ pub(crate) fn transpose_in_place<F: Field>(mat: &mut [F], scratch: &mut [F], row
     mat.copy_from_slice(scratch);
 }
 
+/**********************************
+ * TENSOR IOP BASED PCS UTILITIES *
+ **********************************/
+
+pub trait TensorIOPPCS {
+    fn codeword_len(&self) -> usize;
+
+    fn hamming_weight(&self) -> f64;
+
+    fn row_col_from_variables<F: Field>(num_vars: usize) -> (usize, usize) {
+        let elems_for_smallest_tree = tree::leaf_adic::<F>() * 2;
+
+        let row_num: usize = elems_for_smallest_tree;
+        let msg_size: usize = (1 << num_vars) / row_num;
+
+        (row_num, msg_size)
+    }
+
+    fn query_complexity(&self, soundness_bits: usize) -> usize {
+        // NOTE: use Ligero (AHIV22) or Avg-case dist to a code (BKS18)
+        // version of avg case dist in unique decoding technique.
+        let avg_case_dist = self.hamming_weight() / 3f64;
+        let sec_bits = -(1f64 - avg_case_dist).log2();
+
+        (soundness_bits as f64 / sec_bits).ceil() as usize
+    }
+
+    fn proximity_repetitions<F: Field>(&self, soundness_bits: usize) -> usize {
+        // NOTE: use Ligero (AHIV22) or Avg-case dist to a code (BKS18)
+        // version of avg case dist in unique decoding technique.
+        // Here is the probability union bound
+        let single_run_soundness_bits = F::FIELD_SIZE - self.codeword_len().ilog2() as usize;
+
+        (soundness_bits as f64 / single_run_soundness_bits as f64).ceil() as usize
+    }
+}
+
 /*********************
  * LINEAR OPERATIONS *
  *********************/
 
-#[allow(unused)]
-#[inline]
-pub(crate) fn simd_inner_prod<F0, F1, IPPackF0, IPPackF1>(
-    l: &[F0],
-    r: &[F1],
-    scratch_pl: &mut [IPPackF0],
-    scratch_pr: &mut [IPPackF1],
-) -> F1
-where
-    F0: Field,
-    F1: Field + From<F0> + Mul<F0, Output = F1>,
-    IPPackF0: SimdField<Scalar = F0>,
-    IPPackF1: SimdField<Scalar = F1> + Mul<IPPackF0, Output = IPPackF1>,
-{
-    scratch_pl
-        .iter_mut()
-        .zip(l.chunks(IPPackF0::PACK_SIZE))
-        .for_each(|(pl, ls)| *pl = IPPackF0::pack(ls));
-
-    scratch_pr
-        .iter_mut()
-        .zip(r.chunks(IPPackF1::PACK_SIZE))
-        .for_each(|(pr, rs)| *pr = IPPackF1::pack(rs));
-
-    let simd_sum: IPPackF1 = scratch_pl
-        .iter()
-        .zip(scratch_pr.iter())
-        .map(|(pl, pr)| *pr * *pl)
-        .sum();
-
-    simd_sum.unpack().iter().sum()
-}
-
-pub(crate) struct SubsetSumLUTs<F: Field> {
+pub struct SubsetSumLUTs<F: Field> {
     pub entry_bits: usize,
     pub tables: Vec<Vec<F>>,
 }
diff --git a/pcs/tests/test_orion.rs b/pcs/tests/test_orion.rs
@@ -26,18 +26,16 @@ where
     let mut rng = test_rng();
     let poly = MultiLinearPoly::<F>::random(params.num_vars, &mut rng);
 
-    (0..5).for_each(|_| {
-        let opening_point: Vec<_> = (0..params.num_vars)
-            .map(|_| EvalF::random_unsafe(&mut rng))
-            .collect();
+    let opening_point: Vec<_> = (0..params.num_vars)
+        .map(|_| EvalF::random_unsafe(&mut rng))
+        .collect();
 
-        common::test_pcs_e2e::<OrionPCS<F, EvalF, ComPackF, OpenPackF, T>>(
-            &params,
-            &poly,
-            &opening_point,
-            &mut rng,
-        );
-    })
+    common::test_pcs_e2e::<OrionPCS<F, EvalF, ComPackF, OpenPackF, T>>(
+        &params,
+        &poly,
+        &opening_point,
+        &mut rng,
+    );
 }
 
 #[test]
