winter_math/polynom/mod.rs
1// Copyright (c) Facebook, Inc. and its affiliates.
2//
3// This source code is licensed under the MIT license found in the
4// LICENSE file in the root directory of this source tree.
5
6//! Basic polynomial operations.
7//!
8//! This module provides a set of function for basic polynomial operations, including:
9//! - Polynomial evaluation using Horner method.
10//! - Polynomial interpolation using Lagrange method.
11//! - Polynomial addition, subtraction, multiplication, and division.
12//! - Synthetic polynomial division for efficient division by polynomials of the form `x`^`a` - `b`.
13//!
14//! In the context of this module any slice of field elements is considered to be a polynomial
15//! in reverse coefficient form. A few examples:
16//!
17//! ```
18//! # use winter_math::{fields::{f128::BaseElement}, FieldElement};
19//! // p(x) = 2 * x + 1
20//! let p = vec![BaseElement::new(1), BaseElement::new(2)];
21//!
22//! // p(x) = 4 * x^2 + 3
23//! let p = [BaseElement::new(3), BaseElement::ZERO, BaseElement::new(4)];
24//! ```
25
26use alloc::vec::Vec;
27use core::mem;
28
29use utils::group_slice_elements;
30
31use crate::{field::FieldElement, utils::batch_inversion};
32
33#[cfg(test)]
34mod tests;
35
36// POLYNOMIAL EVALUATION
37// ================================================================================================
38
39/// Evaluates a polynomial at a single point and returns the result.
40///
41/// Evaluates polynomial `p` at coordinate `x` using
42/// [Horner's method](https://en.wikipedia.org/wiki/Horner%27s_method).
43///
44/// # Examples
45/// ```
46/// # use winter_math::polynom::*;
47/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
48/// // define polynomial: f(x) = 3 * x^2 + 2 * x + 1
49/// let p = (1u32..4).map(BaseElement::from).collect::<Vec<_>>();
50///
51/// // evaluate the polynomial at point 4
52/// let x = BaseElement::new(4);
53/// assert_eq!(BaseElement::new(57), eval(&p, x));
54/// ```
55pub fn eval<B, E>(p: &[B], x: E) -> E
56where
57 B: FieldElement,
58 E: FieldElement + From<B>,
59{
60 // Horner evaluation
61 p.iter().rev().fold(E::ZERO, |acc, &coeff| acc * x + E::from(coeff))
62}
63
64/// Evaluates a polynomial at multiple points and returns a vector of results.
65///
66/// Evaluates polynomial `p` at all coordinates in `xs` slice by repeatedly invoking
67/// `polynom::eval()` function.
68///
69/// # Examples
70/// ```
71/// # use winter_math::polynom::*;
72/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
73/// // define polynomial: f(x) = 3 * x^2 + 2 * x + 1
74/// let p = (1_u32..4).map(BaseElement::from).collect::<Vec<_>>();
75/// let xs = (3_u32..6).map(BaseElement::from).collect::<Vec<_>>();
76///
77/// let expected = xs.iter().map(|x| eval(&p, *x)).collect::<Vec<_>>();
78/// assert_eq!(expected, eval_many(&p, &xs));
79/// ```
80pub fn eval_many<B, E>(p: &[B], xs: &[E]) -> Vec<E>
81where
82 B: FieldElement,
83 E: FieldElement + From<B>,
84{
85 xs.iter().map(|x| eval(p, *x)).collect()
86}
87
88// POLYNOMIAL INTERPOLATION
89// ================================================================================================
90
91/// Returns a polynomial in coefficient form interpolated from a set of X and Y coordinates.
92///
93/// Uses [Lagrange interpolation](https://en.wikipedia.org/wiki/Lagrange_polynomial) to build a
94/// polynomial from X and Y coordinates. If `remove_leading_zeros = true`, all leading coefficients
95/// which are ZEROs will be truncated; otherwise, the length of result will be equal to the number
96/// of X coordinates.
97///
98/// # Panics
99/// Panics if number of X and Y coordinates is not the same.
100///
101/// # Example
102/// ```
103/// # use winter_math::polynom::*;
104/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
105/// # use rand_utils::rand_vector;
106/// let xs: Vec<BaseElement> = rand_vector(16);
107/// let ys: Vec<BaseElement> = rand_vector(16);
108///
109/// let p = interpolate(&xs, &ys, false);
110/// assert_eq!(ys, eval_many(&p, &xs));
111/// ```
112pub fn interpolate<E>(xs: &[E], ys: &[E], remove_leading_zeros: bool) -> Vec<E>
113where
114 E: FieldElement,
115{
116 debug_assert!(xs.len() == ys.len(), "number of X and Y coordinates must be the same");
117
118 let roots = poly_from_roots(xs);
119 let numerators: Vec<Vec<E>> = xs.iter().map(|&x| syn_div(&roots, 1, x)).collect();
120
121 let denominators: Vec<E> = numerators.iter().zip(xs).map(|(e, &x)| eval(e, x)).collect();
122 let denominators = batch_inversion(&denominators);
123
124 let mut result = vec![E::ZERO; xs.len()];
125 for i in 0..xs.len() {
126 let y_slice = ys[i] * denominators[i];
127 for (j, res) in result.iter_mut().enumerate() {
128 *res += numerators[i][j] * y_slice;
129 }
130 }
131
132 if remove_leading_zeros {
133 crate::polynom::remove_leading_zeros(&result)
134 } else {
135 result
136 }
137}
138
139/// Returns a vector of polynomials interpolated from the provided X and Y coordinate batches.
140///
141/// Uses [Lagrange interpolation](https://en.wikipedia.org/wiki/Lagrange_polynomial) to build a
142/// vector of polynomial from X and Y coordinate batches (one polynomial per batch).
143///
144/// When the number of batches is larger, this function is significantly faster than using
145/// `polynom::interpolate()` function individually for each batch of coordinates. The speed-up
146/// is primarily due to computing all inversions as a single batch inversion across all
147/// coordinate batches.
148///
149/// # Panics
150/// Panics if the number of X coordinate batches and Y coordinate batches is not the same.
151///
152/// # Examples
153/// ```
154/// # use winter_math::polynom::*;
155/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
156/// # use rand_utils::rand_array;
157/// let x_batches: Vec<[BaseElement; 8]> = vec![rand_array(), rand_array()];
158/// let y_batches: Vec<[BaseElement; 8]> = vec![rand_array(), rand_array()];
159///
160/// let polys = interpolate_batch(&x_batches, &y_batches);
161/// for ((p, xs), ys) in polys.iter().zip(x_batches).zip(y_batches) {
162/// assert_eq!(ys.to_vec(), eval_many(p, &xs));
163/// }
164/// ```
165pub fn interpolate_batch<E, const N: usize>(xs: &[[E; N]], ys: &[[E; N]]) -> Vec<[E; N]>
166where
167 E: FieldElement,
168{
169 debug_assert!(
170 xs.len() == ys.len(),
171 "number of X coordinate batches and Y coordinate batches must be the same"
172 );
173
174 let n = xs.len();
175 let mut equations = vec![[E::ZERO; N]; n * N];
176 let mut inverses = vec![E::ZERO; n * N];
177
178 // TODO: converting this to an array results in about 5% speed-up, but unfortunately, complex
179 // generic constraints are not yet supported: https://github.com/rust-lang/rust/issues/76560
180 let mut roots = vec![E::ZERO; N + 1];
181
182 for (i, xs) in xs.iter().enumerate() {
183 fill_zero_roots(xs, &mut roots);
184 for (j, &x) in xs.iter().enumerate() {
185 let equation = &mut equations[i * N + j];
186 // optimized synthetic division for this context
187 equation[N - 1] = roots[N];
188 for k in (0..N - 1).rev() {
189 equation[k] = roots[k + 1] + equation[k + 1] * x;
190 }
191 inverses[i * N + j] = eval(equation, x);
192 }
193 }
194 let equations = group_slice_elements::<[E; N], N>(&equations);
195 let inverses_vec = batch_inversion(&inverses);
196 let inverses = group_slice_elements::<E, N>(&inverses_vec);
197
198 let mut result = vec![[E::ZERO; N]; n];
199 for (i, poly) in result.iter_mut().enumerate() {
200 for j in 0..N {
201 let inv_y = ys[i][j] * inverses[i][j];
202 for (res_coeff, &eq_coeff) in poly.iter_mut().zip(equations[i][j].iter()) {
203 *res_coeff += eq_coeff * inv_y;
204 }
205 }
206 }
207
208 result
209}
210
211// POLYNOMIAL MATH OPERATIONS
212// ================================================================================================
213
214/// Returns a polynomial resulting from adding two polynomials together.
215///
216/// Polynomials `a` and `b` are expected to be in the coefficient form, and the returned
217/// polynomial will be in the coefficient form as well. The length of the returned vector
218/// will be max(a.len(), b.len()).
219///
220/// # Examples
221/// ```
222/// # use winter_math::polynom::*;
223/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
224/// // p1(x) = 4 * x^2 + 3 * x + 2
225/// let p1 = (2_u32..5).map(BaseElement::from).collect::<Vec<_>>();
226/// // p2(x) = 2 * x + 1
227/// let p2 = (1_u32..3).map(BaseElement::from).collect::<Vec<_>>();
228///
229/// // expected result = 4 * x^2 + 5 * x + 3
230/// let expected = vec![BaseElement::new(3), BaseElement::new(5), BaseElement::new(4)];
231/// assert_eq!(expected, add(&p1, &p2));
232/// ```
233pub fn add<E>(a: &[E], b: &[E]) -> Vec<E>
234where
235 E: FieldElement,
236{
237 let result_len = core::cmp::max(a.len(), b.len());
238 let mut result = Vec::with_capacity(result_len);
239 for i in 0..result_len {
240 let c1 = if i < a.len() { a[i] } else { E::ZERO };
241 let c2 = if i < b.len() { b[i] } else { E::ZERO };
242 result.push(c1 + c2);
243 }
244 result
245}
246
247/// Returns a polynomial resulting from subtracting one polynomial from another.
248///
249/// Specifically, subtracts polynomial `b` from polynomial `a` and returns the result. Both
250/// polynomials are expected to be in the coefficient form, and the returned polynomial will
251/// be in the coefficient form as well. The length of the returned vector will be
252/// max(a.len(), b.len()).
253///
254/// # Examples
255/// ```
256/// # use winter_math::polynom::*;
257/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
258/// // p1(x) = 4 * x^2 + 3 * x + 2
259/// let p1 = (2_u32..5).map(BaseElement::from).collect::<Vec<_>>();
260/// // p2(x) = 2 * x + 1
261/// let p2 = (1_u32..3).map(BaseElement::from).collect::<Vec<_>>();
262///
263/// // expected result = 4 * x^2 + x + 1
264/// let expected = vec![BaseElement::new(1), BaseElement::new(1), BaseElement::new(4)];
265/// assert_eq!(expected, sub(&p1, &p2));
266/// ```
267pub fn sub<E>(a: &[E], b: &[E]) -> Vec<E>
268where
269 E: FieldElement,
270{
271 let result_len = core::cmp::max(a.len(), b.len());
272 let mut result = Vec::with_capacity(result_len);
273 for i in 0..result_len {
274 let c1 = if i < a.len() { a[i] } else { E::ZERO };
275 let c2 = if i < b.len() { b[i] } else { E::ZERO };
276 result.push(c1 - c2);
277 }
278 result
279}
280
281/// Returns a polynomial resulting from multiplying two polynomials together.
282///
283/// Polynomials `a` and `b` are expected to be in the coefficient form, and the returned
284/// polynomial will be in the coefficient form as well. The length of the returned vector
285/// will be a.len() + b.len() - 1.
286///
287/// # Examples
288/// ```
289/// # use winter_math::polynom::*;
290/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
291/// // p1(x) = x + 1
292/// let p1 = [BaseElement::ONE, BaseElement::ONE];
293/// // p2(x) = x^2 + 2
294/// let p2 = [BaseElement::new(2), BaseElement::ZERO, BaseElement::ONE];
295///
296/// // expected result = x^3 + x^2 + 2 * x + 2
297/// let expected = vec![
298/// BaseElement::new(2),
299/// BaseElement::new(2),
300/// BaseElement::new(1),
301/// BaseElement::new(1),
302/// ];
303/// assert_eq!(expected, mul(&p1, &p2));
304/// ```
305pub fn mul<E>(a: &[E], b: &[E]) -> Vec<E>
306where
307 E: FieldElement,
308{
309 let result_len = a.len() + b.len() - 1;
310 let mut result = vec![E::ZERO; result_len];
311 for i in 0..a.len() {
312 for j in 0..b.len() {
313 let s = a[i] * b[j];
314 result[i + j] += s;
315 }
316 }
317 result
318}
319
320/// Returns a polynomial resulting from multiplying a given polynomial by a scalar value.
321///
322/// Specifically, multiplies every coefficient of polynomial `p` by constant `k` and returns
323/// the resulting vector.
324///
325/// # Examples
326/// ```
327/// # use winter_math::polynom::*;
328/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
329/// let p = [BaseElement::new(1), BaseElement::new(2), BaseElement::new(3)];
330/// let k = BaseElement::new(2);
331///
332/// let expected = vec![BaseElement::new(2), BaseElement::new(4), BaseElement::new(6)];
333/// assert_eq!(expected, mul_by_scalar(&p, k));
334/// ```
335pub fn mul_by_scalar<E>(p: &[E], k: E) -> Vec<E>
336where
337 E: FieldElement,
338{
339 let mut result = Vec::with_capacity(p.len());
340 for coeff in p {
341 result.push(*coeff * k);
342 }
343 result
344}
345
346/// Returns a polynomial resulting from dividing one polynomial by another.
347///
348/// Specifically, divides polynomial `a` by polynomial `b` and returns the result. If the
349/// polynomials don't divide evenly, the remainder is ignored. Both polynomials are expected to
350/// be in the coefficient form, and the returned polynomial will be in the coefficient form as
351/// well. The length of the returned vector will be a.len() - b.len() + 1.
352///
353/// # Panics
354/// Panics if:
355/// * Polynomial `b` is empty.
356/// * Degree of polynomial `b` is zero and the constant coefficient is ZERO.
357/// * The degree of polynomial `b` is greater than the degree of polynomial `a`.
358///
359/// # Examples
360/// ```
361/// # use winter_math::polynom::*;
362/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
363/// // p1(x) = x^3 + x^2 + 2 * x + 2
364/// let p1 = [
365/// BaseElement::new(2),
366/// BaseElement::new(2),
367/// BaseElement::new(1),
368/// BaseElement::new(1),
369/// ];
370/// // p2(x) = x^2 + 2
371/// let p2 = [BaseElement::new(2), BaseElement::ZERO, BaseElement::ONE];
372///
373/// // expected result = x + 1
374/// let expected = vec![BaseElement::ONE, BaseElement::ONE];
375/// assert_eq!(expected, div(&p1, &p2));
376/// ```
377pub fn div<E>(a: &[E], b: &[E]) -> Vec<E>
378where
379 E: FieldElement,
380{
381 let mut apos = degree_of(a);
382 let mut a = a.to_vec();
383
384 let bpos = degree_of(b);
385 assert!(apos >= bpos, "cannot divide by polynomial of higher degree");
386 if bpos == 0 {
387 assert!(!b.is_empty(), "cannot divide by empty polynomial");
388 assert!(b[0] != E::ZERO, "cannot divide polynomial by zero");
389 }
390
391 let mut result = vec![E::ZERO; apos - bpos + 1];
392 for i in (0..result.len()).rev() {
393 let quot = a[apos] / b[bpos];
394 result[i] = quot;
395 for j in (0..bpos).rev() {
396 a[i + j] -= b[j] * quot;
397 }
398 apos = apos.wrapping_sub(1);
399 }
400
401 result
402}
403
404/// Returns a polynomial resulting from dividing a polynomial by a polynomial of special form.
405///
406/// Specifically, divides polynomial `p` by polynomial (x^`a` - `b`) using
407/// [synthetic division](https://en.wikipedia.org/wiki/Synthetic_division) method; if the
408/// polynomials don't divide evenly, the remainder is ignored. Polynomial `p` is expected
409/// to be in the coefficient form, and the result will be in the coefficient form as well.
410/// The length of the resulting polynomial will be equal to `p.len()`.
411///
412/// This function is significantly faster than the generic `polynom::div()` function.
413///
414/// # Panics
415/// Panics if:
416/// * `a` is zero;
417/// * `b` is zero;
418/// * `p.len()` is smaller than or equal to `a`.
419///
420/// # Examples
421/// ```
422/// # use winter_math::polynom::*;
423/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
424/// // p(x) = x^3 + x^2 + 2 * x + 2
425/// let p = [
426/// BaseElement::new(2),
427/// BaseElement::new(2),
428/// BaseElement::new(1),
429/// BaseElement::new(1),
430/// ];
431///
432/// // expected result = x^2 + 2
433/// let expected =
434/// vec![BaseElement::new(2), BaseElement::ZERO, BaseElement::new(1), BaseElement::ZERO];
435///
436/// // divide by x + 1
437/// assert_eq!(expected, syn_div(&p, 1, -BaseElement::ONE));
438/// ```
439pub fn syn_div<E>(p: &[E], a: usize, b: E) -> Vec<E>
440where
441 E: FieldElement,
442{
443 let mut result = p.to_vec();
444 syn_div_in_place(&mut result, a, b);
445 result
446}
447
448/// Divides a polynomial by a polynomial of special form and saves the result into the original
449/// polynomial.
450///
451/// Specifically, divides polynomial `p` by polynomial (x^`a` - `b`) using
452/// [synthetic division](https://en.wikipedia.org/wiki/Synthetic_division) method and saves the
453/// result into `p`. If the polynomials don't divide evenly, the remainder is ignored. Polynomial
454/// `p` is expected to be in the coefficient form, and the result will be in coefficient form as
455/// well.
456///
457/// This function is significantly faster than the generic `polynom::div()` function, and as
458/// compared to `polynom::syn_div()` function, this function does not allocate any additional
459/// memory.
460///
461/// # Panics
462/// Panics if:
463/// * `a` is zero;
464/// * `b` is zero;
465/// * `p.len()` is smaller than or equal to `a`.
466///
467/// # Examples
468/// ```
469/// # use winter_math::polynom::*;
470/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
471/// // p(x) = x^3 + x^2 + 2 * x + 2
472/// let mut p = [
473/// BaseElement::new(2),
474/// BaseElement::new(2),
475/// BaseElement::new(1),
476/// BaseElement::new(1),
477/// ];
478///
479/// // divide by x + 1
480/// syn_div_in_place(&mut p, 1, -BaseElement::ONE);
481///
482/// // expected result = x^2 + 2
483/// let expected = [
484/// BaseElement::new(2),
485/// BaseElement::ZERO,
486/// BaseElement::new(1),
487/// BaseElement::ZERO,
488/// ];
489///
490/// assert_eq!(expected, p);
491pub fn syn_div_in_place<E>(p: &mut [E], a: usize, b: E)
492where
493 E: FieldElement,
494{
495 assert!(a != 0, "divisor degree cannot be zero");
496 assert!(b != E::ZERO, "constant cannot be zero");
497 assert!(p.len() > a, "divisor degree cannot be greater than dividend size");
498
499 if a == 1 {
500 // if we are dividing by (x - `b`), we can use a single variable to keep track
501 // of the remainder; this way, we can avoid shifting the values in the slice later
502 let mut c = E::ZERO;
503 for coeff in p.iter_mut().rev() {
504 *coeff += b * c;
505 mem::swap(coeff, &mut c);
506 }
507 } else {
508 // if we are dividing by a polynomial of higher power, we need to keep track of the
509 // full remainder. we do that in place, but then need to shift the values at the end
510 // to discard the remainder
511 let degree_offset = p.len() - a;
512 if b == E::ONE {
513 // if `b` is 1, no need to multiply by `b` in every iteration of the loop
514 for i in (0..degree_offset).rev() {
515 p[i] += p[i + a];
516 }
517 } else {
518 for i in (0..degree_offset).rev() {
519 p[i] += p[i + a] * b;
520 }
521 }
522 // discard the remainder
523 p.copy_within(a.., 0);
524 p[degree_offset..].fill(E::ZERO);
525 }
526}
527
528/// Divides a polynomial by a polynomial given its roots and saves the result into the original
529/// polynomial.
530///
531/// Specifically, divides polynomial `p` by polynomial \prod_{i = 1}^m (x - `x_i`) using
532/// [synthetic division](https://en.wikipedia.org/wiki/Synthetic_division) method and saves the
533/// result into `p`. If the polynomials don't divide evenly, the remainder is ignored. Polynomial
534/// `p` is expected to be in the coefficient form, and the result will be in coefficient form as
535/// well.
536///
537/// This function is significantly faster than the generic `polynom::div()` function, using
538/// the coefficients of the divisor.
539///
540/// # Panics
541/// Panics if:
542/// * `roots.len()` is zero;
543/// * `p.len()` is smaller than or equal to `roots.len()`.
544///
545/// # Examples
546/// ```
547/// # use winter_math::polynom::*;
548/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
549/// // p(x) = x^3 - 7 * x + 6
550/// let mut p = [
551/// BaseElement::new(6),
552/// -BaseElement::new(7),
553/// BaseElement::new(0),
554/// BaseElement::new(1),
555/// ];
556///
557/// // divide by (x - 1) * (x - 2)
558/// let zeros = vec![BaseElement::new(1), BaseElement::new(2)];
559/// syn_div_roots_in_place(&mut p, &zeros);
560///
561/// // expected result = x + 3
562/// let expected = [
563/// BaseElement::new(3),
564/// BaseElement::new(1),
565/// BaseElement::ZERO,
566/// BaseElement::ZERO,
567/// ];
568///
569/// assert_eq!(expected, p);
570pub fn syn_div_roots_in_place<E>(p: &mut [E], roots: &[E])
571where
572 E: FieldElement,
573{
574 assert!(!roots.is_empty(), "divisor should contain at least one linear factor");
575 assert!(p.len() > roots.len(), "divisor degree cannot be greater than dividend size");
576
577 for root in roots {
578 let mut c = E::ZERO;
579 for coeff in p.iter_mut().rev() {
580 *coeff += *root * c;
581 mem::swap(coeff, &mut c);
582 }
583 }
584}
585
586// DEGREE INFERENCE
587// ================================================================================================
588
589/// Returns the degree of the provided polynomial.
590///
591/// If the size of the provided slice is much larger than the degree of the polynomial (i.e.,
592/// a large number of leading coefficients is ZERO), this operation can be quite inefficient.
593///
594/// # Examples
595/// ```
596/// # use winter_math::polynom::*;
597/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
598/// assert_eq!(0, degree_of::<BaseElement>(&[]));
599/// assert_eq!(0, degree_of(&[BaseElement::ONE]));
600/// assert_eq!(1, degree_of(&[BaseElement::ONE, BaseElement::new(2)]));
601/// assert_eq!(1, degree_of(&[BaseElement::ONE, BaseElement::new(2), BaseElement::ZERO]));
602/// assert_eq!(2, degree_of(&[BaseElement::ONE, BaseElement::new(2), BaseElement::new(3)]));
603/// assert_eq!(
604/// 2,
605/// degree_of(&[BaseElement::ONE, BaseElement::new(2), BaseElement::new(3), BaseElement::ZERO])
606/// );
607/// ```
608pub fn degree_of<E>(poly: &[E]) -> usize
609where
610 E: FieldElement,
611{
612 for i in (0..poly.len()).rev() {
613 if poly[i] != E::ZERO {
614 return i;
615 }
616 }
617 0
618}
619
620/// Returns a polynomial with all leading ZERO coefficients removed.
621///
622/// # Examples
623/// ```
624/// # use winter_math::polynom::*;
625/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
626/// let a = vec![1u128, 2, 3, 4, 5, 6, 0, 0]
627/// .into_iter()
628/// .map(BaseElement::new)
629/// .collect::<Vec<_>>();
630/// let b = remove_leading_zeros(&a);
631/// assert_eq!(6, b.len());
632/// assert_eq!(a[..6], b);
633///
634/// let a = vec![0u128, 0, 0, 0].into_iter().map(BaseElement::new).collect::<Vec<_>>();
635/// let b = remove_leading_zeros(&a);
636/// assert_eq!(0, b.len());
637/// ```
638pub fn remove_leading_zeros<E>(values: &[E]) -> Vec<E>
639where
640 E: FieldElement,
641{
642 for i in (0..values.len()).rev() {
643 if values[i] != E::ZERO {
644 return values[..(i + 1)].to_vec();
645 }
646 }
647 vec![]
648}
649
650/// Returns the coefficients of polynomial given its roots.
651///
652/// # Examples
653/// ```
654/// # use winter_math::polynom::*;
655/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
656/// let xs = vec![1u128, 2].into_iter().map(BaseElement::new).collect::<Vec<_>>();
657///
658/// let mut expected_poly = vec![2u128, 3, 1].into_iter().map(BaseElement::new).collect::<Vec<_>>();
659/// expected_poly[1] *= -BaseElement::ONE;
660///
661/// let poly = poly_from_roots(&xs);
662/// assert_eq!(expected_poly, poly);
663/// ```
664pub fn poly_from_roots<E: FieldElement>(xs: &[E]) -> Vec<E> {
665 let mut result = unsafe { utils::uninit_vector(xs.len() + 1) };
666 fill_zero_roots(xs, &mut result);
667 result
668}
669
670// HELPER FUNCTIONS
671// ================================================================================================
672
673fn fill_zero_roots<E: FieldElement>(xs: &[E], result: &mut [E]) {
674 let mut n = result.len();
675 n -= 1;
676 result[n] = E::ONE;
677
678 for i in 0..xs.len() {
679 n -= 1;
680 result[n] = E::ZERO;
681 #[allow(clippy::assign_op_pattern)]
682 for j in n..xs.len() {
683 result[j] = result[j] - result[j + 1] * xs[i];
684 }
685 }
686}