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 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 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 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])); 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])); 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 let poly = Polynomial::new(coefficients![1f32, 2.0, 3.0, 0.0]);
327 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 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 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 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 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 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 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}