poly_cool/
quadratic.rs

1use arrayvec::ArrayVec;
2
3#[derive(Copy, Clone, Debug, PartialEq)]
4pub struct Quadratic {
5    pub c0: f64,
6    pub c1: f64,
7    pub c2: f64,
8}
9
10impl std::ops::Mul<f64> for Quadratic {
11    type Output = Quadratic;
12
13    fn mul(self, rhs: f64) -> Self::Output {
14        Self {
15            c0: self.c0 * rhs,
16            c1: self.c1 * rhs,
17            c2: self.c2 * rhs,
18        }
19    }
20}
21
22impl std::ops::Div<f64> for Quadratic {
23    type Output = Quadratic;
24
25    fn div(self, rhs: f64) -> Self::Output {
26        Self {
27            c0: self.c0 / rhs,
28            c1: self.c1 / rhs,
29            c2: self.c2 / rhs,
30        }
31    }
32}
33
34impl Quadratic {
35    pub fn eval(&self, x: f64) -> f64 {
36        self.c0 + self.c1 * x + self.c2 * x * x
37    }
38
39    pub fn is_finite(&self) -> bool {
40        self.c0.is_finite() && self.c1.is_finite() && self.c2.is_finite()
41    }
42
43    pub fn roots(&self) -> ArrayVec<f64, 2> {
44        let a = self.c2;
45        let b = self.c1;
46        let c = self.c0;
47        let disc = b * b - 4.0 * a * c;
48        if disc.is_finite() {
49            let mut ret = ArrayVec::new();
50            let mut push = |r: f64| {
51                if r.is_finite() {
52                    ret.push(r)
53                }
54            };
55            if disc > 0.0 {
56                let q = -0.5 * (b + disc.sqrt().copysign(b));
57                let r0 = q / a;
58                let r1 = c / q;
59                push(r0.min(r1));
60                push(r0.max(r1));
61            } else if disc == 0.0 {
62                let root = -0.5 * b / a;
63                if root.is_finite() {
64                    push(root);
65                } else if c == 0.0 {
66                    // This is kurbo's behavior: the intention is that if the
67                    // whole thing is zero, return zero as a single root. I'm
68                    // not sure I love it.
69                    //
70                    // Bear in mind that this branch is not *only* for the
71                    // identically zero case: if a == c == 0.0 and b * b
72                    // underflows then we will end up here. In that case,
73                    // zero is the only root.
74                    push(0.0);
75                }
76            } else {
77                // No roots.
78            }
79            ret
80        } else {
81            // At least one of the coefficients was too large and triggered
82            // overflow.
83            //
84            // The exponent of f64 maxes out at 1023, so scaling down by
85            // 2^{-512} is enough to ensure that squaring doesn't overflow. We
86            // do an extra factor of 2^{-3} for some wiggle room. This can't
87            // completely destroy all the coefficients: because of the overflow,
88            // we know that at least one of them was big.
89            let scale = 2.0f64.powi(-515);
90            // TODO: this can stack overflow if we're infinite. How should
91            // we handle that?
92            (*self * scale).roots()
93        }
94    }
95
96    pub fn positive_discriminant_roots(&self) -> Option<(f64, f64)> {
97        let a = self.c2;
98        let b = self.c1;
99        let c = self.c0;
100        let disc = b * b - 4.0 * a * c;
101        if disc.is_finite() {
102            if disc > 0.0 {
103                let q = -0.5 * (b + disc.sqrt().copysign(b));
104                let r0 = q / a;
105                let r1 = c / q;
106                Some((r0.min(r1), r0.max(r1)))
107            } else {
108                None
109            }
110        } else {
111            self.positive_discriminant_roots_scaled()
112        }
113    }
114
115    #[cold]
116    fn positive_discriminant_roots_scaled(&self) -> Option<(f64, f64)> {
117        if self.is_finite() {
118            let scale = 2.0f64.powi(-515);
119            (*self * scale).positive_discriminant_roots()
120        } else {
121            None
122        }
123    }
124
125    pub fn positive_discriminant_roots_no_overflow_check(&self) -> Option<(f64, f64)> {
126        let a = self.c2;
127        let b = self.c1;
128        let c = self.c0;
129        let disc = b * b - 4.0 * a * c;
130        if disc > 0.0 {
131            let q = -0.5 * (b + disc.sqrt().copysign(b));
132            let r0 = q / a;
133            let r1 = c / q;
134            Some((r0.min(r1), r0.max(r1)))
135        } else {
136            None
137        }
138    }
139
140    pub fn positive_discriminant_roots_no_overflow_check_half_b(&self) -> Option<(f64, f64)> {
141        let a = self.c2;
142        let b = self.c1;
143        let c = self.c0;
144        let disc = b * b - a * c;
145        if disc > 0.0 {
146            let q = -(b + disc.sqrt().copysign(b));
147            let r0 = q / a;
148            let r1 = c / q;
149            Some((r0.min(r1), r0.max(r1)))
150        } else {
151            None
152        }
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    #[test]
159    fn root_evaluation() {
160        arbtest::arbtest(|u| {
161            let q = crate::arbitrary::quadratic(u)?;
162            // Arbitrary quadratics can have coefficients with wild magnitudes,
163            // so we need to adjust our error expectations accordingly.
164            let magnitude = q.c0.abs().max(q.c1.abs()).max(q.c2.abs()).max(1.0);
165
166            for r in q.roots() {
167                let y = q.eval(r);
168                // To evaluate the polynomial, we need to square r, so our error
169                // should be relative to the magnitude of r squared.
170                let r_magnitude = r.abs().max(1.0);
171                let threshold = r_magnitude * 1e-14 * r_magnitude * magnitude;
172                if y.is_finite() && threshold.is_finite() {
173                    assert!(y.abs() <= threshold);
174                }
175            }
176            Ok(())
177        })
178        .budget_ms(5_000);
179    }
180}
poly_cool/quadratic.rs

poly_cool/
quadratic.rs
