kmeans: Split k-independent inner loop into its own function

This reduces code size of the k-means implementation by 28%.
Hot instructions are concentrated in a shared block which is
only 5.1% of total k-means code bytes.

Before:
  Size Name
1,277B rav1e::util::kmeans::kmeans
1,229B rav1e::util::kmeans::kmeans
1,163B rav1e::util::kmeans::kmeans
1,085B rav1e::util::kmeans::kmeans
1,080B rav1e::util::kmeans::kmeans
1,004B rav1e::util::kmeans::kmeans
6,838B filtered data size

After:
  Size Name
1,001B rav1e::util::kmeans::kmeans
  891B rav1e::util::kmeans::kmeans
  768B rav1e::util::kmeans::kmeans
  766B rav1e::util::kmeans::kmeans
  669B rav1e::util::kmeans::kmeans
  573B rav1e::util::kmeans::kmeans
  251B rav1e::util::kmeans::scan
4,919B filtered data size

Author: barrbrain

src/util/kmeans.rs

@@ -27,9 +27,6 @@ where
   high[K - 1] = data.len();
   sum[K - 1] = means[K - 1].into();
 
-  let data_to = |n: usize| unsafe { data.get_unchecked(..n) }.iter();
-  let data_from = |n: usize| unsafe { data.get_unchecked(n..) }.iter();
-
   // Constrain complexity to O(n log n)
   let limit = 2 * (usize::BITS - data.len().leading_zeros());
   for _ in 0..limit {
@@ -40,31 +37,9 @@ where
       .zip(low.iter_mut().skip(1).zip(&mut high))
       .enumerate()
     {
-      let mut n = *high;
-      let mut s = sum[i];
-      for &d in data_to(n).rev().take_while(|&d| *d > threshold) {
-        s -= d.into();
-        n -= 1;
-      }
-      for &d in data_from(n).take_while(|&d| *d <= threshold) {
-        s += d.into();
-        n += 1;
-      }
-      *high = n;
-      sum[i] = s;
-
-      let mut n = *low;
-      let mut s = sum[i + 1];
-      for &d in data_from(n).take_while(|&d| *d < threshold) {
-        s -= d.into();
-        n += 1;
-      }
-      for &d in data_to(n).rev().take_while(|&d| *d >= threshold) {
-        s += d.into();
-        n -= 1;
+      unsafe {
+        scan(high, low, sum.get_unchecked_mut(i..=i + 1), data, threshold);
       }
-      *low = n;
-      sum[i + 1] = s;
     }
     let mut changed = false;
     for (((m, sum), high), low) in
@@ -90,6 +65,42 @@ where
   means
 }
 
+#[inline(never)]
+unsafe fn scan<T>(
+  high: &mut usize, low: &mut usize, sum: &mut [i64], data: &[T], t: T,
+) where
+  T: Copy,
+  T: Into<i64>,
+  T: PartialEq,
+  T: PartialOrd,
+{
+  let mut n = *high;
+  let mut s = *sum.get_unchecked(0);
+  for &d in data.get_unchecked(..n).iter().rev().take_while(|&d| *d > t) {
+    s -= d.into();
+    n -= 1;
+  }
+  for &d in data.get_unchecked(n..).iter().take_while(|&d| *d <= t) {
+    s += d.into();
+    n += 1;
+  }
+  *high = n;
+  *sum.get_unchecked_mut(0) = s;
+
+  let mut n = *low;
+  let mut s = *sum.get_unchecked(1);
+  for &d in data.get_unchecked(n..).iter().take_while(|&d| *d < t) {
+    s -= d.into();
+    n += 1;
+  }
+  for &d in data.get_unchecked(..n).iter().rev().take_while(|&d| *d >= t) {
+    s += d.into();
+    n -= 1;
+  }
+  *low = n;
+  *sum.get_unchecked_mut(1) = s;
+}
+
 #[cfg(test)]
 mod test {
   use super::*;