Skip to main content

poly_cool/
quadratic.rs

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