clean up api

2 files changed, 134 insertions, 67 deletions

diff --git a/src/difference.rs b/src/difference.rs
new file mode 100644
index 0000000..56e40a0
--- /dev/null
+++ b/src/difference.rs
@@ -0,0 +1,28 @@
+//! A set of difference functions you can use with [SquasherBuilder::difference]
+
+use rgb::RGB8;
+
+/// A naïve comparison just summing the channel differences
+/// I.E. `|a.red - b.red| + |a.green - b.green| + |a.blue - b.blue|`
+#[allow(clippy::many_single_char_names)]
+#[inline(always)]
+pub fn rgb_difference(a: &RGB8, b: &RGB8) -> f32 {
+    let absdiff = |a: u8, b: u8| (a as f32 - b as f32).abs();
+    absdiff(a.r, b.r) + absdiff(a.g, b.g) + absdiff(a.b, b.b)
+}
+
+// https://en.wikipedia.org/wiki/Color_difference#sRGB
+#[inline(always)]
+pub fn redmean_difference(a: &RGB8, b: &RGB8) -> f32 {
+    let delta_r = a.r as f32 - b.r as f32;
+    let delta_g = a.g as f32 - b.g as f32;
+    let delta_b = a.b as f32 - b.b as f32;
+    // reasonably sure calling it prime is wrong, but
+    let r_prime = 0.5 * (a.r as f32 + b.r as f32);
+
+    let red_part = (2.0 + (r_prime / 256.0)) * (delta_r * delta_r);
+    let green_part = 4.0 * (delta_g * delta_g);
+    let blue_part = (2.0 + (255.0 - r_prime) / 256.0) * (delta_b * delta_b);
+
+    (red_part + green_part + blue_part).sqrt()
+}
diff --git a/src/lib.rs b/src/lib.rs
index c16bfb2..8a709e9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,49 +1,111 @@
 use rgb::RGB8;
 use std::collections::HashMap;
 
+pub mod difference;
+
 type DiffFn = dyn Fn(&RGB8, &RGB8) -> f32;
 
+pub struct SquasherBuilder<T> {
+    max_colours: T,
+    difference_fn: Box<DiffFn>,
+    tolerance: f32,
+}
+
+impl<T: Count> SquasherBuilder<T> {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// The max number of colors selected for the palette, minus one.
+    ///
+    /// `max_colors(255)` will attempt to make a 256 color palette
+    pub fn max_colors(mut self, max_minus_one: T) -> SquasherBuilder<T> {
+        self.max_colours = max_minus_one;
+        self
+    }
+
+    /// The function to use to compare colours.
+    ///
+    /// see the [difference] module for functions included with the crate.
+    pub fn difference(mut self, difference: &'static DiffFn) -> SquasherBuilder<T> {
+        self.difference_fn = Box::new(difference);
+        self
+    }
+
+    /// Percent colours have to differ by to be included into the palette.
+    /// between and including 0.0 to 100.0
+    pub fn tolerance(mut self, percent: f32) -> SquasherBuilder<T> {
+        self.tolerance = percent;
+        self
+    }
+
+    pub fn build(self, image: &[u8]) -> Squasher<T> {
+        let mut squasher =
+            Squasher::from_parts(self.max_colours, self.difference_fn, self.tolerance);
+        squasher.recolor(image);
+
+        squasher
+    }
+}
+
+impl<T: Count> Default for SquasherBuilder<T> {
+    fn default() -> Self {
+        Self {
+            max_colours: T::from_usize(255),
+            difference_fn: Box::new(difference::rgb_difference),
+            tolerance: 2.0,
+        }
+    }
+}
+
 pub struct Squasher<T> {
+    // one less than the max colours as you can't have a zero colour image.
+    max_colours_min1: T,
     palette: Vec<RGB8>,
     map: Vec<T>,
     difference_fn: Box<DiffFn>,
+    tolerance_percent: f32,
 }
 
 impl<T: Count> Squasher<T> {
     /// Creates a new squasher and allocates a new color map. A color map
     /// contains every 24-bit color and ends up with an amount of memory
     /// equal to `16MB * std::mem::size_of(T)`.
-    ///
-    ///
     pub fn new(max_colors_minus_one: T, buffer: &[u8]) -> Self {
-        let sorted = Self::unique_and_sort(buffer);
-        Self::from_sorted(max_colors_minus_one, sorted, Box::new(rgb_difference))
+        let mut this = Self::from_parts(
+            max_colors_minus_one,
+            Box::new(difference::rgb_difference),
+            2.0,
+        );
+        this.recolor(buffer);
+
+        this
     }
 
-    /// Like [Squasher::new] but lets you pass your own difference function
-    /// to compare values while selecting colours. The default difference
-    /// function sums to difference between the RGB channels.
-    pub fn new_with_difference(
-        max_colors_minus_one: T,
-        buffer: &[u8],
-        difference_fn: &'static DiffFn,
-    ) -> Self {
-        let sorted = Self::unique_and_sort(buffer);
-        Self::from_sorted(max_colors_minus_one, sorted, Box::new(difference_fn))
+    pub fn builder() -> SquasherBuilder<T> {
+        SquasherBuilder::new()
     }
 
-    fn from_sorted(max_colors: T, sorted: Vec<(RGB8, usize)>, difference_fn: Box<DiffFn>) -> Self {
-        let selected = Self::select_colors(&sorted, max_colors, difference_fn.as_ref());
+    pub fn set_tolerance(&mut self, percent: f32) {
+        self.tolerance_percent = percent;
+    }
 
-        let mut this = Self {
-            palette: selected,
+    /// Create a new palette from the colours in the given image.
+    pub fn recolor(&mut self, image: &[u8]) {
+        let sorted = Self::unique_and_sort(image);
+        let selected = self.select_colors(sorted);
+        self.palette = selected;
+    }
+
+    /// Create a Squasher from parts. Noteably, this leave your palette empty
+    fn from_parts(max_colours_min1: T, difference_fn: Box<DiffFn>, tolerance: f32) -> Self {
+        Self {
+            max_colours_min1,
+            palette: vec![],
             map: vec![T::zero(); 256 * 256 * 256],
             difference_fn,
-        };
-
-        this.map_selected(&sorted);
-
-        this
+            tolerance_percent: tolerance,
+        }
     }
 
     /// Take an RGB image buffer and an output buffer. The function will fill
@@ -97,7 +159,7 @@ impl<T: Count> Squasher<T> {
     }
 
     /// Takes an image buffer of RGB data and fill the color map
-    fn unique_and_sort(buffer: &[u8]) -> Vec<(RGB8, usize)> {
+    fn unique_and_sort(buffer: &[u8]) -> Vec<RGB8> {
         let mut colors: HashMap<RGB8, usize> = HashMap::default();
 
         //count pixels
@@ -115,7 +177,7 @@ impl<T: Count> Squasher<T> {
         Self::sort(colors)
     }
 
-    fn sort(map: HashMap<RGB8, usize>) -> Vec<(RGB8, usize)> {
+    fn sort(map: HashMap<RGB8, usize>) -> Vec<RGB8> {
         let mut sorted: Vec<(RGB8, usize)> = map.into_iter().collect();
         sorted.sort_by(|(colour1, freq1), (colour2, freq2)| {
             freq2
@@ -125,40 +187,42 @@ impl<T: Count> Squasher<T> {
                 .then(colour2.b.cmp(&colour1.b))
         });
 
-        sorted
+        sorted.into_iter().map(|(color, _count)| color).collect()
     }
 
-    fn select_colors(sorted: &[(RGB8, usize)], max_colors: T, difference: &DiffFn) -> Vec<RGB8> {
+    /// Pick the colors in the palette from a Vec of colors sorted by number
+    /// of times they occur, high to low.
+    fn select_colors(&self, sorted: Vec<RGB8>) -> Vec<RGB8> {
         // I made these numbers up
         #[allow(non_snake_case)]
-        //let RGB_TOLERANCE: f32 = 0.01 * 768.0;
-        let RGB_TOLERANCE: f32 = 36.0;
-        let mut selected_colors: Vec<(RGB8, usize)> = Vec::with_capacity(max_colors.as_usize());
-
-        for (key, count) in sorted.iter() {
-            if max_colors.le(&selected_colors.len().saturating_sub(1)) {
+        //let RGB_TOLERANCE: f32 = 0.01 * 765.0;
+        //let RGB_TOLERANCE: f32 = 36.0;
+        let tolerance = (self.tolerance_percent / 100.0) * 765.0;
+        let max_colours = self.max_colours_min1.as_usize() + 1;
+        let mut selected_colors: Vec<RGB8> = Vec::with_capacity(max_colours);
+
+        for sorted_color in sorted {
+            if max_colours > selected_colors.len() {
                 break;
             } else if selected_colors
                 .iter()
-                .all(|color| difference(key, &color.0) > RGB_TOLERANCE)
+                .all(|color| (self.difference_fn)(&sorted_color, color) > tolerance)
             {
-                selected_colors.push((*key, *count));
+                selected_colors.push(sorted_color);
             }
         }
 
         selected_colors
-            .into_iter()
-            .map(|(color, _count)| color)
-            .collect()
     }
 
-    fn map_selected(&mut self, sorted: &[(RGB8, usize)]) {
-        for (sorted, _) in sorted {
+    /// Pick the closest colour in the palette for each unique color in the image
+    fn map_selected(&mut self, sorted: &[RGB8]) {
+        for colour in sorted {
             let mut min_diff = f32::MAX;
             let mut min_index = usize::MAX;
 
             for (index, selected) in self.palette.iter().enumerate() {
-                let diff = (self.difference_fn)(sorted, selected);
+                let diff = (self.difference_fn)(colour, selected);
 
                 if diff.max(0.0) < min_diff {
                     min_diff = diff;
@@ -166,7 +230,7 @@ impl<T: Count> Squasher<T> {
                 }
             }
 
-            self.map[color_index(sorted)] = T::from_usize(min_index);
+            self.map[color_index(colour)] = T::from_usize(min_index);
         }
     }
 }
@@ -231,28 +295,3 @@ count_impl!(usize);
 fn color_index(c: &RGB8) -> usize {
     c.r as usize * (256 * 256) + c.g as usize * 256 + c.b as usize
 }
-
-/// The default comparison function. Returns a sum of the channel differences.
-/// I.E. `|a.red - b.red| + |a.green - b.green| + |a.blue - b.blue|`
-#[allow(clippy::many_single_char_names)]
-#[inline(always)]
-pub fn rgb_difference(a: &RGB8, b: &RGB8) -> f32 {
-    let absdiff = |a: u8, b: u8| (a as f32 - b as f32).abs();
-    absdiff(a.r, b.r) + absdiff(a.g, b.g) + absdiff(a.b, b.b)
-}
-
-// https://en.wikipedia.org/wiki/Color_difference#sRGB
-#[inline(always)]
-pub fn redmean_difference(a: &RGB8, b: &RGB8) -> f32 {
-    let delta_r = a.r as f32 - b.r as f32;
-    let delta_g = a.g as f32 - b.g as f32;
-    let delta_b = a.b as f32 - b.b as f32;
-    // reasonably sure calling it prime is wrong, but
-    let r_prime = 0.5 * (a.r as f32 + b.r as f32);
-
-    let red_part = (2.0 + (r_prime / 256.0)) * (delta_r * delta_r);
-    let green_part = 4.0 * (delta_g * delta_g);
-    let blue_part = (2.0 + (255.0 - r_prime) / 256.0) * (delta_b * delta_b);
-
-    (red_part + green_part + blue_part).sqrt()
-}