inkr/src/handwriting/canvas_rasterizer.rs

use std::sync::Arc;

use egui::{
    Color32, ColorImage, CornerRadius, Painter, Pos2, Rect, Stroke, StrokeKind, TextureHandle,
    Vec2,
    ahash::HashMap,
    emath::TSTransform,
    epaint::{Brush, RectShape, Vertex},
    load::SizedTexture,
};

use crate::rasterizer::{PxBoundingBox, rasterize_triangle_onto, triangle_bounding_box};

use super::StrokeBlendMode;

const CHUNK_SIZE: usize = 64;

/// Rasterize onto a resizeable canvas.
#[derive(Default)]
pub struct CanvasRasterizer {
    image_size: [usize; 2],
    tiles: HashMap<[usize; 2], Tile>,
}

struct Tile {
    bounding_box: PxBoundingBox,
    image: ColorImage,
    texture: Option<TextureHandle>,
    texture_is_dirty: bool,
}

impl Tile {
    fn new(xi: usize, yi: usize) -> Self {
        let x_from = xi * CHUNK_SIZE;
        let y_from = yi * CHUNK_SIZE;
        let bounding_box = PxBoundingBox {
            x_from,
            y_from,
            x_to: x_from + CHUNK_SIZE,
            y_to: y_from + CHUNK_SIZE,
        };

        Self {
            bounding_box,
            image: ColorImage::new([CHUNK_SIZE, CHUNK_SIZE], Color32::TRANSPARENT),
            texture: None,
            texture_is_dirty: false,
        }
    }
}

impl CanvasRasterizer {
    pub fn set_size(&mut self, width: usize, height: usize) {
        self.image_size = [width, height];
        self.populate_tiles();
    }

    pub fn clear(&mut self) {
        log::debug!("clearing all tiles");
        self.tiles.clear();
        self.populate_tiles();
    }

    fn populate_tiles(&mut self) {
        let [width, height] = self.image_size;

        // discard tiles that are out of bounds
        self.tiles.retain(|_, tile| {
            tile.bounding_box.x_from <= width && tile.bounding_box.y_from <= height
        });

        let chunk = |max: usize| {
            (0..)
                .step_by(CHUNK_SIZE)
                .take_while(move |n| n <= &max)
                .enumerate()
        };

        // create new tiles where we need them
        for (xi, _x) in chunk(width) {
            for (yi, _y) in chunk(height) {
                self.tiles
                    .entry([xi, yi])
                    .or_insert_with(|| Tile::new(xi, yi));
            }
        }
    }

    pub fn rasterize<'a>(
        &mut self,
        point_to_pixel: TSTransform,
        triangles: impl Iterator<Item = [&'a Vertex; 3]> + Clone,
    ) {
        for triangle in triangles {
            let triangle_bounding_box = triangle_bounding_box(&triangle, point_to_pixel);
            for chunk in chunks_from_bounding_box(triangle_bounding_box) {
                let Some(tile) = self.tiles.get_mut(&chunk) else {
                    continue;
                };

                let mut point_to_tile_pixel = point_to_pixel;
                point_to_tile_pixel.translation -= Vec2::new(
                    tile.bounding_box.x_from as f32,
                    tile.bounding_box.y_from as f32,
                );

                tile.texture_is_dirty = true;
                rasterize_triangle_onto::<StrokeBlendMode>(
                    &mut tile.image,
                    point_to_tile_pixel,
                    triangle,
                );
            }
        }
    }

    /// `at` defines the location in screen-coordinates where the canvas should be drawn.
    pub fn show(&mut self, ctx: &egui::Context, painter: &Painter, at: Rect) {
        let pixels_per_point = ctx.pixels_per_point();
        let chunk_vec = Vec2::splat(CHUNK_SIZE as f32) / pixels_per_point;

        for ([xi, yi], tile) in &mut self.tiles {
            if tile.texture_is_dirty {
                tile.texture_is_dirty = false;
                if let Some(texture) = &mut tile.texture {
                    texture.set(tile.image.clone(), Default::default());
                } else {
                    tile.texture = Some(ctx.load_texture(
                        "handwriting",
                        tile.image.clone(),
                        Default::default(),
                    ));
                }
            }

            if let Some(texture) = &mut tile.texture {
                let texture = SizedTexture::new(texture.id(), texture.size_vec2());
                let shape = RectShape {
                    rect: Rect::from_min_size(
                        at.min + Vec2::new(*xi as f32, *yi as f32) * chunk_vec,
                        chunk_vec,
                    ),
                    corner_radius: CornerRadius::ZERO,
                    fill: Color32::WHITE,
                    stroke: Stroke::NONE,
                    stroke_kind: StrokeKind::Inside,
                    round_to_pixels: None,
                    blur_width: 0.0,
                    brush: Some(Arc::new(Brush {
                        fill_texture_id: texture.id,
                        uv: Rect {
                            min: Pos2::ZERO,
                            max: Pos2::new(1.0, 1.0),
                        },
                    })),
                };
                painter.add(shape);
            }
        }
    }
}

/// Get all chunk indices that overlaps with a PxBoundingBox.
fn chunks_from_bounding_box(
    triangle_bounding_box: PxBoundingBox,
) -> impl Iterator<Item = [usize; 2]> {
    let x_from_chunk = triangle_bounding_box.x_from / CHUNK_SIZE;
    let y_from_chunk = triangle_bounding_box.y_from / CHUNK_SIZE;
    let x_to_chunk = triangle_bounding_box.x_to.saturating_sub(1) / CHUNK_SIZE;
    let y_to_chunk = triangle_bounding_box.y_to.saturating_sub(1) / CHUNK_SIZE;

    let xs = x_from_chunk..=x_to_chunk;
    let ys = y_from_chunk..=y_to_chunk;

    ys.flat_map(move |yi| xs.clone().map(move |xi| [xi, yi]))
}