Skip to main content

sp1_primitives/
polynomial.rs

1use 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/// A polynomial represented as a vector of coefficients.
12#[derive(Debug, Clone)]
13pub struct Polynomial<T> {
14    coefficients: Vec<T>,
15}
16
17impl<T> Polynomial<T> {
18    /// Create a new polynomial from a vector of coefficients.
19    #[must_use]
20    pub const fn new(coefficients: Vec<T>) -> Self {
21        Self { coefficients }
22    }
23
24    /// Create a new polynomial from a slice of coefficients.
25    pub fn from_coefficients(coefficients: &[T]) -> Self
26    where
27        T: Clone,
28    {
29        Self { coefficients: coefficients.to_vec() }
30    }
31
32    /// Gets the coefficients of the polynomial.
33    #[must_use]
34    pub fn as_coefficients(self) -> Vec<T> {
35        self.coefficients
36    }
37
38    /// Gets the coefficients of the polynomial.
39    #[must_use]
40    pub fn coefficients(&self) -> &[T] {
41        &self.coefficients
42    }
43
44    /// Gets the degree of the polynomial.
45    #[must_use]
46    pub fn degree(&self) -> usize {
47        self.coefficients.len() - 1
48    }
49
50    /// Evaluates the polynomial at a given point.
51    #[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    /// Computes the root quotient of the polynomial.
61    #[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    /// Converts the polynomial to a field polynomial.
265    #[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}