Merge pull request #1 from tavianator/rust

Rewrite in rust

3 files changed, 364 insertions, 0 deletions

diff --git a/src/frontier/image.rs b/src/frontier/image.rs
new file mode 100644
index 0000000..3655580
--- /dev/null
+++ b/src/frontier/image.rs
@@ -0,0 +1,74 @@
+//! Frontier that targets an image.
+
+use super::{Frontier, Pixel};
+
+use crate::color::{ColorSpace, Rgb8};
+use crate::metric::soft::SoftKdTree;
+use crate::metric::NearestNeighbors;
+
+use image::RgbImage;
+
+/// A [Frontier] that places colors on the closest pixel of a target image.
+#[derive(Debug)]
+pub struct ImageFrontier<C: ColorSpace> {
+    nodes: SoftKdTree<Pixel<C>>,
+    width: u32,
+    height: u32,
+    len: usize,
+    deleted: usize,
+}
+
+impl<C: ColorSpace> ImageFrontier<C> {
+    /// Create an ImageFrontier from an image.
+    pub fn new(img: &RgbImage) -> Self {
+        let width = img.width();
+        let height = img.height();
+        let len = (width as usize) * (height as usize);
+
+        Self {
+            nodes: img
+                .enumerate_pixels()
+                .map(|(x, y, p)| Pixel::new(x, y, C::from(*p)))
+                .collect(),
+            width,
+            height,
+            len,
+            deleted: 0,
+        }
+    }
+}
+
+impl<C: ColorSpace> Frontier for ImageFrontier<C> {
+    fn width(&self) -> u32 {
+        self.width
+    }
+
+    fn height(&self) -> u32 {
+        self.height
+    }
+
+    fn len(&self) -> usize {
+        self.len - self.deleted
+    }
+
+    fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)> {
+        let color = C::from(rgb8);
+
+        if let Some(node) = self.nodes.nearest(&color).map(|n| n.item) {
+            let pos = node.pos;
+
+            node.delete();
+            self.deleted += 1;
+
+            if 32 * self.deleted >= self.len {
+                self.nodes.rebuild();
+                self.len -= self.deleted;
+                self.deleted = 0;
+            }
+
+            Some(pos)
+        } else {
+            None
+        }
+    }
+}
diff --git a/src/frontier/mean.rs b/src/frontier/mean.rs
new file mode 100644
index 0000000..889c5ba
--- /dev/null
+++ b/src/frontier/mean.rs
@@ -0,0 +1,140 @@
+//! Mean selection frontier.
+
+use super::{neighbors, Frontier, Pixel};
+
+use crate::color::{ColorSpace, Rgb8};
+use crate::metric::soft::SoftKdForest;
+use crate::metric::NearestNeighbors;
+
+use std::iter;
+use std::rc::Rc;
+
+/// A pixel on a mean frontier.
+#[derive(Debug)]
+enum MeanPixel<C> {
+    Empty,
+    Fillable(Rc<Pixel<C>>),
+    Filled(C),
+}
+
+impl<C: ColorSpace> MeanPixel<C> {
+    fn filled_color(&self) -> Option<C> {
+        match self {
+            Self::Filled(color) => Some(*color),
+            _ => None,
+        }
+    }
+}
+
+/// A [Frontier] that looks at the average color of each pixel's neighbors.
+pub struct MeanFrontier<C> {
+    pixels: Vec<MeanPixel<C>>,
+    forest: SoftKdForest<Rc<Pixel<C>>>,
+    width: u32,
+    height: u32,
+    len: usize,
+    deleted: usize,
+}
+
+impl<C: ColorSpace> MeanFrontier<C> {
+    /// Create a MeanFrontier with the given dimensions and initial pixel location.
+    pub fn new(width: u32, height: u32, x0: u32, y0: u32) -> Self {
+        let size = (width as usize) * (height as usize);
+        let mut pixels = Vec::with_capacity(size);
+        for _ in 0..size {
+            pixels.push(MeanPixel::Empty);
+        }
+
+        let pixel0 = Rc::new(Pixel::new(x0, y0, C::from(Rgb8::from([0, 0, 0]))));
+        let i = (x0 + y0 * width) as usize;
+        pixels[i] = MeanPixel::Fillable(pixel0.clone());
+
+        Self {
+            pixels,
+            forest: iter::once(pixel0).collect(),
+            width,
+            height,
+            len: 1,
+            deleted: 0,
+        }
+    }
+
+    fn pixel_index(&self, x: u32, y: u32) -> usize {
+        debug_assert!(x < self.width);
+        debug_assert!(y < self.height);
+
+        (x + y * self.width) as usize
+    }
+
+    fn fill(&mut self, x: u32, y: u32, color: C) {
+        let i = self.pixel_index(x, y);
+        match &self.pixels[i] {
+            MeanPixel::Empty => {}
+            MeanPixel::Fillable(pixel) => {
+                pixel.delete();
+                self.deleted += 1;
+            }
+            _ => unreachable!(),
+        }
+        self.pixels[i] = MeanPixel::Filled(color);
+
+        let mut pixels = Vec::new();
+        for &(x, y) in &neighbors(x, y) {
+            if x < self.width && y < self.height {
+                let i = self.pixel_index(x, y);
+                match &self.pixels[i] {
+                    MeanPixel::Empty => {}
+                    MeanPixel::Fillable(pixel) => {
+                        pixel.delete();
+                        self.deleted += 1;
+                    }
+                    MeanPixel::Filled(_) => continue,
+                }
+                let color = C::average(
+                    neighbors(x, y)
+                        .iter()
+                        .filter(|(x, y)| *x < self.width && *y < self.height)
+                        .map(|(x, y)| self.pixel_index(*x, *y))
+                        .map(|i| &self.pixels[i])
+                        .map(MeanPixel::filled_color)
+                        .flatten(),
+                );
+                let pixel = Rc::new(Pixel::new(x, y, color));
+                self.pixels[i] = MeanPixel::Fillable(pixel.clone());
+                pixels.push(pixel);
+            }
+        }
+
+        self.len += pixels.len();
+        self.forest.extend(pixels);
+
+        if 2 * self.deleted >= self.len {
+            self.forest.rebuild();
+            self.len -= self.deleted;
+            self.deleted = 0;
+        }
+    }
+}
+
+impl<C: ColorSpace> Frontier for MeanFrontier<C> {
+    fn width(&self) -> u32 {
+        self.width
+    }
+
+    fn height(&self) -> u32 {
+        self.height
+    }
+
+    fn len(&self) -> usize {
+        self.len - self.deleted
+    }
+
+    fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)> {
+        let color = C::from(rgb8);
+        let (x, y) = self.forest.nearest(&color).map(|n| n.item.pos)?;
+
+        self.fill(x, y, color);
+
+        Some((x, y))
+    }
+}
diff --git a/src/frontier/min.rs b/src/frontier/min.rs
new file mode 100644
index 0000000..b22b290
--- /dev/null
+++ b/src/frontier/min.rs
@@ -0,0 +1,150 @@
+//! Minimum selection frontier.
+
+use super::{neighbors, Frontier, Pixel};
+
+use crate::color::{ColorSpace, Rgb8};
+use crate::metric::soft::SoftKdForest;
+use crate::metric::NearestNeighbors;
+
+use rand::Rng;
+
+use std::rc::Rc;
+
+/// A pixel on a min frontier.
+#[derive(Debug)]
+struct MinPixel<C> {
+    pixel: Option<Rc<Pixel<C>>>,
+    filled: bool,
+}
+
+impl<C: ColorSpace> MinPixel<C> {
+    fn new() -> Self {
+        Self {
+            pixel: None,
+            filled: false,
+        }
+    }
+}
+
+/// A [Frontier] that places colors on a neighbor of the closest pixel so far.
+#[derive(Debug)]
+pub struct MinFrontier<C, R> {
+    rng: R,
+    pixels: Vec<MinPixel<C>>,
+    forest: SoftKdForest<Rc<Pixel<C>>>,
+    width: u32,
+    height: u32,
+    x0: u32,
+    y0: u32,
+    len: usize,
+    deleted: usize,
+}
+
+impl<C: ColorSpace, R: Rng> MinFrontier<C, R> {
+    /// Create a MinFrontier with the given dimensions and initial pixel location.
+    pub fn new(rng: R, width: u32, height: u32, x0: u32, y0: u32) -> Self {
+        let size = (width as usize) * (height as usize);
+        let mut pixels = Vec::with_capacity(size);
+        for _ in 0..size {
+            pixels.push(MinPixel::new());
+        }
+
+        Self {
+            rng,
+            pixels,
+            forest: SoftKdForest::new(),
+            width,
+            height,
+            x0,
+            y0,
+            len: 0,
+            deleted: 0,
+        }
+    }
+
+    fn pixel_index(&self, x: u32, y: u32) -> usize {
+        debug_assert!(x < self.width);
+        debug_assert!(y < self.height);
+
+        (x + y * self.width) as usize
+    }
+
+    fn free_neighbor(&mut self, x: u32, y: u32) -> Option<(u32, u32)> {
+        // Pick a pseudo-random neighbor
+        let offset: usize = self.rng.gen();
+
+        let neighbors = neighbors(x, y);
+        for i in 0..8 {
+            let (x, y) = neighbors[(i + offset) % 8];
+            if x < self.width && y < self.height {
+                let i = self.pixel_index(x, y);
+                if !self.pixels[i].filled {
+                    return Some((x, y));
+                }
+            }
+        }
+
+        None
+    }
+
+    fn fill(&mut self, x: u32, y: u32, color: C) -> Option<(u32, u32)> {
+        let i = self.pixel_index(x, y);
+        let pixel = &mut self.pixels[i];
+        if pixel.filled {
+            return None;
+        }
+
+        let rc = Rc::new(Pixel::new(x, y, color));
+        pixel.pixel = Some(rc.clone());
+        pixel.filled = true;
+
+        if self.free_neighbor(x, y).is_some() {
+            self.forest.push(rc);
+            self.len += 1;
+        }
+
+        for &(x, y) in &neighbors(x, y) {
+            if x < self.width && y < self.height && self.free_neighbor(x, y).is_none() {
+                let i = self.pixel_index(x, y);
+                if let Some(pixel) = self.pixels[i].pixel.take() {
+                    pixel.delete();
+                    self.deleted += 1;
+                }
+            }
+        }
+
+        if 2 * self.deleted >= self.len {
+            self.forest.rebuild();
+            self.len -= self.deleted;
+            self.deleted = 0;
+        }
+
+        Some((x, y))
+    }
+}
+
+impl<C: ColorSpace, R: Rng> Frontier for MinFrontier<C, R> {
+    fn width(&self) -> u32 {
+        self.width
+    }
+
+    fn height(&self) -> u32 {
+        self.height
+    }
+
+    fn len(&self) -> usize {
+        self.len - self.deleted
+    }
+
+    fn place(&mut self, rgb8: Rgb8) -> Option<(u32, u32)> {
+        let color = C::from(rgb8);
+        let (x, y) = self
+            .forest
+            .nearest(&color)
+            .map(|n| n.item.pos)
+            .map(|(x, y)| self.free_neighbor(x, y).unwrap())
+            .unwrap_or((self.x0, self.y0));
+
+        self.fill(x, y, color)
+    }
+}