collect colours multiple times with colorcollector

1 files changed, 51 insertions, 1 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 43aec69..3af73f1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,6 +3,37 @@ use std::ops::Deref;
 
 use ahash::RandomState;
 
+pub struct ColorCollector {
+    colors: HashMap<Rgb, usize, RandomState>,
+}
+
+impl ColorCollector {
+    pub fn new() -> Self {
+        Self {
+            colors: HashMap::default(),
+        }
+    }
+
+    /// Wants an RGB buffer
+    pub fn add(&mut self, buffer: &[u8]) {
+        for pixel in buffer.chunks(3) {
+            let rgb = Rgb([pixel[0], pixel[1], pixel[2]]);
+
+            match self.colors.get_mut(&rgb) {
+                None => {
+                    self.colors.insert(rgb, 1);
+                }
+                Some(n) => *n += 1,
+            }
+        }
+    }
+
+    pub fn as_squasher<T: Count>(self, max_colors: T) -> Squasher<T> {
+        let sorted = Squasher::<T>::sort(self.colors);
+        Squasher::from_sorted(max_colors, sorted)
+    }
+}
+
 pub struct Squasher<T> {
     palette: Vec<(Rgb, usize)>,
     larget_count: usize,
@@ -15,6 +46,10 @@ impl<T: Count> Squasher<T> {
     /// equal to `16MB * std::mem::size_of(T)`
     pub fn new(max_colors: T, buffer: &[u8]) -> Self {
         let sorted = Self::unique_and_sort(buffer);
+        Self::from_sorted(max_colors, sorted)
+    }
+
+    fn from_sorted(max_colors: T, sorted: Vec<(Rgb, usize)>) -> Self {
         let selected = Self::select_colors(&sorted, max_colors);
 
         let mut this = Self {
@@ -43,6 +78,17 @@ impl<T: Count> Squasher<T> {
         }
     }
 
+    /// Take an RGB image buffer and an output buffer. The function will fill
+    /// the output buffer with indexes into the Palette.
+    //TODO: gen- Better name?
+    pub fn map_unsafe(&self, image: &[u8], buffer: &mut [T]) {
+        for (idx, color) in image.chunks(3).enumerate() {
+            let index = self.map[color_index(&Rgb([color[0], color[1], color[2]]))];
+
+            buffer[idx] = index;
+        }
+    }
+
     /// Retrieve the palette this squasher is working from
     pub fn palette(&self) -> Vec<Rgb> {
         self.palette.iter().map(|ahh| ahh.0).collect()
@@ -74,7 +120,11 @@ impl<T: Count> Squasher<T> {
             }
         }
 
-        let mut sorted: Vec<(Rgb, usize)> = colors.into_iter().collect();
+        Self::sort(colors)
+    }
+
+    fn sort(map: HashMap<Rgb, usize, RandomState>) -> Vec<(Rgb, usize)> {
+        let mut sorted: Vec<(Rgb, usize)> = map.into_iter().collect();
         sorted.sort_by(|(colour1, freq1), (colour2, freq2)| {
             freq2
                 .cmp(freq1)