1use crate::rs_config::ecfft::ECFFTConfig;
  2use crate::tree::BaseOpening;
  3use crate::FieldExt;
  4use ecfft::extend;
  5use serde::{Deserialize, Serialize};
  6
  7use crate::fft::fft;
  8use crate::polynomial::eq_poly::EqPoly;
  9use crate::tensor_code::TensorCode;
 10use crate::transcript::Transcript;
 11use crate::utils::{det_num_cols, det_num_rows, dot_prod, hash_all, rlc_rows, sample_indices};
 12
 13use super::tensor_code::CommittedTensorCode;
 14
 15#[derive(Clone)]
 16pub struct TensorRSMultilinearPCSConfig<F: FieldExt> {
 17    pub expansion_factor: usize,
 18    pub domain_powers: Option<Vec<Vec<F>>>,
 19    pub fft_domain: Option<Vec<F>>,
 20    pub ecfft_config: Option<ECFFTConfig<F>>,
 21    pub l: usize,
 22}
 23
 24impl<F: FieldExt> TensorRSMultilinearPCSConfig<F> {
 25    pub fn num_cols(&self, num_entries: usize) -> usize {
 26        det_num_cols(num_entries, self.l)
 27    }
 28
 29    pub fn num_rows(&self, num_entries: usize) -> usize {
 30        det_num_rows(num_entries, self.l)
 31    }
 32}
 33
 34#[derive(Clone)]
 35pub struct TensorMultilinearPCS<F: FieldExt> {
 36    config: TensorRSMultilinearPCSConfig<F>,
 37}
 38
 39#[derive(Clone, Serialize, Deserialize)]
 40pub struct TensorMLOpening<F: FieldExt> {
 41    pub x: Vec<F>,
 42    pub y: F,
 43    pub base_opening: BaseOpening,
 44    pub test_query_leaves: Vec<Vec<F>>,
 45    pub eval_query_leaves: Vec<Vec<F>>,
 46    pub u_hat_comm: [u8; 32],
 47    pub test_u_prime: Vec<F>,
 48    pub test_r_prime: Vec<F>,
 49    pub eval_r_prime: Vec<F>,
 50    pub eval_u_prime: Vec<F>,
 51    pub poly_num_vars: usize,
 52}
 53
 54impl<F: FieldExt> TensorMultilinearPCS<F> {
 55    pub fn new(config: TensorRSMultilinearPCSConfig<F>) -> Self {
 56        Self { config }
 57    }
 58
 59    pub fn commit(&self, ml_poly_evals: &[F]) -> CommittedTensorCode<F> {
 60        // Merkle commit to the evaluations of the polynomial
 61        let n = ml_poly_evals.len();
 62        assert!(n.is_power_of_two());
 63        let tensor_code = self.encode_zk(ml_poly_evals);
 64        let tree = tensor_code.commit(self.config.num_cols(n), self.config.num_rows(n));
 65        tree
 66    }
 67
 68    pub fn open(
 69        &self,
 70        u_hat_comm: &CommittedTensorCode<F>,
 71        // TODO: Remove poly and use u_hat_comm
 72        ml_poly_evals: &[F],
 73        point: &[F],
 74        eval: F,
 75        transcript: &mut Transcript<F>,
 76    ) -> TensorMLOpening<F> {
 77        let n = ml_poly_evals.len();
 78        assert!(n.is_power_of_two());
 79        let num_vars = (n as f64).log2() as usize;
 80
 81        let num_cols = self.config.num_cols(n);
 82        let num_rows = self.config.num_rows(n);
 83
 84        debug_assert_eq!(num_vars, point.len());
 85
 86        // ########################################
 87        // Testing phase
 88        // Prove the consistency between the random linear combination of the evaluation tensor (u_prime)
 89        // and the tensor codeword (u_hat)
 90        // ########################################
 91
 92        // Derive the challenge vector;
 93        let r_u = transcript.challenge_vec(num_rows);
 94
 95        let u = (0..num_rows)
 96            .map(|i| ml_poly_evals[(i * num_cols)..((i + 1) * num_cols)].to_vec())
 97            .collect::<Vec<Vec<F>>>();
 98
 99        // Random linear combination of the rows of the polynomial in a tensor structure
100        let test_u_prime = rlc_rows(u.clone(), &r_u);
101
102        // Random linear combination of the blinder
103        let blinder = u_hat_comm
104            .tensor_codeword
105            .0
106            .iter()
107            .map(|row| row[(row.len() / 2)..].to_vec())
108            .collect::<Vec<Vec<F>>>();
109
110        debug_assert_eq!(blinder[0].len(), u_hat_comm.tensor_codeword.0[0].len() / 2);
111
112        let test_r_prime = rlc_rows(blinder.clone(), &r_u);
113
114        let num_indices = self.config.l;
115        let indices = sample_indices(num_indices, num_cols * 2, transcript);
116
117        let test_queries = self.test_phase(&indices, &u_hat_comm);
118
119        // ########################################
120        // Evaluation phase
121        // Prove the consistency
122        // ########################################
123
124        // Get the evaluation point
125
126        let log2_num_rows = (num_rows as f64).log2() as usize;
127        let q1 = EqPoly::new(point[0..log2_num_rows].to_vec()).evals();
128
129        let eval_r_prime = rlc_rows(blinder, &q1);
130
131        let eval_u_prime = rlc_rows(u.clone(), &q1);
132
133        let eval_queries = self.test_phase(&indices, &u_hat_comm);
134
135        TensorMLOpening {
136            x: point.to_vec(),
137            y: eval,
138            eval_query_leaves: eval_queries,
139            test_query_leaves: test_queries,
140            u_hat_comm: u_hat_comm.committed_tree.root(),
141            test_u_prime,
142            test_r_prime,
143            eval_r_prime,
144            eval_u_prime,
145            base_opening: BaseOpening {
146                hashes: u_hat_comm.committed_tree.column_roots.clone(),
147            },
148            poly_num_vars: num_vars,
149        }
150    }
151}
152
153impl<F: FieldExt> TensorMultilinearPCS<F> {
154    pub fn verify(&self, opening: &TensorMLOpening<F>, transcript: &mut Transcript<F>) {
155        let poly_num_entries = 2usize.pow(opening.poly_num_vars as u32);
156        let num_rows = self.config.num_rows(poly_num_entries);
157        let num_cols = self.config.num_cols(poly_num_entries);
158
159        // Verify the base opening
160        let base_opening = &opening.base_opening;
161        base_opening.verify(opening.u_hat_comm);
162
163        // ########################################
164        // Verify test phase
165        // ########################################
166
167        let r_u = transcript.challenge_vec(num_rows);
168
169        let test_u_prime_rs_codeword = self
170            .rs_encode(&opening.test_u_prime)
171            .iter()
172            .zip(opening.test_r_prime.iter())
173            .map(|(c, r)| *c + *r)
174            .collect::<Vec<F>>();
175
176        let num_indices = self.config.l;
177        let indices = sample_indices(num_indices, num_cols * 2, transcript);
178
179        debug_assert_eq!(indices.len(), opening.test_query_leaves.len());
180        for (expected_index, leaves) in indices.iter().zip(opening.test_query_leaves.iter()) {
181            // Verify that the hashes of the leaves equals the corresponding column root
182            let leaf_bytes = leaves
183                .iter()
184                .map(|x| x.to_repr())
185                .collect::<Vec<[u8; 32]>>();
186            let column_root = hash_all(&leaf_bytes);
187            let expected_column_root = base_opening.hashes[*expected_index];
188            assert_eq!(column_root, expected_column_root);
189
190            let mut sum = F::ZERO;
191            for (leaf, r_i) in leaves.iter().zip(r_u.iter()) {
192                sum += *r_i * *leaf;
193            }
194            assert_eq!(sum, test_u_prime_rs_codeword[*expected_index]);
195        }
196
197        // ########################################
198        // Verify evaluation phase
199        // ########################################
200
201        let log2_num_rows = (num_rows as f64).log2() as usize;
202        let q1 = EqPoly::new(opening.x[0..log2_num_rows].to_vec()).evals();
203        let q2 = EqPoly::new(opening.x[log2_num_rows..].to_vec()).evals();
204
205        let eval_u_prime_rs_codeword = self
206            .rs_encode(&opening.eval_u_prime)
207            .iter()
208            .zip(opening.eval_r_prime.iter())
209            .map(|(c, r)| *c + *r)
210            .collect::<Vec<F>>();
211
212        debug_assert_eq!(q1.len(), opening.eval_query_leaves[0].len());
213        debug_assert_eq!(indices.len(), opening.test_query_leaves.len());
214        for (expected_index, leaves) in indices.iter().zip(opening.eval_query_leaves.iter()) {
215            // TODO: Don't need to check the leaves again?
216            // Verify that the hashes of the leaves equals the corresponding column root
217            let leaf_bytes = leaves
218                .iter()
219                .map(|x| x.to_repr())
220                .collect::<Vec<[u8; 32]>>();
221            let column_root = hash_all(&leaf_bytes);
222            let expected_column_root = base_opening.hashes[*expected_index];
223            assert_eq!(column_root, expected_column_root);
224
225            let mut sum = F::ZERO;
226            for (leaf, q1_i) in leaves.iter().zip(q1.iter()) {
227                sum += *q1_i * *leaf;
228            }
229            assert_eq!(sum, eval_u_prime_rs_codeword[*expected_index]);
230        }
231
232        let expected_eval = dot_prod(&opening.eval_u_prime, &q2);
233        assert_eq!(expected_eval, opening.y);
234    }
235
236    fn split_encode(&self, message: &[F]) -> Vec<F> {
237        let codeword = self.rs_encode(message);
238
239        let mut rng = rand::thread_rng();
240        let blinder = (0..codeword.len())
241            .map(|_| F::random(&mut rng))
242            .collect::<Vec<F>>();
243
244        let mut randomized_codeword = codeword
245            .iter()
246            .zip(blinder.clone().iter())
247            .map(|(c, b)| *b + *c)
248            .collect::<Vec<F>>();
249
250        randomized_codeword.extend_from_slice(&blinder);
251        debug_assert_eq!(randomized_codeword.len(), codeword.len() * 2);
252        randomized_codeword
253    }
254
255    fn rs_encode(&self, message: &[F]) -> Vec<F> {
256        let mut padded_message = message.to_vec();
257        padded_message.resize(message.len().next_power_of_two(), F::ZERO);
258        let codeword_len = padded_message.len() * self.config.expansion_factor;
259
260        let codeword_len_log2 = (codeword_len as f64).log2() as usize;
261        let codeword = if self.config.fft_domain.is_some() {
262            // TODO: Resize the domain according to the message length
263            let fft_domain = self.config.fft_domain.as_ref().unwrap();
264            let mut padded_coeffs = message.clone().to_vec();
265            padded_coeffs.resize(fft_domain.len(), F::ZERO);
266            let codeword = fft(&padded_coeffs, &fft_domain);
267
268            codeword
269        } else if self.config.ecfft_config.is_some() {
270            let mut ecfft_config = self.config.ecfft_config.clone().unwrap();
271
272            // Resize the domain to the correct size
273            let config_domain_size = ecfft_config.domain.len();
274
275            assert!(config_domain_size >= codeword_len_log2 - 1);
276            ecfft_config.domain =
277                ecfft_config.domain[(config_domain_size - (codeword_len_log2 - 1))..].to_vec();
278            ecfft_config.matrices =
279                ecfft_config.matrices[(config_domain_size - (codeword_len_log2 - 1))..].to_vec();
280            ecfft_config.inverse_matrices = ecfft_config.inverse_matrices
281                [(config_domain_size - (codeword_len_log2 - 1))..]
282                .to_vec();
283
284            assert_eq!(
285                padded_message.len() * self.config.expansion_factor,
286                ecfft_config.domain[0].len()
287            );
288
289            let extended_evals = extend(
290                &padded_message,
291                &ecfft_config.domain,
292                &ecfft_config.matrices,
293                &ecfft_config.inverse_matrices,
294                0,
295            );
296
297            let codeword = [message.to_vec(), extended_evals].concat();
298            codeword
299        } else {
300            // TODO: Resize the domain according to the message length
301            let domain_powers = self.config.domain_powers.as_ref().unwrap();
302            assert_eq!(message.len(), domain_powers[0].len());
303            assert_eq!(
304                message.len() * self.config.expansion_factor,
305                domain_powers.len()
306            );
307
308            let codeword = domain_powers
309                .iter()
310                .map(|powers| {
311                    message
312                        .iter()
313                        .zip(powers.iter())
314                        .fold(F::ZERO, |acc, (m, p)| acc + *m * *p)
315                })
316                .collect::<Vec<F>>();
317
318            codeword
319        };
320
321        codeword
322    }
323
324    fn test_phase(&self, indices: &[usize], u_hat_comm: &CommittedTensorCode<F>) -> Vec<Vec<F>> {
325        // Query the columns of u_hat
326        let num_indices = self.config.l;
327
328        let u_hat_openings = indices
329            .iter()
330            .map(|index| u_hat_comm.query_column(*index))
331            .collect::<Vec<Vec<F>>>();
332
333        debug_assert_eq!(u_hat_openings.len(), num_indices);
334
335        u_hat_openings
336    }
337
338    fn encode_zk(&self, ml_poly_evals: &[F]) -> TensorCode<F> {
339        let n = ml_poly_evals.len();
340        assert!(n.is_power_of_two());
341
342        let num_rows = self.config.num_rows(n);
343        let num_cols = self.config.num_cols(n);
344        debug_assert_eq!(n, num_cols * num_rows);
345
346        let codewords = (0..num_rows)
347            .map(|i| &ml_poly_evals[i * num_cols..(i + 1) * num_cols])
348            .map(|row| self.split_encode(&row))
349            .collect::<Vec<Vec<F>>>();
350
351        TensorCode(codewords)
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358    use crate::polynomial::ml_poly::MlPoly;
359    use crate::rs_config::{ecfft, good_curves::secp256k1::secp256k1_good_curve, naive, smooth};
360
361    const TEST_NUM_VARS: usize = 8;
362    const TEST_L: usize = 10;
363
364    fn test_poly_evals<F: FieldExt>() -> MlPoly<F> {
365        let num_entries: usize = 2usize.pow(TEST_NUM_VARS as u32);
366
367        let evals = (0..num_entries)
368            .map(|i| F::from((i + 1) as u64))
369            .collect::<Vec<F>>();
370
371        MlPoly::new(evals)
372    }
373
374    fn prove_and_verify<F: FieldExt>(ml_poly: &MlPoly<F>, pcs: TensorMultilinearPCS<F>) {
375        let ml_poly_evals = &ml_poly.evals;
376
377        let comm = pcs.commit(ml_poly_evals);
378
379        let ml_poly_num_vars = (ml_poly_evals.len() as f64).log2() as usize;
380        let open_at = (0..ml_poly_num_vars)
381            .map(|i| F::from(i as u64))
382            .collect::<Vec<F>>();
383        let y = ml_poly.eval(&open_at);
384
385        let mut prover_transcript = Transcript::<F>::new(b"test");
386        prover_transcript.append_bytes(&comm.committed_tree.root);
387        let opening = pcs.open(&comm, ml_poly_evals, &open_at, y, &mut prover_transcript);
388
389        let mut verifier_transcript = Transcript::<F>::new(b"test");
390        verifier_transcript.append_bytes(&comm.committed_tree.root);
391        pcs.verify(&opening, &mut verifier_transcript);
392    }
393
394    fn config_base<F: FieldExt>() -> TensorRSMultilinearPCSConfig<F> {
395        let expansion_factor = 2;
396
397        TensorRSMultilinearPCSConfig::<F> {
398            expansion_factor,
399            domain_powers: None,
400            fft_domain: None,
401            ecfft_config: None,
402            l: TEST_L,
403        }
404    }
405
406    #[test]
407    fn test_tensor_pcs_fft() {
408        type F = halo2curves::pasta::Fp;
409        // FFT config
410        let ml_poly = test_poly_evals();
411        let mut config = config_base();
412
413        // The test polynomial has 2^k non-zero entries
414        let num_entries = ml_poly.evals.len();
415        config.fft_domain = Some(smooth::gen_config(config.num_cols(num_entries)));
416
417        // Test FFT PCS
418        let tensor_pcs_fft = TensorMultilinearPCS::<F>::new(config);
419        prove_and_verify(&ml_poly, tensor_pcs_fft);
420    }
421
422    #[test]
423    fn test_tensor_pcs_ecfft() {
424        type F = halo2curves::secp256k1::Fp;
425        let ml_poly = test_poly_evals();
426
427        let mut config = config_base();
428
429        let n = ml_poly.evals.len();
430        let num_cols = config.num_cols(n);
431        let k = ((num_cols * config.expansion_factor).next_power_of_two() as f64).log2() as usize;
432
433        let (curve, coset_offset) = secp256k1_good_curve(k);
434        config.ecfft_config = Some(ecfft::gen_config_form_curve(curve, coset_offset));
435
436        // Test FFT PCS
437        let tensor_pcs_ecf = TensorMultilinearPCS::<F>::new(config);
438        prove_and_verify(&ml_poly, tensor_pcs_ecf);
439    }
440
441    #[test]
442    fn test_tensor_pcs_naive() {
443        type F = halo2curves::secp256k1::Fp;
444        // FFT config
445        let ml_poly = test_poly_evals();
446        let n = ml_poly.evals.len();
447
448        // Naive config
449        let mut config = config_base();
450        config.domain_powers = Some(naive::gen_config(config.num_cols(n)));
451
452        // Test FFT PCS
453        let tensor_pcs_naive = TensorMultilinearPCS::<F>::new(config);
454        prove_and_verify(&ml_poly, tensor_pcs_naive);
455    }
456}