poly_cool/
quadratic.rs

1use arrayvec::ArrayVec;
2
3use crate::Quadratic;
4
5impl Quadratic {
6    /// This is like [`Quadratic::eval`] but faster.
7    ///
8    /// It would be nice if we could just make `eval` like this, but I couldn't
9    /// figure out how, given the lack of specialization.
10    #[doc(hidden)]
11    pub fn eval_opt(&self, x: f64) -> f64 {
12        let [c0, c1, c2] = self.coeffs;
13        c0 + c1 * x + c2 * x * x
14    }
15
16    /// Returns the roots of this quadratic, in increasing order.
17    ///
18    /// Double-roots are only counted once.
19    pub fn roots(&self) -> ArrayVec<f64, 2> {
20        let &[c, b, a] = self.coeffs();
21        let disc = b * b - 4.0 * a * c;
22        if disc.is_finite() {
23            let mut ret = ArrayVec::new();
24            let mut push = |r: f64| {
25                if r.is_finite() {
26                    ret.push(r)
27                }
28            };
29            if disc > 0.0 {
30                let q = -0.5 * (b + disc.sqrt().copysign(b));
31                let r0 = q / a;
32                let r1 = c / q;
33                push(r0.min(r1));
34                push(r0.max(r1));
35            } else if disc == 0.0 {
36                let root = -0.5 * b / a;
37                if root.is_finite() {
38                    push(root);
39                } else if c == 0.0 {
40                    // This is kurbo's behavior: the intention is that if the
41                    // whole thing is zero, return zero as a single root. I'm
42                    // not sure I love it.
43                    //
44                    // Bear in mind that this branch is not *only* for the
45                    // identically zero case: if a == c == 0.0 and b * b
46                    // underflows then we will end up here. In that case,
47                    // zero is the only root.
48                    push(0.0);
49                }
50            } else {
51                // No roots.
52            }
53            ret
54        } else {
55            // At least one of the coefficients was too large and triggered
56            // overflow.
57            //
58            // The exponent of f64 maxes out at 1023, so scaling down by
59            // 2^{-512} is enough to ensure that squaring doesn't overflow. We
60            // do an extra factor of 2^{-3} for some wiggle room. This can't
61            // completely destroy all the coefficients: because of the overflow,
62            // we know that at least one of them was big.
63            let scale = 2.0f64.powi(-515);
64            // If we're infinite, just give up. (Otherwise, we'd stack overflow
65            // by repeatedly trying to rescale.)
66            if self.is_finite() {
67                (*self * scale).roots()
68            } else {
69                ArrayVec::new()
70            }
71        }
72    }
73
74    /// Returns the two distinct roots of this quadratic, but only if the
75    /// discriminant is positive.
76    pub fn positive_discriminant_roots(&self) -> Option<(f64, f64)> {
77        let &[c, b, a] = self.coeffs();
78        let disc = b * b - 4.0 * a * c;
79        if disc.is_finite() {
80            if disc > 0.0 {
81                let q = -0.5 * (b + disc.sqrt().copysign(b));
82                let r0 = q / a;
83                let r1 = c / q;
84                Some((r0.min(r1), r0.max(r1)))
85            } else {
86                None
87            }
88        } else {
89            self.positive_discriminant_roots_scaled()
90        }
91    }
92
93    #[cold]
94    fn positive_discriminant_roots_scaled(&self) -> Option<(f64, f64)> {
95        if self.is_finite() {
96            let scale = 2.0f64.powi(-515);
97            (*self * scale).positive_discriminant_roots()
98        } else {
99            None
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    #[test]
107    fn root_evaluation() {
108        arbtest::arbtest(|u| {
109            let q = crate::arbitrary::quadratic(u)?;
110            // Arbitrary quadratics can have coefficients with wild magnitudes,
111            // so we need to adjust our error expectations accordingly.
112            let magnitude = q.magnitude().max(1.0);
113
114            for r in q.roots() {
115                let y = q.eval(r);
116                // To evaluate the polynomial, we need to square r, so our error
117                // should be relative to the magnitude of r squared.
118                let r_magnitude = r.abs().max(1.0);
119                let threshold = r_magnitude * 1e-14 * r_magnitude * magnitude;
120                if y.is_finite() && threshold.is_finite() {
121                    assert!(y.abs() <= threshold);
122                }
123            }
124            Ok(())
125        })
126        .budget_ms(5_000);
127    }
128}