sp1_primitives/
polynomial.rs1use alloc::{vec, vec::Vec};
2use core::{
3 fmt::Debug,
4 ops::{Add, AddAssign, Mul, Neg, Sub},
5 slice::Iter,
6};
7
8use itertools::Itertools;
9use slop_algebra::{AbstractExtensionField, AbstractField, Field};
10
11#[derive(Debug, Clone)]
13pub struct Polynomial<T> {
14 coefficients: Vec<T>,
15}
16
17impl<T> Polynomial<T> {
18 #[must_use]
20 pub const fn new(coefficients: Vec<T>) -> Self {
21 Self { coefficients }
22 }
23
24 pub fn from_coefficients(coefficients: &[T]) -> Self
26 where
27 T: Clone,
28 {
29 Self { coefficients: coefficients.to_vec() }
30 }
31
32 #[must_use]
34 pub fn as_coefficients(self) -> Vec<T> {
35 self.coefficients
36 }
37
38 #[must_use]
40 pub fn coefficients(&self) -> &[T] {
41 &self.coefficients
42 }
43
44 #[must_use]
46 pub fn degree(&self) -> usize {
47 self.coefficients.len() - 1
48 }
49
50 #[allow(clippy::needless_pass_by_value)]
52 pub fn eval<S: AbstractExtensionField<T>>(&self, x: S) -> S
53 where
54 T: AbstractField,
55 {
56 let powers = x.powers();
57 self.coefficients.iter().zip(powers).map(|(c, x)| x * c.clone()).sum()
58 }
59
60 #[must_use]
62 pub fn root_quotient(&self, r: T) -> Self
63 where
64 T: Field,
65 {
66 let len = self.coefficients.len();
67 let mut result = Vec::with_capacity(len - 1);
68 let r_inv = r.inverse();
69
70 result.push(-self.coefficients[0] * r_inv);
71 for i in 1..len - 1 {
72 let element = result[i - 1] - self.coefficients[i];
73 result.push(element * r_inv);
74 }
75 Self { coefficients: result }
76 }
77}
78
79impl<T> FromIterator<T> for Polynomial<T> {
80 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
81 Self { coefficients: iter.into_iter().collect() }
82 }
83}
84
85impl<T: Add<Output = T> + Clone> Add for Polynomial<T> {
86 type Output = Self;
87
88 fn add(self, other: Self) -> Self {
89 self + &other
90 }
91}
92
93impl<T: Add<Output = T> + Clone> Add for &Polynomial<T> {
94 type Output = Polynomial<T>;
95
96 fn add(self, other: Self) -> Polynomial<T> {
97 self.coefficients
98 .iter()
99 .zip_longest(other.coefficients.iter())
100 .map(|x| match x {
101 itertools::EitherOrBoth::Both(a, b) => a.clone() + b.clone(),
102 itertools::EitherOrBoth::Left(a) => a.clone(),
103 itertools::EitherOrBoth::Right(b) => b.clone(),
104 })
105 .collect()
106 }
107}
108
109impl<T: Add<Output = T> + Clone> Add<&Polynomial<T>> for Polynomial<T> {
110 type Output = Polynomial<T>;
111
112 fn add(self, other: &Polynomial<T>) -> Polynomial<T> {
113 self.coefficients
114 .iter()
115 .zip_longest(other.coefficients.iter())
116 .map(|x| match x {
117 itertools::EitherOrBoth::Both(a, b) => a.clone() + b.clone(),
118 itertools::EitherOrBoth::Left(a) => a.clone(),
119 itertools::EitherOrBoth::Right(b) => b.clone(),
120 })
121 .collect()
122 }
123}
124
125impl<T: Mul<Output = T> + Add<Output = T> + AddAssign + Clone> Add<T> for Polynomial<T> {
126 type Output = Polynomial<T>;
127
128 fn add(self, other: T) -> Polynomial<T> {
129 let mut coefficients = self.coefficients;
130 coefficients[0] = coefficients[0].clone() + other;
131 Self::new(coefficients)
132 }
133}
134
135impl<T: Mul<Output = T> + Add<Output = T> + Add + Clone> Add<T> for &Polynomial<T> {
136 type Output = Polynomial<T>;
137
138 fn add(self, other: T) -> Polynomial<T> {
139 let mut coefficients = self.coefficients.clone();
140 coefficients[0] = coefficients[0].clone() + other;
141 Polynomial::new(coefficients)
142 }
143}
144
145impl<T: Neg<Output = T>> Neg for Polynomial<T> {
146 type Output = Self;
147
148 fn neg(self) -> Self {
149 Self::new(self.coefficients.into_iter().map(|x| -x).collect())
150 }
151}
152
153impl<T: Sub<Output = T> + Neg<Output = T> + Clone> Sub for Polynomial<T> {
154 type Output = Self;
155
156 fn sub(self, other: Self) -> Self {
157 self - &other
158 }
159}
160
161impl<T: Sub<Output = T> + Neg<Output = T> + Clone> Sub<&Polynomial<T>> for Polynomial<T> {
162 type Output = Polynomial<T>;
163
164 fn sub(self, other: &Polynomial<T>) -> Polynomial<T> {
165 Polynomial::new(
166 self.coefficients
167 .iter()
168 .zip_longest(other.coefficients.iter())
169 .map(|x| match x {
170 itertools::EitherOrBoth::Both(a, b) => a.clone() - b.clone(),
171 itertools::EitherOrBoth::Left(a) => a.clone(),
172 itertools::EitherOrBoth::Right(b) => -b.clone(),
173 })
174 .collect(),
175 )
176 }
177}
178
179impl<T: Sub<Output = T> + Neg<Output = T> + Clone> Sub for &Polynomial<T> {
180 type Output = Polynomial<T>;
181
182 fn sub(self, other: Self) -> Polynomial<T> {
183 Polynomial::new(
184 self.coefficients
185 .iter()
186 .zip_longest(other.coefficients.iter())
187 .map(|x| match x {
188 itertools::EitherOrBoth::Both(a, b) => a.clone() - b.clone(),
189 itertools::EitherOrBoth::Left(a) => a.clone(),
190 itertools::EitherOrBoth::Right(b) => -b.clone(),
191 })
192 .collect(),
193 )
194 }
195}
196
197impl<T: AbstractField> Mul for Polynomial<T> {
198 type Output = Self;
199
200 fn mul(self, other: Self) -> Self {
201 let mut result = vec![T::zero(); self.coefficients.len() + other.coefficients.len() - 1];
202 for (i, a) in self.coefficients.into_iter().enumerate() {
203 for (j, b) in other.coefficients.iter().enumerate() {
204 result[i + j] = result[i + j].clone() + a.clone() * b.clone();
205 }
206 }
207 Self::new(result)
208 }
209}
210
211impl<T: AbstractField> Mul for &Polynomial<T> {
212 type Output = Polynomial<T>;
213
214 fn mul(self, other: Self) -> Polynomial<T> {
215 let mut result = vec![T::zero(); self.coefficients.len() + other.coefficients.len() - 1];
216 for (i, a) in self.coefficients.iter().enumerate() {
217 for (j, b) in other.coefficients.iter().enumerate() {
218 result[i + j] = result[i + j].clone() + a.clone() * b.clone();
219 }
220 }
221 Polynomial::new(result)
222 }
223}
224
225impl<T: AbstractField> Mul<T> for Polynomial<T> {
226 type Output = Self;
227
228 fn mul(self, other: T) -> Self {
229 Self::new(self.coefficients.into_iter().map(|x| x * other.clone()).collect())
230 }
231}
232
233impl<T: AbstractField> Mul<T> for &Polynomial<T> {
234 type Output = Polynomial<T>;
235
236 fn mul(self, other: T) -> Polynomial<T> {
237 Polynomial::new(self.coefficients.iter().cloned().map(|x| x * other.clone()).collect())
238 }
239}
240
241impl<T: Eq + AbstractField> PartialEq<Polynomial<T>> for Polynomial<T> {
242 fn eq(&self, other: &Polynomial<T>) -> bool {
243 if self.coefficients.len() != other.coefficients.len() {
244 let (shorter, longer) = if self.coefficients.len() < other.coefficients.len() {
245 (self, other)
246 } else {
247 (other, self)
248 };
249 for i in 0..longer.coefficients.len() {
250 if (i < shorter.coefficients.len()
251 && shorter.coefficients[i] != longer.coefficients[i])
252 || (i >= shorter.coefficients.len() && longer.coefficients[i] != T::zero())
253 {
254 return false;
255 }
256 }
257 return true;
258 }
259 self.coefficients == other.coefficients
260 }
261}
262
263impl Polynomial<u8> {
264 #[must_use]
266 pub fn as_field<F: Field>(self) -> Polynomial<F> {
267 Polynomial {
268 coefficients: self.coefficients.iter().map(|x| F::from_canonical_u8(*x)).collect(),
269 }
270 }
271}
272
273impl<'a, Var: Into<Expr> + Clone, Expr: Clone> From<Iter<'a, Var>> for Polynomial<Expr> {
274 fn from(value: Iter<'a, Var>) -> Self {
275 Polynomial::from_coefficients(&value.map(|x| (*x).clone().into()).collect::<Vec<_>>())
276 }
277}