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