Merge pull request #16 from JulianKnodt/fix_bug

check 0 and refactor search netindex

src/lib.rs

@@ -89,6 +89,21 @@ const BETAGAMMA: f64 = BETA * GAMMA;
 // that it is divisible by all four primes
 const PRIMES: [usize; 4] = [499, 491, 487, 503];
 
+pub enum ControlFlow {
+    Break,
+    Continue,
+}
+
+impl ControlFlow {
+    fn is_break(self) -> bool {
+        if let ControlFlow::Break = self {
+            true
+        } else {
+            false
+        }
+    }
+}
+
 #[derive(Clone, Copy)]
 struct Quad<T> {
     r: T,
@@ -416,65 +431,70 @@ impl NeuQuant {
             self.netindex[j] = max_netpos
         } // really 256
     }
-
     /// Search for best matching color
     fn search_netindex(&self, b: u8, g: u8, r: u8, a: u8) -> usize {
-        let mut bestd = 1 << 30; // ~ 1_000_000
-        let mut best = 0;
-        // start at netindex[g] and work outwards
-        let mut i = self.netindex[g as usize];
-        let mut j = if i > 0 { i - 1 } else { 0 };
-
-        while (i < self.netsize) || (j > 0) {
-            if i < self.netsize {
-                let p = self.colormap[i];
-                let mut e = p.g - g as i32;
-                let mut dist = e * e; // inx key
-                if dist >= bestd {
-                    break;
-                } else {
-                    e = p.b - b as i32;
-                    dist += e * e;
-                    if dist < bestd {
-                        e = p.r - r as i32;
-                        dist += e * e;
-                        if dist < bestd {
-                            e = p.a - a as i32;
-                            dist += e * e;
-                            if dist < bestd {
-                                bestd = dist;
-                                best = i;
-                            }
-                        }
-                    }
-                    i += 1;
+        let mut best_dist = std::i32::MAX;
+        let first_guess = self.netindex[g as usize];
+        let mut best_pos = first_guess;
+        let mut i = best_pos;
+
+        #[inline]
+        fn sqr_dist(a: i32, b: u8) -> i32 {
+            let dist = a - b as i32;
+            dist * dist
+        }
+
+        {
+            let mut cmp = |i| {
+                let Quad {
+                    r: pr,
+                    g: pg,
+                    b: pb,
+                    a: pa,
+                } = self.colormap[i];
+                let mut dist = sqr_dist(pg, g);
+                if dist > best_dist {
+                    // If the green is less than optimal, break.
+                    // Seems to be arbitrary choice from the original implementation,
+                    // but can't change this w/o invasive changes since
+                    // we also sort by green.
+                    return ControlFlow::Break;
+                }
+                // otherwise, continue searching through the colormap.
+                dist += sqr_dist(pr, r);
+                if dist >= best_dist {
+                    return ControlFlow::Continue;
+                }
+                dist += sqr_dist(pb, b);
+                if dist >= best_dist {
+                    return ControlFlow::Continue;
+                }
+                dist += sqr_dist(pa, a);
+                if dist >= best_dist {
+                    return ControlFlow::Continue;
                 }
+                best_dist = dist;
+                best_pos = i;
+                ControlFlow::Continue
+            };
+            while i < self.netsize {
+                i = if cmp(i).is_break() { break } else { i + 1 };
             }
-            if j > 0 {
-                let p = self.colormap[j];
-                let mut e = p.g - g as i32;
-                let mut dist = e * e; // inx key
-                if dist >= bestd {
+            // this j < C is a cheat to avoid the bounds check, as when the loop reaches 0
+            // it will wrap to usize::MAX. Assume that self.netsize < usize::MAX, otherwise
+            // using a lot of memory. This also must be true since netsize < isize::MAX <
+            // usize::MAX, otherwise it would fail while allocating a Vec since they can be at
+            // most isize::MAX elements.
+            let mut j = first_guess.wrapping_sub(1);
+            while j < self.netsize {
+                j = if cmp(j).is_break() {
                     break;
                 } else {
-                    e = p.b - b as i32;
-                    dist += e * e;
-                    if dist < bestd {
-                        e = p.r - r as i32;
-                        dist += e * e;
-                        if dist < bestd {
-                            e = p.a - a as i32;
-                            dist += e * e;
-                            if dist < bestd {
-                                bestd = dist;
-                                best = j;
-                            }
-                        }
-                    }
-                    j -= 1;
-                }
+                    j.wrapping_sub(1)
+                };
             }
         }
-        best
+
+        best_pos
     }
 }