poly/
lib.rs

1use core::convert::TryFrom;
2use core::ops::{AddAssign, Div, Mul, MulAssign, SubAssign};
3use num::traits::{MulAddAssign, Pow};
4use num::Zero;
5use smallvec::SmallVec;
6
7mod add;
8mod display;
9mod mul;
10mod sub;
11pub use add::*;
12pub use display::*;
13pub use mul::*;
14pub use sub::*;
15
16#[cfg(debug_assertions)]
17macro_rules! assert_assume {
18	($cond:expr) => {
19		assert!($cond);
20	};
21	($cond:expr, $($arg:tt)+) => {
22		assert!($cond, $($arg)+);
23	};
24}
25
26#[cfg(not(debug_assertions))]
27macro_rules! assert_assume {
28	($cond:expr) => {
29		if !($cond) {
30			unsafe {
31				std::hint::unreachable_unchecked();
32				}
33			}
34	};
35	($cond:expr, $($arg:tt)+) => {
36		if !($cond) {
37			unsafe {
38				std::hint::unreachable_unchecked();
39				}
40			}
41	};
42}
43
44#[macro_export]
45macro_rules! coefficients {
46	($elem:expr; $n:expr) => ({
47		use smallvec::{smallvec, SmallVec};
48		let ret : SmallVec<[_; 8]> = smallvec![$elem; $n];
49		ret
50	});
51	($($x:expr),*$(,)*) => ({
52		use smallvec::{smallvec, SmallVec};
53		let ret : SmallVec<[_; 8]> = smallvec![$($x,)*];
54		ret
55	});
56}
57
58#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
59pub struct Polynomial<T> {
60	rev_coeffs: SmallVec<[T; 8]>,
61}
62
63impl<T> Polynomial<T> {
64	pub fn order(&self) -> i32 {
65		assert_assume!(!self.rev_coeffs.is_empty());
66		assert_assume!(i32::try_from(self.rev_coeffs.len() - 1).is_ok());
67		(self.rev_coeffs.len() - 1) as i32
68	}
69
70	pub fn reverse_coeffs(&self) -> &SmallVec<[T; 8]> {
71		&self.rev_coeffs
72	}
73
74	pub fn into_coeffs(self) -> SmallVec<[T; 8]> {
75		let mut coeffs = self.rev_coeffs;
76		coeffs.reverse();
77		coeffs
78	}
79
80	pub fn into_reverse_coeffs(self) -> SmallVec<[T; 8]> {
81		self.rev_coeffs
82	}
83
84	pub fn coeffs(&self) -> SmallVec<[T; 8]>
85	where
86		T: Clone,
87	{
88		let mut coeffs = self.rev_coeffs.clone();
89		coeffs.reverse();
90		coeffs
91	}
92
93	fn fixup_coefficients(&mut self)
94	where
95		T: Zero,
96	{
97		while {
98			if let Some(x) = self.rev_coeffs.last() {
99				x.is_zero()
100			} else {
101				false
102			}
103		} {
104			self.rev_coeffs.pop();
105		}
106		if self.rev_coeffs.is_empty() {
107			self.rev_coeffs.push(T::zero());
108		}
109		assert!(i32::try_from(self.rev_coeffs.len() - 1).is_ok());
110	}
111
112	pub fn new(coefficients: SmallVec<[T; 8]>) -> Self
113	where
114		T: Zero,
115	{
116		let mut rev_coeffs = coefficients;
117		rev_coeffs.reverse();
118		Self::new_reversed(rev_coeffs)
119	}
120
121	pub fn new_reversed(rev_coeffs: SmallVec<[T; 8]>) -> Self
122	where
123		T: Zero,
124	{
125		let mut ret = Self { rev_coeffs };
126		ret.fixup_coefficients();
127		ret
128	}
129
130	pub fn eval<X, Y>(&self, x: X) -> Y
131	where
132		T: Zero,
133		for<'l, 'r> &'l T: Mul<&'r X, Output = Y>,
134		X: for<'r> MulAssign<&'r X> + num::One,
135		Y: AddAssign + Zero,
136	{
137		let mut xn = X::one();
138		let mut y = Y::zero();
139		for a in self.rev_coeffs.iter() {
140			y += a * &xn;
141			xn *= &x;
142		}
143		y
144	}
145
146	pub fn eval_precise<X, Y>(&self, x: X) -> Y
147	where
148		T: Zero + Clone + Into<Y>,
149		for<'l> &'l X: Pow<i32, Output = X>,
150		Y: MulAddAssign<X, Y> + Zero,
151		// this trait is only required to guide type inference
152		for<'l> &'l T: Mul<X, Output = Y>,
153	{
154		let mut y = Y::zero();
155		for (a, e) in self.rev_coeffs.iter().zip(0i32..) {
156			// in num, MulAddAssign is only defined for values not references so we clone a
157			let mut ay = a.clone().into();
158			let mut y_old = Y::zero();
159			core::mem::swap(&mut y, &mut y_old);
160			ay.mul_add_assign(x.pow(e), y_old);
161			core::mem::swap(&mut y, &mut ay);
162		}
163		y
164	}
165
166	pub fn eval_der<X, Y>(&self, x: X, n: i32) -> Y
167	where
168		T: Zero + num::FromPrimitive + for<'r> Mul<&'r X, Output = Y>,
169		for<'l> &'l T: Mul<T, Output = T>,
170		X: for<'r> MulAssign<&'r X> + num::One,
171		Y: AddAssign + Zero,
172	{
173		assert!(n > 0);
174		let mut xn = X::one();
175		let mut y = Y::zero();
176		for ((a, e_old), e_new) in self
177			.rev_coeffs
178			.iter()
179			.zip(0i32..)
180			.skip(n as usize)
181			.zip(0i32..)
182		{
183			let mul = ((e_new + 1)..=e_old).product();
184			y += (a * T::from_i32(mul).unwrap()) * &xn;
185			xn *= &x;
186		}
187		y
188	}
189
190	pub fn eval_der_precise<X, Y>(&self, x: X, n: i32) -> Y
191	where
192		T: Zero + num::FromPrimitive + Into<Y>,
193		for<'l> &'l T: Mul<T, Output = T>,
194		for<'l> &'l X: Pow<i32, Output = X>,
195		Y: MulAddAssign<X, Y> + Zero + Clone,
196		// this trait is only required to guide type inference
197		for<'l> &'l T: Mul<X, Output = Y>,
198	{
199		let mut y = Y::zero();
200		for ((a, e_old), e_new) in self
201			.rev_coeffs
202			.iter()
203			.zip(0i32..)
204			.skip(n as usize)
205			.zip(0i32..)
206		{
207			let mul = ((e_new + 1)..=e_old).product();
208			let mut ay = (a * T::from_i32(mul).unwrap()).into();
209			let mut y_old = Y::zero();
210			core::mem::swap(&mut y, &mut y_old);
211			ay.mul_add_assign(x.pow(e_new), y_old);
212			core::mem::swap(&mut y, &mut ay);
213		}
214		y
215	}
216
217	pub fn div_rem(&self, rhs: &Self) -> (Self, Self)
218	where
219		T: Zero + Clone + for<'r> AddAssign<&'r T> + SubAssign,
220		for<'l, 'r> &'l T: Mul<&'r T, Output = T> + Div<&'r T, Output = T>,
221	{
222		assert!(!rhs.is_zero());
223
224		let order_l = self.order() as usize;
225		let order_r = rhs.order() as usize;
226		if order_l < order_r {
227			return (Self::zero(), self.clone());
228		}
229		let order_o = order_l - order_r;
230
231		let rhs = &rhs.rev_coeffs;
232		let mut remainder = self.rev_coeffs.clone();
233		let mut quotient = coefficients![T::zero(); order_o + 1];
234
235		for el in (order_r..=order_l).rev() {
236			let v = &remainder[el] / &rhs[order_r];
237			remainder[el] = T::zero();
238			for k in 1..=order_r {
239				remainder[el - k] -= &v * &rhs[order_r - k];
240			}
241			quotient[el - order_r] = v;
242		}
243
244		(Self::new_reversed(quotient), Self::new_reversed(remainder))
245	}
246}
247
248#[cfg(test)]
249mod tests {
250	use crate::*;
251
252	#[cfg(not(debug_assertions))]
253	use smallvec::SmallVec;
254
255	#[test]
256	fn test_new() {
257		let poly = Polynomial::new(coefficients![0f32, 0.0, 1.0, 2.0, 3.0, 0.0]);
258		assert_eq!(poly.order(), 3);
259		let mut coeffs = coefficients![1f32, 2.0, 3.0, 0.0];
260		assert_eq!(poly.coeffs(), coeffs);
261		coeffs.reverse();
262		assert_eq!(*poly.reverse_coeffs(), coeffs);
263	}
264
265	#[cfg(not(debug_assertions))]
266	#[test]
267	fn test_new_max_size() {
268		let mut coefficients = SmallVec::new();
269		coefficients.resize(
270			usize::try_from(core::i32::MAX)
271				.unwrap()
272				.checked_add(1)
273				.unwrap(),
274			1u8,
275		);
276		let poly = Polynomial::new_reversed(coefficients);
277		assert_eq!(poly.order(), core::i32::MAX);
278	}
279
280	#[cfg(not(debug_assertions))]
281	#[test]
282	#[should_panic(expected = "assertion failed: i32::try_from(self.rev_coeffs.len() - 1).is_ok()")]
283	fn test_new_oversized() {
284		let mut coefficients = SmallVec::new();
285		coefficients.resize(
286			usize::try_from(core::i32::MAX)
287				.unwrap()
288				.checked_add(2)
289				.unwrap(),
290			1u8,
291		);
292		Polynomial::new_reversed(coefficients);
293	}
294
295	#[test]
296	fn test_eval() {
297		let poly = Polynomial::new(coefficients![1f32, 2.0, 3.0, 0.0]);
298		assert_eq!(poly.eval(2f32), 22.0);
299		assert_eq!(
300			poly.eval(num::Complex::new(0f32, 1.0)),
301			num::Complex::new(-2f32, 2.0)
302		);
303		assert_eq!(poly.eval_precise(2f32), 22.0);
304	}
305
306	#[test]
307	fn test_eval_substitution() {
308		let mut poly = Polynomial::new(coefficients![1f32, 2.0, 3.0, 0.0]);
309		poly = poly.eval(Polynomial::new(coefficients![1f32, 0.0, 0.0])); // x → x²
310		assert_eq!(poly.order(), 6);
311		assert_eq!(
312			poly.into_coeffs(),
313			coefficients![1f32, 0.0, 2.0, 0.0, 3.0, 0.0, 0.0]
314		);
315		poly = Polynomial::new(coefficients![1f32, 2.0, 3.0, 0.0]);
316		poly = poly.eval(Polynomial::new(coefficients![1f32, 0.0, 1.0])); // x → x² + 1
317		assert_eq!(
318			poly.into_coeffs(),
319			coefficients![1f32, 0.0, 5.0, 0.0, 10.0, 0.0, 6.0]
320		);
321	}
322
323	#[test]
324	fn test_eval_der() {
325		// x³ + 2x² + 3x + 0
326		let poly = Polynomial::new(coefficients![1f32, 2.0, 3.0, 0.0]);
327		// 3x² + 4x + 3
328		assert_eq!(poly.eval_der(0f32, 1), 3.0);
329		assert_eq!(poly.eval_der(1f32, 1), 10.0);
330		assert_eq!(poly.eval_der(2f32, 1), 23.0);
331		assert_eq!(poly.eval_der_precise(0f32, 1), 3.0);
332		assert_eq!(poly.eval_der_precise(1f32, 1), 10.0);
333		assert_eq!(poly.eval_der_precise(2f32, 1), 23.0);
334		// 6x + 4
335		assert_eq!(poly.eval_der(0f32, 2), 4.0);
336		assert_eq!(poly.eval_der(1f32, 2), 10.0);
337		assert_eq!(poly.eval_der_precise(0f32, 2), 4.0);
338		assert_eq!(poly.eval_der_precise(1f32, 2), 10.0);
339	}
340
341	#[test]
342	fn test_div_rem() {
343		let poly_a = Polynomial::new(coefficients![1f32, 2.0, 3.0, 0.0]);
344
345		// (x³ + 2x² + 3x) / x = x² + 2x + 3 | 0
346		let poly_b = Polynomial::new(coefficients![1f32, 0.0]);
347		let (q, r) = poly_a.div_rem(&poly_b);
348		assert_eq!(q.coeffs(), coefficients![1f32, 2.0, 3.0]);
349		assert!(r.is_zero());
350
351		// (x³ + 2x² + 3x) / 2x = 0.5x² + x + 1.5 | 0
352		let poly_b = Polynomial::new(coefficients![2f32, 0.0]);
353		let (q, r) = poly_a.div_rem(&poly_b);
354		assert_eq!(q.coeffs(), coefficients![0.5f32, 1.0, 1.5]);
355		assert!(r.is_zero());
356
357		// (x³ + 2x² + 3x) / x² = x + 2 | 3x
358		let poly_b = Polynomial::new(coefficients![1f32, 0.0, 0.0]);
359		let (q, r) = poly_a.div_rem(&poly_b);
360		assert_eq!(q.coeffs(), coefficients![1f32, 2.0]);
361		assert_eq!(r.coeffs(), coefficients![3f32, 0.0]);
362
363		// (x³ + 2x² + 3x) / (x² + 1) = x + 2 | 2x - 2
364		let poly_b = Polynomial::new(coefficients![1f32, 0.0, 1.0]);
365		let (q, r) = poly_a.div_rem(&poly_b);
366		assert_eq!(q.coeffs(), coefficients![1f32, 2.0]);
367		assert_eq!(r.coeffs(), coefficients![2f32, -2.0]);
368
369		// (x² + 1) / (x³ + 2x² + 3x) = 0 | x² + 1
370		let (q, r) = poly_b.div_rem(&poly_a);
371		assert!(q.is_zero());
372		assert_eq!(r, poly_b);
373	}
374
375	#[test]
376	#[should_panic(expected = "assertion failed: !rhs.is_zero()")]
377	fn test_div_rem_zero() {
378		let poly_a = Polynomial::new(coefficients![1f32, 2.0, 3.0, 0.0]);
379		let poly_b = Polynomial::zero();
380		let _ = poly_a.div_rem(&poly_b);
381	}
382}