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