Morph: pick the next curve to split from a heap, not a full scan

crates/geom/src/morph.rs

@@ -126,6 +126,21 @@ fn split_curve(c: &[Point], n: usize, out: &mut Vec<Point>) {
     out.extend([p1, p2, p3]);
 }
 
+/// A piece length ordered totally, so it can sit in a heap.
+#[derive(PartialEq)]
+struct Share(f64);
+impl Eq for Share {}
+impl PartialOrd for Share {
+    fn partial_cmp(&self, o: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(o))
+    }
+}
+impl Ord for Share {
+    fn cmp(&self, o: &Self) -> std::cmp::Ordering {
+        self.0.total_cmp(&o.0)
+    }
+}
+
 /// Brings a contour to `n` curves. Unlike manim (which splits every curve equally often), the extra
 /// splits go to the longest pieces, so a long side gets more points than a short one.
 fn refine(c: &Contour, n: usize) -> Contour {
@@ -136,11 +151,15 @@ fn refine(c: &Contour, n: usize) -> Contour {
     let curves: Vec<&[Point]> = (0..have).filter_map(|k| c.pts.get(3 * k..3 * k + 4)).collect();
     let len: Vec<f64> = curves.iter().map(|q| curve_length(q) + 1e-9).collect();
     let mut parts = vec![1usize; curves.len()];
+    // The curve with the longest piece gets the next split (the lowest index on a tie); a heap keeps
+    // this O(n log n) instead of rescanning every curve for each split.
+    let mut heap: std::collections::BinaryHeap<(Share, std::cmp::Reverse<usize>)> =
+        len.iter().enumerate().map(|(k, l)| (Share(*l), std::cmp::Reverse(k))).collect();
     for _ in 0..n - have {
-        let share = |k: usize| len.get(k).copied().unwrap_or(0.0) / parts.get(k).copied().unwrap_or(1) as f64;
-        let best = (1..curves.len()).fold(0, |b, k| if share(k) > share(b) { k } else { b });
+        let Some((_, std::cmp::Reverse(best))) = heap.pop() else { break };
         if let Some(p) = parts.get_mut(best) {
             *p += 1;
+            heap.push((Share(len.get(best).copied().unwrap_or(0.0) / *p as f64), std::cmp::Reverse(best)));
         }
     }
     let mut out = Vec::with_capacity(1 + 3 * n);
@@ -420,6 +439,16 @@ mod tests {
         assert_eq!(contours(&to_path(&cs)), cs);
     }
 
+    #[test]
+    fn refine_many_curves_finishes_at_the_asked_count() {
+        let mut c = Contour { closed: false, pts: vec![Point::ZERO] };
+        for i in 1..=20_000 {
+            line_to(&mut c.pts, Point::new(f64::from(i), f64::from(i % 7)));
+        }
+        let r = refine(&c, 40_000);
+        assert_eq!(r.curves(), 40_000);
+    }
+
     #[test]
     fn refine_splits_the_longest_curves_and_keeps_the_shape() {
         let c = Contour { closed: false, pts: vec![Point::ZERO, Point::new(1.0, 0.0), Point::new(2.0, 0.0), Point::new(3.0, 0.0)] };