inkr/src/rasterizer.rs

use core::f32;
use egui::{Color32, ColorImage, Pos2, Rect, Vec2, emath::TSTransform, epaint::Vertex};
use std::ops::Range;

pub trait BlendFn {
    fn blend(a: Color32, b: Color32) -> Color32;
}

pub mod blend {
    pub struct Normal;
    pub struct Add;
    pub struct Multiply;
}

/// Rasterize some triangles onto a new image,
///
/// Triangle positions must be in image-local point-coords.
/// `width` and `height` are in pixel coords.
pub fn rasterize<'a, Blend: BlendFn>(
    width: usize,
    height: usize,
    point_to_pixel: TSTransform,
    triangles: impl Iterator<Item = [&'a Vertex; 3]>,
) -> ColorImage {
    let mut image = ColorImage::new([width, height], Color32::TRANSPARENT);
    rasterize_onto::<Blend>(&mut image, point_to_pixel, triangles);
    image
}

/// Rasterize some triangles onto an image,
///
/// Triangle positions must be in image-local point-coords.
pub fn rasterize_onto<'a, Blend: BlendFn>(
    image: &mut ColorImage,
    point_to_pixel: TSTransform,
    triangles: impl Iterator<Item = [&'a Vertex; 3]>,
) {
    let width = image.width();
    let height = image.height();

    let mut set_pixel = |x: usize, y: usize, color| {
        let pixel = &mut image.pixels[y * width + x];
        *pixel = Blend::blend(*pixel, color);
    };

    let image_box = PxBoundingBox {
        x_from: 0,
        y_from: 0,
        x_to: width,
        y_to: height,
    };

    let pixel_to_point = point_to_pixel.inverse();

    for triangle in triangles {
        let [a, b, c] = triangle;
        if triangle_area(a.pos, b.pos, c.pos) == 0.0 {
            continue;
        }

        // Check all pixels within the triangle's bounding box.
        let bounding_box =
            triangle_bounding_box(&triangle, point_to_pixel).intersection(&image_box);

        // TODO: consider subdividing the triangle if it's very large.
        let pixels = pixels_in_box(bounding_box);

        for [x, y] in pixels {
            // Calculate point-coordinate of the pixel
            let pt_pos = pixel_to_point * Pos2::new(x as f32, y as f32);

            let point_in_triangle = point_in_triangle(pt_pos, triangle);

            // If the pixel is within the triangle, fill it in.
            if point_in_triangle.inside {
                let [c0, c1, c2] = [0, 1, 2].map(|i| {
                    triangle[i]
                        .color
                        .linear_multiply(point_in_triangle.weights[i])
                });

                let color = c0 + c1 + c2;

                set_pixel(x, y, color);
            }
        }
    }
}

/// Rasterize a single triangles onto an image,
///
/// Triangle positions must be in image-local point-coords.
pub fn rasterize_triangle_onto<'a, Blend: BlendFn>(
    image: &mut ColorImage,
    point_to_pixel: TSTransform,
    triangle: [&'a Vertex; 3],
) {
    rasterize_onto::<Blend>(image, point_to_pixel, [triangle].into_iter());
}

/// A bounding box, measured in pixels.
#[derive(Debug, PartialEq, Eq)]
pub struct PxBoundingBox {
    pub x_from: usize,
    pub y_from: usize,
    pub x_to: usize,
    pub y_to: usize,
}

impl PxBoundingBox {
    pub fn intersection(&self, other: &PxBoundingBox) -> PxBoundingBox {
        PxBoundingBox {
            x_from: self.x_from.max(other.x_from),
            y_from: self.y_from.max(other.y_from),
            x_to: self.x_to.min(other.x_to),
            y_to: self.y_to.min(other.y_to),
        }
    }

    pub fn union(&self, other: &PxBoundingBox) -> PxBoundingBox {
        PxBoundingBox {
            x_from: self.x_from.min(other.x_from),
            y_from: self.y_from.min(other.y_from),
            x_to: self.x_to.max(other.x_to),
            y_to: self.y_to.max(other.y_to),
        }
    }

    pub fn x_range(&self) -> Range<usize> {
        self.x_from..self.x_to
    }

    pub fn y_range(&self) -> Range<usize> {
        self.y_from..self.y_to
    }

    /// Test whether two boxes do NOT overlap
    pub fn overlaps_with(&self, other: &PxBoundingBox) -> bool {
        !self.is_disjoint_from(other)
    }

    pub fn is_disjoint_from(&self, other: &PxBoundingBox) -> bool {
        false
            || self.x_from > other.x_to
            || self.y_from > other.y_to
            || other.x_from > self.x_to
            || other.y_from > self.y_to
    }
}

pub fn triangle_bounding_box(
    triangle: &[&Vertex; 3],
    point_to_pixel: TSTransform,
) -> PxBoundingBox {
    // calculate bounding box in point coords
    let mut rect = Rect::NOTHING;
    for vertex in triangle {
        rect.min = rect.min.min(vertex.pos);
        rect.max = rect.max.max(vertex.pos);
    }

    // convert bounding box to pixel coords
    let rect = point_to_pixel.mul_rect(rect);
    PxBoundingBox {
        x_from: rect.min.x.floor() as usize,
        y_from: rect.min.y.floor() as usize,
        x_to: rect.max.x.ceil() as usize,
        y_to: rect.max.y.ceil() as usize,
    }
}

/// Calculate the perpendicular vector (90 degrees from the given vector)
fn perpendicular(v: Vec2) -> Vec2 {
    Vec2::new(v.y, -v.x)
}

#[derive(Clone, Debug)]
struct PointInTriangle {
    /// Is the point inside the triangle?
    inside: bool,

    /// Normalized weights describing the vicinity between the point and the three verticies of
    /// thre triangle.
    weights: [f32; 3],
}

/// Calculate whether a point is within a triangle, and the relative vicinities between the point
/// and each triangle vertex. The triangle must have a non-zero area.
fn point_in_triangle(point: Pos2, triangle: [&Vertex; 3]) -> PointInTriangle {
    let [a, b, c] = triangle;

    let sides = [[b, c], [c, a], [a, b]];

    // For each side of the triangle, imagine a new triangle consisting of the side and `point`.
    // Calculate the areas of those triangles.
    let areas = sides.map(|[start, end]| signed_triangle_area(start.pos, end.pos, point));

    // Use the areas to determine the side of the line at which the point exists.
    // If the area is positive, the point is on the right side of the triangle line.
    let [side_ab, side_bc, side_ca] = areas.map(|area| area >= 0.0);

    // Total area of the traingle.
    let triangle_area: f32 = areas.iter().sum();

    // egui does not wind the triangles in a consistent order, otherwise we might check if the
    // point is on a *specific* side of each line. As it is, we just check if the point is on the
    // same side of each line.
    let inside = side_ab == side_bc && side_bc == side_ca;

    // Normalize the weights.
    let weights = areas.map(|area| area / triangle_area);

    if cfg!(debug_assertions) && weights.into_iter().any(f32::is_nan) {
        panic!("weights must not be NaN! {weights:?} {triangle_area:?} {areas:?} {sides:?}");
    }

    PointInTriangle { inside, weights }
}

/// Calculate the area of a triangle.
fn triangle_area(a: Pos2, b: Pos2, c: Pos2) -> f32 {
    signed_triangle_area(a, b, c).abs()
}

/// Calculate the area of a triangle.
///
/// The area will be positive if the triangle is wound clockwise, and negative otherwise.
fn signed_triangle_area(a: Pos2, b: Pos2, c: Pos2) -> f32 {
    // Vector of an arbitrary "base" side of the triangle.
    let base = c - a;
    let base_perp = perpendicular(base);
    let diagonal = c - b;
    base_perp.dot(diagonal) / 2.0
}

/// Iterate over every pixel coordinate in a box.
#[inline(always)]
fn pixels_in_box(
    PxBoundingBox {
        x_from,
        y_from,
        x_to,
        y_to,
    }: PxBoundingBox,
) -> impl ExactSizeIterator<Item = [usize; 2]> {
    struct IterWithLen<I>(I, usize);
    impl<I: Iterator> ExactSizeIterator for IterWithLen<I> {}
    impl<I: Iterator> Iterator for IterWithLen<I> {
        type Item = I::Item;

        fn size_hint(&self) -> (usize, Option<usize>) {
            (self.1, Some(self.1))
        }

        fn next(&mut self) -> Option<Self::Item> {
            self.0.next()
        }
    }

    let len = (x_from..x_to).len() * (y_from..y_to).len();
    let iter = (x_from..x_to).flat_map(move |x| (y_from..y_to).map(move |y| [x, y]));

    debug_assert_eq!(len, iter.clone().count());

    IterWithLen(iter, len)
}

#[cfg(test)]
mod test {
    use egui::{Color32, Pos2, Vec2, emath::TSTransform, epaint::Vertex};

    use super::triangle_bounding_box;

    #[test]
    fn px_bounding_box() {
        let triangle = [
            Vertex {
                pos: Pos2::new(56.3, 18.9),
                uv: Default::default(),
                color: Color32::WHITE,
            },
            Vertex {
                pos: Pos2::new(56.4, 19.6),
                uv: Default::default(),
                color: Color32::WHITE,
            },
            Vertex {
                pos: Pos2::new(55.8, 20.5),
                uv: Default::default(),
                color: Color32::WHITE,
            },
        ];

        let pixels_per_point = 2.0;
        let point_to_pixel = TSTransform {
            scaling: pixels_per_point,
            translation: Vec2::new(-55.8, -18.9) * pixels_per_point,
        };

        let bounding_box = triangle_bounding_box(&triangle.each_ref(), point_to_pixel);

        insta::assert_debug_snapshot!((triangle, point_to_pixel, bounding_box));
    }
}

impl BlendFn for blend::Normal {
    fn blend(a: Color32, b: Color32) -> Color32 {
        a.blend(b)
    }
}

impl BlendFn for blend::Add {
    fn blend(a: Color32, b: Color32) -> Color32 {
        a + b
    }
}

impl BlendFn for blend::Multiply {
    fn blend(a: Color32, b: Color32) -> Color32 {
        a * b
    }
}