1use crate::algebra::{
2 Additive, CryptoGroup, Field, Object, Random, Ring, Space, msm_naive, powers,
3};
4#[cfg(not(feature = "std"))]
5use alloc::{borrow::Cow, vec, vec::Vec};
6use commonware_codec::{
7 Buf, EncodeSize, MaxEncodeSize, RangeCfg, Read, Write, util::max_items_size,
8};
9use commonware_parallel::Strategy;
10use commonware_utils::{TryCollect, non_empty_vec, ordered::Map, vec::NonEmptyVec};
11use core::{
12 fmt::Debug,
13 iter,
14 num::NonZeroU32,
15 ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
16};
17use rand_core::CryptoRng;
18#[cfg(feature = "std")]
19use std::borrow::Cow;
20
21const MIN_POINTS_FOR_MSM: usize = 2;
23
24#[derive(Clone)]
26pub struct Poly<K> {
27 coeffs: NonEmptyVec<K>,
29}
30
31impl<K> Poly<K> {
32 fn len(&self) -> NonZeroU32 {
33 self.coeffs
34 .len()
35 .try_into()
36 .expect("Impossible: polynomial length not in 1..=u32::MAX")
37 }
38
39 const fn len_usize(&self) -> usize {
40 self.coeffs.len().get()
41 }
42
43 fn from_iter_unchecked(iter: impl IntoIterator<Item = K>) -> Self {
48 let coeffs = iter
49 .into_iter()
50 .try_collect::<NonEmptyVec<_>>()
51 .expect("polynomial must have a least 1 coefficient");
52 Self { coeffs }
53 }
54
55 pub fn degree(&self) -> u32 {
66 self.len().get() - 1
67 }
68
69 pub fn required(&self) -> NonZeroU32 {
73 self.len()
74 }
75
76 pub fn constant(&self) -> &K {
80 &self.coeffs[0]
81 }
82
83 pub fn translate<L>(&self, f: impl Fn(&K) -> L) -> Poly<L> {
88 Poly {
89 coeffs: self.coeffs.map(f),
90 }
91 }
92
93 pub fn eval<R>(&self, r: &R) -> K
109 where
110 K: Space<R>,
111 {
112 let mut iter = self.coeffs.iter().rev();
113 let mut acc = iter
120 .next()
121 .expect("Impossible: Polynomial has no coefficients")
122 .clone();
123 for coeff in iter {
124 acc *= r;
125 acc += coeff;
126 }
127 acc
128 }
129
130 pub fn eval_msm<R: Ring>(&self, r: &R, strategy: &impl Strategy) -> K
136 where
137 K: Space<R>,
138 {
139 let weights = powers(R::one(), r)
141 .take(self.len_usize())
142 .collect::<Vec<_>>();
143 K::msm(&self.coeffs, &weights, strategy)
144 }
145
146 pub fn lin_comb_eval<'a, R: Ring + 'a>(
153 &self,
154 into_iter: impl IntoIterator<Item = (R, Cow<'a, R>)>,
155 strategy: &impl Strategy,
156 ) -> K
157 where
158 K: Space<R>,
159 {
160 let weights = {
162 let mut iter = into_iter.into_iter();
163 let Some((a0, b0)) = iter.next() else {
164 return K::zero();
165 };
166
167 let len = self.len_usize();
168 let mut out: Vec<_> = powers(a0, b0.as_ref()).take(len).collect();
169 for (ai, bi) in iter {
170 powers(ai, bi.as_ref())
171 .take(len)
172 .zip(out.iter_mut())
173 .for_each(|(c_j, o_j)| *o_j += &c_j);
174 }
175 out
176 };
177 K::msm(&self.coeffs, &weights, strategy)
178 }
179}
180
181impl<K: Debug> Debug for Poly<K> {
182 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
183 write!(f, "Poly(")?;
184 for (i, c) in self.coeffs.iter().enumerate() {
185 if i > 0 {
186 write!(f, " + {c:?} X^{i}")?;
187 } else {
188 write!(f, "{c:?}")?;
189 }
190 }
191 write!(f, ")")?;
192 Ok(())
193 }
194}
195
196impl<K: EncodeSize> EncodeSize for Poly<K> {
197 fn encode_size(&self) -> usize {
198 self.coeffs.encode_size()
199 }
200}
201
202impl<K: MaxEncodeSize> Poly<K> {
203 pub fn max_encode_size(coefficients: usize) -> usize {
206 max_items_size(coefficients, K::MAX_ENCODE_SIZE)
207 }
208}
209
210impl<K: Write> Write for Poly<K> {
211 fn write(&self, buf: &mut impl bytes::BufMut) {
212 self.coeffs.write(buf);
213 }
214}
215
216impl<K: Read> Read for Poly<K> {
217 type Cfg = (RangeCfg<NonZeroU32>, <K as Read>::Cfg);
218
219 fn read_cfg(buf: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
220 Ok(Self {
221 coeffs: NonEmptyVec::<K>::read_cfg(buf, &(cfg.0.into(), cfg.1.clone()))?,
222 })
223 }
224}
225
226impl<K: Random> Poly<K> {
227 pub fn new(mut rng: impl CryptoRng, degree: u32) -> Self {
230 Self::from_iter_unchecked((0..=degree).map(|_| K::random(&mut rng)))
231 }
232
233 pub fn new_with_constant(mut rng: impl CryptoRng, degree: u32, constant: K) -> Self {
237 Self::from_iter_unchecked(
238 iter::once(constant).chain((0..=degree).skip(1).map(|_| K::random(&mut rng))),
239 )
240 }
241}
242
243impl<K: Additive> PartialEq for Poly<K> {
248 fn eq(&self, other: &Self) -> bool {
249 let zero = K::zero();
250 let max_len = self.len().max(other.len());
251 let self_then_zeros = self.coeffs.iter().chain(iter::repeat(&zero));
252 let other_then_zeros = other.coeffs.iter().chain(iter::repeat(&zero));
253 self_then_zeros
254 .zip(other_then_zeros)
255 .take(max_len.get() as usize)
256 .all(|(a, b)| a == b)
257 }
258}
259
260impl<K: Additive> Eq for Poly<K> {}
261
262impl<K: Additive> Poly<K> {
263 fn merge_with(&mut self, rhs: &Self, f: impl Fn(&mut K, &K)) {
264 self.coeffs
265 .resize(self.coeffs.len().max(rhs.coeffs.len()), K::zero());
266 self.coeffs
267 .iter_mut()
268 .zip(&rhs.coeffs)
269 .for_each(|(a, b)| f(a, b));
270 }
271
272 pub fn degree_exact(&self) -> u32 {
278 let zero = K::zero();
279 let leading_zeroes = self.coeffs.iter().rev().take_while(|&x| x == &zero).count();
280 let lz_u32 =
281 u32::try_from(leading_zeroes).expect("Impossible: Polynomial has >= 2^32 coefficients");
282 self.degree().saturating_sub(lz_u32)
285 }
286}
287
288impl<K: Additive> Object for Poly<K> {}
289
290impl<'a, K: Additive> AddAssign<&'a Self> for Poly<K> {
293 fn add_assign(&mut self, rhs: &'a Self) {
294 self.merge_with(rhs, |a, b| *a += b);
295 }
296}
297
298impl<'a, K: Additive> Add<&'a Self> for Poly<K> {
299 type Output = Self;
300
301 fn add(mut self, rhs: &'a Self) -> Self::Output {
302 self += rhs;
303 self
304 }
305}
306
307impl<'a, K: Additive> SubAssign<&'a Self> for Poly<K> {
308 fn sub_assign(&mut self, rhs: &'a Self) {
309 self.merge_with(rhs, |a, b| *a -= b);
310 }
311}
312
313impl<'a, K: Additive> Sub<&'a Self> for Poly<K> {
314 type Output = Self;
315
316 fn sub(mut self, rhs: &'a Self) -> Self::Output {
317 self -= rhs;
318 self
319 }
320}
321
322impl<K: Additive> Neg for Poly<K> {
323 type Output = Self;
324
325 fn neg(self) -> Self::Output {
326 Self {
327 coeffs: self.coeffs.map_into(Neg::neg),
328 }
329 }
330}
331
332impl<K: Additive> Additive for Poly<K> {
333 fn zero() -> Self {
334 Self {
335 coeffs: non_empty_vec![K::zero()],
336 }
337 }
338}
339
340impl<'a, R, K: Space<R>> MulAssign<&'a R> for Poly<K> {
343 fn mul_assign(&mut self, rhs: &'a R) {
344 self.coeffs.iter_mut().for_each(|c| *c *= rhs);
345 }
346}
347
348impl<'a, R, K: Space<R>> Mul<&'a R> for Poly<K> {
349 type Output = Self;
350
351 fn mul(mut self, rhs: &'a R) -> Self::Output {
352 self *= rhs;
353 self
354 }
355}
356
357impl<R: Sync, K: Space<R>> Space<R> for Poly<K> {
358 fn msm(polys: &[Self], scalars: &[R], strategy: &impl Strategy) -> Self {
359 if polys.len() < MIN_POINTS_FOR_MSM {
360 return msm_naive(polys, scalars);
361 }
362
363 let cols = polys.len().min(scalars.len());
364 let polys = &polys[..cols];
365 let scalars = &scalars[..cols];
366
367 let rows = polys
368 .iter()
369 .map(|x| x.len_usize())
370 .max()
371 .expect("at least 1 point");
372
373 let coeffs = strategy.map_init_collect_vec(
374 0..rows,
375 || Vec::with_capacity(cols),
376 |row, i| {
377 row.clear();
378 for p in polys {
379 row.push(p.coeffs.get(i).cloned().unwrap_or_else(K::zero));
380 }
381 K::msm(row, scalars, strategy)
382 },
383 );
384 Self::from_iter_unchecked(coeffs)
385 }
386}
387
388impl<G: CryptoGroup> Poly<G> {
389 pub fn commit(p: Poly<G::Scalar>) -> Self {
391 p.translate(|c| G::generator() * c)
392 }
393}
394
395pub struct Interpolator<I, F> {
426 weights: Map<I, F>,
427}
428
429impl<I: PartialEq, F: Ring> Interpolator<I, F> {
430 pub fn interpolate<K: Space<F>>(
435 &self,
436 evals: &Map<I, K>,
437 strategy: &impl Strategy,
438 ) -> Option<K> {
439 if evals.keys() != self.weights.keys() {
440 return None;
441 }
442 Some(K::msm(evals.values(), self.weights.values(), strategy))
443 }
444}
445
446impl<I: Ord, F: Field> Interpolator<I, F> {
447 pub fn new(points: impl IntoIterator<Item = (I, F)>) -> Self {
454 let points = Map::from_iter_dedup(points);
455 let n = points.len();
456 if n == 0 {
457 return Self { weights: points };
458 }
459
460 let values = points.values();
463 let zero = F::zero();
464 let mut total_product = F::one();
465 let mut c = Vec::with_capacity(n);
466 for (i, w_i) in values.iter().enumerate() {
467 if w_i == &zero {
469 let mut out = points;
470 for (j, w) in out.values_mut().iter_mut().enumerate() {
471 *w = if j == i { F::one() } else { F::zero() };
472 }
473 return Self { weights: out };
474 }
475
476 total_product *= w_i;
478 let mut c_i = w_i.clone();
479 for w_j in values
480 .iter()
481 .enumerate()
482 .filter_map(|(j, v)| (j != i).then_some(v))
483 {
484 c_i *= &(w_j.clone() - w_i);
485 }
486 c.push(c_i);
487 }
488
489 let mut prefix = Vec::with_capacity(n + 1);
492 prefix.push(F::one());
493 let mut acc = F::one();
494 for c_i in &c {
495 acc *= c_i;
496 prefix.push(acc.clone());
497 }
498
499 let mut inv_acc = total_product * &prefix[n].inv();
501
502 let mut out = points;
504 let out_vals = out.values_mut();
505 for i in (0..n).rev() {
506 out_vals[i] = inv_acc.clone() * &prefix[i];
507 inv_acc *= &c[i];
508 }
509 Self { weights: out }
510 }
511}
512
513#[commonware_macros::stability(ALPHA)]
514impl<I: Clone + Ord, F: crate::algebra::FieldNTT> Interpolator<I, F> {
515 pub fn roots_of_unity(
524 total: NonZeroU32,
525 points: commonware_utils::ordered::BiMap<I, u32>,
526 ) -> Self {
527 let weights = <Map<I, F> as commonware_utils::TryFromIterator<(I, F)>>::try_from_iter(
528 crate::ntt::lagrange_coefficients(total, points.values().iter().copied())
529 .into_iter()
530 .filter_map(|(k, coeff)| Some((points.get_key(&k)?.clone(), coeff))),
531 )
532 .expect("points has already been deduped");
533 Self { weights }
534 }
535
536 #[cfg(any(test, feature = "fuzz"))]
543 fn roots_of_unity_naive(
544 total: NonZeroU32,
545 points: commonware_utils::ordered::BiMap<I, u32>,
546 ) -> Self {
547 use crate::algebra::powers;
548
549 let total_u32 = total.get();
550 let size = (total_u32 as u64).next_power_of_two();
551 let lg_size = size.ilog2() as u8;
552 let w = F::root_of_unity(lg_size).expect("domain too large for NTT");
553
554 let points: Vec<(I, u32)> = points.into_iter().filter(|(_, k)| *k < total_u32).collect();
555 let max_k = points.iter().map(|(_, k)| *k).max().unwrap_or(0) as usize;
556 let powers: Vec<_> = powers(F::one(), &w).take(max_k + 1).collect();
557
558 let eval_points = points
559 .into_iter()
560 .map(|(i, k)| (i, powers[k as usize].clone()));
561 Self::new(eval_points)
562 }
563}
564
565#[cfg(any(test, feature = "arbitrary"))]
566mod impl_arbitrary {
567 use super::*;
568 use arbitrary::Arbitrary;
569
570 impl<'a, F: Arbitrary<'a>> Arbitrary<'a> for Poly<F> {
571 fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
572 let first = u.arbitrary()?;
573 let rest: Vec<F> = u.arbitrary()?;
574 let mut coeffs = NonEmptyVec::new(first);
575 coeffs.extend(rest);
576 Ok(Self { coeffs })
577 }
578 }
579}
580
581#[commonware_macros::stability(ALPHA)]
582#[cfg(any(test, feature = "fuzz"))]
583pub mod fuzz {
584 use super::*;
585 use crate::{
586 algebra::test_suites,
587 test::{F, G},
588 };
589 use arbitrary::{Arbitrary, Unstructured};
590 use commonware_codec::Encode as _;
591 use commonware_parallel::Sequential;
592 use commonware_utils::{
593 TryFromIterator,
594 ordered::{BiMap, Map},
595 };
596
597 #[derive(Debug, Arbitrary)]
598 pub enum Plan {
599 Codec(Poly<F>),
600 EvalAdd(Poly<F>, Poly<F>, F),
601 EvalScale(Poly<F>, F, F),
602 EvalZero(Poly<F>),
603 EvalMsm(Poly<F>, F),
604 LinCombEval(Poly<F>, Vec<(F, F)>),
605 Interpolate(Poly<F>),
606 InterpolateWithZeroPoint(Poly<F>),
607 InterpolateWithZeroPointMiddle(Poly<F>),
608 TranslateScale(Poly<F>, F),
609 CommitEval(Poly<F>, F),
610 RootsOfUnityEqNaive(u16),
611 FuzzAdditive,
612 FuzzSpaceRing,
613 }
614
615 impl Plan {
616 pub fn run(self, u: &mut Unstructured<'_>) -> arbitrary::Result<()> {
617 match self {
618 Self::Codec(f) => {
619 assert_eq!(
620 &f,
621 &Poly::<F>::read_cfg(&mut f.encode(), &(RangeCfg::exact(f.required()), ()))
622 .unwrap()
623 );
624 }
625 Self::EvalAdd(f, g, x) => {
626 assert_eq!(f.eval(&x) + &g.eval(&x), (f + &g).eval(&x));
627 }
628 Self::EvalScale(f, x, w) => {
629 assert_eq!(f.eval(&x) * &w, (f * &w).eval(&x));
630 }
631 Self::EvalZero(f) => {
632 assert_eq!(&f.eval(&F::zero()), f.constant());
633 }
634 Self::EvalMsm(f, x) => {
635 assert_eq!(f.eval(&x), f.eval_msm(&x, &Sequential));
636 }
637 Self::LinCombEval(f, pairs) => {
638 let naive_eval = pairs.iter().fold(F::zero(), |mut acc, (a, b)| {
639 acc += &(*a * &f.eval(b));
640 acc
641 });
642 let lin_comb = f.lin_comb_eval(
643 pairs.iter().map(|(a, b)| (*a, Cow::Borrowed(b))),
644 &Sequential,
645 );
646 assert_eq!(naive_eval, lin_comb);
647 }
648 Self::Interpolate(f) => {
649 if f == Poly::zero() || u64::from(f.required().get()) >= F::MAX {
650 return Ok(());
651 }
652 let mut points = (0..f.required().get())
653 .map(|i| F::from(u64::from(i) + 1))
654 .collect::<Vec<_>>();
655 let interpolator = Interpolator::new(points.iter().copied().enumerate());
656 let evals = Map::from_iter_dedup(points.iter().map(|p| f.eval(p)).enumerate());
657 let recovered = interpolator.interpolate(&evals, &Sequential);
658 assert_eq!(recovered.as_ref(), Some(f.constant()));
659 points.pop();
660 assert_eq!(
661 interpolator.interpolate(
662 &Map::from_iter_dedup(points.iter().map(|p| f.eval(p)).enumerate()),
663 &Sequential
664 ),
665 None
666 );
667 }
668 Self::InterpolateWithZeroPoint(f) => {
669 if f == Poly::zero() || u64::from(f.required().get()) >= F::MAX {
670 return Ok(());
671 }
672 let points: Vec<_> = (0..f.required().get())
673 .map(|i| F::from(u64::from(i)))
674 .collect();
675 let interpolator = Interpolator::new(points.iter().copied().enumerate());
676 let evals = Map::from_iter_dedup(points.iter().map(|p| f.eval(p)).enumerate());
677 let recovered = interpolator.interpolate(&evals, &Sequential);
678 assert_eq!(recovered.as_ref(), Some(f.constant()));
679 }
680 Self::InterpolateWithZeroPointMiddle(f) => {
681 if f == Poly::zero()
682 || f.required().get() < 2
683 || u64::from(f.required().get()) >= F::MAX
684 {
685 return Ok(());
686 }
687 let n = f.required().get();
688 let points: Vec<_> = (1..n)
689 .map(|i| F::from(u64::from(i)))
690 .chain(core::iter::once(F::zero()))
691 .collect();
692 let interpolator = Interpolator::new(points.iter().copied().enumerate());
693 let evals = Map::from_iter_dedup(points.iter().map(|p| f.eval(p)).enumerate());
694 let recovered = interpolator.interpolate(&evals, &Sequential);
695 assert_eq!(recovered.as_ref(), Some(f.constant()));
696 }
697 Self::TranslateScale(f, x) => {
698 assert_eq!(f.translate(|c| x * c), f * &x);
699 }
700 Self::CommitEval(f, x) => {
701 assert_eq!(G::generator() * &f.eval(&x), Poly::<G>::commit(f).eval(&x));
702 }
703 Self::RootsOfUnityEqNaive(n) => {
704 let n = (u32::from(n) % 256) + 1;
705 let total = NonZeroU32::new(n).expect("n is in 1..=256");
706 let points = BiMap::try_from_iter((0..n as usize).map(|i| (i, i as u32)))
707 .expect("interpolation points should be bijective");
708 let fast = Interpolator::<usize, crate::fields::goldilocks::F>::roots_of_unity(
709 total,
710 points.clone(),
711 );
712 let naive =
713 Interpolator::<usize, crate::fields::goldilocks::F>::roots_of_unity_naive(
714 total, points,
715 );
716 assert_eq!(fast.weights, naive.weights);
717 }
718 Self::FuzzAdditive => {
719 test_suites::fuzz_additive::<Poly<F>>(u)?;
720 }
721 Self::FuzzSpaceRing => {
722 test_suites::fuzz_space_ring::<F, Poly<F>>(u)?;
723 }
724 }
725 Ok(())
726 }
727 }
728
729 #[test]
730 fn test_fuzz() {
731 commonware_invariants::minifuzz::test(|u| u.arbitrary::<Plan>()?.run(u));
732 }
733}
734#[cfg(test)]
735mod test {
736 use super::{fuzz::Plan, *};
737 use crate::test::F;
738 use arbitrary::Unstructured;
739
740 #[test]
741 fn test_eq() {
742 fn eq(a: &[u8], b: &[u8]) -> bool {
743 Poly {
744 coeffs: a.iter().copied().map(F::from).try_collect().unwrap(),
745 } == Poly {
746 coeffs: b.iter().copied().map(F::from).try_collect().unwrap(),
747 }
748 }
749 assert!(eq(&[1, 2], &[1, 2]));
750 assert!(!eq(&[1, 2], &[2, 3]));
751 assert!(!eq(&[1, 2], &[1, 2, 3]));
752 assert!(!eq(&[1, 2, 3], &[1, 2]));
753 assert!(eq(&[1, 2], &[1, 2, 0, 0]));
754 assert!(eq(&[1, 2, 0, 0], &[1, 2]));
755 assert!(!eq(&[1, 2, 0], &[2, 3]));
756 assert!(!eq(&[2, 3], &[1, 2, 0]));
757 }
758
759 #[test]
760 fn lin_comb_eval_edge_cases() {
761 fn poly(coeffs: &[u8]) -> Poly<F> {
762 Poly {
763 coeffs: coeffs.iter().copied().map(F::from).try_collect().unwrap(),
764 }
765 }
766
767 fn pairs(values: &[(u8, u8)]) -> Vec<(F, F)> {
768 values
769 .iter()
770 .map(|(a, b)| (F::from(*a), F::from(*b)))
771 .collect()
772 }
773
774 let cases = [
775 Plan::LinCombEval(poly(&[3, 5, 7]), vec![]),
776 Plan::LinCombEval(poly(&[11]), pairs(&[(2, 0), (3, 1), (5, 8)])),
777 Plan::LinCombEval(poly(&[4, 6, 8]), pairs(&[(2, 5), (7, 5), (3, 5)])),
778 Plan::LinCombEval(poly(&[9, 2, 3, 4]), pairs(&[(6, 0), (1, 0), (5, 7)])),
779 Plan::LinCombEval(poly(&[1, 2, 4, 8]), pairs(&[(3, 1), (7, 1), (2, 6)])),
780 ];
781 let mut u = Unstructured::new(&[]);
782 for case in cases {
783 case.run(&mut u).unwrap();
784 }
785 }
786
787 #[test]
788 fn interpolation_plans_cover_large_test_field() {
789 let poly = Poly {
790 coeffs: [3u8, 5, 7, 11]
791 .into_iter()
792 .map(F::from)
793 .try_collect()
794 .unwrap(),
795 };
796 let mut u = Unstructured::new(&[]);
797 for make_plan in [
798 Plan::Interpolate as fn(Poly<F>) -> Plan,
799 Plan::InterpolateWithZeroPoint,
800 Plan::InterpolateWithZeroPointMiddle,
801 ] {
802 make_plan(poly.clone()).run(&mut u).unwrap();
803 }
804 }
805
806 #[cfg(feature = "arbitrary")]
807 mod conformance {
808 use super::*;
809 use commonware_codec::conformance::CodecConformance;
810
811 commonware_conformance::conformance_tests! {
812 CodecConformance<Poly<F>>
813 }
814 }
815}