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