1use std::{
145 collections::HashMap,
146 fmt::{Debug, Formatter, Result},
147 ops::{Add, Mul, Neg, Sub},
148};
149
150use ark_ff::{One, PrimeField, Zero};
151use kimchi::circuits::{expr::Variable, gate::CurrOrNext};
152use num_integer::binomial;
153use o1_utils::FieldHelpers;
154use rand::{Rng, RngCore};
155use std::ops::{Index, IndexMut};
156
157use crate::{
158 utils::{compute_all_two_factors_decomposition, naive_prime_factors, PrimeNumberGenerator},
159 MVPoly,
160};
161
162#[derive(Clone)]
169pub struct Dense<F: PrimeField, const N: usize, const D: usize> {
170 coeff: Vec<F>,
171 normalized_indices: Vec<usize>,
176}
177
178impl<F: PrimeField, const N: usize, const D: usize> Index<usize> for Dense<F, N, D> {
179 type Output = F;
180
181 fn index(&self, index: usize) -> &Self::Output {
182 &self.coeff[index]
183 }
184}
185
186impl<F: PrimeField, const N: usize, const D: usize> IndexMut<usize> for Dense<F, N, D> {
187 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
188 &mut self.coeff[index]
189 }
190}
191
192impl<F: PrimeField, const N: usize, const D: usize> Zero for Dense<F, N, D> {
193 fn is_zero(&self) -> bool {
194 self.coeff.iter().all(|c| c.is_zero())
195 }
196
197 fn zero() -> Self {
198 Dense {
199 coeff: vec![F::zero(); Self::dimension()],
200 normalized_indices: Self::compute_normalized_indices(),
201 }
202 }
203}
204
205impl<F: PrimeField, const N: usize, const D: usize> One for Dense<F, N, D> {
206 fn one() -> Self {
207 let mut result = Dense::zero();
208 result.coeff[0] = F::one();
209 result
210 }
211}
212
213impl<F: PrimeField, const N: usize, const D: usize> MVPoly<F, N, D> for Dense<F, N, D> {
214 unsafe fn random<RNG: RngCore>(rng: &mut RNG, max_degree: Option<usize>) -> Self {
231 let mut prime_gen = PrimeNumberGenerator::new();
232 let normalized_indices = Self::compute_normalized_indices();
233 let coeff = if let Some(max_degree) = max_degree {
236 normalized_indices
237 .iter()
238 .map(|idx| {
239 let degree = naive_prime_factors(*idx, &mut prime_gen)
240 .iter()
241 .fold(0, |acc, (_, d)| acc + d);
242 if degree > max_degree || rng.gen_range(0..10) == 0 {
243 F::zero()
245 } else {
246 F::rand(rng)
247 }
248 })
249 .collect::<Vec<F>>()
250 } else {
251 normalized_indices
252 .iter()
253 .map(|_| {
254 if rng.gen_range(0..10) == 0 {
255 F::zero()
257 } else {
258 F::rand(rng)
259 }
260 })
261 .collect()
262 };
263 Self {
264 coeff,
265 normalized_indices,
266 }
267 }
268
269 fn is_constant(&self) -> bool {
270 self.coeff.iter().skip(1).all(|c| c.is_zero())
271 }
272
273 unsafe fn degree(&self) -> usize {
284 if self.is_constant() {
285 return 0;
286 }
287 let mut prime_gen = PrimeNumberGenerator::new();
288 self.coeff.iter().enumerate().fold(1, |acc, (i, c)| {
289 if *c != F::zero() {
290 let decomposition_of_i =
291 naive_prime_factors(self.normalized_indices[i], &mut prime_gen);
292 let monomial_degree = decomposition_of_i.iter().fold(0, |acc, (_, d)| acc + d);
293 acc.max(monomial_degree)
294 } else {
295 acc
296 }
297 })
298 }
299
300 fn double(&self) -> Self {
301 let coeffs = self.coeff.iter().map(|c| c.double()).collect();
302 Self::from_coeffs(coeffs)
303 }
304
305 fn mul_by_scalar(&self, c: F) -> Self {
306 let coeffs = self.coeff.iter().map(|coef| *coef * c).collect();
307 Self::from_coeffs(coeffs)
308 }
309
310 fn eval(&self, x: &[F; N]) -> F {
315 let mut prime_gen = PrimeNumberGenerator::new();
316 let primes = prime_gen.get_first_nth_primes(N);
317 self.coeff
318 .iter()
319 .enumerate()
320 .fold(F::zero(), |acc, (i, c)| {
321 if i == 0 {
322 acc + c
323 } else {
324 let normalized_index = self.normalized_indices[i];
325 let prime_decomposition = naive_prime_factors(normalized_index, &mut prime_gen);
329 let mut monomial = F::one();
330 prime_decomposition.iter().for_each(|(p, d)| {
331 let inv_p = primes.iter().position(|&x| x == *p).unwrap();
333 let x_p = x[inv_p].pow([*d as u64]);
334 monomial *= x_p;
335 });
336 acc + *c * monomial
337 }
338 })
339 }
340
341 fn from_variable<Column: Into<usize>>(
342 var: Variable<Column>,
343 offset_next_row: Option<usize>,
344 ) -> Self {
345 let Variable { col, row } = var;
346 if row == CurrOrNext::Next {
347 assert!(
348 offset_next_row.is_some(),
349 "The offset for the next row must be provided"
350 );
351 }
352 let offset = if row == CurrOrNext::Curr {
353 0
354 } else {
355 offset_next_row.unwrap()
356 };
357 let var_usize: usize = col.into();
358
359 let mut prime_gen = PrimeNumberGenerator::new();
360 let primes = prime_gen.get_first_nth_primes(N);
361 assert!(primes.contains(&var_usize), "The usize representation of the variable must be a prime number, and unique for each variable");
362
363 let prime_idx = primes.iter().position(|&x| x == var_usize).unwrap();
364 let idx = prime_gen.get_nth_prime(prime_idx + offset + 1);
365
366 let mut res = Self::zero();
367 let inv_idx = res
368 .normalized_indices
369 .iter()
370 .position(|&x| x == idx)
371 .unwrap();
372 res[inv_idx] = F::one();
373 res
374 }
375
376 fn is_homogeneous(&self) -> bool {
377 let normalized_indices = self.normalized_indices.clone();
378 let mut prime_gen = PrimeNumberGenerator::new();
379 let is_homogeneous = normalized_indices
380 .iter()
381 .zip(self.coeff.clone())
382 .all(|(idx, c)| {
383 let decomposition_of_i = naive_prime_factors(*idx, &mut prime_gen);
384 let monomial_degree = decomposition_of_i.iter().fold(0, |acc, (_, d)| acc + d);
385 monomial_degree == D || c == F::zero()
386 });
387 is_homogeneous
388 }
389
390 fn homogeneous_eval(&self, x: &[F; N], u: F) -> F {
391 let normalized_indices = self.normalized_indices.clone();
392 let mut prime_gen = PrimeNumberGenerator::new();
393 let primes = prime_gen.get_first_nth_primes(N);
394 normalized_indices
395 .iter()
396 .zip(self.coeff.clone())
397 .fold(F::zero(), |acc, (idx, c)| {
398 let decomposition_of_i = naive_prime_factors(*idx, &mut prime_gen);
399 let monomial_degree = decomposition_of_i.iter().fold(0, |acc, (_, d)| acc + d);
400 let u_power = D - monomial_degree;
401 let monomial = decomposition_of_i.iter().fold(F::one(), |acc, (p, d)| {
402 let inv_p = primes.iter().position(|&x| x == *p).unwrap();
403 let x_p = x[inv_p].pow([*d as u64]);
404 acc * x_p
405 });
406 acc + c * monomial * u.pow([u_power as u64])
407 })
408 }
409
410 fn add_monomial(&mut self, exponents: [usize; N], coeff: F) {
411 let mut prime_gen = PrimeNumberGenerator::new();
412 let primes = prime_gen.get_first_nth_primes(N);
413 let normalized_index = exponents
414 .iter()
415 .zip(primes.iter())
416 .fold(1, |acc, (d, p)| acc * p.pow(*d as u32));
417 let inv_idx = self
418 .normalized_indices
419 .iter()
420 .position(|&x| x == normalized_index)
421 .unwrap();
422 self.coeff[inv_idx] += coeff;
423 }
424
425 fn compute_cross_terms(
426 &self,
427 _eval1: &[F; N],
428 _eval2: &[F; N],
429 _u1: F,
430 _u2: F,
431 ) -> HashMap<usize, F> {
432 unimplemented!()
433 }
434
435 fn compute_cross_terms_scaled(
436 &self,
437 _eval1: &[F; N],
438 _eval2: &[F; N],
439 _u1: F,
440 _u2: F,
441 _scalar1: F,
442 _scalar2: F,
443 ) -> HashMap<usize, F> {
444 unimplemented!()
445 }
446
447 fn modify_monomial(&mut self, exponents: [usize; N], coeff: F) {
448 let mut prime_gen = PrimeNumberGenerator::new();
449 let primes = prime_gen.get_first_nth_primes(N);
450 let index = exponents
451 .iter()
452 .zip(primes.iter())
453 .fold(1, |acc, (exp, &prime)| acc * prime.pow(*exp as u32));
454 if let Some(pos) = self.normalized_indices.iter().position(|&x| x == index) {
455 self.coeff[pos] = coeff;
456 } else {
457 panic!("Exponent combination out of bounds for the given polynomial degree and number of variables.");
458 }
459 }
460
461 fn is_multilinear(&self) -> bool {
462 if self.is_zero() {
463 return true;
464 }
465 let normalized_indices = self.normalized_indices.clone();
466 let mut prime_gen = PrimeNumberGenerator::new();
467 normalized_indices
468 .iter()
469 .zip(self.coeff.iter())
470 .all(|(idx, c)| {
471 if c.is_zero() {
472 true
473 } else {
474 let decomposition_of_i = naive_prime_factors(*idx, &mut prime_gen);
475 decomposition_of_i.iter().all(|(_p, d)| *d <= 1)
477 }
478 })
479 }
480}
481
482impl<F: PrimeField, const N: usize, const D: usize> Dense<F, N, D> {
483 pub fn new() -> Self {
484 let normalized_indices = Self::compute_normalized_indices();
485 Self {
486 coeff: vec![F::zero(); Self::dimension()],
487 normalized_indices,
488 }
489 }
490 pub fn iter(&self) -> impl Iterator<Item = &F> {
491 self.coeff.iter()
492 }
493
494 pub fn dimension() -> usize {
495 binomial(N + D, D)
496 }
497
498 pub fn from_coeffs(coeff: Vec<F>) -> Self {
499 let normalized_indices = Self::compute_normalized_indices();
500 Dense {
501 coeff,
502 normalized_indices,
503 }
504 }
505
506 pub fn number_of_variables(&self) -> usize {
507 N
508 }
509
510 pub fn maximum_degree(&self) -> usize {
511 D
512 }
513
514 pub fn compute_normalized_indices() -> Vec<usize> {
524 let mut normalized_indices = vec![1; Self::dimension()];
525 let mut prime_gen = PrimeNumberGenerator::new();
526 let primes = prime_gen.get_first_nth_primes(N);
527 let max_index = primes[N - 1].checked_pow(D as u32);
528 let max_index = max_index.expect("Overflow in computing the maximum index");
529 let mut j = 0;
530 (1..=max_index).for_each(|i| {
531 let prime_decomposition_of_index = naive_prime_factors(i, &mut prime_gen);
532 let is_valid_decomposition = prime_decomposition_of_index
533 .iter()
534 .all(|(p, _)| primes.contains(p));
535 let monomial_degree = prime_decomposition_of_index
536 .iter()
537 .fold(0, |acc, (_, d)| acc + d);
538 let is_valid_decomposition = is_valid_decomposition && monomial_degree <= D;
539 if is_valid_decomposition {
540 normalized_indices[j] = i;
541 j += 1;
542 }
543 });
544 normalized_indices
545 }
546
547 pub fn increase_degree<const D_PRIME: usize>(&self) -> Dense<F, N, D_PRIME> {
548 assert!(D <= D_PRIME, "The degree of the target polynomial must be greater or equal to the degree of the source polynomial");
549 let mut result: Dense<F, N, D_PRIME> = Dense::zero();
550 let dst_normalized_indices = Dense::<F, N, D_PRIME>::compute_normalized_indices();
551 let src_normalized_indices = Dense::<F, N, D>::compute_normalized_indices();
552 src_normalized_indices
553 .iter()
554 .enumerate()
555 .for_each(|(i, idx)| {
556 let inv_idx = dst_normalized_indices
558 .iter()
559 .position(|&x| x == *idx)
560 .unwrap();
561 result[inv_idx] = self[i];
562 });
563 result
564 }
565}
566
567impl<F: PrimeField, const N: usize, const D: usize> Default for Dense<F, N, D> {
568 fn default() -> Self {
569 Dense::new()
570 }
571}
572
573impl<F: PrimeField, const N: usize, const D: usize> Add for Dense<F, N, D> {
575 type Output = Self;
576
577 fn add(self, other: Self) -> Self {
578 &self + &other
579 }
580}
581
582impl<F: PrimeField, const N: usize, const D: usize> Add<&Dense<F, N, D>> for Dense<F, N, D> {
583 type Output = Dense<F, N, D>;
584
585 fn add(self, other: &Dense<F, N, D>) -> Dense<F, N, D> {
586 &self + other
587 }
588}
589
590impl<F: PrimeField, const N: usize, const D: usize> Add<Dense<F, N, D>> for &Dense<F, N, D> {
591 type Output = Dense<F, N, D>;
592
593 fn add(self, other: Dense<F, N, D>) -> Dense<F, N, D> {
594 self + &other
595 }
596}
597
598impl<F: PrimeField, const N: usize, const D: usize> Add<&Dense<F, N, D>> for &Dense<F, N, D> {
599 type Output = Dense<F, N, D>;
600
601 fn add(self, other: &Dense<F, N, D>) -> Dense<F, N, D> {
602 let coeffs = self
603 .coeff
604 .iter()
605 .zip(other.coeff.iter())
606 .map(|(a, b)| *a + *b)
607 .collect();
608 Dense::from_coeffs(coeffs)
609 }
610}
611
612impl<F: PrimeField, const N: usize, const D: usize> Sub for Dense<F, N, D> {
614 type Output = Self;
615
616 fn sub(self, other: Self) -> Self {
617 self + (-other)
618 }
619}
620
621impl<F: PrimeField, const N: usize, const D: usize> Sub<&Dense<F, N, D>> for Dense<F, N, D> {
622 type Output = Dense<F, N, D>;
623
624 fn sub(self, other: &Dense<F, N, D>) -> Dense<F, N, D> {
625 self + (-other)
626 }
627}
628
629impl<F: PrimeField, const N: usize, const D: usize> Sub<Dense<F, N, D>> for &Dense<F, N, D> {
630 type Output = Dense<F, N, D>;
631
632 fn sub(self, other: Dense<F, N, D>) -> Dense<F, N, D> {
633 self + (-other)
634 }
635}
636
637impl<F: PrimeField, const N: usize, const D: usize> Sub<&Dense<F, N, D>> for &Dense<F, N, D> {
638 type Output = Dense<F, N, D>;
639
640 fn sub(self, other: &Dense<F, N, D>) -> Dense<F, N, D> {
641 self + (-other)
642 }
643}
644
645impl<F: PrimeField, const N: usize, const D: usize> Neg for Dense<F, N, D> {
647 type Output = Self;
648
649 fn neg(self) -> Self::Output {
650 -&self
651 }
652}
653
654impl<F: PrimeField, const N: usize, const D: usize> Neg for &Dense<F, N, D> {
655 type Output = Dense<F, N, D>;
656
657 fn neg(self) -> Self::Output {
658 let coeffs = self.coeff.iter().map(|c| -*c).collect();
659 Dense::from_coeffs(coeffs)
660 }
661}
662
663impl<F: PrimeField, const N: usize, const D: usize> Mul<Dense<F, N, D>> for Dense<F, N, D> {
665 type Output = Self;
666
667 fn mul(self, other: Self) -> Self {
668 let mut cache = HashMap::new();
669 let mut prime_gen = PrimeNumberGenerator::new();
670 let mut result = vec![];
671 (0..self.coeff.len()).for_each(|i| {
672 let mut sum = F::zero();
673 let normalized_index = self.normalized_indices[i];
674 let two_factors_decomposition =
675 compute_all_two_factors_decomposition(normalized_index, &mut cache, &mut prime_gen);
676 two_factors_decomposition.iter().for_each(|(a, b)| {
677 let inv_a = self
679 .normalized_indices
680 .iter()
681 .position(|&x| x == *a)
682 .unwrap();
683 let inv_b = self
684 .normalized_indices
685 .iter()
686 .position(|&x| x == *b)
687 .unwrap();
688 let a_coeff = self.coeff[inv_a];
689 let b_coeff = other.coeff[inv_b];
690 let product = a_coeff * b_coeff;
691 sum += product;
692 });
693 result.push(sum);
694 });
695 Self::from_coeffs(result)
696 }
697}
698
699impl<F: PrimeField, const N: usize, const D: usize> PartialEq for Dense<F, N, D> {
700 fn eq(&self, other: &Self) -> bool {
701 self.coeff == other.coeff
702 }
703}
704
705impl<F: PrimeField, const N: usize, const D: usize> Eq for Dense<F, N, D> {}
706
707impl<F: PrimeField, const N: usize, const D: usize> Debug for Dense<F, N, D> {
708 fn fmt(&self, f: &mut Formatter<'_>) -> Result {
709 let mut prime_gen = PrimeNumberGenerator::new();
710 let primes = prime_gen.get_first_nth_primes(N);
711 let coeff: Vec<_> = self
712 .coeff
713 .iter()
714 .enumerate()
715 .filter(|(_i, c)| *c != &F::zero())
716 .collect();
717 if coeff.is_empty() {
719 write!(f, "0").unwrap();
720 return Ok(());
721 }
722 let l = coeff.len();
723 coeff.into_iter().for_each(|(i, c)| {
724 let normalized_idx = self.normalized_indices[i];
725 if normalized_idx == 1 && *c != F::one() {
726 write!(f, "{}", c.to_biguint()).unwrap();
727 } else {
728 let prime_decomposition = naive_prime_factors(normalized_idx, &mut prime_gen);
729 if *c != F::one() {
730 write!(f, "{}", c.to_biguint()).unwrap();
731 }
732 prime_decomposition.iter().for_each(|(p, d)| {
733 let inv_p = primes.iter().position(|&x| x == *p).unwrap();
734 if *d > 1 {
735 write!(f, "x_{}^{}", inv_p, d).unwrap();
736 } else {
737 write!(f, "x_{}", inv_p).unwrap();
738 }
739 });
740 }
741 if i != l - 1 && l != 1 {
744 write!(f, " + ").unwrap();
745 }
746 });
747 Ok(())
748 }
749}
750
751impl<F: PrimeField, const N: usize, const D: usize> From<F> for Dense<F, N, D> {
752 fn from(value: F) -> Self {
753 let mut result = Self::zero();
754 result.coeff[0] = value;
755 result
756 }
757}
758
759